mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Move model code from Rust monorepo to improve Rust compile times
This will need work done to support w2l.
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,11 @@
|
||||
README
|
||||
======
|
||||
|
||||
These models were moved out of the Rust monorepo on 2022-11-13.
|
||||
TTS is already accounted for here (under `tts/`), but we need to get W2L working again from this repository.
|
||||
|
||||
These are being removed to improve the build times of the Rust monorepo.
|
||||
|
||||
The idea is that these Docker images will build independently, move less frequently, and can be mounted into a
|
||||
k8s container as a volume or sidecar.
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
*.pkl
|
||||
*.jpg
|
||||
*.mp4
|
||||
*.pth
|
||||
*.pyc
|
||||
__pycache__
|
||||
*.h5
|
||||
*.avi
|
||||
*.wav
|
||||
filelists/*.txt
|
||||
evaluation/test_filelists/lr*.txt
|
||||
*.pyc
|
||||
*.mkv
|
||||
*.gif
|
||||
*.webm
|
||||
*.mp3
|
||||
@@ -0,0 +1,145 @@
|
||||
# **Wav2Lip**: *Accurately Lip-syncing Videos In The Wild*
|
||||
|
||||
This code is part of the paper: _A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild_ published at ACM Multimedia 2020.
|
||||
|
||||
[](https://paperswithcode.com/sota/lip-sync-on-lrs2?p=a-lip-sync-expert-is-all-you-need-for-speech)
|
||||
[](https://paperswithcode.com/sota/lip-sync-on-lrs3?p=a-lip-sync-expert-is-all-you-need-for-speech)
|
||||
[](https://paperswithcode.com/sota/lip-sync-on-lrw?p=a-lip-sync-expert-is-all-you-need-for-speech)
|
||||
|
||||
[[Paper]](http://arxiv.org/abs/2008.10010) | [[Project Page]](http://cvit.iiit.ac.in/research/projects/cvit-projects/a-lip-sync-expert-is-all-you-need-for-speech-to-lip-generation-in-the-wild/) | [[Demo Video]](https://youtu.be/0fXaDCZNOJc) | [[Interactive Demo]](https://bhaasha.iiit.ac.in/lipsync) | [[Colab Notebook]](https://colab.research.google.com/drive/1tZpDWXz49W6wDcTprANRGLo2D_EbD5J8?usp=sharing) | [[ReSyncED] (coming soon)](#)
|
||||
|
||||
<img src="https://drive.google.com/uc?export=view&id=1Wn0hPmpo4GRbCIJR8Tf20Akzdi1qjjG9"/>
|
||||
|
||||
----------
|
||||
**Highlights**
|
||||
----------
|
||||
- Lip-sync videos to any target speech with high accuracy. Try our [interactive demo](https://bhaasha.iiit.ac.in/lipsync).
|
||||
- Works for any identity, voice, and language. Also works for CGI faces and synthetic voices.
|
||||
- Complete training code, inference code, and pretrained models are available.
|
||||
- Or, quick-start with the Google Colab Notebook: [Link](https://colab.research.google.com/drive/1tZpDWXz49W6wDcTprANRGLo2D_EbD5J8?usp=sharing). Checkpoints and samples are available in a Google Drive [folder](https://drive.google.com/drive/folders/1I-0dNLfFOSFwrfqjNa-SXuwaURHE5K4k?usp=sharing) as well.
|
||||
- Several new, reliable evaluation benchmarks and metrics [[`evaluation/` folder of this repo]](https://github.com/Rudrabha/Wav2Lip/tree/master/evaluation) released.
|
||||
- Code to calculate metrics reported in the paper is also made available.
|
||||
|
||||
--------
|
||||
**Disclaimer**
|
||||
--------
|
||||
All results from this open-source code or our [demo website](https://bhaasha.iiit.ac.in/lipsync) should only be used for research/academic/personal purposes only. As the models are trained on the <a href="http://www.robots.ox.ac.uk/~vgg/data/lip_reading/lrs2.html">LRS2 dataset</a>, any form of commercial use is strictly prohibhited. Please contact us for all further queries.
|
||||
|
||||
Prerequisites
|
||||
-------------
|
||||
- `Python 3.5.2` (code has been tested with this version at our end, but several other users say that `3.6+` is the one that works instead.)
|
||||
- ffmpeg: `sudo apt-get install ffmpeg`
|
||||
- Install necessary packages using `pip install -r requirements.txt`
|
||||
- Face detection [pre-trained model](https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth) should be downloaded to `face_detection/detection/sfd/s3fd.pth`. Alternative [link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/prajwal_k_research_iiit_ac_in/EZsy6qWuivtDnANIG73iHjIBjMSoojcIV0NULXV-yiuiIg?e=qTasa8) if the above does not work.
|
||||
|
||||
Getting the weights
|
||||
----------
|
||||
| Model | Description | Link to the model |
|
||||
| :-------------: | :---------------: | :---------------: |
|
||||
| Wav2Lip | Highly accurate lip-sync | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/Eb3LEzbfuKlJiR600lQWRxgBIY27JZg80f7V9jtMfbNDaQ?e=TBFBVW) |
|
||||
| Wav2Lip + GAN | Slightly inferior lip-sync, but better visual quality | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EdjI7bZlgApMqsVoEUUXpLsBxqXbn5z8VTmoxp55YNDcIA?e=n9ljGW) |
|
||||
| Expert Discriminator | Weights of the expert discriminator | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EQRvmiZg-HRAjvI6zqN9eTEBP74KefynCwPWVmF57l-AYA?e=ZRPHKP) |
|
||||
|
||||
Lip-syncing videos using the pre-trained models (Inference)
|
||||
-------
|
||||
You can lip-sync any video to any audio:
|
||||
```bash
|
||||
python inference.py --checkpoint_path <ckpt> --face <video.mp4> --audio <an-audio-source>
|
||||
```
|
||||
The result is saved (by default) in `results/result_voice.mp4`. You can specify it as an argument, similar to several other available options. The audio source can be any file supported by `FFMPEG` containing audio data: `*.wav`, `*.mp3` or even a video file, from which the code will automatically extract the audio.
|
||||
|
||||
##### Tips for better results:
|
||||
- Experiment with the `--pads` argument to adjust the detected face bounding box. Often leads to improved results. You might need to increase the bottom padding to include the chin region. E.g. `--pads 0 20 0 0`.
|
||||
- If you see the mouth position dislocated or some weird artifacts such as two mouths, then it can be because of over-smoothing the face detections. Use the `--nosmooth` argument and give another try.
|
||||
- Experiment with the `--resize_factor` argument, to get a lower resolution video. Why? The models are trained on faces which were at a lower resolution. You might get better, visually pleasing results for 720p videos than for 1080p videos (in many cases, the latter works well too).
|
||||
- The Wav2Lip model without GAN usually needs more experimenting with the above two to get the most ideal results, and sometimes, can give you a better result as well.
|
||||
|
||||
Preparing LRS2 for training
|
||||
----------
|
||||
Our models are trained on LRS2. See [here](#training-on-datasets-other-than-lrs2) for a few suggestions regarding training on other datasets.
|
||||
##### LRS2 dataset folder structure
|
||||
|
||||
```
|
||||
data_root (mvlrs_v1)
|
||||
├── main, pretrain (we use only main folder in this work)
|
||||
| ├── list of folders
|
||||
| │ ├── five-digit numbered video IDs ending with (.mp4)
|
||||
```
|
||||
|
||||
Place the LRS2 filelists (train, val, test) `.txt` files in the `filelists/` folder.
|
||||
|
||||
##### Preprocess the dataset for fast training
|
||||
|
||||
```bash
|
||||
python preprocess.py --data_root data_root/main --preprocessed_root lrs2_preprocessed/
|
||||
```
|
||||
Additional options like `batch_size` and number of GPUs to use in parallel to use can also be set.
|
||||
|
||||
##### Preprocessed LRS2 folder structure
|
||||
```
|
||||
preprocessed_root (lrs2_preprocessed)
|
||||
├── list of folders
|
||||
| ├── Folders with five-digit numbered video IDs
|
||||
| │ ├── *.jpg
|
||||
| │ ├── audio.wav
|
||||
```
|
||||
|
||||
Train!
|
||||
----------
|
||||
There are two major steps: (i) Train the expert lip-sync discriminator, (ii) Train the Wav2Lip model(s).
|
||||
|
||||
##### Training the expert discriminator
|
||||
You can download [the pre-trained weights](#getting-the-weights) if you want to skip this step. To train it:
|
||||
```bash
|
||||
python color_syncnet_train.py --data_root lrs2_preprocessed/ --checkpoint_dir <folder_to_save_checkpoints>
|
||||
```
|
||||
##### Training the Wav2Lip models
|
||||
You can either train the model without the additional visual quality disriminator (< 1 day of training) or use the discriminator (~2 days). For the former, run:
|
||||
```bash
|
||||
python wav2lip_train.py --data_root lrs2_preprocessed/ --checkpoint_dir <folder_to_save_checkpoints> --syncnet_checkpoint_path <path_to_expert_disc_checkpoint>
|
||||
```
|
||||
|
||||
To train with the visual quality discriminator, you should run `hq_wav2lip_train.py` instead. The arguments for both the files are similar. In both the cases, you can resume training as well. Look at `python wav2lip_train.py --help` for more details. You can also set additional less commonly-used hyper-parameters at the bottom of the `hparams.py` file.
|
||||
|
||||
Training on datasets other than LRS2
|
||||
------------------------------------
|
||||
Training on other datasets might require modifications to the code. Please read the following before you raise an issue:
|
||||
|
||||
- You might not get good results by training/fine-tuning on a few minutes of a single speaker. This is a separate research problem, to which we do not have a solution yet. Thus, we would most likely not be able to resolve your issue.
|
||||
- You must train the expert discriminator for your own dataset before training Wav2Lip.
|
||||
- If it is your own dataset downloaded from the web, in most cases, needs to be sync-corrected.
|
||||
- Be mindful of the FPS of the videos of your dataset. Changes to FPS would need significant code changes.
|
||||
- The expert discriminator's eval loss should go down to ~0.25 and the Wav2Lip eval sync loss should go down to ~0.2 to get good results.
|
||||
|
||||
When raising an issue on this topic, please let us know that you are aware of all these points.
|
||||
|
||||
Evaluation
|
||||
----------
|
||||
Will be updated.
|
||||
|
||||
License and Citation
|
||||
----------
|
||||
The software can only be used for personal/research/non-commercial purposes. Please cite the following paper if you have use this code:
|
||||
```
|
||||
@inproceedings{10.1145/3394171.3413532,
|
||||
author = {Prajwal, K R and Mukhopadhyay, Rudrabha and Namboodiri, Vinay P. and Jawahar, C.V.},
|
||||
title = {A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild},
|
||||
year = {2020},
|
||||
isbn = {9781450379885},
|
||||
publisher = {Association for Computing Machinery},
|
||||
address = {New York, NY, USA},
|
||||
url = {https://doi.org/10.1145/3394171.3413532},
|
||||
doi = {10.1145/3394171.3413532},
|
||||
booktitle = {Proceedings of the 28th ACM International Conference on Multimedia},
|
||||
pages = {484–492},
|
||||
numpages = {9},
|
||||
keywords = {lip sync, talking face generation, video generation},
|
||||
location = {Seattle, WA, USA},
|
||||
series = {MM '20}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
Acknowledgements
|
||||
----------
|
||||
Parts of the code structure is inspired by this [TTS repository](https://github.com/r9y9/deepvoice3_pytorch). We thank the author for this wonderful code. The code for Face Detection has been taken from the [face_alignment](https://github.com/1adrianb/face-alignment) repository. We thank the authors for releasing their code and models.
|
||||
@@ -0,0 +1,136 @@
|
||||
import librosa
|
||||
import librosa.filters
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from scipy import signal
|
||||
from scipy.io import wavfile
|
||||
from hparams import hparams as hp
|
||||
|
||||
def load_wav(path, sr):
|
||||
return librosa.core.load(path, sr=sr)[0]
|
||||
|
||||
def save_wav(wav, path, sr):
|
||||
wav *= 32767 / max(0.01, np.max(np.abs(wav)))
|
||||
#proposed by @dsmiller
|
||||
wavfile.write(path, sr, wav.astype(np.int16))
|
||||
|
||||
def save_wavenet_wav(wav, path, sr):
|
||||
librosa.output.write_wav(path, wav, sr=sr)
|
||||
|
||||
def preemphasis(wav, k, preemphasize=True):
|
||||
if preemphasize:
|
||||
return signal.lfilter([1, -k], [1], wav)
|
||||
return wav
|
||||
|
||||
def inv_preemphasis(wav, k, inv_preemphasize=True):
|
||||
if inv_preemphasize:
|
||||
return signal.lfilter([1], [1, -k], wav)
|
||||
return wav
|
||||
|
||||
def get_hop_size():
|
||||
hop_size = hp.hop_size
|
||||
if hop_size is None:
|
||||
assert hp.frame_shift_ms is not None
|
||||
hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate)
|
||||
return hop_size
|
||||
|
||||
def linearspectrogram(wav):
|
||||
D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
|
||||
S = _amp_to_db(np.abs(D)) - hp.ref_level_db
|
||||
|
||||
if hp.signal_normalization:
|
||||
return _normalize(S)
|
||||
return S
|
||||
|
||||
def melspectrogram(wav):
|
||||
D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
|
||||
S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db
|
||||
|
||||
if hp.signal_normalization:
|
||||
return _normalize(S)
|
||||
return S
|
||||
|
||||
def _lws_processor():
|
||||
import lws
|
||||
return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech")
|
||||
|
||||
def _stft(y):
|
||||
if hp.use_lws:
|
||||
return _lws_processor(hp).stft(y).T
|
||||
else:
|
||||
return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=get_hop_size(), win_length=hp.win_size)
|
||||
|
||||
##########################################################
|
||||
#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
|
||||
def num_frames(length, fsize, fshift):
|
||||
"""Compute number of time frames of spectrogram
|
||||
"""
|
||||
pad = (fsize - fshift)
|
||||
if length % fshift == 0:
|
||||
M = (length + pad * 2 - fsize) // fshift + 1
|
||||
else:
|
||||
M = (length + pad * 2 - fsize) // fshift + 2
|
||||
return M
|
||||
|
||||
|
||||
def pad_lr(x, fsize, fshift):
|
||||
"""Compute left and right padding
|
||||
"""
|
||||
M = num_frames(len(x), fsize, fshift)
|
||||
pad = (fsize - fshift)
|
||||
T = len(x) + 2 * pad
|
||||
r = (M - 1) * fshift + fsize - T
|
||||
return pad, pad + r
|
||||
##########################################################
|
||||
#Librosa correct padding
|
||||
def librosa_pad_lr(x, fsize, fshift):
|
||||
return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
|
||||
|
||||
# Conversions
|
||||
_mel_basis = None
|
||||
|
||||
def _linear_to_mel(spectogram):
|
||||
global _mel_basis
|
||||
if _mel_basis is None:
|
||||
_mel_basis = _build_mel_basis()
|
||||
return np.dot(_mel_basis, spectogram)
|
||||
|
||||
def _build_mel_basis():
|
||||
assert hp.fmax <= hp.sample_rate // 2
|
||||
return librosa.filters.mel(hp.sample_rate, hp.n_fft, n_mels=hp.num_mels,
|
||||
fmin=hp.fmin, fmax=hp.fmax)
|
||||
|
||||
def _amp_to_db(x):
|
||||
min_level = np.exp(hp.min_level_db / 20 * np.log(10))
|
||||
return 20 * np.log10(np.maximum(min_level, x))
|
||||
|
||||
def _db_to_amp(x):
|
||||
return np.power(10.0, (x) * 0.05)
|
||||
|
||||
def _normalize(S):
|
||||
if hp.allow_clipping_in_normalization:
|
||||
if hp.symmetric_mels:
|
||||
return np.clip((2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value,
|
||||
-hp.max_abs_value, hp.max_abs_value)
|
||||
else:
|
||||
return np.clip(hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)), 0, hp.max_abs_value)
|
||||
|
||||
assert S.max() <= 0 and S.min() - hp.min_level_db >= 0
|
||||
if hp.symmetric_mels:
|
||||
return (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value
|
||||
else:
|
||||
return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db))
|
||||
|
||||
def _denormalize(D):
|
||||
if hp.allow_clipping_in_normalization:
|
||||
if hp.symmetric_mels:
|
||||
return (((np.clip(D, -hp.max_abs_value,
|
||||
hp.max_abs_value) + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value))
|
||||
+ hp.min_level_db)
|
||||
else:
|
||||
return ((np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
|
||||
|
||||
if hp.symmetric_mels:
|
||||
return (((D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + hp.min_level_db)
|
||||
else:
|
||||
return ((D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
|
||||
@@ -0,0 +1 @@
|
||||
Place all your checkpoints (.pth files) here.
|
||||
@@ -0,0 +1,279 @@
|
||||
from os.path import dirname, join, basename, isfile
|
||||
from tqdm import tqdm
|
||||
|
||||
from models import SyncNet_color as SyncNet
|
||||
import audio
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch import optim
|
||||
import torch.backends.cudnn as cudnn
|
||||
from torch.utils import data as data_utils
|
||||
import numpy as np
|
||||
|
||||
from glob import glob
|
||||
|
||||
import os, random, cv2, argparse
|
||||
from hparams import hparams, get_image_list
|
||||
|
||||
parser = argparse.ArgumentParser(description='Code to train the expert lip-sync discriminator')
|
||||
|
||||
parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True)
|
||||
|
||||
parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str)
|
||||
parser.add_argument('--checkpoint_path', help='Resumed from this checkpoint', default=None, type=str)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
global_step = 0
|
||||
global_epoch = 0
|
||||
use_cuda = torch.cuda.is_available()
|
||||
print('use_cuda: {}'.format(use_cuda))
|
||||
|
||||
syncnet_T = 5
|
||||
syncnet_mel_step_size = 16
|
||||
|
||||
class Dataset(object):
|
||||
def __init__(self, split):
|
||||
self.all_videos = get_image_list(args.data_root, split)
|
||||
|
||||
def get_frame_id(self, frame):
|
||||
return int(basename(frame).split('.')[0])
|
||||
|
||||
def get_window(self, start_frame):
|
||||
start_id = self.get_frame_id(start_frame)
|
||||
vidname = dirname(start_frame)
|
||||
|
||||
window_fnames = []
|
||||
for frame_id in range(start_id, start_id + syncnet_T):
|
||||
frame = join(vidname, '{}.jpg'.format(frame_id))
|
||||
if not isfile(frame):
|
||||
return None
|
||||
window_fnames.append(frame)
|
||||
return window_fnames
|
||||
|
||||
def crop_audio_window(self, spec, start_frame):
|
||||
# num_frames = (T x hop_size * fps) / sample_rate
|
||||
start_frame_num = self.get_frame_id(start_frame)
|
||||
start_idx = int(80. * (start_frame_num / float(hparams.fps)))
|
||||
|
||||
end_idx = start_idx + syncnet_mel_step_size
|
||||
|
||||
return spec[start_idx : end_idx, :]
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_videos)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
while 1:
|
||||
idx = random.randint(0, len(self.all_videos) - 1)
|
||||
vidname = self.all_videos[idx]
|
||||
|
||||
img_names = list(glob(join(vidname, '*.jpg')))
|
||||
if len(img_names) <= 3 * syncnet_T:
|
||||
continue
|
||||
img_name = random.choice(img_names)
|
||||
wrong_img_name = random.choice(img_names)
|
||||
while wrong_img_name == img_name:
|
||||
wrong_img_name = random.choice(img_names)
|
||||
|
||||
if random.choice([True, False]):
|
||||
y = torch.ones(1).float()
|
||||
chosen = img_name
|
||||
else:
|
||||
y = torch.zeros(1).float()
|
||||
chosen = wrong_img_name
|
||||
|
||||
window_fnames = self.get_window(chosen)
|
||||
if window_fnames is None:
|
||||
continue
|
||||
|
||||
window = []
|
||||
all_read = True
|
||||
for fname in window_fnames:
|
||||
img = cv2.imread(fname)
|
||||
if img is None:
|
||||
all_read = False
|
||||
break
|
||||
try:
|
||||
img = cv2.resize(img, (hparams.img_size, hparams.img_size))
|
||||
except Exception as e:
|
||||
all_read = False
|
||||
break
|
||||
|
||||
window.append(img)
|
||||
|
||||
if not all_read: continue
|
||||
|
||||
try:
|
||||
wavpath = join(vidname, "audio.wav")
|
||||
wav = audio.load_wav(wavpath, hparams.sample_rate)
|
||||
|
||||
orig_mel = audio.melspectrogram(wav).T
|
||||
except Exception as e:
|
||||
continue
|
||||
|
||||
mel = self.crop_audio_window(orig_mel.copy(), img_name)
|
||||
|
||||
if (mel.shape[0] != syncnet_mel_step_size):
|
||||
continue
|
||||
|
||||
# H x W x 3 * T
|
||||
x = np.concatenate(window, axis=2) / 255.
|
||||
x = x.transpose(2, 0, 1)
|
||||
x = x[:, x.shape[1]//2:]
|
||||
|
||||
x = torch.FloatTensor(x)
|
||||
mel = torch.FloatTensor(mel.T).unsqueeze(0)
|
||||
|
||||
return x, mel, y
|
||||
|
||||
logloss = nn.BCELoss()
|
||||
def cosine_loss(a, v, y):
|
||||
d = nn.functional.cosine_similarity(a, v)
|
||||
loss = logloss(d.unsqueeze(1), y)
|
||||
|
||||
return loss
|
||||
|
||||
def train(device, model, train_data_loader, test_data_loader, optimizer,
|
||||
checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
|
||||
|
||||
global global_step, global_epoch
|
||||
resumed_step = global_step
|
||||
|
||||
while global_epoch < nepochs:
|
||||
running_loss = 0.
|
||||
prog_bar = tqdm(enumerate(train_data_loader))
|
||||
for step, (x, mel, y) in prog_bar:
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Transform data to CUDA device
|
||||
x = x.to(device)
|
||||
|
||||
mel = mel.to(device)
|
||||
|
||||
a, v = model(mel, x)
|
||||
y = y.to(device)
|
||||
|
||||
loss = cosine_loss(a, v, y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
global_step += 1
|
||||
cur_session_steps = global_step - resumed_step
|
||||
running_loss += loss.item()
|
||||
|
||||
if global_step == 1 or global_step % checkpoint_interval == 0:
|
||||
save_checkpoint(
|
||||
model, optimizer, global_step, checkpoint_dir, global_epoch)
|
||||
|
||||
if global_step % hparams.syncnet_eval_interval == 0:
|
||||
with torch.no_grad():
|
||||
eval_model(test_data_loader, global_step, device, model, checkpoint_dir)
|
||||
|
||||
prog_bar.set_description('Loss: {}'.format(running_loss / (step + 1)))
|
||||
|
||||
global_epoch += 1
|
||||
|
||||
def eval_model(test_data_loader, global_step, device, model, checkpoint_dir):
|
||||
eval_steps = 1400
|
||||
print('Evaluating for {} steps'.format(eval_steps))
|
||||
losses = []
|
||||
while 1:
|
||||
for step, (x, mel, y) in enumerate(test_data_loader):
|
||||
|
||||
model.eval()
|
||||
|
||||
# Transform data to CUDA device
|
||||
x = x.to(device)
|
||||
|
||||
mel = mel.to(device)
|
||||
|
||||
a, v = model(mel, x)
|
||||
y = y.to(device)
|
||||
|
||||
loss = cosine_loss(a, v, y)
|
||||
losses.append(loss.item())
|
||||
|
||||
if step > eval_steps: break
|
||||
|
||||
averaged_loss = sum(losses) / len(losses)
|
||||
print(averaged_loss)
|
||||
|
||||
return
|
||||
|
||||
def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch):
|
||||
|
||||
checkpoint_path = join(
|
||||
checkpoint_dir, "checkpoint_step{:09d}.pth".format(global_step))
|
||||
optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None
|
||||
torch.save({
|
||||
"state_dict": model.state_dict(),
|
||||
"optimizer": optimizer_state,
|
||||
"global_step": step,
|
||||
"global_epoch": epoch,
|
||||
}, checkpoint_path)
|
||||
print("Saved checkpoint:", checkpoint_path)
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if use_cuda:
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_checkpoint(path, model, optimizer, reset_optimizer=False):
|
||||
global global_step
|
||||
global global_epoch
|
||||
|
||||
print("Load checkpoint from: {}".format(path))
|
||||
checkpoint = _load(path)
|
||||
model.load_state_dict(checkpoint["state_dict"])
|
||||
if not reset_optimizer:
|
||||
optimizer_state = checkpoint["optimizer"]
|
||||
if optimizer_state is not None:
|
||||
print("Load optimizer state from {}".format(path))
|
||||
optimizer.load_state_dict(checkpoint["optimizer"])
|
||||
global_step = checkpoint["global_step"]
|
||||
global_epoch = checkpoint["global_epoch"]
|
||||
|
||||
return model
|
||||
|
||||
if __name__ == "__main__":
|
||||
checkpoint_dir = args.checkpoint_dir
|
||||
checkpoint_path = args.checkpoint_path
|
||||
|
||||
if not os.path.exists(checkpoint_dir): os.mkdir(checkpoint_dir)
|
||||
|
||||
# Dataset and Dataloader setup
|
||||
train_dataset = Dataset('train')
|
||||
test_dataset = Dataset('val')
|
||||
|
||||
train_data_loader = data_utils.DataLoader(
|
||||
train_dataset, batch_size=hparams.syncnet_batch_size, shuffle=True,
|
||||
num_workers=hparams.num_workers)
|
||||
|
||||
test_data_loader = data_utils.DataLoader(
|
||||
test_dataset, batch_size=hparams.syncnet_batch_size,
|
||||
num_workers=8)
|
||||
|
||||
device = torch.device("cuda" if use_cuda else "cpu")
|
||||
|
||||
# Model
|
||||
model = SyncNet().to(device)
|
||||
print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad)))
|
||||
|
||||
optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
|
||||
lr=hparams.syncnet_lr)
|
||||
|
||||
if checkpoint_path is not None:
|
||||
load_checkpoint(checkpoint_path, model, optimizer, reset_optimizer=False)
|
||||
|
||||
train(device, model, train_data_loader, test_data_loader, optimizer,
|
||||
checkpoint_dir=checkpoint_dir,
|
||||
checkpoint_interval=hparams.syncnet_checkpoint_interval,
|
||||
nepochs=hparams.nepochs)
|
||||
@@ -0,0 +1,59 @@
|
||||
# Evaluation of Lip-sync using LSE-D and LSE-C metric.
|
||||
|
||||
We use the pre-trained syncnet model available in this [repository](https://github.com/joonson/syncnet_python).
|
||||
|
||||
### Steps to set-up the evaluation repository for LSE-D and LSE-C metric:
|
||||
* Clone the SyncNet repository.
|
||||
```
|
||||
git clone https://github.com/joonson/syncnet_python.git
|
||||
```
|
||||
* Follow the procedure given in the above linked [repository](https://github.com/joonson/syncnet_python) to download the pretrained models and set up the dependencies.
|
||||
* **Note: Please install a separate virtual environment for the evaluation scripts. The versions used by Wav2Lip and the publicly released code of SyncNet is different and can cause version mis-match issues. To avoid this, we suggest the users to install a separate virtual environment for the evaluation scripts**
|
||||
```
|
||||
cd syncnet_python
|
||||
pip install -r requirements.txt
|
||||
sh download_model.sh
|
||||
```
|
||||
* The above step should ensure that all the dependencies required by the repository is installed and the pre-trained models are downloaded.
|
||||
|
||||
### Running the evaluation scripts:
|
||||
* Copy our evaluation scripts given in this folder to the cloned repository.
|
||||
```
|
||||
cd Wav2Lip/evaluation/scores_LSE/
|
||||
cp *.py syncnet_python/
|
||||
cp *.sh syncnet_python/
|
||||
```
|
||||
**Note: We will release the test filelists for LRW, LRS2 and LRS3 shortly once we receive permission from the dataset creators. We will also release the Real World Dataset we have collected shortly.**
|
||||
|
||||
* Our evaluation technique does not require ground-truth of any sorts. Given lip-synced videos we can directly calculate the scores from only the generated videos. Please store the generated videos (from our test sets or your own generated videos) in the following folder structure.
|
||||
```
|
||||
video data root (Folder containing all videos)
|
||||
├── All .mp4 files
|
||||
```
|
||||
* Change the folder back to the cloned repository.
|
||||
```
|
||||
cd syncnet_python
|
||||
```
|
||||
* To run evaluation on the LRW, LRS2 and LRS3 test files, please run the following command:
|
||||
```
|
||||
python calculate_scores_LRS.py --data_root /path/to/video/data/root --tmp_dir tmp_dir/
|
||||
```
|
||||
|
||||
* To run evaluation on the ReSynced dataset or your own generated videos, please run the following command:
|
||||
```
|
||||
sh calculate_scores_real_videos.sh /path/to/video/data/root
|
||||
```
|
||||
* The generated scores will be present in the all_scores.txt generated in the ```syncnet_python/``` folder
|
||||
|
||||
# Evaluation of image quality using FID metric.
|
||||
We will update the scripts for the FID metric shortly
|
||||
|
||||
|
||||
# Opening issues related to evaluation scripts
|
||||
* Please open the issues with the "Evaluation" label if you face any issues in the evaluation scripts.
|
||||
|
||||
# Acknowledgements
|
||||
Our evaluation pipeline in based on two existing repositories. LSE metrics are based on the [syncnet_python](https://github.com/joonson/syncnet_python) repository and the FID score is based on [pytorch-fid](https://github.com/mseitzer/pytorch-fid) repository. We thank the authors of both the repositories for releasing their wonderful code.
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
from os import listdir, path
|
||||
import numpy as np
|
||||
import scipy, cv2, os, sys, argparse
|
||||
import dlib, json, subprocess
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch
|
||||
|
||||
sys.path.append('../')
|
||||
import audio
|
||||
import face_detection
|
||||
from models import Wav2Lip
|
||||
|
||||
parser = argparse.ArgumentParser(description='Code to generate results for test filelists')
|
||||
|
||||
parser.add_argument('--filelist', type=str,
|
||||
help='Filepath of filelist file to read', required=True)
|
||||
parser.add_argument('--results_dir', type=str, help='Folder to save all results into',
|
||||
required=True)
|
||||
parser.add_argument('--data_root', type=str, required=True)
|
||||
parser.add_argument('--checkpoint_path', type=str,
|
||||
help='Name of saved checkpoint to load weights from', required=True)
|
||||
|
||||
parser.add_argument('--pads', nargs='+', type=int, default=[0, 0, 0, 0],
|
||||
help='Padding (top, bottom, left, right)')
|
||||
parser.add_argument('--face_det_batch_size', type=int,
|
||||
help='Single GPU batch size for face detection', default=64)
|
||||
parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip', default=128)
|
||||
|
||||
# parser.add_argument('--resize_factor', default=1, type=int)
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = 96
|
||||
|
||||
def get_smoothened_boxes(boxes, T):
|
||||
for i in range(len(boxes)):
|
||||
if i + T > len(boxes):
|
||||
window = boxes[len(boxes) - T:]
|
||||
else:
|
||||
window = boxes[i : i + T]
|
||||
boxes[i] = np.mean(window, axis=0)
|
||||
return boxes
|
||||
|
||||
def face_detect(images):
|
||||
batch_size = args.face_det_batch_size
|
||||
|
||||
while 1:
|
||||
predictions = []
|
||||
try:
|
||||
for i in range(0, len(images), batch_size):
|
||||
predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
|
||||
except RuntimeError:
|
||||
if batch_size == 1:
|
||||
raise RuntimeError('Image too big to run face detection on GPU')
|
||||
batch_size //= 2
|
||||
args.face_det_batch_size = batch_size
|
||||
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
||||
continue
|
||||
break
|
||||
|
||||
results = []
|
||||
pady1, pady2, padx1, padx2 = args.pads
|
||||
for rect, image in zip(predictions, images):
|
||||
if rect is None:
|
||||
raise ValueError('Face not detected!')
|
||||
|
||||
y1 = max(0, rect[1] - pady1)
|
||||
y2 = min(image.shape[0], rect[3] + pady2)
|
||||
x1 = max(0, rect[0] - padx1)
|
||||
x2 = min(image.shape[1], rect[2] + padx2)
|
||||
|
||||
results.append([x1, y1, x2, y2])
|
||||
|
||||
boxes = get_smoothened_boxes(np.array(results), T=5)
|
||||
results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
|
||||
|
||||
return results
|
||||
|
||||
def datagen(frames, face_det_results, mels):
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
for i, m in enumerate(mels):
|
||||
if i >= len(frames): raise ValueError('Equal or less lengths only')
|
||||
|
||||
frame_to_save = frames[i].copy()
|
||||
face, coords, valid_frame = face_det_results[i].copy()
|
||||
if not valid_frame:
|
||||
continue
|
||||
|
||||
face = cv2.resize(face, (args.img_size, args.img_size))
|
||||
|
||||
img_batch.append(face)
|
||||
mel_batch.append(m)
|
||||
frame_batch.append(frame_to_save)
|
||||
coords_batch.append(coords)
|
||||
|
||||
if len(img_batch) >= args.wav2lip_batch_size:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if len(img_batch) > 0:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
|
||||
fps = 25
|
||||
mel_step_size = 16
|
||||
mel_idx_multiplier = 80./fps
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
print('Using {} for inference.'.format(device))
|
||||
|
||||
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
||||
flip_input=False, device=device)
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if device == 'cuda':
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_model(path):
|
||||
model = Wav2Lip()
|
||||
print("Load checkpoint from: {}".format(path))
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
|
||||
model = model.to(device)
|
||||
return model.eval()
|
||||
|
||||
model = load_model(args.checkpoint_path)
|
||||
|
||||
def main():
|
||||
assert args.data_root is not None
|
||||
data_root = args.data_root
|
||||
|
||||
if not os.path.isdir(args.results_dir): os.makedirs(args.results_dir)
|
||||
|
||||
with open(args.filelist, 'r') as filelist:
|
||||
lines = filelist.readlines()
|
||||
|
||||
for idx, line in enumerate(tqdm(lines)):
|
||||
audio_src, video = line.strip().split()
|
||||
|
||||
audio_src = os.path.join(data_root, audio_src) + '.mp4'
|
||||
video = os.path.join(data_root, video) + '.mp4'
|
||||
|
||||
command = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'.format(audio_src, '../temp/temp.wav')
|
||||
subprocess.call(command, shell=True)
|
||||
temp_audio = '../temp/temp.wav'
|
||||
|
||||
wav = audio.load_wav(temp_audio, 16000)
|
||||
mel = audio.melspectrogram(wav)
|
||||
if np.isnan(mel.reshape(-1)).sum() > 0:
|
||||
continue
|
||||
|
||||
mel_chunks = []
|
||||
i = 0
|
||||
while 1:
|
||||
start_idx = int(i * mel_idx_multiplier)
|
||||
if start_idx + mel_step_size > len(mel[0]):
|
||||
break
|
||||
mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
|
||||
i += 1
|
||||
|
||||
video_stream = cv2.VideoCapture(video)
|
||||
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading or len(full_frames) > len(mel_chunks):
|
||||
video_stream.release()
|
||||
break
|
||||
full_frames.append(frame)
|
||||
|
||||
if len(full_frames) < len(mel_chunks):
|
||||
continue
|
||||
|
||||
full_frames = full_frames[:len(mel_chunks)]
|
||||
|
||||
try:
|
||||
face_det_results = face_detect(full_frames.copy())
|
||||
except ValueError as e:
|
||||
continue
|
||||
|
||||
batch_size = args.wav2lip_batch_size
|
||||
gen = datagen(full_frames.copy(), face_det_results, mel_chunks)
|
||||
|
||||
for i, (img_batch, mel_batch, frames, coords) in enumerate(gen):
|
||||
if i == 0:
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
out = cv2.VideoWriter('../temp/result.avi',
|
||||
cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
|
||||
|
||||
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
|
||||
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(mel_batch, img_batch)
|
||||
|
||||
|
||||
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
|
||||
|
||||
for pl, f, c in zip(pred, frames, coords):
|
||||
y1, y2, x1, x2 = c
|
||||
pl = cv2.resize(pl.astype(np.uint8), (x2 - x1, y2 - y1))
|
||||
f[y1:y2, x1:x2] = pl
|
||||
out.write(f)
|
||||
|
||||
out.release()
|
||||
|
||||
vid = os.path.join(args.results_dir, '{}.mp4'.format(idx))
|
||||
|
||||
command = 'ffmpeg -loglevel panic -y -i {} -i {} -strict -2 -q:v 1 {}'.format(temp_audio,
|
||||
'../temp/result.avi', vid)
|
||||
subprocess.call(command, shell=True)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,305 @@
|
||||
from os import listdir, path
|
||||
import numpy as np
|
||||
import scipy, cv2, os, sys, argparse
|
||||
import dlib, json, subprocess
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch
|
||||
|
||||
sys.path.append('../')
|
||||
import audio
|
||||
import face_detection
|
||||
from models import Wav2Lip
|
||||
|
||||
parser = argparse.ArgumentParser(description='Code to generate results on ReSyncED evaluation set')
|
||||
|
||||
parser.add_argument('--mode', type=str,
|
||||
help='random | dubbed | tts', required=True)
|
||||
|
||||
parser.add_argument('--filelist', type=str,
|
||||
help='Filepath of filelist file to read', default=None)
|
||||
|
||||
parser.add_argument('--results_dir', type=str, help='Folder to save all results into',
|
||||
required=True)
|
||||
parser.add_argument('--data_root', type=str, required=True)
|
||||
parser.add_argument('--checkpoint_path', type=str,
|
||||
help='Name of saved checkpoint to load weights from', required=True)
|
||||
parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0],
|
||||
help='Padding (top, bottom, left, right)')
|
||||
|
||||
parser.add_argument('--face_det_batch_size', type=int,
|
||||
help='Single GPU batch size for face detection', default=16)
|
||||
|
||||
parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip', default=128)
|
||||
parser.add_argument('--face_res', help='Approximate resolution of the face at which to test', default=180)
|
||||
parser.add_argument('--min_frame_res', help='Do not downsample further below this frame resolution', default=480)
|
||||
parser.add_argument('--max_frame_res', help='Downsample to at least this frame resolution', default=720)
|
||||
# parser.add_argument('--resize_factor', default=1, type=int)
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = 96
|
||||
|
||||
def get_smoothened_boxes(boxes, T):
|
||||
for i in range(len(boxes)):
|
||||
if i + T > len(boxes):
|
||||
window = boxes[len(boxes) - T:]
|
||||
else:
|
||||
window = boxes[i : i + T]
|
||||
boxes[i] = np.mean(window, axis=0)
|
||||
return boxes
|
||||
|
||||
def rescale_frames(images):
|
||||
rect = detector.get_detections_for_batch(np.array([images[0]]))[0]
|
||||
if rect is None:
|
||||
raise ValueError('Face not detected!')
|
||||
h, w = images[0].shape[:-1]
|
||||
|
||||
x1, y1, x2, y2 = rect
|
||||
|
||||
face_size = max(np.abs(y1 - y2), np.abs(x1 - x2))
|
||||
|
||||
diff = np.abs(face_size - args.face_res)
|
||||
for factor in range(2, 16):
|
||||
downsampled_res = face_size // factor
|
||||
if min(h//factor, w//factor) < args.min_frame_res: break
|
||||
if np.abs(downsampled_res - args.face_res) >= diff: break
|
||||
|
||||
factor -= 1
|
||||
if factor == 1: return images
|
||||
|
||||
return [cv2.resize(im, (im.shape[1]//(factor), im.shape[0]//(factor))) for im in images]
|
||||
|
||||
|
||||
def face_detect(images):
|
||||
batch_size = args.face_det_batch_size
|
||||
images = rescale_frames(images)
|
||||
|
||||
while 1:
|
||||
predictions = []
|
||||
try:
|
||||
for i in range(0, len(images), batch_size):
|
||||
predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
|
||||
except RuntimeError:
|
||||
if batch_size == 1:
|
||||
raise RuntimeError('Image too big to run face detection on GPU')
|
||||
batch_size //= 2
|
||||
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
||||
continue
|
||||
break
|
||||
|
||||
results = []
|
||||
pady1, pady2, padx1, padx2 = args.pads
|
||||
for rect, image in zip(predictions, images):
|
||||
if rect is None:
|
||||
raise ValueError('Face not detected!')
|
||||
|
||||
y1 = max(0, rect[1] - pady1)
|
||||
y2 = min(image.shape[0], rect[3] + pady2)
|
||||
x1 = max(0, rect[0] - padx1)
|
||||
x2 = min(image.shape[1], rect[2] + padx2)
|
||||
|
||||
results.append([x1, y1, x2, y2])
|
||||
|
||||
boxes = get_smoothened_boxes(np.array(results), T=5)
|
||||
results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)]
|
||||
|
||||
return results, images
|
||||
|
||||
def datagen(frames, face_det_results, mels):
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
for i, m in enumerate(mels):
|
||||
if i >= len(frames): raise ValueError('Equal or less lengths only')
|
||||
|
||||
frame_to_save = frames[i].copy()
|
||||
face, coords, valid_frame = face_det_results[i].copy()
|
||||
if not valid_frame:
|
||||
continue
|
||||
|
||||
face = cv2.resize(face, (args.img_size, args.img_size))
|
||||
|
||||
img_batch.append(face)
|
||||
mel_batch.append(m)
|
||||
frame_batch.append(frame_to_save)
|
||||
coords_batch.append(coords)
|
||||
|
||||
if len(img_batch) >= args.wav2lip_batch_size:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if len(img_batch) > 0:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
|
||||
def increase_frames(frames, l):
|
||||
## evenly duplicating frames to increase length of video
|
||||
while len(frames) < l:
|
||||
dup_every = float(l) / len(frames)
|
||||
|
||||
final_frames = []
|
||||
next_duplicate = 0.
|
||||
|
||||
for i, f in enumerate(frames):
|
||||
final_frames.append(f)
|
||||
|
||||
if int(np.ceil(next_duplicate)) == i:
|
||||
final_frames.append(f)
|
||||
|
||||
next_duplicate += dup_every
|
||||
|
||||
frames = final_frames
|
||||
|
||||
return frames[:l]
|
||||
|
||||
mel_step_size = 16
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
print('Using {} for inference.'.format(device))
|
||||
|
||||
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
||||
flip_input=False, device=device)
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if device == 'cuda':
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_model(path):
|
||||
model = Wav2Lip()
|
||||
print("Load checkpoint from: {}".format(path))
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
|
||||
model = model.to(device)
|
||||
return model.eval()
|
||||
|
||||
model = load_model(args.checkpoint_path)
|
||||
|
||||
def main():
|
||||
if not os.path.isdir(args.results_dir): os.makedirs(args.results_dir)
|
||||
|
||||
if args.mode == 'dubbed':
|
||||
files = listdir(args.data_root)
|
||||
lines = ['{} {}'.format(f, f) for f in files]
|
||||
|
||||
else:
|
||||
assert args.filelist is not None
|
||||
with open(args.filelist, 'r') as filelist:
|
||||
lines = filelist.readlines()
|
||||
|
||||
for idx, line in enumerate(tqdm(lines)):
|
||||
video, audio_src = line.strip().split()
|
||||
|
||||
audio_src = os.path.join(args.data_root, audio_src)
|
||||
video = os.path.join(args.data_root, video)
|
||||
|
||||
command = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'.format(audio_src, '../temp/temp.wav')
|
||||
subprocess.call(command, shell=True)
|
||||
temp_audio = '../temp/temp.wav'
|
||||
|
||||
wav = audio.load_wav(temp_audio, 16000)
|
||||
mel = audio.melspectrogram(wav)
|
||||
|
||||
if np.isnan(mel.reshape(-1)).sum() > 0:
|
||||
raise ValueError('Mel contains nan!')
|
||||
|
||||
video_stream = cv2.VideoCapture(video)
|
||||
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
mel_idx_multiplier = 80./fps
|
||||
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
|
||||
if min(frame.shape[:-1]) > args.max_frame_res:
|
||||
h, w = frame.shape[:-1]
|
||||
scale_factor = min(h, w) / float(args.max_frame_res)
|
||||
h = int(h/scale_factor)
|
||||
w = int(w/scale_factor)
|
||||
|
||||
frame = cv2.resize(frame, (w, h))
|
||||
full_frames.append(frame)
|
||||
|
||||
mel_chunks = []
|
||||
i = 0
|
||||
while 1:
|
||||
start_idx = int(i * mel_idx_multiplier)
|
||||
if start_idx + mel_step_size > len(mel[0]):
|
||||
break
|
||||
mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
|
||||
i += 1
|
||||
|
||||
if len(full_frames) < len(mel_chunks):
|
||||
if args.mode == 'tts':
|
||||
full_frames = increase_frames(full_frames, len(mel_chunks))
|
||||
else:
|
||||
raise ValueError('#Frames, audio length mismatch')
|
||||
|
||||
else:
|
||||
full_frames = full_frames[:len(mel_chunks)]
|
||||
|
||||
try:
|
||||
face_det_results, full_frames = face_detect(full_frames.copy())
|
||||
except ValueError as e:
|
||||
continue
|
||||
|
||||
batch_size = args.wav2lip_batch_size
|
||||
gen = datagen(full_frames.copy(), face_det_results, mel_chunks)
|
||||
|
||||
for i, (img_batch, mel_batch, frames, coords) in enumerate(gen):
|
||||
if i == 0:
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
|
||||
out = cv2.VideoWriter('../temp/result.avi',
|
||||
cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
|
||||
|
||||
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
|
||||
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(mel_batch, img_batch)
|
||||
|
||||
|
||||
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
|
||||
|
||||
for pl, f, c in zip(pred, frames, coords):
|
||||
y1, y2, x1, x2 = c
|
||||
pl = cv2.resize(pl.astype(np.uint8), (x2 - x1, y2 - y1))
|
||||
f[y1:y2, x1:x2] = pl
|
||||
out.write(f)
|
||||
|
||||
out.release()
|
||||
|
||||
vid = os.path.join(args.results_dir, '{}.mp4'.format(idx))
|
||||
command = 'ffmpeg -loglevel panic -y -i {} -i {} -strict -2 -q:v 1 {}'.format('../temp/temp.wav',
|
||||
'../temp/result.avi', vid)
|
||||
subprocess.call(command, shell=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
+210
@@ -0,0 +1,210 @@
|
||||
#!/usr/bin/python
|
||||
#-*- coding: utf-8 -*-
|
||||
# Video 25 FPS, Audio 16000HZ
|
||||
|
||||
import torch
|
||||
import numpy
|
||||
import time, pdb, argparse, subprocess, os, math, glob
|
||||
import cv2
|
||||
import python_speech_features
|
||||
|
||||
from scipy import signal
|
||||
from scipy.io import wavfile
|
||||
from SyncNetModel import *
|
||||
from shutil import rmtree
|
||||
|
||||
|
||||
# ==================== Get OFFSET ====================
|
||||
|
||||
def calc_pdist(feat1, feat2, vshift=10):
|
||||
|
||||
win_size = vshift*2+1
|
||||
|
||||
feat2p = torch.nn.functional.pad(feat2,(0,0,vshift,vshift))
|
||||
|
||||
dists = []
|
||||
|
||||
for i in range(0,len(feat1)):
|
||||
|
||||
dists.append(torch.nn.functional.pairwise_distance(feat1[[i],:].repeat(win_size, 1), feat2p[i:i+win_size,:]))
|
||||
|
||||
return dists
|
||||
|
||||
# ==================== MAIN DEF ====================
|
||||
|
||||
class SyncNetInstance(torch.nn.Module):
|
||||
|
||||
def __init__(self, dropout = 0, num_layers_in_fc_layers = 1024):
|
||||
super(SyncNetInstance, self).__init__();
|
||||
|
||||
self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers).cuda();
|
||||
|
||||
def evaluate(self, opt, videofile):
|
||||
|
||||
self.__S__.eval();
|
||||
|
||||
# ========== ==========
|
||||
# Convert files
|
||||
# ========== ==========
|
||||
|
||||
if os.path.exists(os.path.join(opt.tmp_dir,opt.reference)):
|
||||
rmtree(os.path.join(opt.tmp_dir,opt.reference))
|
||||
|
||||
os.makedirs(os.path.join(opt.tmp_dir,opt.reference))
|
||||
|
||||
command = ("ffmpeg -loglevel error -y -i %s -threads 1 -f image2 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'%06d.jpg')))
|
||||
output = subprocess.call(command, shell=True, stdout=None)
|
||||
|
||||
command = ("ffmpeg -loglevel error -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'audio.wav')))
|
||||
output = subprocess.call(command, shell=True, stdout=None)
|
||||
|
||||
# ========== ==========
|
||||
# Load video
|
||||
# ========== ==========
|
||||
|
||||
images = []
|
||||
|
||||
flist = glob.glob(os.path.join(opt.tmp_dir,opt.reference,'*.jpg'))
|
||||
flist.sort()
|
||||
|
||||
for fname in flist:
|
||||
img_input = cv2.imread(fname)
|
||||
img_input = cv2.resize(img_input, (224,224)) #HARD CODED, CHANGE BEFORE RELEASE
|
||||
images.append(img_input)
|
||||
|
||||
im = numpy.stack(images,axis=3)
|
||||
im = numpy.expand_dims(im,axis=0)
|
||||
im = numpy.transpose(im,(0,3,4,1,2))
|
||||
|
||||
imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
|
||||
|
||||
# ========== ==========
|
||||
# Load audio
|
||||
# ========== ==========
|
||||
|
||||
sample_rate, audio = wavfile.read(os.path.join(opt.tmp_dir,opt.reference,'audio.wav'))
|
||||
mfcc = zip(*python_speech_features.mfcc(audio,sample_rate))
|
||||
mfcc = numpy.stack([numpy.array(i) for i in mfcc])
|
||||
|
||||
cc = numpy.expand_dims(numpy.expand_dims(mfcc,axis=0),axis=0)
|
||||
cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float())
|
||||
|
||||
# ========== ==========
|
||||
# Check audio and video input length
|
||||
# ========== ==========
|
||||
|
||||
#if (float(len(audio))/16000) != (float(len(images))/25) :
|
||||
# print("WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."%(float(len(audio))/16000,float(len(images))/25))
|
||||
|
||||
min_length = min(len(images),math.floor(len(audio)/640))
|
||||
|
||||
# ========== ==========
|
||||
# Generate video and audio feats
|
||||
# ========== ==========
|
||||
|
||||
lastframe = min_length-5
|
||||
im_feat = []
|
||||
cc_feat = []
|
||||
|
||||
tS = time.time()
|
||||
for i in range(0,lastframe,opt.batch_size):
|
||||
|
||||
im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
|
||||
im_in = torch.cat(im_batch,0)
|
||||
im_out = self.__S__.forward_lip(im_in.cuda());
|
||||
im_feat.append(im_out.data.cpu())
|
||||
|
||||
cc_batch = [ cct[:,:,:,vframe*4:vframe*4+20] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
|
||||
cc_in = torch.cat(cc_batch,0)
|
||||
cc_out = self.__S__.forward_aud(cc_in.cuda())
|
||||
cc_feat.append(cc_out.data.cpu())
|
||||
|
||||
im_feat = torch.cat(im_feat,0)
|
||||
cc_feat = torch.cat(cc_feat,0)
|
||||
|
||||
# ========== ==========
|
||||
# Compute offset
|
||||
# ========== ==========
|
||||
|
||||
#print('Compute time %.3f sec.' % (time.time()-tS))
|
||||
|
||||
dists = calc_pdist(im_feat,cc_feat,vshift=opt.vshift)
|
||||
mdist = torch.mean(torch.stack(dists,1),1)
|
||||
|
||||
minval, minidx = torch.min(mdist,0)
|
||||
|
||||
offset = opt.vshift-minidx
|
||||
conf = torch.median(mdist) - minval
|
||||
|
||||
fdist = numpy.stack([dist[minidx].numpy() for dist in dists])
|
||||
# fdist = numpy.pad(fdist, (3,3), 'constant', constant_values=15)
|
||||
fconf = torch.median(mdist).numpy() - fdist
|
||||
fconfm = signal.medfilt(fconf,kernel_size=9)
|
||||
|
||||
numpy.set_printoptions(formatter={'float': '{: 0.3f}'.format})
|
||||
#print('Framewise conf: ')
|
||||
#print(fconfm)
|
||||
#print('AV offset: \t%d \nMin dist: \t%.3f\nConfidence: \t%.3f' % (offset,minval,conf))
|
||||
|
||||
dists_npy = numpy.array([ dist.numpy() for dist in dists ])
|
||||
return offset.numpy(), conf.numpy(), minval.numpy()
|
||||
|
||||
def extract_feature(self, opt, videofile):
|
||||
|
||||
self.__S__.eval();
|
||||
|
||||
# ========== ==========
|
||||
# Load video
|
||||
# ========== ==========
|
||||
cap = cv2.VideoCapture(videofile)
|
||||
|
||||
frame_num = 1;
|
||||
images = []
|
||||
while frame_num:
|
||||
frame_num += 1
|
||||
ret, image = cap.read()
|
||||
if ret == 0:
|
||||
break
|
||||
|
||||
images.append(image)
|
||||
|
||||
im = numpy.stack(images,axis=3)
|
||||
im = numpy.expand_dims(im,axis=0)
|
||||
im = numpy.transpose(im,(0,3,4,1,2))
|
||||
|
||||
imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
|
||||
|
||||
# ========== ==========
|
||||
# Generate video feats
|
||||
# ========== ==========
|
||||
|
||||
lastframe = len(images)-4
|
||||
im_feat = []
|
||||
|
||||
tS = time.time()
|
||||
for i in range(0,lastframe,opt.batch_size):
|
||||
|
||||
im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
|
||||
im_in = torch.cat(im_batch,0)
|
||||
im_out = self.__S__.forward_lipfeat(im_in.cuda());
|
||||
im_feat.append(im_out.data.cpu())
|
||||
|
||||
im_feat = torch.cat(im_feat,0)
|
||||
|
||||
# ========== ==========
|
||||
# Compute offset
|
||||
# ========== ==========
|
||||
|
||||
print('Compute time %.3f sec.' % (time.time()-tS))
|
||||
|
||||
return im_feat
|
||||
|
||||
|
||||
def loadParameters(self, path):
|
||||
loaded_state = torch.load(path, map_location=lambda storage, loc: storage);
|
||||
|
||||
self_state = self.__S__.state_dict();
|
||||
|
||||
for name, param in loaded_state.items():
|
||||
|
||||
self_state[name].copy_(param);
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/python
|
||||
#-*- coding: utf-8 -*-
|
||||
|
||||
import time, pdb, argparse, subprocess
|
||||
import glob
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
|
||||
from SyncNetInstance_calc_scores import *
|
||||
|
||||
# ==================== LOAD PARAMS ====================
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description = "SyncNet");
|
||||
|
||||
parser.add_argument('--initial_model', type=str, default="data/syncnet_v2.model", help='');
|
||||
parser.add_argument('--batch_size', type=int, default='20', help='');
|
||||
parser.add_argument('--vshift', type=int, default='15', help='');
|
||||
parser.add_argument('--data_root', type=str, required=True, help='');
|
||||
parser.add_argument('--tmp_dir', type=str, default="data/work/pytmp", help='');
|
||||
parser.add_argument('--reference', type=str, default="demo", help='');
|
||||
|
||||
opt = parser.parse_args();
|
||||
|
||||
|
||||
# ==================== RUN EVALUATION ====================
|
||||
|
||||
s = SyncNetInstance();
|
||||
|
||||
s.loadParameters(opt.initial_model);
|
||||
#print("Model %s loaded."%opt.initial_model);
|
||||
path = os.path.join(opt.data_root, "*.mp4")
|
||||
|
||||
all_videos = glob.glob(path)
|
||||
|
||||
prog_bar = tqdm(range(len(all_videos)))
|
||||
avg_confidence = 0.
|
||||
avg_min_distance = 0.
|
||||
|
||||
|
||||
for videofile_idx in prog_bar:
|
||||
videofile = all_videos[videofile_idx]
|
||||
offset, confidence, min_distance = s.evaluate(opt, videofile=videofile)
|
||||
avg_confidence += confidence
|
||||
avg_min_distance += min_distance
|
||||
prog_bar.set_description('Avg Confidence: {}, Avg Minimum Dist: {}'.format(round(avg_confidence / (videofile_idx + 1), 3), round(avg_min_distance / (videofile_idx + 1), 3)))
|
||||
prog_bar.refresh()
|
||||
|
||||
print ('Average Confidence: {}'.format(avg_confidence/len(all_videos)))
|
||||
print ('Average Minimum Distance: {}'.format(avg_min_distance/len(all_videos)))
|
||||
|
||||
|
||||
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/python
|
||||
#-*- coding: utf-8 -*-
|
||||
|
||||
import time, pdb, argparse, subprocess, pickle, os, gzip, glob
|
||||
|
||||
from SyncNetInstance_calc_scores import *
|
||||
|
||||
# ==================== PARSE ARGUMENT ====================
|
||||
|
||||
parser = argparse.ArgumentParser(description = "SyncNet");
|
||||
parser.add_argument('--initial_model', type=str, default="data/syncnet_v2.model", help='');
|
||||
parser.add_argument('--batch_size', type=int, default='20', help='');
|
||||
parser.add_argument('--vshift', type=int, default='15', help='');
|
||||
parser.add_argument('--data_dir', type=str, default='data/work', help='');
|
||||
parser.add_argument('--videofile', type=str, default='', help='');
|
||||
parser.add_argument('--reference', type=str, default='', help='');
|
||||
opt = parser.parse_args();
|
||||
|
||||
setattr(opt,'avi_dir',os.path.join(opt.data_dir,'pyavi'))
|
||||
setattr(opt,'tmp_dir',os.path.join(opt.data_dir,'pytmp'))
|
||||
setattr(opt,'work_dir',os.path.join(opt.data_dir,'pywork'))
|
||||
setattr(opt,'crop_dir',os.path.join(opt.data_dir,'pycrop'))
|
||||
|
||||
|
||||
# ==================== LOAD MODEL AND FILE LIST ====================
|
||||
|
||||
s = SyncNetInstance();
|
||||
|
||||
s.loadParameters(opt.initial_model);
|
||||
#print("Model %s loaded."%opt.initial_model);
|
||||
|
||||
flist = glob.glob(os.path.join(opt.crop_dir,opt.reference,'0*.avi'))
|
||||
flist.sort()
|
||||
|
||||
# ==================== GET OFFSETS ====================
|
||||
|
||||
dists = []
|
||||
for idx, fname in enumerate(flist):
|
||||
offset, conf, dist = s.evaluate(opt,videofile=fname)
|
||||
print (str(dist)+" "+str(conf))
|
||||
|
||||
# ==================== PRINT RESULTS TO FILE ====================
|
||||
|
||||
#with open(os.path.join(opt.work_dir,opt.reference,'activesd.pckl'), 'wb') as fil:
|
||||
# pickle.dump(dists, fil)
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
rm all_scores.txt
|
||||
yourfilenames=`ls $1`
|
||||
for eachfile in $yourfilenames
|
||||
do
|
||||
python run_pipeline.py --videofile $eachfile --reference wav2lip --data_dir tmp_dir
|
||||
python calculate_scores_real_videos.py --videofile $eachfile --reference wav2lip --data_dir /ssd_scratch/cvit/prajwalkr_rudra/tmp_dir >> all_scores.txt
|
||||
done
|
||||
@@ -0,0 +1 @@
|
||||
Place LRS2, LRW, and LRS3 (and any other) test set filelists here.
|
||||
+160
@@ -0,0 +1,160 @@
|
||||
sachin.mp4 emma_cropped.mp4
|
||||
sachin.mp4 mourinho.mp4
|
||||
sachin.mp4 elon.mp4
|
||||
sachin.mp4 messi2.mp4
|
||||
sachin.mp4 cr1.mp4
|
||||
sachin.mp4 sachin.mp4
|
||||
sachin.mp4 sg.mp4
|
||||
sachin.mp4 fergi.mp4
|
||||
sachin.mp4 spanish_lec1.mp4
|
||||
sachin.mp4 bush_small.mp4
|
||||
sachin.mp4 macca_cut.mp4
|
||||
sachin.mp4 ca_cropped.mp4
|
||||
sachin.mp4 lecun.mp4
|
||||
sachin.mp4 spanish_lec0.mp4
|
||||
srk.mp4 emma_cropped.mp4
|
||||
srk.mp4 mourinho.mp4
|
||||
srk.mp4 elon.mp4
|
||||
srk.mp4 messi2.mp4
|
||||
srk.mp4 cr1.mp4
|
||||
srk.mp4 srk.mp4
|
||||
srk.mp4 sachin.mp4
|
||||
srk.mp4 sg.mp4
|
||||
srk.mp4 fergi.mp4
|
||||
srk.mp4 spanish_lec1.mp4
|
||||
srk.mp4 bush_small.mp4
|
||||
srk.mp4 macca_cut.mp4
|
||||
srk.mp4 ca_cropped.mp4
|
||||
srk.mp4 guardiola.mp4
|
||||
srk.mp4 lecun.mp4
|
||||
srk.mp4 spanish_lec0.mp4
|
||||
cr1.mp4 emma_cropped.mp4
|
||||
cr1.mp4 elon.mp4
|
||||
cr1.mp4 messi2.mp4
|
||||
cr1.mp4 cr1.mp4
|
||||
cr1.mp4 spanish_lec1.mp4
|
||||
cr1.mp4 bush_small.mp4
|
||||
cr1.mp4 macca_cut.mp4
|
||||
cr1.mp4 ca_cropped.mp4
|
||||
cr1.mp4 lecun.mp4
|
||||
cr1.mp4 spanish_lec0.mp4
|
||||
macca_cut.mp4 emma_cropped.mp4
|
||||
macca_cut.mp4 elon.mp4
|
||||
macca_cut.mp4 messi2.mp4
|
||||
macca_cut.mp4 spanish_lec1.mp4
|
||||
macca_cut.mp4 macca_cut.mp4
|
||||
macca_cut.mp4 ca_cropped.mp4
|
||||
macca_cut.mp4 spanish_lec0.mp4
|
||||
lecun.mp4 emma_cropped.mp4
|
||||
lecun.mp4 elon.mp4
|
||||
lecun.mp4 messi2.mp4
|
||||
lecun.mp4 spanish_lec1.mp4
|
||||
lecun.mp4 macca_cut.mp4
|
||||
lecun.mp4 ca_cropped.mp4
|
||||
lecun.mp4 lecun.mp4
|
||||
lecun.mp4 spanish_lec0.mp4
|
||||
messi2.mp4 emma_cropped.mp4
|
||||
messi2.mp4 elon.mp4
|
||||
messi2.mp4 messi2.mp4
|
||||
messi2.mp4 spanish_lec1.mp4
|
||||
messi2.mp4 macca_cut.mp4
|
||||
messi2.mp4 ca_cropped.mp4
|
||||
messi2.mp4 spanish_lec0.mp4
|
||||
ca_cropped.mp4 emma_cropped.mp4
|
||||
ca_cropped.mp4 elon.mp4
|
||||
ca_cropped.mp4 spanish_lec1.mp4
|
||||
ca_cropped.mp4 ca_cropped.mp4
|
||||
ca_cropped.mp4 spanish_lec0.mp4
|
||||
spanish_lec1.mp4 spanish_lec1.mp4
|
||||
spanish_lec1.mp4 spanish_lec0.mp4
|
||||
elon.mp4 elon.mp4
|
||||
elon.mp4 spanish_lec1.mp4
|
||||
elon.mp4 spanish_lec0.mp4
|
||||
guardiola.mp4 emma_cropped.mp4
|
||||
guardiola.mp4 mourinho.mp4
|
||||
guardiola.mp4 elon.mp4
|
||||
guardiola.mp4 messi2.mp4
|
||||
guardiola.mp4 cr1.mp4
|
||||
guardiola.mp4 sachin.mp4
|
||||
guardiola.mp4 sg.mp4
|
||||
guardiola.mp4 fergi.mp4
|
||||
guardiola.mp4 spanish_lec1.mp4
|
||||
guardiola.mp4 bush_small.mp4
|
||||
guardiola.mp4 macca_cut.mp4
|
||||
guardiola.mp4 ca_cropped.mp4
|
||||
guardiola.mp4 guardiola.mp4
|
||||
guardiola.mp4 lecun.mp4
|
||||
guardiola.mp4 spanish_lec0.mp4
|
||||
fergi.mp4 emma_cropped.mp4
|
||||
fergi.mp4 mourinho.mp4
|
||||
fergi.mp4 elon.mp4
|
||||
fergi.mp4 messi2.mp4
|
||||
fergi.mp4 cr1.mp4
|
||||
fergi.mp4 sachin.mp4
|
||||
fergi.mp4 sg.mp4
|
||||
fergi.mp4 fergi.mp4
|
||||
fergi.mp4 spanish_lec1.mp4
|
||||
fergi.mp4 bush_small.mp4
|
||||
fergi.mp4 macca_cut.mp4
|
||||
fergi.mp4 ca_cropped.mp4
|
||||
fergi.mp4 lecun.mp4
|
||||
fergi.mp4 spanish_lec0.mp4
|
||||
spanish.mp4 emma_cropped.mp4
|
||||
spanish.mp4 spanish.mp4
|
||||
spanish.mp4 mourinho.mp4
|
||||
spanish.mp4 elon.mp4
|
||||
spanish.mp4 messi2.mp4
|
||||
spanish.mp4 cr1.mp4
|
||||
spanish.mp4 srk.mp4
|
||||
spanish.mp4 sachin.mp4
|
||||
spanish.mp4 sg.mp4
|
||||
spanish.mp4 fergi.mp4
|
||||
spanish.mp4 spanish_lec1.mp4
|
||||
spanish.mp4 bush_small.mp4
|
||||
spanish.mp4 macca_cut.mp4
|
||||
spanish.mp4 ca_cropped.mp4
|
||||
spanish.mp4 guardiola.mp4
|
||||
spanish.mp4 lecun.mp4
|
||||
spanish.mp4 spanish_lec0.mp4
|
||||
bush_small.mp4 emma_cropped.mp4
|
||||
bush_small.mp4 elon.mp4
|
||||
bush_small.mp4 messi2.mp4
|
||||
bush_small.mp4 spanish_lec1.mp4
|
||||
bush_small.mp4 bush_small.mp4
|
||||
bush_small.mp4 macca_cut.mp4
|
||||
bush_small.mp4 ca_cropped.mp4
|
||||
bush_small.mp4 lecun.mp4
|
||||
bush_small.mp4 spanish_lec0.mp4
|
||||
emma_cropped.mp4 emma_cropped.mp4
|
||||
emma_cropped.mp4 elon.mp4
|
||||
emma_cropped.mp4 spanish_lec1.mp4
|
||||
emma_cropped.mp4 spanish_lec0.mp4
|
||||
sg.mp4 emma_cropped.mp4
|
||||
sg.mp4 mourinho.mp4
|
||||
sg.mp4 elon.mp4
|
||||
sg.mp4 messi2.mp4
|
||||
sg.mp4 cr1.mp4
|
||||
sg.mp4 sachin.mp4
|
||||
sg.mp4 sg.mp4
|
||||
sg.mp4 fergi.mp4
|
||||
sg.mp4 spanish_lec1.mp4
|
||||
sg.mp4 bush_small.mp4
|
||||
sg.mp4 macca_cut.mp4
|
||||
sg.mp4 ca_cropped.mp4
|
||||
sg.mp4 lecun.mp4
|
||||
sg.mp4 spanish_lec0.mp4
|
||||
spanish_lec0.mp4 spanish_lec0.mp4
|
||||
mourinho.mp4 emma_cropped.mp4
|
||||
mourinho.mp4 mourinho.mp4
|
||||
mourinho.mp4 elon.mp4
|
||||
mourinho.mp4 messi2.mp4
|
||||
mourinho.mp4 cr1.mp4
|
||||
mourinho.mp4 sachin.mp4
|
||||
mourinho.mp4 sg.mp4
|
||||
mourinho.mp4 fergi.mp4
|
||||
mourinho.mp4 spanish_lec1.mp4
|
||||
mourinho.mp4 bush_small.mp4
|
||||
mourinho.mp4 macca_cut.mp4
|
||||
mourinho.mp4 ca_cropped.mp4
|
||||
mourinho.mp4 lecun.mp4
|
||||
mourinho.mp4 spanish_lec0.mp4
|
||||
@@ -0,0 +1,18 @@
|
||||
adam_1.mp4 andreng_optimization.wav
|
||||
agad_2.mp4 agad_2.wav
|
||||
agad_1.mp4 agad_1.wav
|
||||
agad_3.mp4 agad_3.wav
|
||||
rms_prop_1.mp4 rms_prop_tts.wav
|
||||
tf_1.mp4 tf_1.wav
|
||||
tf_2.mp4 tf_2.wav
|
||||
andrew_ng_ai_business.mp4 andrewng_business_tts.wav
|
||||
covid_autopsy_1.mp4 autopsy_tts.wav
|
||||
news_1.mp4 news_tts.wav
|
||||
andrew_ng_fund_1.mp4 andrewng_ai_fund.wav
|
||||
covid_treatments_1.mp4 covid_tts.wav
|
||||
pytorch_v_tf.mp4 pytorch_vs_tf_eng.wav
|
||||
pytorch_1.mp4 pytorch.wav
|
||||
pkb_1.mp4 pkb_1.wav
|
||||
ss_1.mp4 ss_1.wav
|
||||
carlsen_1.mp4 carlsen_eng.wav
|
||||
french.mp4 french.wav
|
||||
@@ -0,0 +1 @@
|
||||
The code for Face Detection in this folder has been taken from the wonderful [face_alignment](https://github.com/1adrianb/face-alignment) repository. This has been modified to take batches of faces at a time.
|
||||
@@ -0,0 +1,7 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
__author__ = """Adrian Bulat"""
|
||||
__email__ = 'adrian.bulat@nottingham.ac.uk'
|
||||
__version__ = '1.0.1'
|
||||
|
||||
from .api import FaceAlignment, LandmarksType, NetworkSize
|
||||
@@ -0,0 +1,79 @@
|
||||
from __future__ import print_function
|
||||
import os
|
||||
import torch
|
||||
from torch.utils.model_zoo import load_url
|
||||
from enum import Enum
|
||||
import numpy as np
|
||||
import cv2
|
||||
try:
|
||||
import urllib.request as request_file
|
||||
except BaseException:
|
||||
import urllib as request_file
|
||||
|
||||
from .models import FAN, ResNetDepth
|
||||
from .utils import *
|
||||
|
||||
|
||||
class LandmarksType(Enum):
|
||||
"""Enum class defining the type of landmarks to detect.
|
||||
|
||||
``_2D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face
|
||||
``_2halfD`` - this points represent the projection of the 3D points into 3D
|
||||
``_3D`` - detect the points ``(x,y,z)``` in a 3D space
|
||||
|
||||
"""
|
||||
_2D = 1
|
||||
_2halfD = 2
|
||||
_3D = 3
|
||||
|
||||
|
||||
class NetworkSize(Enum):
|
||||
# TINY = 1
|
||||
# SMALL = 2
|
||||
# MEDIUM = 3
|
||||
LARGE = 4
|
||||
|
||||
def __new__(cls, value):
|
||||
member = object.__new__(cls)
|
||||
member._value_ = value
|
||||
return member
|
||||
|
||||
def __int__(self):
|
||||
return self.value
|
||||
|
||||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
class FaceAlignment:
|
||||
def __init__(self, landmarks_type, network_size=NetworkSize.LARGE,
|
||||
device='cuda', flip_input=False, face_detector='sfd', verbose=False):
|
||||
self.device = device
|
||||
self.flip_input = flip_input
|
||||
self.landmarks_type = landmarks_type
|
||||
self.verbose = verbose
|
||||
|
||||
network_size = int(network_size)
|
||||
|
||||
if 'cuda' in device:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
# Get the face detector
|
||||
face_detector_module = __import__('face_detection.detection.' + face_detector,
|
||||
globals(), locals(), [face_detector], 0)
|
||||
self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose)
|
||||
|
||||
def get_detections_for_batch(self, images):
|
||||
images = images[..., ::-1]
|
||||
detected_faces = self.face_detector.detect_from_batch(images.copy())
|
||||
results = []
|
||||
|
||||
for i, d in enumerate(detected_faces):
|
||||
if len(d) == 0:
|
||||
results.append(None)
|
||||
continue
|
||||
d = d[0]
|
||||
d = np.clip(d, 0, None)
|
||||
|
||||
x1, y1, x2, y2 = map(int, d[:-1])
|
||||
results.append((x1, y1, x2, y2))
|
||||
|
||||
return results
|
||||
@@ -0,0 +1 @@
|
||||
from .core import FaceDetector
|
||||
@@ -0,0 +1,130 @@
|
||||
import logging
|
||||
import glob
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
|
||||
|
||||
class FaceDetector(object):
|
||||
"""An abstract class representing a face detector.
|
||||
|
||||
Any other face detection implementation must subclass it. All subclasses
|
||||
must implement ``detect_from_image``, that return a list of detected
|
||||
bounding boxes. Optionally, for speed considerations detect from path is
|
||||
recommended.
|
||||
"""
|
||||
|
||||
def __init__(self, device, verbose):
|
||||
self.device = device
|
||||
self.verbose = verbose
|
||||
|
||||
if verbose:
|
||||
if 'cpu' in device:
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning("Detection running on CPU, this may be potentially slow.")
|
||||
|
||||
if 'cpu' not in device and 'cuda' not in device:
|
||||
if verbose:
|
||||
logger.error("Expected values for device are: {cpu, cuda} but got: %s", device)
|
||||
raise ValueError
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
"""Detects faces in a given image.
|
||||
|
||||
This function detects the faces present in a provided BGR(usually)
|
||||
image. The input can be either the image itself or the path to it.
|
||||
|
||||
Arguments:
|
||||
tensor_or_path {numpy.ndarray, torch.tensor or string} -- the path
|
||||
to an image or the image itself.
|
||||
|
||||
Example::
|
||||
|
||||
>>> path_to_image = 'data/image_01.jpg'
|
||||
... detected_faces = detect_from_image(path_to_image)
|
||||
[A list of bounding boxes (x1, y1, x2, y2)]
|
||||
>>> image = cv2.imread(path_to_image)
|
||||
... detected_faces = detect_from_image(image)
|
||||
[A list of bounding boxes (x1, y1, x2, y2)]
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def detect_from_directory(self, path, extensions=['.jpg', '.png'], recursive=False, show_progress_bar=True):
|
||||
"""Detects faces from all the images present in a given directory.
|
||||
|
||||
Arguments:
|
||||
path {string} -- a string containing a path that points to the folder containing the images
|
||||
|
||||
Keyword Arguments:
|
||||
extensions {list} -- list of string containing the extensions to be
|
||||
consider in the following format: ``.extension_name`` (default:
|
||||
{['.jpg', '.png']}) recursive {bool} -- option wherever to scan the
|
||||
folder recursively (default: {False}) show_progress_bar {bool} --
|
||||
display a progressbar (default: {True})
|
||||
|
||||
Example:
|
||||
>>> directory = 'data'
|
||||
... detected_faces = detect_from_directory(directory)
|
||||
{A dictionary of [lists containing bounding boxes(x1, y1, x2, y2)]}
|
||||
|
||||
"""
|
||||
if self.verbose:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if len(extensions) == 0:
|
||||
if self.verbose:
|
||||
logger.error("Expected at list one extension, but none was received.")
|
||||
raise ValueError
|
||||
|
||||
if self.verbose:
|
||||
logger.info("Constructing the list of images.")
|
||||
additional_pattern = '/**/*' if recursive else '/*'
|
||||
files = []
|
||||
for extension in extensions:
|
||||
files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
|
||||
|
||||
if self.verbose:
|
||||
logger.info("Finished searching for images. %s images found", len(files))
|
||||
logger.info("Preparing to run the detection.")
|
||||
|
||||
predictions = {}
|
||||
for image_path in tqdm(files, disable=not show_progress_bar):
|
||||
if self.verbose:
|
||||
logger.info("Running the face detector on image: %s", image_path)
|
||||
predictions[image_path] = self.detect_from_image(image_path)
|
||||
|
||||
if self.verbose:
|
||||
logger.info("The detector was successfully run on all %s images", len(files))
|
||||
|
||||
return predictions
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def tensor_or_path_to_ndarray(tensor_or_path, rgb=True):
|
||||
"""Convert path (represented as a string) or torch.tensor to a numpy.ndarray
|
||||
|
||||
Arguments:
|
||||
tensor_or_path {numpy.ndarray, torch.tensor or string} -- path to the image, or the image itself
|
||||
"""
|
||||
if isinstance(tensor_or_path, str):
|
||||
return cv2.imread(tensor_or_path) if not rgb else cv2.imread(tensor_or_path)[..., ::-1]
|
||||
elif torch.is_tensor(tensor_or_path):
|
||||
# Call cpu in case its coming from cuda
|
||||
return tensor_or_path.cpu().numpy()[..., ::-1].copy() if not rgb else tensor_or_path.cpu().numpy()
|
||||
elif isinstance(tensor_or_path, np.ndarray):
|
||||
return tensor_or_path[..., ::-1].copy() if not rgb else tensor_or_path
|
||||
else:
|
||||
raise TypeError
|
||||
@@ -0,0 +1 @@
|
||||
from .sfd_detector import SFDDetector as FaceDetector
|
||||
@@ -0,0 +1,129 @@
|
||||
from __future__ import print_function
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import random
|
||||
import datetime
|
||||
import time
|
||||
import math
|
||||
import argparse
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
try:
|
||||
from iou import IOU
|
||||
except BaseException:
|
||||
# IOU cython speedup 10x
|
||||
def IOU(ax1, ay1, ax2, ay2, bx1, by1, bx2, by2):
|
||||
sa = abs((ax2 - ax1) * (ay2 - ay1))
|
||||
sb = abs((bx2 - bx1) * (by2 - by1))
|
||||
x1, y1 = max(ax1, bx1), max(ay1, by1)
|
||||
x2, y2 = min(ax2, bx2), min(ay2, by2)
|
||||
w = x2 - x1
|
||||
h = y2 - y1
|
||||
if w < 0 or h < 0:
|
||||
return 0.0
|
||||
else:
|
||||
return 1.0 * w * h / (sa + sb - w * h)
|
||||
|
||||
|
||||
def bboxlog(x1, y1, x2, y2, axc, ayc, aww, ahh):
|
||||
xc, yc, ww, hh = (x2 + x1) / 2, (y2 + y1) / 2, x2 - x1, y2 - y1
|
||||
dx, dy = (xc - axc) / aww, (yc - ayc) / ahh
|
||||
dw, dh = math.log(ww / aww), math.log(hh / ahh)
|
||||
return dx, dy, dw, dh
|
||||
|
||||
|
||||
def bboxloginv(dx, dy, dw, dh, axc, ayc, aww, ahh):
|
||||
xc, yc = dx * aww + axc, dy * ahh + ayc
|
||||
ww, hh = math.exp(dw) * aww, math.exp(dh) * ahh
|
||||
x1, x2, y1, y2 = xc - ww / 2, xc + ww / 2, yc - hh / 2, yc + hh / 2
|
||||
return x1, y1, x2, y2
|
||||
|
||||
|
||||
def nms(dets, thresh):
|
||||
if 0 == len(dets):
|
||||
return []
|
||||
x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1, yy1 = np.maximum(x1[i], x1[order[1:]]), np.maximum(y1[i], y1[order[1:]])
|
||||
xx2, yy2 = np.minimum(x2[i], x2[order[1:]]), np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w, h = np.maximum(0.0, xx2 - xx1 + 1), np.maximum(0.0, yy2 - yy1 + 1)
|
||||
ovr = w * h / (areas[i] + areas[order[1:]] - w * h)
|
||||
|
||||
inds = np.where(ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def encode(matched, priors, variances):
|
||||
"""Encode the variances from the priorbox layers into the ground truth boxes
|
||||
we have matched (based on jaccard overlap) with the prior boxes.
|
||||
Args:
|
||||
matched: (tensor) Coords of ground truth for each prior in point-form
|
||||
Shape: [num_priors, 4].
|
||||
priors: (tensor) Prior boxes in center-offset form
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
encoded boxes (tensor), Shape: [num_priors, 4]
|
||||
"""
|
||||
|
||||
# dist b/t match center and prior's center
|
||||
g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
|
||||
# encode variance
|
||||
g_cxcy /= (variances[0] * priors[:, 2:])
|
||||
# match wh / prior wh
|
||||
g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
|
||||
g_wh = torch.log(g_wh) / variances[1]
|
||||
# return target for smooth_l1_loss
|
||||
return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4]
|
||||
|
||||
|
||||
def decode(loc, priors, variances):
|
||||
"""Decode locations from predictions using priors to undo
|
||||
the encoding we did for offset regression at train time.
|
||||
Args:
|
||||
loc (tensor): location predictions for loc layers,
|
||||
Shape: [num_priors,4]
|
||||
priors (tensor): Prior boxes in center-offset form.
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
decoded bounding box predictions
|
||||
"""
|
||||
|
||||
boxes = torch.cat((
|
||||
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
|
||||
priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
|
||||
boxes[:, :2] -= boxes[:, 2:] / 2
|
||||
boxes[:, 2:] += boxes[:, :2]
|
||||
return boxes
|
||||
|
||||
def batch_decode(loc, priors, variances):
|
||||
"""Decode locations from predictions using priors to undo
|
||||
the encoding we did for offset regression at train time.
|
||||
Args:
|
||||
loc (tensor): location predictions for loc layers,
|
||||
Shape: [num_priors,4]
|
||||
priors (tensor): Prior boxes in center-offset form.
|
||||
Shape: [num_priors,4].
|
||||
variances: (list[float]) Variances of priorboxes
|
||||
Return:
|
||||
decoded bounding box predictions
|
||||
"""
|
||||
|
||||
boxes = torch.cat((
|
||||
priors[:, :, :2] + loc[:, :, :2] * variances[0] * priors[:, :, 2:],
|
||||
priors[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1])), 2)
|
||||
boxes[:, :, :2] -= boxes[:, :, 2:] / 2
|
||||
boxes[:, :, 2:] += boxes[:, :, :2]
|
||||
return boxes
|
||||
@@ -0,0 +1,112 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import random
|
||||
import datetime
|
||||
import math
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import scipy.io as sio
|
||||
import zipfile
|
||||
from .net_s3fd import s3fd
|
||||
from .bbox import *
|
||||
|
||||
|
||||
def detect(net, img, device):
|
||||
img = img - np.array([104, 117, 123])
|
||||
img = img.transpose(2, 0, 1)
|
||||
img = img.reshape((1,) + img.shape)
|
||||
|
||||
if 'cuda' in device:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
img = torch.from_numpy(img).float().to(device)
|
||||
BB, CC, HH, WW = img.size()
|
||||
with torch.no_grad():
|
||||
olist = net(img)
|
||||
|
||||
bboxlist = []
|
||||
for i in range(len(olist) // 2):
|
||||
olist[i * 2] = F.softmax(olist[i * 2], dim=1)
|
||||
olist = [oelem.data.cpu() for oelem in olist]
|
||||
for i in range(len(olist) // 2):
|
||||
ocls, oreg = olist[i * 2], olist[i * 2 + 1]
|
||||
FB, FC, FH, FW = ocls.size() # feature map size
|
||||
stride = 2**(i + 2) # 4,8,16,32,64,128
|
||||
anchor = stride * 4
|
||||
poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
|
||||
for Iindex, hindex, windex in poss:
|
||||
axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
|
||||
score = ocls[0, 1, hindex, windex]
|
||||
loc = oreg[0, :, hindex, windex].contiguous().view(1, 4)
|
||||
priors = torch.Tensor([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]])
|
||||
variances = [0.1, 0.2]
|
||||
box = decode(loc, priors, variances)
|
||||
x1, y1, x2, y2 = box[0] * 1.0
|
||||
# cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1)
|
||||
bboxlist.append([x1, y1, x2, y2, score])
|
||||
bboxlist = np.array(bboxlist)
|
||||
if 0 == len(bboxlist):
|
||||
bboxlist = np.zeros((1, 5))
|
||||
|
||||
return bboxlist
|
||||
|
||||
def batch_detect(net, imgs, device):
|
||||
imgs = imgs - np.array([104, 117, 123])
|
||||
imgs = imgs.transpose(0, 3, 1, 2)
|
||||
|
||||
if 'cuda' in device:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
imgs = torch.from_numpy(imgs).float().to(device)
|
||||
BB, CC, HH, WW = imgs.size()
|
||||
with torch.no_grad():
|
||||
olist = net(imgs)
|
||||
|
||||
bboxlist = []
|
||||
for i in range(len(olist) // 2):
|
||||
olist[i * 2] = F.softmax(olist[i * 2], dim=1)
|
||||
olist = [oelem.data.cpu() for oelem in olist]
|
||||
for i in range(len(olist) // 2):
|
||||
ocls, oreg = olist[i * 2], olist[i * 2 + 1]
|
||||
FB, FC, FH, FW = ocls.size() # feature map size
|
||||
stride = 2**(i + 2) # 4,8,16,32,64,128
|
||||
anchor = stride * 4
|
||||
poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
|
||||
for Iindex, hindex, windex in poss:
|
||||
axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
|
||||
score = ocls[:, 1, hindex, windex]
|
||||
loc = oreg[:, :, hindex, windex].contiguous().view(BB, 1, 4)
|
||||
priors = torch.Tensor([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]]).view(1, 1, 4)
|
||||
variances = [0.1, 0.2]
|
||||
box = batch_decode(loc, priors, variances)
|
||||
box = box[:, 0] * 1.0
|
||||
# cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1)
|
||||
bboxlist.append(torch.cat([box, score.unsqueeze(1)], 1).cpu().numpy())
|
||||
bboxlist = np.array(bboxlist)
|
||||
if 0 == len(bboxlist):
|
||||
bboxlist = np.zeros((1, BB, 5))
|
||||
|
||||
return bboxlist
|
||||
|
||||
def flip_detect(net, img, device):
|
||||
img = cv2.flip(img, 1)
|
||||
b = detect(net, img, device)
|
||||
|
||||
bboxlist = np.zeros(b.shape)
|
||||
bboxlist[:, 0] = img.shape[1] - b[:, 2]
|
||||
bboxlist[:, 1] = b[:, 1]
|
||||
bboxlist[:, 2] = img.shape[1] - b[:, 0]
|
||||
bboxlist[:, 3] = b[:, 3]
|
||||
bboxlist[:, 4] = b[:, 4]
|
||||
return bboxlist
|
||||
|
||||
|
||||
def pts_to_bb(pts):
|
||||
min_x, min_y = np.min(pts, axis=0)
|
||||
max_x, max_y = np.max(pts, axis=0)
|
||||
return np.array([min_x, min_y, max_x, max_y])
|
||||
@@ -0,0 +1,129 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class L2Norm(nn.Module):
|
||||
def __init__(self, n_channels, scale=1.0):
|
||||
super(L2Norm, self).__init__()
|
||||
self.n_channels = n_channels
|
||||
self.scale = scale
|
||||
self.eps = 1e-10
|
||||
self.weight = nn.Parameter(torch.Tensor(self.n_channels))
|
||||
self.weight.data *= 0.0
|
||||
self.weight.data += self.scale
|
||||
|
||||
def forward(self, x):
|
||||
norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
|
||||
x = x / norm * self.weight.view(1, -1, 1, 1)
|
||||
return x
|
||||
|
||||
|
||||
class s3fd(nn.Module):
|
||||
def __init__(self):
|
||||
super(s3fd, self).__init__()
|
||||
self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
|
||||
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
|
||||
self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
|
||||
self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
|
||||
self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.fc6 = nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=3)
|
||||
self.fc7 = nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
self.conv6_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
|
||||
self.conv6_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
|
||||
|
||||
self.conv7_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
|
||||
self.conv7_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
|
||||
|
||||
self.conv3_3_norm = L2Norm(256, scale=10)
|
||||
self.conv4_3_norm = L2Norm(512, scale=8)
|
||||
self.conv5_3_norm = L2Norm(512, scale=5)
|
||||
|
||||
self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.fc7_mbox_conf = nn.Conv2d(1024, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.fc7_mbox_loc = nn.Conv2d(1024, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv6_2_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv6_2_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1)
|
||||
self.conv7_2_mbox_conf = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1)
|
||||
self.conv7_2_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
h = F.relu(self.conv1_1(x))
|
||||
h = F.relu(self.conv1_2(h))
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv2_1(h))
|
||||
h = F.relu(self.conv2_2(h))
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv3_1(h))
|
||||
h = F.relu(self.conv3_2(h))
|
||||
h = F.relu(self.conv3_3(h))
|
||||
f3_3 = h
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv4_1(h))
|
||||
h = F.relu(self.conv4_2(h))
|
||||
h = F.relu(self.conv4_3(h))
|
||||
f4_3 = h
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.conv5_1(h))
|
||||
h = F.relu(self.conv5_2(h))
|
||||
h = F.relu(self.conv5_3(h))
|
||||
f5_3 = h
|
||||
h = F.max_pool2d(h, 2, 2)
|
||||
|
||||
h = F.relu(self.fc6(h))
|
||||
h = F.relu(self.fc7(h))
|
||||
ffc7 = h
|
||||
h = F.relu(self.conv6_1(h))
|
||||
h = F.relu(self.conv6_2(h))
|
||||
f6_2 = h
|
||||
h = F.relu(self.conv7_1(h))
|
||||
h = F.relu(self.conv7_2(h))
|
||||
f7_2 = h
|
||||
|
||||
f3_3 = self.conv3_3_norm(f3_3)
|
||||
f4_3 = self.conv4_3_norm(f4_3)
|
||||
f5_3 = self.conv5_3_norm(f5_3)
|
||||
|
||||
cls1 = self.conv3_3_norm_mbox_conf(f3_3)
|
||||
reg1 = self.conv3_3_norm_mbox_loc(f3_3)
|
||||
cls2 = self.conv4_3_norm_mbox_conf(f4_3)
|
||||
reg2 = self.conv4_3_norm_mbox_loc(f4_3)
|
||||
cls3 = self.conv5_3_norm_mbox_conf(f5_3)
|
||||
reg3 = self.conv5_3_norm_mbox_loc(f5_3)
|
||||
cls4 = self.fc7_mbox_conf(ffc7)
|
||||
reg4 = self.fc7_mbox_loc(ffc7)
|
||||
cls5 = self.conv6_2_mbox_conf(f6_2)
|
||||
reg5 = self.conv6_2_mbox_loc(f6_2)
|
||||
cls6 = self.conv7_2_mbox_conf(f7_2)
|
||||
reg6 = self.conv7_2_mbox_loc(f7_2)
|
||||
|
||||
# max-out background label
|
||||
chunk = torch.chunk(cls1, 4, 1)
|
||||
bmax = torch.max(torch.max(chunk[0], chunk[1]), chunk[2])
|
||||
cls1 = torch.cat([bmax, chunk[3]], dim=1)
|
||||
|
||||
return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]
|
||||
@@ -0,0 +1,59 @@
|
||||
import os
|
||||
import cv2
|
||||
from torch.utils.model_zoo import load_url
|
||||
|
||||
from ..core import FaceDetector
|
||||
|
||||
from .net_s3fd import s3fd
|
||||
from .bbox import *
|
||||
from .detect import *
|
||||
|
||||
models_urls = {
|
||||
's3fd': 'https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth',
|
||||
}
|
||||
|
||||
|
||||
class SFDDetector(FaceDetector):
|
||||
def __init__(self, device, path_to_detector=os.path.join(os.path.dirname(os.path.abspath(__file__)), 's3fd.pth'), verbose=False):
|
||||
super(SFDDetector, self).__init__(device, verbose)
|
||||
|
||||
# Initialise the face detector
|
||||
if not os.path.isfile(path_to_detector):
|
||||
model_weights = load_url(models_urls['s3fd'])
|
||||
else:
|
||||
model_weights = torch.load(path_to_detector)
|
||||
|
||||
self.face_detector = s3fd()
|
||||
self.face_detector.load_state_dict(model_weights)
|
||||
self.face_detector.to(device)
|
||||
self.face_detector.eval()
|
||||
|
||||
def detect_from_image(self, tensor_or_path):
|
||||
image = self.tensor_or_path_to_ndarray(tensor_or_path)
|
||||
|
||||
bboxlist = detect(self.face_detector, image, device=self.device)
|
||||
keep = nms(bboxlist, 0.3)
|
||||
bboxlist = bboxlist[keep, :]
|
||||
bboxlist = [x for x in bboxlist if x[-1] > 0.5]
|
||||
|
||||
return bboxlist
|
||||
|
||||
def detect_from_batch(self, images):
|
||||
bboxlists = batch_detect(self.face_detector, images, device=self.device)
|
||||
keeps = [nms(bboxlists[:, i, :], 0.3) for i in range(bboxlists.shape[1])]
|
||||
bboxlists = [bboxlists[keep, i, :] for i, keep in enumerate(keeps)]
|
||||
bboxlists = [[x for x in bboxlist if x[-1] > 0.5] for bboxlist in bboxlists]
|
||||
|
||||
return bboxlists
|
||||
|
||||
@property
|
||||
def reference_scale(self):
|
||||
return 195
|
||||
|
||||
@property
|
||||
def reference_x_shift(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def reference_y_shift(self):
|
||||
return 0
|
||||
@@ -0,0 +1,261 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False):
|
||||
"3x3 convolution with padding"
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3,
|
||||
stride=strd, padding=padding, bias=bias)
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_planes, out_planes):
|
||||
super(ConvBlock, self).__init__()
|
||||
self.bn1 = nn.BatchNorm2d(in_planes)
|
||||
self.conv1 = conv3x3(in_planes, int(out_planes / 2))
|
||||
self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
|
||||
self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4))
|
||||
self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
|
||||
self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4))
|
||||
|
||||
if in_planes != out_planes:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.BatchNorm2d(in_planes),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(in_planes, out_planes,
|
||||
kernel_size=1, stride=1, bias=False),
|
||||
)
|
||||
else:
|
||||
self.downsample = None
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out1 = self.bn1(x)
|
||||
out1 = F.relu(out1, True)
|
||||
out1 = self.conv1(out1)
|
||||
|
||||
out2 = self.bn2(out1)
|
||||
out2 = F.relu(out2, True)
|
||||
out2 = self.conv2(out2)
|
||||
|
||||
out3 = self.bn3(out2)
|
||||
out3 = F.relu(out3, True)
|
||||
out3 = self.conv3(out3)
|
||||
|
||||
out3 = torch.cat((out1, out2, out3), 1)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(residual)
|
||||
|
||||
out3 += residual
|
||||
|
||||
return out3
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * 4)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class HourGlass(nn.Module):
|
||||
def __init__(self, num_modules, depth, num_features):
|
||||
super(HourGlass, self).__init__()
|
||||
self.num_modules = num_modules
|
||||
self.depth = depth
|
||||
self.features = num_features
|
||||
|
||||
self._generate_network(self.depth)
|
||||
|
||||
def _generate_network(self, level):
|
||||
self.add_module('b1_' + str(level), ConvBlock(self.features, self.features))
|
||||
|
||||
self.add_module('b2_' + str(level), ConvBlock(self.features, self.features))
|
||||
|
||||
if level > 1:
|
||||
self._generate_network(level - 1)
|
||||
else:
|
||||
self.add_module('b2_plus_' + str(level), ConvBlock(self.features, self.features))
|
||||
|
||||
self.add_module('b3_' + str(level), ConvBlock(self.features, self.features))
|
||||
|
||||
def _forward(self, level, inp):
|
||||
# Upper branch
|
||||
up1 = inp
|
||||
up1 = self._modules['b1_' + str(level)](up1)
|
||||
|
||||
# Lower branch
|
||||
low1 = F.avg_pool2d(inp, 2, stride=2)
|
||||
low1 = self._modules['b2_' + str(level)](low1)
|
||||
|
||||
if level > 1:
|
||||
low2 = self._forward(level - 1, low1)
|
||||
else:
|
||||
low2 = low1
|
||||
low2 = self._modules['b2_plus_' + str(level)](low2)
|
||||
|
||||
low3 = low2
|
||||
low3 = self._modules['b3_' + str(level)](low3)
|
||||
|
||||
up2 = F.interpolate(low3, scale_factor=2, mode='nearest')
|
||||
|
||||
return up1 + up2
|
||||
|
||||
def forward(self, x):
|
||||
return self._forward(self.depth, x)
|
||||
|
||||
|
||||
class FAN(nn.Module):
|
||||
|
||||
def __init__(self, num_modules=1):
|
||||
super(FAN, self).__init__()
|
||||
self.num_modules = num_modules
|
||||
|
||||
# Base part
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.conv2 = ConvBlock(64, 128)
|
||||
self.conv3 = ConvBlock(128, 128)
|
||||
self.conv4 = ConvBlock(128, 256)
|
||||
|
||||
# Stacking part
|
||||
for hg_module in range(self.num_modules):
|
||||
self.add_module('m' + str(hg_module), HourGlass(1, 4, 256))
|
||||
self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256))
|
||||
self.add_module('conv_last' + str(hg_module),
|
||||
nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
|
||||
self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256))
|
||||
self.add_module('l' + str(hg_module), nn.Conv2d(256,
|
||||
68, kernel_size=1, stride=1, padding=0))
|
||||
|
||||
if hg_module < self.num_modules - 1:
|
||||
self.add_module(
|
||||
'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
|
||||
self.add_module('al' + str(hg_module), nn.Conv2d(68,
|
||||
256, kernel_size=1, stride=1, padding=0))
|
||||
|
||||
def forward(self, x):
|
||||
x = F.relu(self.bn1(self.conv1(x)), True)
|
||||
x = F.avg_pool2d(self.conv2(x), 2, stride=2)
|
||||
x = self.conv3(x)
|
||||
x = self.conv4(x)
|
||||
|
||||
previous = x
|
||||
|
||||
outputs = []
|
||||
for i in range(self.num_modules):
|
||||
hg = self._modules['m' + str(i)](previous)
|
||||
|
||||
ll = hg
|
||||
ll = self._modules['top_m_' + str(i)](ll)
|
||||
|
||||
ll = F.relu(self._modules['bn_end' + str(i)]
|
||||
(self._modules['conv_last' + str(i)](ll)), True)
|
||||
|
||||
# Predict heatmaps
|
||||
tmp_out = self._modules['l' + str(i)](ll)
|
||||
outputs.append(tmp_out)
|
||||
|
||||
if i < self.num_modules - 1:
|
||||
ll = self._modules['bl' + str(i)](ll)
|
||||
tmp_out_ = self._modules['al' + str(i)](tmp_out)
|
||||
previous = previous + ll + tmp_out_
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class ResNetDepth(nn.Module):
|
||||
|
||||
def __init__(self, block=Bottleneck, layers=[3, 8, 36, 3], num_classes=68):
|
||||
self.inplanes = 64
|
||||
super(ResNetDepth, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3 + 68, 64, kernel_size=7, stride=2, padding=3,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
self.avgpool = nn.AvgPool2d(7)
|
||||
self.fc = nn.Linear(512 * block.expansion, num_classes)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,313 @@
|
||||
from __future__ import print_function
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import torch
|
||||
import math
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
def _gaussian(
|
||||
size=3, sigma=0.25, amplitude=1, normalize=False, width=None,
|
||||
height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5,
|
||||
mean_vert=0.5):
|
||||
# handle some defaults
|
||||
if width is None:
|
||||
width = size
|
||||
if height is None:
|
||||
height = size
|
||||
if sigma_horz is None:
|
||||
sigma_horz = sigma
|
||||
if sigma_vert is None:
|
||||
sigma_vert = sigma
|
||||
center_x = mean_horz * width + 0.5
|
||||
center_y = mean_vert * height + 0.5
|
||||
gauss = np.empty((height, width), dtype=np.float32)
|
||||
# generate kernel
|
||||
for i in range(height):
|
||||
for j in range(width):
|
||||
gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / (
|
||||
sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0))
|
||||
if normalize:
|
||||
gauss = gauss / np.sum(gauss)
|
||||
return gauss
|
||||
|
||||
|
||||
def draw_gaussian(image, point, sigma):
|
||||
# Check if the gaussian is inside
|
||||
ul = [math.floor(point[0] - 3 * sigma), math.floor(point[1] - 3 * sigma)]
|
||||
br = [math.floor(point[0] + 3 * sigma), math.floor(point[1] + 3 * sigma)]
|
||||
if (ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1):
|
||||
return image
|
||||
size = 6 * sigma + 1
|
||||
g = _gaussian(size)
|
||||
g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]
|
||||
g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]
|
||||
img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
|
||||
img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
|
||||
assert (g_x[0] > 0 and g_y[1] > 0)
|
||||
image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]
|
||||
] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]]
|
||||
image[image > 1] = 1
|
||||
return image
|
||||
|
||||
|
||||
def transform(point, center, scale, resolution, invert=False):
|
||||
"""Generate and affine transformation matrix.
|
||||
|
||||
Given a set of points, a center, a scale and a targer resolution, the
|
||||
function generates and affine transformation matrix. If invert is ``True``
|
||||
it will produce the inverse transformation.
|
||||
|
||||
Arguments:
|
||||
point {torch.tensor} -- the input 2D point
|
||||
center {torch.tensor or numpy.array} -- the center around which to perform the transformations
|
||||
scale {float} -- the scale of the face/object
|
||||
resolution {float} -- the output resolution
|
||||
|
||||
Keyword Arguments:
|
||||
invert {bool} -- define wherever the function should produce the direct or the
|
||||
inverse transformation matrix (default: {False})
|
||||
"""
|
||||
_pt = torch.ones(3)
|
||||
_pt[0] = point[0]
|
||||
_pt[1] = point[1]
|
||||
|
||||
h = 200.0 * scale
|
||||
t = torch.eye(3)
|
||||
t[0, 0] = resolution / h
|
||||
t[1, 1] = resolution / h
|
||||
t[0, 2] = resolution * (-center[0] / h + 0.5)
|
||||
t[1, 2] = resolution * (-center[1] / h + 0.5)
|
||||
|
||||
if invert:
|
||||
t = torch.inverse(t)
|
||||
|
||||
new_point = (torch.matmul(t, _pt))[0:2]
|
||||
|
||||
return new_point.int()
|
||||
|
||||
|
||||
def crop(image, center, scale, resolution=256.0):
|
||||
"""Center crops an image or set of heatmaps
|
||||
|
||||
Arguments:
|
||||
image {numpy.array} -- an rgb image
|
||||
center {numpy.array} -- the center of the object, usually the same as of the bounding box
|
||||
scale {float} -- scale of the face
|
||||
|
||||
Keyword Arguments:
|
||||
resolution {float} -- the size of the output cropped image (default: {256.0})
|
||||
|
||||
Returns:
|
||||
[type] -- [description]
|
||||
""" # Crop around the center point
|
||||
""" Crops the image around the center. Input is expected to be an np.ndarray """
|
||||
ul = transform([1, 1], center, scale, resolution, True)
|
||||
br = transform([resolution, resolution], center, scale, resolution, True)
|
||||
# pad = math.ceil(torch.norm((ul - br).float()) / 2.0 - (br[0] - ul[0]) / 2.0)
|
||||
if image.ndim > 2:
|
||||
newDim = np.array([br[1] - ul[1], br[0] - ul[0],
|
||||
image.shape[2]], dtype=np.int32)
|
||||
newImg = np.zeros(newDim, dtype=np.uint8)
|
||||
else:
|
||||
newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)
|
||||
newImg = np.zeros(newDim, dtype=np.uint8)
|
||||
ht = image.shape[0]
|
||||
wd = image.shape[1]
|
||||
newX = np.array(
|
||||
[max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
|
||||
newY = np.array(
|
||||
[max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
|
||||
oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
|
||||
oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
|
||||
newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1]
|
||||
] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
|
||||
newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)),
|
||||
interpolation=cv2.INTER_LINEAR)
|
||||
return newImg
|
||||
|
||||
|
||||
def get_preds_fromhm(hm, center=None, scale=None):
|
||||
"""Obtain (x,y) coordinates given a set of N heatmaps. If the center
|
||||
and the scale is provided the function will return the points also in
|
||||
the original coordinate frame.
|
||||
|
||||
Arguments:
|
||||
hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
|
||||
|
||||
Keyword Arguments:
|
||||
center {torch.tensor} -- the center of the bounding box (default: {None})
|
||||
scale {float} -- face scale (default: {None})
|
||||
"""
|
||||
max, idx = torch.max(
|
||||
hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
|
||||
idx += 1
|
||||
preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
|
||||
preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
|
||||
preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
|
||||
|
||||
for i in range(preds.size(0)):
|
||||
for j in range(preds.size(1)):
|
||||
hm_ = hm[i, j, :]
|
||||
pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
|
||||
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
|
||||
diff = torch.FloatTensor(
|
||||
[hm_[pY, pX + 1] - hm_[pY, pX - 1],
|
||||
hm_[pY + 1, pX] - hm_[pY - 1, pX]])
|
||||
preds[i, j].add_(diff.sign_().mul_(.25))
|
||||
|
||||
preds.add_(-.5)
|
||||
|
||||
preds_orig = torch.zeros(preds.size())
|
||||
if center is not None and scale is not None:
|
||||
for i in range(hm.size(0)):
|
||||
for j in range(hm.size(1)):
|
||||
preds_orig[i, j] = transform(
|
||||
preds[i, j], center, scale, hm.size(2), True)
|
||||
|
||||
return preds, preds_orig
|
||||
|
||||
def get_preds_fromhm_batch(hm, centers=None, scales=None):
|
||||
"""Obtain (x,y) coordinates given a set of N heatmaps. If the centers
|
||||
and the scales is provided the function will return the points also in
|
||||
the original coordinate frame.
|
||||
|
||||
Arguments:
|
||||
hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H]
|
||||
|
||||
Keyword Arguments:
|
||||
centers {torch.tensor} -- the centers of the bounding box (default: {None})
|
||||
scales {float} -- face scales (default: {None})
|
||||
"""
|
||||
max, idx = torch.max(
|
||||
hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
|
||||
idx += 1
|
||||
preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
|
||||
preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
|
||||
preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
|
||||
|
||||
for i in range(preds.size(0)):
|
||||
for j in range(preds.size(1)):
|
||||
hm_ = hm[i, j, :]
|
||||
pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
|
||||
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
|
||||
diff = torch.FloatTensor(
|
||||
[hm_[pY, pX + 1] - hm_[pY, pX - 1],
|
||||
hm_[pY + 1, pX] - hm_[pY - 1, pX]])
|
||||
preds[i, j].add_(diff.sign_().mul_(.25))
|
||||
|
||||
preds.add_(-.5)
|
||||
|
||||
preds_orig = torch.zeros(preds.size())
|
||||
if centers is not None and scales is not None:
|
||||
for i in range(hm.size(0)):
|
||||
for j in range(hm.size(1)):
|
||||
preds_orig[i, j] = transform(
|
||||
preds[i, j], centers[i], scales[i], hm.size(2), True)
|
||||
|
||||
return preds, preds_orig
|
||||
|
||||
def shuffle_lr(parts, pairs=None):
|
||||
"""Shuffle the points left-right according to the axis of symmetry
|
||||
of the object.
|
||||
|
||||
Arguments:
|
||||
parts {torch.tensor} -- a 3D or 4D object containing the
|
||||
heatmaps.
|
||||
|
||||
Keyword Arguments:
|
||||
pairs {list of integers} -- [order of the flipped points] (default: {None})
|
||||
"""
|
||||
if pairs is None:
|
||||
pairs = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
|
||||
26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35,
|
||||
34, 33, 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41,
|
||||
40, 54, 53, 52, 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63,
|
||||
62, 61, 60, 67, 66, 65]
|
||||
if parts.ndimension() == 3:
|
||||
parts = parts[pairs, ...]
|
||||
else:
|
||||
parts = parts[:, pairs, ...]
|
||||
|
||||
return parts
|
||||
|
||||
|
||||
def flip(tensor, is_label=False):
|
||||
"""Flip an image or a set of heatmaps left-right
|
||||
|
||||
Arguments:
|
||||
tensor {numpy.array or torch.tensor} -- [the input image or heatmaps]
|
||||
|
||||
Keyword Arguments:
|
||||
is_label {bool} -- [denote wherever the input is an image or a set of heatmaps ] (default: {False})
|
||||
"""
|
||||
if not torch.is_tensor(tensor):
|
||||
tensor = torch.from_numpy(tensor)
|
||||
|
||||
if is_label:
|
||||
tensor = shuffle_lr(tensor).flip(tensor.ndimension() - 1)
|
||||
else:
|
||||
tensor = tensor.flip(tensor.ndimension() - 1)
|
||||
|
||||
return tensor
|
||||
|
||||
# From pyzolib/paths.py (https://bitbucket.org/pyzo/pyzolib/src/tip/paths.py)
|
||||
|
||||
|
||||
def appdata_dir(appname=None, roaming=False):
|
||||
""" appdata_dir(appname=None, roaming=False)
|
||||
|
||||
Get the path to the application directory, where applications are allowed
|
||||
to write user specific files (e.g. configurations). For non-user specific
|
||||
data, consider using common_appdata_dir().
|
||||
If appname is given, a subdir is appended (and created if necessary).
|
||||
If roaming is True, will prefer a roaming directory (Windows Vista/7).
|
||||
"""
|
||||
|
||||
# Define default user directory
|
||||
userDir = os.getenv('FACEALIGNMENT_USERDIR', None)
|
||||
if userDir is None:
|
||||
userDir = os.path.expanduser('~')
|
||||
if not os.path.isdir(userDir): # pragma: no cover
|
||||
userDir = '/var/tmp' # issue #54
|
||||
|
||||
# Get system app data dir
|
||||
path = None
|
||||
if sys.platform.startswith('win'):
|
||||
path1, path2 = os.getenv('LOCALAPPDATA'), os.getenv('APPDATA')
|
||||
path = (path2 or path1) if roaming else (path1 or path2)
|
||||
elif sys.platform.startswith('darwin'):
|
||||
path = os.path.join(userDir, 'Library', 'Application Support')
|
||||
# On Linux and as fallback
|
||||
if not (path and os.path.isdir(path)):
|
||||
path = userDir
|
||||
|
||||
# Maybe we should store things local to the executable (in case of a
|
||||
# portable distro or a frozen application that wants to be portable)
|
||||
prefix = sys.prefix
|
||||
if getattr(sys, 'frozen', None):
|
||||
prefix = os.path.abspath(os.path.dirname(sys.executable))
|
||||
for reldir in ('settings', '../settings'):
|
||||
localpath = os.path.abspath(os.path.join(prefix, reldir))
|
||||
if os.path.isdir(localpath): # pragma: no cover
|
||||
try:
|
||||
open(os.path.join(localpath, 'test.write'), 'wb').close()
|
||||
os.remove(os.path.join(localpath, 'test.write'))
|
||||
except IOError:
|
||||
pass # We cannot write in this directory
|
||||
else:
|
||||
path = localpath
|
||||
break
|
||||
|
||||
# Get path specific for this app
|
||||
if appname:
|
||||
if path == userDir:
|
||||
appname = '.' + appname.lstrip('.') # Make it a hidden directory
|
||||
path = os.path.join(path, appname)
|
||||
if not os.path.isdir(path): # pragma: no cover
|
||||
os.mkdir(path)
|
||||
|
||||
# Done
|
||||
return path
|
||||
@@ -0,0 +1 @@
|
||||
Place LRS2 (and any other) filelists here for training.
|
||||
@@ -0,0 +1,86 @@
|
||||
from tensorflow.contrib.training import HParams
|
||||
from glob import glob
|
||||
import os
|
||||
|
||||
def get_image_list(data_root, split):
|
||||
filelist = []
|
||||
|
||||
with open('filelists/{}.txt'.format(split)) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if ' ' in line: line = line.split()[0]
|
||||
filelist.append(os.path.join(data_root, line))
|
||||
|
||||
return filelist
|
||||
|
||||
# Default hyperparameters
|
||||
hparams = HParams(
|
||||
num_mels=80, # Number of mel-spectrogram channels and local conditioning dimensionality
|
||||
# network
|
||||
rescale=True, # Whether to rescale audio prior to preprocessing
|
||||
rescaling_max=0.9, # Rescaling value
|
||||
|
||||
# Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
|
||||
# It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
|
||||
# Does not work if n_ffit is not multiple of hop_size!!
|
||||
use_lws=False,
|
||||
|
||||
n_fft=800, # Extra window size is filled with 0 paddings to match this parameter
|
||||
hop_size=200, # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
|
||||
win_size=800, # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
|
||||
sample_rate=16000, # 16000Hz (corresponding to librispeech) (sox --i <filename>)
|
||||
|
||||
frame_shift_ms=None, # Can replace hop_size parameter. (Recommended: 12.5)
|
||||
|
||||
# Mel and Linear spectrograms normalization/scaling and clipping
|
||||
signal_normalization=True,
|
||||
# Whether to normalize mel spectrograms to some predefined range (following below parameters)
|
||||
allow_clipping_in_normalization=True, # Only relevant if mel_normalization = True
|
||||
symmetric_mels=True,
|
||||
# Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2,
|
||||
# faster and cleaner convergence)
|
||||
max_abs_value=4.,
|
||||
# max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not
|
||||
# be too big to avoid gradient explosion,
|
||||
# not too small for fast convergence)
|
||||
# Contribution by @begeekmyfriend
|
||||
# Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude
|
||||
# levels. Also allows for better G&L phase reconstruction)
|
||||
preemphasize=True, # whether to apply filter
|
||||
preemphasis=0.97, # filter coefficient.
|
||||
|
||||
# Limits
|
||||
min_level_db=-100,
|
||||
ref_level_db=20,
|
||||
fmin=55,
|
||||
# Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To
|
||||
# test depending on dataset. Pitch info: male~[65, 260], female~[100, 525])
|
||||
fmax=7600, # To be increased/reduced depending on data.
|
||||
|
||||
###################### Our training parameters #################################
|
||||
img_size=96,
|
||||
fps=25,
|
||||
|
||||
batch_size=16,
|
||||
initial_learning_rate=1e-4,
|
||||
nepochs=200000000000000000, ### ctrl + c, stop whenever eval loss is consistently greater than train loss for ~10 epochs
|
||||
num_workers=16,
|
||||
checkpoint_interval=3000,
|
||||
eval_interval=3000,
|
||||
save_optimizer_state=True,
|
||||
|
||||
syncnet_wt=0.0, # is initially zero, will be set automatically to 0.03 later. Leads to faster convergence.
|
||||
syncnet_batch_size=64,
|
||||
syncnet_lr=1e-4,
|
||||
syncnet_eval_interval=10000,
|
||||
syncnet_checkpoint_interval=10000,
|
||||
|
||||
disc_wt=0.07,
|
||||
disc_initial_learning_rate=1e-4,
|
||||
)
|
||||
|
||||
|
||||
def hparams_debug_string():
|
||||
values = hparams.values()
|
||||
hp = [" %s: %s" % (name, values[name]) for name in sorted(values) if name != "sentences"]
|
||||
return "Hyperparameters:\n" + "\n".join(hp)
|
||||
@@ -0,0 +1,443 @@
|
||||
from os.path import dirname, join, basename, isfile
|
||||
from tqdm import tqdm
|
||||
|
||||
from models import SyncNet_color as SyncNet
|
||||
from models import Wav2Lip, Wav2Lip_disc_qual
|
||||
import audio
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from torch import optim
|
||||
import torch.backends.cudnn as cudnn
|
||||
from torch.utils import data as data_utils
|
||||
import numpy as np
|
||||
|
||||
from glob import glob
|
||||
|
||||
import os, random, cv2, argparse
|
||||
from hparams import hparams, get_image_list
|
||||
|
||||
parser = argparse.ArgumentParser(description='Code to train the Wav2Lip model WITH the visual quality discriminator')
|
||||
|
||||
parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True, type=str)
|
||||
|
||||
parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str)
|
||||
parser.add_argument('--syncnet_checkpoint_path', help='Load the pre-trained Expert discriminator', required=True, type=str)
|
||||
|
||||
parser.add_argument('--checkpoint_path', help='Resume generator from this checkpoint', default=None, type=str)
|
||||
parser.add_argument('--disc_checkpoint_path', help='Resume quality disc from this checkpoint', default=None, type=str)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
global_step = 0
|
||||
global_epoch = 0
|
||||
use_cuda = torch.cuda.is_available()
|
||||
print('use_cuda: {}'.format(use_cuda))
|
||||
|
||||
syncnet_T = 5
|
||||
syncnet_mel_step_size = 16
|
||||
|
||||
class Dataset(object):
|
||||
def __init__(self, split):
|
||||
self.all_videos = get_image_list(args.data_root, split)
|
||||
|
||||
def get_frame_id(self, frame):
|
||||
return int(basename(frame).split('.')[0])
|
||||
|
||||
def get_window(self, start_frame):
|
||||
start_id = self.get_frame_id(start_frame)
|
||||
vidname = dirname(start_frame)
|
||||
|
||||
window_fnames = []
|
||||
for frame_id in range(start_id, start_id + syncnet_T):
|
||||
frame = join(vidname, '{}.jpg'.format(frame_id))
|
||||
if not isfile(frame):
|
||||
return None
|
||||
window_fnames.append(frame)
|
||||
return window_fnames
|
||||
|
||||
def read_window(self, window_fnames):
|
||||
if window_fnames is None: return None
|
||||
window = []
|
||||
for fname in window_fnames:
|
||||
img = cv2.imread(fname)
|
||||
if img is None:
|
||||
return None
|
||||
try:
|
||||
img = cv2.resize(img, (hparams.img_size, hparams.img_size))
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
window.append(img)
|
||||
|
||||
return window
|
||||
|
||||
def crop_audio_window(self, spec, start_frame):
|
||||
if type(start_frame) == int:
|
||||
start_frame_num = start_frame
|
||||
else:
|
||||
start_frame_num = self.get_frame_id(start_frame)
|
||||
start_idx = int(80. * (start_frame_num / float(hparams.fps)))
|
||||
|
||||
end_idx = start_idx + syncnet_mel_step_size
|
||||
|
||||
return spec[start_idx : end_idx, :]
|
||||
|
||||
def get_segmented_mels(self, spec, start_frame):
|
||||
mels = []
|
||||
assert syncnet_T == 5
|
||||
start_frame_num = self.get_frame_id(start_frame) + 1 # 0-indexing ---> 1-indexing
|
||||
if start_frame_num - 2 < 0: return None
|
||||
for i in range(start_frame_num, start_frame_num + syncnet_T):
|
||||
m = self.crop_audio_window(spec, i - 2)
|
||||
if m.shape[0] != syncnet_mel_step_size:
|
||||
return None
|
||||
mels.append(m.T)
|
||||
|
||||
mels = np.asarray(mels)
|
||||
|
||||
return mels
|
||||
|
||||
def prepare_window(self, window):
|
||||
# 3 x T x H x W
|
||||
x = np.asarray(window) / 255.
|
||||
x = np.transpose(x, (3, 0, 1, 2))
|
||||
|
||||
return x
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_videos)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
while 1:
|
||||
idx = random.randint(0, len(self.all_videos) - 1)
|
||||
vidname = self.all_videos[idx]
|
||||
img_names = list(glob(join(vidname, '*.jpg')))
|
||||
if len(img_names) <= 3 * syncnet_T:
|
||||
continue
|
||||
|
||||
img_name = random.choice(img_names)
|
||||
wrong_img_name = random.choice(img_names)
|
||||
while wrong_img_name == img_name:
|
||||
wrong_img_name = random.choice(img_names)
|
||||
|
||||
window_fnames = self.get_window(img_name)
|
||||
wrong_window_fnames = self.get_window(wrong_img_name)
|
||||
if window_fnames is None or wrong_window_fnames is None:
|
||||
continue
|
||||
|
||||
window = self.read_window(window_fnames)
|
||||
if window is None:
|
||||
continue
|
||||
|
||||
wrong_window = self.read_window(wrong_window_fnames)
|
||||
if wrong_window is None:
|
||||
continue
|
||||
|
||||
try:
|
||||
wavpath = join(vidname, "audio.wav")
|
||||
wav = audio.load_wav(wavpath, hparams.sample_rate)
|
||||
|
||||
orig_mel = audio.melspectrogram(wav).T
|
||||
except Exception as e:
|
||||
continue
|
||||
|
||||
mel = self.crop_audio_window(orig_mel.copy(), img_name)
|
||||
|
||||
if (mel.shape[0] != syncnet_mel_step_size):
|
||||
continue
|
||||
|
||||
indiv_mels = self.get_segmented_mels(orig_mel.copy(), img_name)
|
||||
if indiv_mels is None: continue
|
||||
|
||||
window = self.prepare_window(window)
|
||||
y = window.copy()
|
||||
window[:, :, window.shape[2]//2:] = 0.
|
||||
|
||||
wrong_window = self.prepare_window(wrong_window)
|
||||
x = np.concatenate([window, wrong_window], axis=0)
|
||||
|
||||
x = torch.FloatTensor(x)
|
||||
mel = torch.FloatTensor(mel.T).unsqueeze(0)
|
||||
indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1)
|
||||
y = torch.FloatTensor(y)
|
||||
return x, indiv_mels, mel, y
|
||||
|
||||
def save_sample_images(x, g, gt, global_step, checkpoint_dir):
|
||||
x = (x.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
|
||||
g = (g.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
|
||||
gt = (gt.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
|
||||
|
||||
refs, inps = x[..., 3:], x[..., :3]
|
||||
folder = join(checkpoint_dir, "samples_step{:09d}".format(global_step))
|
||||
if not os.path.exists(folder): os.mkdir(folder)
|
||||
collage = np.concatenate((refs, inps, g, gt), axis=-2)
|
||||
for batch_idx, c in enumerate(collage):
|
||||
for t in range(len(c)):
|
||||
cv2.imwrite('{}/{}_{}.jpg'.format(folder, batch_idx, t), c[t])
|
||||
|
||||
logloss = nn.BCELoss()
|
||||
def cosine_loss(a, v, y):
|
||||
d = nn.functional.cosine_similarity(a, v)
|
||||
loss = logloss(d.unsqueeze(1), y)
|
||||
|
||||
return loss
|
||||
|
||||
device = torch.device("cuda" if use_cuda else "cpu")
|
||||
syncnet = SyncNet().to(device)
|
||||
for p in syncnet.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
recon_loss = nn.L1Loss()
|
||||
def get_sync_loss(mel, g):
|
||||
g = g[:, :, :, g.size(3)//2:]
|
||||
g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1)
|
||||
# B, 3 * T, H//2, W
|
||||
a, v = syncnet(mel, g)
|
||||
y = torch.ones(g.size(0), 1).float().to(device)
|
||||
return cosine_loss(a, v, y)
|
||||
|
||||
def train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer,
|
||||
checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
|
||||
global global_step, global_epoch
|
||||
resumed_step = global_step
|
||||
|
||||
while global_epoch < nepochs:
|
||||
print('Starting Epoch: {}'.format(global_epoch))
|
||||
running_sync_loss, running_l1_loss, disc_loss, running_perceptual_loss = 0., 0., 0., 0.
|
||||
running_disc_real_loss, running_disc_fake_loss = 0., 0.
|
||||
prog_bar = tqdm(enumerate(train_data_loader))
|
||||
for step, (x, indiv_mels, mel, gt) in prog_bar:
|
||||
disc.train()
|
||||
model.train()
|
||||
|
||||
x = x.to(device)
|
||||
mel = mel.to(device)
|
||||
indiv_mels = indiv_mels.to(device)
|
||||
gt = gt.to(device)
|
||||
|
||||
### Train generator now. Remove ALL grads.
|
||||
optimizer.zero_grad()
|
||||
disc_optimizer.zero_grad()
|
||||
|
||||
g = model(indiv_mels, x)
|
||||
|
||||
if hparams.syncnet_wt > 0.:
|
||||
sync_loss = get_sync_loss(mel, g)
|
||||
else:
|
||||
sync_loss = 0.
|
||||
|
||||
if hparams.disc_wt > 0.:
|
||||
perceptual_loss = disc.perceptual_forward(g)
|
||||
else:
|
||||
perceptual_loss = 0.
|
||||
|
||||
l1loss = recon_loss(g, gt)
|
||||
|
||||
loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \
|
||||
(1. - hparams.syncnet_wt - hparams.disc_wt) * l1loss
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
### Remove all gradients before Training disc
|
||||
disc_optimizer.zero_grad()
|
||||
|
||||
pred = disc(gt)
|
||||
disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device))
|
||||
disc_real_loss.backward()
|
||||
|
||||
pred = disc(g.detach())
|
||||
disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device))
|
||||
disc_fake_loss.backward()
|
||||
|
||||
disc_optimizer.step()
|
||||
|
||||
running_disc_real_loss += disc_real_loss.item()
|
||||
running_disc_fake_loss += disc_fake_loss.item()
|
||||
|
||||
if global_step % checkpoint_interval == 0:
|
||||
save_sample_images(x, g, gt, global_step, checkpoint_dir)
|
||||
|
||||
# Logs
|
||||
global_step += 1
|
||||
cur_session_steps = global_step - resumed_step
|
||||
|
||||
running_l1_loss += l1loss.item()
|
||||
if hparams.syncnet_wt > 0.:
|
||||
running_sync_loss += sync_loss.item()
|
||||
else:
|
||||
running_sync_loss += 0.
|
||||
|
||||
if hparams.disc_wt > 0.:
|
||||
running_perceptual_loss += perceptual_loss.item()
|
||||
else:
|
||||
running_perceptual_loss += 0.
|
||||
|
||||
if global_step == 1 or global_step % checkpoint_interval == 0:
|
||||
save_checkpoint(
|
||||
model, optimizer, global_step, checkpoint_dir, global_epoch)
|
||||
save_checkpoint(disc, disc_optimizer, global_step, checkpoint_dir, global_epoch, prefix='disc_')
|
||||
|
||||
|
||||
if global_step % hparams.eval_interval == 0:
|
||||
with torch.no_grad():
|
||||
average_sync_loss = eval_model(test_data_loader, global_step, device, model, disc)
|
||||
|
||||
if average_sync_loss < .75:
|
||||
hparams.set_hparam('syncnet_wt', 0.03)
|
||||
|
||||
prog_bar.set_description('L1: {}, Sync: {}, Percep: {} | Fake: {}, Real: {}'.format(running_l1_loss / (step + 1),
|
||||
running_sync_loss / (step + 1),
|
||||
running_perceptual_loss / (step + 1),
|
||||
running_disc_fake_loss / (step + 1),
|
||||
running_disc_real_loss / (step + 1)))
|
||||
|
||||
global_epoch += 1
|
||||
|
||||
def eval_model(test_data_loader, global_step, device, model, disc):
|
||||
eval_steps = 300
|
||||
print('Evaluating for {} steps'.format(eval_steps))
|
||||
running_sync_loss, running_l1_loss, running_disc_real_loss, running_disc_fake_loss, running_perceptual_loss = [], [], [], [], []
|
||||
while 1:
|
||||
for step, (x, indiv_mels, mel, gt) in enumerate((test_data_loader)):
|
||||
model.eval()
|
||||
disc.eval()
|
||||
|
||||
x = x.to(device)
|
||||
mel = mel.to(device)
|
||||
indiv_mels = indiv_mels.to(device)
|
||||
gt = gt.to(device)
|
||||
|
||||
pred = disc(gt)
|
||||
disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device))
|
||||
|
||||
g = model(indiv_mels, x)
|
||||
pred = disc(g)
|
||||
disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device))
|
||||
|
||||
running_disc_real_loss.append(disc_real_loss.item())
|
||||
running_disc_fake_loss.append(disc_fake_loss.item())
|
||||
|
||||
sync_loss = get_sync_loss(mel, g)
|
||||
|
||||
if hparams.disc_wt > 0.:
|
||||
perceptual_loss = disc.perceptual_forward(g)
|
||||
else:
|
||||
perceptual_loss = 0.
|
||||
|
||||
l1loss = recon_loss(g, gt)
|
||||
|
||||
loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \
|
||||
(1. - hparams.syncnet_wt - hparams.disc_wt) * l1loss
|
||||
|
||||
running_l1_loss.append(l1loss.item())
|
||||
running_sync_loss.append(sync_loss.item())
|
||||
|
||||
if hparams.disc_wt > 0.:
|
||||
running_perceptual_loss.append(perceptual_loss.item())
|
||||
else:
|
||||
running_perceptual_loss.append(0.)
|
||||
|
||||
if step > eval_steps: break
|
||||
|
||||
print('L1: {}, Sync: {}, Percep: {} | Fake: {}, Real: {}'.format(sum(running_l1_loss) / len(running_l1_loss),
|
||||
sum(running_sync_loss) / len(running_sync_loss),
|
||||
sum(running_perceptual_loss) / len(running_perceptual_loss),
|
||||
sum(running_disc_fake_loss) / len(running_disc_fake_loss),
|
||||
sum(running_disc_real_loss) / len(running_disc_real_loss)))
|
||||
return sum(running_sync_loss) / len(running_sync_loss)
|
||||
|
||||
|
||||
def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch, prefix=''):
|
||||
checkpoint_path = join(
|
||||
checkpoint_dir, "{}checkpoint_step{:09d}.pth".format(prefix, global_step))
|
||||
optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None
|
||||
torch.save({
|
||||
"state_dict": model.state_dict(),
|
||||
"optimizer": optimizer_state,
|
||||
"global_step": step,
|
||||
"global_epoch": epoch,
|
||||
}, checkpoint_path)
|
||||
print("Saved checkpoint:", checkpoint_path)
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if use_cuda:
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
|
||||
def load_checkpoint(path, model, optimizer, reset_optimizer=False, overwrite_global_states=True):
|
||||
global global_step
|
||||
global global_epoch
|
||||
|
||||
print("Load checkpoint from: {}".format(path))
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
if not reset_optimizer:
|
||||
optimizer_state = checkpoint["optimizer"]
|
||||
if optimizer_state is not None:
|
||||
print("Load optimizer state from {}".format(path))
|
||||
optimizer.load_state_dict(checkpoint["optimizer"])
|
||||
if overwrite_global_states:
|
||||
global_step = checkpoint["global_step"]
|
||||
global_epoch = checkpoint["global_epoch"]
|
||||
|
||||
return model
|
||||
|
||||
if __name__ == "__main__":
|
||||
checkpoint_dir = args.checkpoint_dir
|
||||
|
||||
# Dataset and Dataloader setup
|
||||
train_dataset = Dataset('train')
|
||||
test_dataset = Dataset('val')
|
||||
|
||||
train_data_loader = data_utils.DataLoader(
|
||||
train_dataset, batch_size=hparams.batch_size, shuffle=True,
|
||||
num_workers=hparams.num_workers)
|
||||
|
||||
test_data_loader = data_utils.DataLoader(
|
||||
test_dataset, batch_size=hparams.batch_size,
|
||||
num_workers=4)
|
||||
|
||||
device = torch.device("cuda" if use_cuda else "cpu")
|
||||
|
||||
# Model
|
||||
model = Wav2Lip().to(device)
|
||||
disc = Wav2Lip_disc_qual().to(device)
|
||||
|
||||
print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad)))
|
||||
print('total DISC trainable params {}'.format(sum(p.numel() for p in disc.parameters() if p.requires_grad)))
|
||||
|
||||
optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
|
||||
lr=hparams.initial_learning_rate, betas=(0.5, 0.999))
|
||||
disc_optimizer = optim.Adam([p for p in disc.parameters() if p.requires_grad],
|
||||
lr=hparams.disc_initial_learning_rate, betas=(0.5, 0.999))
|
||||
|
||||
if args.checkpoint_path is not None:
|
||||
load_checkpoint(args.checkpoint_path, model, optimizer, reset_optimizer=False)
|
||||
|
||||
if args.disc_checkpoint_path is not None:
|
||||
load_checkpoint(args.disc_checkpoint_path, disc, disc_optimizer,
|
||||
reset_optimizer=False, overwrite_global_states=False)
|
||||
|
||||
load_checkpoint(args.syncnet_checkpoint_path, syncnet, None, reset_optimizer=True,
|
||||
overwrite_global_states=False)
|
||||
|
||||
if not os.path.exists(checkpoint_dir):
|
||||
os.mkdir(checkpoint_dir)
|
||||
|
||||
# Train!
|
||||
train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer,
|
||||
checkpoint_dir=checkpoint_dir,
|
||||
checkpoint_interval=hparams.checkpoint_interval,
|
||||
nepochs=hparams.nepochs)
|
||||
@@ -0,0 +1,286 @@
|
||||
from os import listdir, path
|
||||
import numpy as np
|
||||
import scipy, cv2, os, sys, argparse, audio
|
||||
import json, subprocess, random, string
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch, face_detection
|
||||
from models import Wav2Lip
|
||||
|
||||
for k, v in os.environ.items():
|
||||
print('{}={}'.format(k, v))
|
||||
|
||||
print('PyTorch version', torch.__version__)
|
||||
print('CUDA Available?', torch.cuda.is_available())
|
||||
print('CUDA Device count', torch.cuda.device_count())
|
||||
|
||||
parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
|
||||
|
||||
parser.add_argument('--checkpoint_path', type=str,
|
||||
help='Name of saved checkpoint to load weights from', required=True)
|
||||
|
||||
parser.add_argument('--face', type=str,
|
||||
help='Filepath of video/image that contains faces to use', required=True)
|
||||
parser.add_argument('--audio', type=str,
|
||||
help='Filepath of video/audio file to use as raw audio source', required=True)
|
||||
parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.',
|
||||
default='results/result_voice.mp4')
|
||||
|
||||
parser.add_argument('--static', type=bool,
|
||||
help='If True, then use only first video frame for inference', default=False)
|
||||
parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)',
|
||||
default=25., required=False)
|
||||
|
||||
parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0],
|
||||
help='Padding (top, bottom, left, right). Please adjust to include chin at least')
|
||||
|
||||
parser.add_argument('--face_det_batch_size', type=int,
|
||||
help='Batch size for face detection', default=16)
|
||||
parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
|
||||
|
||||
parser.add_argument('--resize_factor', default=1, type=int,
|
||||
help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
|
||||
|
||||
parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1],
|
||||
help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. '
|
||||
'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
|
||||
|
||||
parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1],
|
||||
help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
|
||||
'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
|
||||
|
||||
parser.add_argument('--rotate', default=False, action='store_true',
|
||||
help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
|
||||
'Use if you get a flipped result, despite feeding a normal looking video')
|
||||
|
||||
parser.add_argument('--nosmooth', default=False, action='store_true',
|
||||
help='Prevent smoothing face detections over a short temporal window')
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = 96
|
||||
|
||||
if os.path.isfile(args.face) and args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
|
||||
args.static = True
|
||||
|
||||
def get_smoothened_boxes(boxes, T):
|
||||
for i in range(len(boxes)):
|
||||
if i + T > len(boxes):
|
||||
window = boxes[len(boxes) - T:]
|
||||
else:
|
||||
window = boxes[i : i + T]
|
||||
boxes[i] = np.mean(window, axis=0)
|
||||
return boxes
|
||||
|
||||
def face_detect(images):
|
||||
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
||||
flip_input=False, device=device)
|
||||
|
||||
batch_size = args.face_det_batch_size
|
||||
|
||||
while 1:
|
||||
predictions = []
|
||||
try:
|
||||
for i in tqdm(range(0, len(images), batch_size)):
|
||||
predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
|
||||
except RuntimeError:
|
||||
if batch_size == 1:
|
||||
raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
|
||||
batch_size //= 2
|
||||
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
||||
continue
|
||||
break
|
||||
|
||||
results = []
|
||||
pady1, pady2, padx1, padx2 = args.pads
|
||||
for rect, image in zip(predictions, images):
|
||||
if rect is None:
|
||||
cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
|
||||
raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
|
||||
|
||||
y1 = max(0, rect[1] - pady1)
|
||||
y2 = min(image.shape[0], rect[3] + pady2)
|
||||
x1 = max(0, rect[0] - padx1)
|
||||
x2 = min(image.shape[1], rect[2] + padx2)
|
||||
|
||||
results.append([x1, y1, x2, y2])
|
||||
|
||||
boxes = np.array(results)
|
||||
if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
|
||||
results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
|
||||
|
||||
del detector
|
||||
return results
|
||||
|
||||
def datagen(frames, mels):
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if args.box[0] == -1:
|
||||
if not args.static:
|
||||
face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
|
||||
else:
|
||||
face_det_results = face_detect([frames[0]])
|
||||
else:
|
||||
print('Using the specified bounding box instead of face detection...')
|
||||
y1, y2, x1, x2 = args.box
|
||||
face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
|
||||
|
||||
for i, m in enumerate(mels):
|
||||
idx = 0 if args.static else i%len(frames)
|
||||
frame_to_save = frames[idx].copy()
|
||||
face, coords = face_det_results[idx].copy()
|
||||
|
||||
face = cv2.resize(face, (args.img_size, args.img_size))
|
||||
|
||||
img_batch.append(face)
|
||||
mel_batch.append(m)
|
||||
frame_batch.append(frame_to_save)
|
||||
coords_batch.append(coords)
|
||||
|
||||
if len(img_batch) >= args.wav2lip_batch_size:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if len(img_batch) > 0:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
|
||||
mel_step_size = 16
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
print('Using {} for inference.'.format(device))
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if device == 'cuda':
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_model(path):
|
||||
model = Wav2Lip()
|
||||
print("Load checkpoint from: {}".format(path))
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
|
||||
model = model.to(device)
|
||||
return model.eval()
|
||||
|
||||
def main():
|
||||
if not os.path.isfile(args.face):
|
||||
raise ValueError('--face argument must be a valid path to video/image file')
|
||||
|
||||
elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
|
||||
full_frames = [cv2.imread(args.face)]
|
||||
fps = args.fps
|
||||
|
||||
else:
|
||||
video_stream = cv2.VideoCapture(args.face)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
|
||||
print('Reading video frames...')
|
||||
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
if args.resize_factor > 1:
|
||||
frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
|
||||
|
||||
if args.rotate:
|
||||
frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
|
||||
|
||||
y1, y2, x1, x2 = args.crop
|
||||
if x2 == -1: x2 = frame.shape[1]
|
||||
if y2 == -1: y2 = frame.shape[0]
|
||||
|
||||
frame = frame[y1:y2, x1:x2]
|
||||
|
||||
full_frames.append(frame)
|
||||
|
||||
print ("Number of frames available for inference: "+str(len(full_frames)))
|
||||
|
||||
if not args.audio.endswith('.wav'):
|
||||
print('Extracting raw audio...')
|
||||
command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, 'temp/temp.wav')
|
||||
|
||||
subprocess.call(command, shell=True)
|
||||
args.audio = 'temp/temp.wav'
|
||||
|
||||
wav = audio.load_wav(args.audio, 16000)
|
||||
mel = audio.melspectrogram(wav)
|
||||
print(mel.shape)
|
||||
|
||||
if np.isnan(mel.reshape(-1)).sum() > 0:
|
||||
raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
|
||||
|
||||
mel_chunks = []
|
||||
mel_idx_multiplier = 80./fps
|
||||
i = 0
|
||||
while 1:
|
||||
start_idx = int(i * mel_idx_multiplier)
|
||||
if start_idx + mel_step_size > len(mel[0]):
|
||||
mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
|
||||
break
|
||||
mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
|
||||
i += 1
|
||||
|
||||
print("Length of mel chunks: {}".format(len(mel_chunks)))
|
||||
|
||||
full_frames = full_frames[:len(mel_chunks)]
|
||||
|
||||
batch_size = args.wav2lip_batch_size
|
||||
gen = datagen(full_frames.copy(), mel_chunks)
|
||||
|
||||
for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen,
|
||||
total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
|
||||
if i == 0:
|
||||
model = load_model(args.checkpoint_path)
|
||||
print ("Model loaded")
|
||||
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
out = cv2.VideoWriter('temp/result.avi',
|
||||
cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
|
||||
|
||||
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
|
||||
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(mel_batch, img_batch)
|
||||
|
||||
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
|
||||
|
||||
for p, f, c in zip(pred, frames, coords):
|
||||
y1, y2, x1, x2 = c
|
||||
p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
|
||||
|
||||
f[y1:y2, x1:x2] = p
|
||||
out.write(f)
|
||||
|
||||
out.release()
|
||||
|
||||
command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/result.avi', args.outfile)
|
||||
subprocess.call(command, shell=True)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,2 @@
|
||||
from .wav2lip import Wav2Lip, Wav2Lip_disc_qual
|
||||
from .syncnet import SyncNet_color
|
||||
@@ -0,0 +1,44 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
class Conv2d(nn.Module):
|
||||
def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(cin, cout, kernel_size, stride, padding),
|
||||
nn.BatchNorm2d(cout)
|
||||
)
|
||||
self.act = nn.ReLU()
|
||||
self.residual = residual
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv_block(x)
|
||||
if self.residual:
|
||||
out += x
|
||||
return self.act(out)
|
||||
|
||||
class nonorm_Conv2d(nn.Module):
|
||||
def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(cin, cout, kernel_size, stride, padding),
|
||||
)
|
||||
self.act = nn.LeakyReLU(0.01, inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv_block(x)
|
||||
return self.act(out)
|
||||
|
||||
class Conv2dTranspose(nn.Module):
|
||||
def __init__(self, cin, cout, kernel_size, stride, padding, output_padding=0, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.ConvTranspose2d(cin, cout, kernel_size, stride, padding, output_padding),
|
||||
nn.BatchNorm2d(cout)
|
||||
)
|
||||
self.act = nn.ReLU()
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv_block(x)
|
||||
return self.act(out)
|
||||
@@ -0,0 +1,66 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from .conv import Conv2d
|
||||
|
||||
class SyncNet_color(nn.Module):
|
||||
def __init__(self):
|
||||
super(SyncNet_color, self).__init__()
|
||||
|
||||
self.face_encoder = nn.Sequential(
|
||||
Conv2d(15, 32, kernel_size=(7, 7), stride=1, padding=3),
|
||||
|
||||
Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=1),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=0),
|
||||
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
|
||||
|
||||
self.audio_encoder = nn.Sequential(
|
||||
Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
|
||||
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
|
||||
|
||||
def forward(self, audio_sequences, face_sequences): # audio_sequences := (B, dim, T)
|
||||
face_embedding = self.face_encoder(face_sequences)
|
||||
audio_embedding = self.audio_encoder(audio_sequences)
|
||||
|
||||
audio_embedding = audio_embedding.view(audio_embedding.size(0), -1)
|
||||
face_embedding = face_embedding.view(face_embedding.size(0), -1)
|
||||
|
||||
audio_embedding = F.normalize(audio_embedding, p=2, dim=1)
|
||||
face_embedding = F.normalize(face_embedding, p=2, dim=1)
|
||||
|
||||
|
||||
return audio_embedding, face_embedding
|
||||
@@ -0,0 +1,184 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
import math
|
||||
|
||||
from .conv import Conv2dTranspose, Conv2d, nonorm_Conv2d
|
||||
|
||||
class Wav2Lip(nn.Module):
|
||||
def __init__(self):
|
||||
super(Wav2Lip, self).__init__()
|
||||
|
||||
self.face_encoder_blocks = nn.ModuleList([
|
||||
nn.Sequential(Conv2d(6, 16, kernel_size=7, stride=1, padding=3)), # 96,96
|
||||
|
||||
nn.Sequential(Conv2d(16, 32, kernel_size=3, stride=2, padding=1), # 48,48
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True)),
|
||||
|
||||
nn.Sequential(Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 24,24
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True)),
|
||||
|
||||
nn.Sequential(Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 12,12
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True)),
|
||||
|
||||
nn.Sequential(Conv2d(128, 256, kernel_size=3, stride=2, padding=1), # 6,6
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True)),
|
||||
|
||||
nn.Sequential(Conv2d(256, 512, kernel_size=3, stride=2, padding=1), # 3,3
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),),
|
||||
|
||||
nn.Sequential(Conv2d(512, 512, kernel_size=3, stride=1, padding=0), # 1, 1
|
||||
Conv2d(512, 512, kernel_size=1, stride=1, padding=0)),])
|
||||
|
||||
self.audio_encoder = nn.Sequential(
|
||||
Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
|
||||
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
|
||||
|
||||
self.face_decoder_blocks = nn.ModuleList([
|
||||
nn.Sequential(Conv2d(512, 512, kernel_size=1, stride=1, padding=0),),
|
||||
|
||||
nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=1, padding=0), # 3,3
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),),
|
||||
|
||||
nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=2, padding=1, output_padding=1),
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),), # 6, 6
|
||||
|
||||
nn.Sequential(Conv2dTranspose(768, 384, kernel_size=3, stride=2, padding=1, output_padding=1),
|
||||
Conv2d(384, 384, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(384, 384, kernel_size=3, stride=1, padding=1, residual=True),), # 12, 12
|
||||
|
||||
nn.Sequential(Conv2dTranspose(512, 256, kernel_size=3, stride=2, padding=1, output_padding=1),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),), # 24, 24
|
||||
|
||||
nn.Sequential(Conv2dTranspose(320, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),), # 48, 48
|
||||
|
||||
nn.Sequential(Conv2dTranspose(160, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),),]) # 96,96
|
||||
|
||||
self.output_block = nn.Sequential(Conv2d(80, 32, kernel_size=3, stride=1, padding=1),
|
||||
nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0),
|
||||
nn.Sigmoid())
|
||||
|
||||
def forward(self, audio_sequences, face_sequences):
|
||||
# audio_sequences = (B, T, 1, 80, 16)
|
||||
B = audio_sequences.size(0)
|
||||
|
||||
input_dim_size = len(face_sequences.size())
|
||||
if input_dim_size > 4:
|
||||
audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
|
||||
face_sequences = torch.cat([face_sequences[:, :, i] for i in range(face_sequences.size(2))], dim=0)
|
||||
|
||||
audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
|
||||
|
||||
feats = []
|
||||
x = face_sequences
|
||||
for f in self.face_encoder_blocks:
|
||||
x = f(x)
|
||||
feats.append(x)
|
||||
|
||||
x = audio_embedding
|
||||
for f in self.face_decoder_blocks:
|
||||
x = f(x)
|
||||
try:
|
||||
x = torch.cat((x, feats[-1]), dim=1)
|
||||
except Exception as e:
|
||||
print(x.size())
|
||||
print(feats[-1].size())
|
||||
raise e
|
||||
|
||||
feats.pop()
|
||||
|
||||
x = self.output_block(x)
|
||||
|
||||
if input_dim_size > 4:
|
||||
x = torch.split(x, B, dim=0) # [(B, C, H, W)]
|
||||
outputs = torch.stack(x, dim=2) # (B, C, T, H, W)
|
||||
|
||||
else:
|
||||
outputs = x
|
||||
|
||||
return outputs
|
||||
|
||||
class Wav2Lip_disc_qual(nn.Module):
|
||||
def __init__(self):
|
||||
super(Wav2Lip_disc_qual, self).__init__()
|
||||
|
||||
self.face_encoder_blocks = nn.ModuleList([
|
||||
nn.Sequential(nonorm_Conv2d(3, 32, kernel_size=7, stride=1, padding=3)), # 48,96
|
||||
|
||||
nn.Sequential(nonorm_Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=2), # 48,48
|
||||
nonorm_Conv2d(64, 64, kernel_size=5, stride=1, padding=2)),
|
||||
|
||||
nn.Sequential(nonorm_Conv2d(64, 128, kernel_size=5, stride=2, padding=2), # 24,24
|
||||
nonorm_Conv2d(128, 128, kernel_size=5, stride=1, padding=2)),
|
||||
|
||||
nn.Sequential(nonorm_Conv2d(128, 256, kernel_size=5, stride=2, padding=2), # 12,12
|
||||
nonorm_Conv2d(256, 256, kernel_size=5, stride=1, padding=2)),
|
||||
|
||||
nn.Sequential(nonorm_Conv2d(256, 512, kernel_size=3, stride=2, padding=1), # 6,6
|
||||
nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1)),
|
||||
|
||||
nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=2, padding=1), # 3,3
|
||||
nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1),),
|
||||
|
||||
nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=0), # 1, 1
|
||||
nonorm_Conv2d(512, 512, kernel_size=1, stride=1, padding=0)),])
|
||||
|
||||
self.binary_pred = nn.Sequential(nn.Conv2d(512, 1, kernel_size=1, stride=1, padding=0), nn.Sigmoid())
|
||||
self.label_noise = .0
|
||||
|
||||
def get_lower_half(self, face_sequences):
|
||||
return face_sequences[:, :, face_sequences.size(2)//2:]
|
||||
|
||||
def to_2d(self, face_sequences):
|
||||
B = face_sequences.size(0)
|
||||
face_sequences = torch.cat([face_sequences[:, :, i] for i in range(face_sequences.size(2))], dim=0)
|
||||
return face_sequences
|
||||
|
||||
def perceptual_forward(self, false_face_sequences):
|
||||
false_face_sequences = self.to_2d(false_face_sequences)
|
||||
false_face_sequences = self.get_lower_half(false_face_sequences)
|
||||
|
||||
false_feats = false_face_sequences
|
||||
for f in self.face_encoder_blocks:
|
||||
false_feats = f(false_feats)
|
||||
|
||||
false_pred_loss = F.binary_cross_entropy(self.binary_pred(false_feats).view(len(false_feats), -1),
|
||||
torch.ones((len(false_feats), 1)).cuda())
|
||||
|
||||
return false_pred_loss
|
||||
|
||||
def forward(self, face_sequences):
|
||||
face_sequences = self.to_2d(face_sequences)
|
||||
face_sequences = self.get_lower_half(face_sequences)
|
||||
|
||||
x = face_sequences
|
||||
for f in self.face_encoder_blocks:
|
||||
x = f(x)
|
||||
|
||||
return self.binary_pred(x).view(len(x), -1)
|
||||
@@ -0,0 +1,113 @@
|
||||
import sys
|
||||
|
||||
if sys.version_info[0] < 3 and sys.version_info[1] < 2:
|
||||
raise Exception("Must be using >= Python 3.2")
|
||||
|
||||
from os import listdir, path
|
||||
|
||||
if not path.isfile('face_detection/detection/sfd/s3fd.pth'):
|
||||
raise FileNotFoundError('Save the s3fd model to face_detection/detection/sfd/s3fd.pth \
|
||||
before running this script!')
|
||||
|
||||
import multiprocessing as mp
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
import numpy as np
|
||||
import argparse, os, cv2, traceback, subprocess
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import audio
|
||||
from hparams import hparams as hp
|
||||
|
||||
import face_detection
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('--ngpu', help='Number of GPUs across which to run in parallel', default=1, type=int)
|
||||
parser.add_argument('--batch_size', help='Single GPU Face detection batch size', default=32, type=int)
|
||||
parser.add_argument("--data_root", help="Root folder of the LRS2 dataset", required=True)
|
||||
parser.add_argument("--preprocessed_root", help="Root folder of the preprocessed dataset", required=True)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
fa = [face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False,
|
||||
device='cuda:{}'.format(id)) for id in range(args.ngpu)]
|
||||
|
||||
template = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'
|
||||
# template2 = 'ffmpeg -hide_banner -loglevel panic -threads 1 -y -i {} -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 {}'
|
||||
|
||||
def process_video_file(vfile, args, gpu_id):
|
||||
video_stream = cv2.VideoCapture(vfile)
|
||||
|
||||
frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
frames.append(frame)
|
||||
|
||||
vidname = os.path.basename(vfile).split('.')[0]
|
||||
dirname = vfile.split('/')[-2]
|
||||
|
||||
fulldir = path.join(args.preprocessed_root, dirname, vidname)
|
||||
os.makedirs(fulldir, exist_ok=True)
|
||||
|
||||
batches = [frames[i:i + args.batch_size] for i in range(0, len(frames), args.batch_size)]
|
||||
|
||||
i = -1
|
||||
for fb in batches:
|
||||
preds = fa[gpu_id].get_detections_for_batch(np.asarray(fb))
|
||||
|
||||
for j, f in enumerate(preds):
|
||||
i += 1
|
||||
if f is None:
|
||||
continue
|
||||
|
||||
x1, y1, x2, y2 = f
|
||||
cv2.imwrite(path.join(fulldir, '{}.jpg'.format(i)), fb[j][y1:y2, x1:x2])
|
||||
|
||||
def process_audio_file(vfile, args):
|
||||
vidname = os.path.basename(vfile).split('.')[0]
|
||||
dirname = vfile.split('/')[-2]
|
||||
|
||||
fulldir = path.join(args.preprocessed_root, dirname, vidname)
|
||||
os.makedirs(fulldir, exist_ok=True)
|
||||
|
||||
wavpath = path.join(fulldir, 'audio.wav')
|
||||
|
||||
command = template.format(vfile, wavpath)
|
||||
subprocess.call(command, shell=True)
|
||||
|
||||
|
||||
def mp_handler(job):
|
||||
vfile, args, gpu_id = job
|
||||
try:
|
||||
process_video_file(vfile, args, gpu_id)
|
||||
except KeyboardInterrupt:
|
||||
exit(0)
|
||||
except:
|
||||
traceback.print_exc()
|
||||
|
||||
def main(args):
|
||||
print('Started processing for {} with {} GPUs'.format(args.data_root, args.ngpu))
|
||||
|
||||
filelist = glob(path.join(args.data_root, '*/*.mp4'))
|
||||
|
||||
jobs = [(vfile, args, i%args.ngpu) for i, vfile in enumerate(filelist)]
|
||||
p = ThreadPoolExecutor(args.ngpu)
|
||||
futures = [p.submit(mp_handler, j) for j in jobs]
|
||||
_ = [r.result() for r in tqdm(as_completed(futures), total=len(futures))]
|
||||
|
||||
print('Dumping audios...')
|
||||
|
||||
for vfile in tqdm(filelist):
|
||||
try:
|
||||
process_audio_file(vfile, args)
|
||||
except KeyboardInterrupt:
|
||||
exit(0)
|
||||
except:
|
||||
traceback.print_exc()
|
||||
continue
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(args)
|
||||
@@ -0,0 +1,16 @@
|
||||
librosa==0.7.0
|
||||
numba==0.48
|
||||
numpy==1.17.1
|
||||
opencv-contrib-python==4.2.0.34
|
||||
opencv-python==4.1.0.25
|
||||
python-magic==0.4.22
|
||||
tensorflow==1.12.0
|
||||
torch==1.1.0
|
||||
torchvision==0.3.0
|
||||
tqdm==4.45.0
|
||||
|
||||
# To download (w2l-download-job)
|
||||
# Note: these will need frequent updates!
|
||||
gdown==3.13.0
|
||||
youtube-dl==2021.5.16
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
Generated results will be placed in this folder by default.
|
||||
@@ -0,0 +1 @@
|
||||
Temporary files at the time of inference/testing will be saved here. You can ignore them.
|
||||
@@ -0,0 +1,458 @@
|
||||
from os import listdir, path
|
||||
import numpy as np
|
||||
import scipy, cv2, os, sys, argparse, audio
|
||||
import json, subprocess, random, string
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch, face_detection
|
||||
from models import Wav2Lip
|
||||
import tempfile
|
||||
import shutil
|
||||
import datetime
|
||||
import pickle
|
||||
from PIL import Image
|
||||
|
||||
if 'PRINT_ENV' in os.environ and os.environ['PRINT_ENV']:
|
||||
for k, v in os.environ.items():
|
||||
print('{}={}'.format(k, v))
|
||||
|
||||
print('========================================')
|
||||
print('Python interpreter', sys.executable)
|
||||
print('PyTorch version', torch.__version__)
|
||||
print('CUDA Available?', torch.cuda.is_available())
|
||||
print('CUDA Device count', torch.cuda.device_count())
|
||||
print('========================================', flush=True)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
|
||||
|
||||
parser.add_argument('--checkpoint_path', type=str,
|
||||
help='Name of saved checkpoint to load weights from', required=True)
|
||||
|
||||
parser.add_argument('--face', type=str,
|
||||
help='Filepath of video/image that contains faces to use', required=True)
|
||||
parser.add_argument('--audio', type=str,
|
||||
help='Filepath of video/audio file to use as raw audio source', required=True)
|
||||
parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.',
|
||||
default='results/result_voice.mp4')
|
||||
|
||||
parser.add_argument('--static', type=bool,
|
||||
help='If True, then use only first video frame for inference', default=False)
|
||||
|
||||
# NB: I'm adjusting the FPS to match that of the end bump I made. The old default FPS was "25." (float 25)
|
||||
parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 29.97)',
|
||||
default=29.97, required=False)
|
||||
|
||||
parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0],
|
||||
help='Padding (top, bottom, left, right). Please adjust to include chin at least')
|
||||
|
||||
parser.add_argument('--face_det_batch_size', type=int,
|
||||
help='Batch size for face detection', default=16)
|
||||
parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
|
||||
|
||||
parser.add_argument('--resize_factor', default=1, type=int,
|
||||
help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
|
||||
|
||||
parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1],
|
||||
help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. '
|
||||
'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
|
||||
|
||||
parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1],
|
||||
help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
|
||||
'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
|
||||
|
||||
parser.add_argument('--rotate', default=False, action='store_true',
|
||||
help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
|
||||
'Use if you get a flipped result, despite feeding a normal looking video')
|
||||
|
||||
parser.add_argument('--nosmooth', default=False, action='store_true',
|
||||
help='Prevent smoothing face detections over a short temporal window')
|
||||
|
||||
parser.add_argument('--audio_start_pad_millis', default=0, type=int,
|
||||
help='Seconds to pad the start of the audio')
|
||||
|
||||
parser.add_argument('--audio_end_pad_millis', default=0, type=int,
|
||||
help='Seconds to pad the end of the audio')
|
||||
|
||||
parser.add_argument('--end_bump_file', type=str,
|
||||
help='Video file to concatenate at the end')
|
||||
|
||||
# This is a little more deliberate than "static" and causes us to pick an FPS, etc.
|
||||
# This is useful when the filename doesn't have an extension.
|
||||
parser.add_argument('--is_image', default=False, action='store_true',
|
||||
help='Denote that the input is a single-frame image, not a video')
|
||||
|
||||
# Purely for debugging on the host machine:
|
||||
parser.add_argument('--preserve_tempdir', default=False, action='store_true',
|
||||
help='Keep the tempdir arround for debugging')
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = 96
|
||||
|
||||
if os.path.isfile(args.face) and os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg']:
|
||||
args.static = True
|
||||
elif args.is_image:
|
||||
args.static = True
|
||||
|
||||
def get_smoothened_boxes(boxes, T):
|
||||
for i in range(len(boxes)):
|
||||
if i + T > len(boxes):
|
||||
window = boxes[len(boxes) - T:]
|
||||
else:
|
||||
window = boxes[i : i + T]
|
||||
boxes[i] = np.mean(window, axis=0)
|
||||
return boxes
|
||||
|
||||
def face_detect(images):
|
||||
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
||||
flip_input=False, device=device)
|
||||
|
||||
batch_size = args.face_det_batch_size
|
||||
|
||||
while 1:
|
||||
predictions = []
|
||||
try:
|
||||
for i in tqdm(range(0, len(images), batch_size)):
|
||||
predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
|
||||
except RuntimeError:
|
||||
if batch_size == 1:
|
||||
raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
|
||||
batch_size //= 2
|
||||
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
||||
continue
|
||||
break
|
||||
|
||||
results = []
|
||||
pady1, pady2, padx1, padx2 = args.pads
|
||||
for rect, image in zip(predictions, images):
|
||||
if rect is None:
|
||||
cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
|
||||
raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
|
||||
|
||||
y1 = max(0, rect[1] - pady1)
|
||||
y2 = min(image.shape[0], rect[3] + pady2)
|
||||
x1 = max(0, rect[0] - padx1)
|
||||
x2 = min(image.shape[1], rect[2] + padx2)
|
||||
|
||||
results.append([x1, y1, x2, y2])
|
||||
|
||||
boxes = np.array(results)
|
||||
if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
|
||||
results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
|
||||
|
||||
del detector
|
||||
return results
|
||||
|
||||
def detect_faces_in_frames(frames, video_faces_pickle_file):
|
||||
t0 = datetime.datetime.now()
|
||||
if args.box[0] == -1:
|
||||
if not args.static:
|
||||
face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
|
||||
else:
|
||||
face_det_results = face_detect([frames[0]])
|
||||
else:
|
||||
print('Using the specified bounding box instead of face detection...')
|
||||
y1, y2, x1, x2 = args.box
|
||||
face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
|
||||
t1 = datetime.datetime.now()
|
||||
|
||||
d = t1 - t0
|
||||
print('Total seconds to detect faces in {} frames: {}'.format(len(frames), d.total_seconds()))
|
||||
|
||||
with open(video_faces_pickle_file, 'wb') as f:
|
||||
print('Saving pickle file: {}'.format(video_faces_pickle_file))
|
||||
pickle.dump(face_det_results, f)
|
||||
|
||||
return face_det_results
|
||||
|
||||
#def datagen(frames, mels):
|
||||
def datagen(frames, face_det_results, mels):
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
#face_det_results = detect_faces_in_frames(frames)
|
||||
|
||||
for i, m in enumerate(mels):
|
||||
idx = 0 if args.static else i%len(frames)
|
||||
frame_to_save = frames[idx].copy()
|
||||
face, coords = face_det_results[idx].copy()
|
||||
|
||||
face = cv2.resize(face, (args.img_size, args.img_size))
|
||||
|
||||
img_batch.append(face)
|
||||
mel_batch.append(m)
|
||||
frame_batch.append(frame_to_save)
|
||||
coords_batch.append(coords)
|
||||
|
||||
if len(img_batch) >= args.wav2lip_batch_size:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if len(img_batch) > 0:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
|
||||
mel_step_size = 16
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
print('Using {} for inference.'.format(device))
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if device == 'cuda':
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_model(path):
|
||||
model = Wav2Lip()
|
||||
print("Load checkpoint from: {}".format(path), flush=True)
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
|
||||
model = model.to(device)
|
||||
return model.eval()
|
||||
|
||||
def maybe_pad_audio_file(tempdir, args):
|
||||
if not args.audio_start_pad_millis and not args.audio_end_pad_millis:
|
||||
return
|
||||
|
||||
print('Padding audio with {} millis at start, {} millis at end.'.format(
|
||||
args.audio_start_pad_millis, args.audio_end_pad_millis), flush=True)
|
||||
|
||||
padded_wav_filename = os.path.join(tempdir, 'padded_audio.wav')
|
||||
# ffmpeg padding: https://superuser.com/a/579110
|
||||
# Structure of filters: https://stackoverflow.com/a/55463101
|
||||
command = ' '.join([
|
||||
'ffmpeg',
|
||||
'-i {}'.format(args.audio),
|
||||
'-filter_complex "[0] adelay={}ms:all=true [a] ; [a] apad=pad_dur={}ms"'.format(args.audio_start_pad_millis, args.audio_end_pad_millis),
|
||||
'{}'.format(padded_wav_filename)
|
||||
])
|
||||
print('command:', command, flush=True)
|
||||
subprocess.call(command, shell=True)
|
||||
args.audio = padded_wav_filename
|
||||
|
||||
|
||||
def maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h):
|
||||
if not args.end_bump_file:
|
||||
print('No end bump file.', flush=True)
|
||||
return
|
||||
|
||||
print('Adding end bump file: {}'.format(args.end_bump_file), flush=True)
|
||||
|
||||
bumpfile = args.end_bump_file
|
||||
|
||||
# May need to resize our bump file...
|
||||
# After many attempts to get this into the same ffmpeg pipeline,
|
||||
# I'm giving up and just adding another step.
|
||||
#
|
||||
# NOTE: The FPS of the bump needs to match the video. Especially
|
||||
# important for videos generated from still photos, because that defaults
|
||||
# to the odd choice of "25 fps".
|
||||
if frame_w is not 1920 or frame_h is not 1080:
|
||||
print('Need to resize bump.', flush=True)
|
||||
bumpfile = os.path.join(tempdir, 'end_bump_resized.mp4')
|
||||
|
||||
# Scale / aspect ratio / sar / resizing / padding:
|
||||
# https://stackoverflow.com/a/51946719
|
||||
command = ' '.join([
|
||||
'ffmpeg',
|
||||
'-i {}'.format(args.end_bump_file),
|
||||
'-vf "scale={}:{}:force_original_aspect_ratio=decrease,pad={}:{}:(ow-iw)/2:(oh-ih)/2,setsar=1"'.format(frame_w, frame_h, frame_w, frame_h),
|
||||
'{}'.format(bumpfile)
|
||||
])
|
||||
print('command:', command, flush=True)
|
||||
failure_code = subprocess.call(command, shell=True)
|
||||
|
||||
if failure_code:
|
||||
print('Failed to resize end bump')
|
||||
return
|
||||
|
||||
padded_wav_filename = os.path.join(tempdir, 'output_with_end_bump.mp4')
|
||||
# ffmpeg video concatenation: https://stackoverflow.com/a/11175851
|
||||
# Structure of filters: https://stackoverflow.com/a/55463101
|
||||
# More info on filters: https://stackoverflow.com/a/22958746
|
||||
command = ' '.join([
|
||||
'ffmpeg',
|
||||
'-i {}'.format(args.outfile),
|
||||
'-i {}'.format(bumpfile),
|
||||
'-filter_complex "concat=n=2:v=1:a=1"',
|
||||
'{}'.format(padded_wav_filename)
|
||||
])
|
||||
print('command:', command, flush=True)
|
||||
failure_code = subprocess.call(command, shell=True)
|
||||
|
||||
if failure_code:
|
||||
print('Failed to concatenate end bump')
|
||||
else:
|
||||
print('Done concatenating end bump; renaming file...', flush=True)
|
||||
os.rename(padded_wav_filename, args.outfile)
|
||||
|
||||
|
||||
def main(tempdir):
|
||||
video_faces_pickle_file = args.face + '.faces'
|
||||
print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True)
|
||||
|
||||
frame_w = 0
|
||||
frame_h = 0
|
||||
|
||||
if not os.path.isfile(args.face):
|
||||
raise ValueError('--face argument must be a valid path to video/image file')
|
||||
|
||||
elif os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg'] \
|
||||
or args.is_image:
|
||||
full_frames = [cv2.imread(args.face)]
|
||||
fps = args.fps
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
|
||||
else:
|
||||
video_stream = cv2.VideoCapture(args.face)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
|
||||
print('Reading video frames...')
|
||||
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
if args.resize_factor > 1:
|
||||
frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
|
||||
|
||||
if args.rotate:
|
||||
frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
|
||||
|
||||
y1, y2, x1, x2 = args.crop
|
||||
if x2 == -1: x2 = frame.shape[1]
|
||||
if y2 == -1: y2 = frame.shape[0]
|
||||
|
||||
frame = frame[y1:y2, x1:x2]
|
||||
|
||||
full_frames.append(frame)
|
||||
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
|
||||
print("Number of frames available for inference: "+str(len(full_frames)), flush=True)
|
||||
print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True)
|
||||
|
||||
if not args.audio.endswith('.wav'):
|
||||
print('Extracting raw audio...')
|
||||
temp_wav_filename = os.path.join(tempdir, 'temp.wav') # previously 'temp/temp.wav'
|
||||
command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, temp_wav_filename)
|
||||
|
||||
print('command:', command, flush=True)
|
||||
subprocess.call(command, shell=True)
|
||||
args.audio = temp_wav_filename
|
||||
|
||||
maybe_pad_audio_file(tempdir, args)
|
||||
|
||||
wav = audio.load_wav(args.audio, 16000)
|
||||
mel = audio.melspectrogram(wav)
|
||||
print(mel.shape)
|
||||
|
||||
if np.isnan(mel.reshape(-1)).sum() > 0:
|
||||
raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
|
||||
|
||||
mel_chunks = []
|
||||
mel_idx_multiplier = 80./fps
|
||||
i = 0
|
||||
while 1:
|
||||
start_idx = int(i * mel_idx_multiplier)
|
||||
if start_idx + mel_step_size > len(mel[0]):
|
||||
mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
|
||||
break
|
||||
mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
|
||||
i += 1
|
||||
|
||||
print("Length of mel chunks: {}".format(len(mel_chunks)), flush=True)
|
||||
|
||||
if len(mel_chunks) > 2000:
|
||||
# TODO: Catch that the file is too long earlier on.
|
||||
raise Exception("Too many mel chunks: {}".format(len(mel_chunks)))
|
||||
|
||||
full_frames = full_frames[:len(mel_chunks)]
|
||||
|
||||
# TODO: Do this offline.
|
||||
print('Face detection time...', flush=True)
|
||||
if not os.path.isfile(video_faces_pickle_file):
|
||||
print('Detecting faces...', flush=True)
|
||||
face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
|
||||
else:
|
||||
with open(video_faces_pickle_file, 'rb') as f:
|
||||
face_det_results = pickle.load(f)
|
||||
if len(face_det_results) < len(full_frames):
|
||||
print("Face detect results have {} frames, but we need {}".format(len(face_det_results), len(full_frames)))
|
||||
print('Detecting faces again...', flush=True)
|
||||
face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
|
||||
else:
|
||||
print("We don't need to find faces! Woo!", flush=True)
|
||||
|
||||
face_det_results = face_det_results[:len(mel_chunks)]
|
||||
|
||||
batch_size = args.wav2lip_batch_size
|
||||
#gen = datagen(full_frames.copy(), mel_chunks)
|
||||
gen = datagen(full_frames.copy(), face_det_results.copy(), mel_chunks)
|
||||
|
||||
output_video_filename = os.path.join(tempdir, 'result.avi') # previously 'temp/result.avi'
|
||||
|
||||
for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen,
|
||||
total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
|
||||
if i == 0:
|
||||
model = load_model(args.checkpoint_path)
|
||||
print ("Model loaded")
|
||||
|
||||
#frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
out = cv2.VideoWriter(output_video_filename,
|
||||
cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
|
||||
|
||||
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
|
||||
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(mel_batch, img_batch)
|
||||
|
||||
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
|
||||
|
||||
for p, f, c in zip(pred, frames, coords):
|
||||
y1, y2, x1, x2 = c
|
||||
p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
|
||||
|
||||
f[y1:y2, x1:x2] = p
|
||||
out.write(f)
|
||||
|
||||
out.release()
|
||||
|
||||
command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, output_video_filename, args.outfile)
|
||||
print('command:', command, flush=True)
|
||||
subprocess.call(command, shell=True)
|
||||
|
||||
maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h)
|
||||
|
||||
if __name__ == '__main__':
|
||||
tempdir = tempfile.mkdtemp()
|
||||
print('Using Python tempdir: {}'.format(tempdir), flush=True)
|
||||
try:
|
||||
main(tempdir)
|
||||
finally:
|
||||
if not args.preserve_tempdir:
|
||||
shutil.rmtree(tempdir)
|
||||
|
||||
@@ -0,0 +1,433 @@
|
||||
from os import listdir, path
|
||||
import numpy as np
|
||||
import scipy, cv2, os, sys, argparse, audio
|
||||
import json, subprocess, random, string
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch, face_detection
|
||||
from models import Wav2Lip
|
||||
import tempfile
|
||||
import shutil
|
||||
import datetime
|
||||
import pickle
|
||||
import magic
|
||||
from PIL import Image
|
||||
|
||||
if 'PRINT_ENV' in os.environ and os.environ['PRINT_ENV']:
|
||||
for k, v in os.environ.items():
|
||||
print('{}={}'.format(k, v))
|
||||
|
||||
print('========================================')
|
||||
print('Python interpreter', sys.executable)
|
||||
print('PyTorch version', torch.__version__)
|
||||
print('CUDA Available?', torch.cuda.is_available())
|
||||
print('CUDA Device count', torch.cuda.device_count())
|
||||
print('========================================', flush=True)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
|
||||
|
||||
parser.add_argument('--checkpoint_path', type=str,
|
||||
help='Name of saved checkpoint to load weights from', required=True)
|
||||
|
||||
parser.add_argument('--image_or_video_filename', type=str,
|
||||
help='Filepath of video/image that contains faces to use', required=True)
|
||||
|
||||
#parser.add_argument('--image_or_video_type', type=str,
|
||||
# help='either "image" or "video', required=True)
|
||||
|
||||
# This is a little more deliberate than "static" and causes us to pick an FPS, etc.
|
||||
# This is useful when the filename doesn't have an extension.
|
||||
parser.add_argument('--is_image', default=False, action='store_true',
|
||||
help='Denote that the input is a single-frame image, not a video')
|
||||
|
||||
parser.add_argument('--cached_faces_filename', type=str,
|
||||
help='Filename for the cached faces file to use', required=True)
|
||||
|
||||
#parser.add_argument('--face', type=str,
|
||||
# help='Filepath of video/image that contains faces to use', required=True)
|
||||
parser.add_argument('--audio_filename', type=str,
|
||||
help='Filepath of video/audio file to use as raw audio source', required=True)
|
||||
|
||||
parser.add_argument('--end_bump_file', type=str,
|
||||
help='Video file to concatenate at the end')
|
||||
|
||||
parser.add_argument('--output_video_filename', type=str,
|
||||
help='Output filename for the final video', required=True)
|
||||
|
||||
parser.add_argument('--output_metadata_filename', type=str,
|
||||
help='Output filename for the JSON containing width, height, etc.', required=True)
|
||||
|
||||
#parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.',
|
||||
# default='results/result_voice.mp4')
|
||||
|
||||
parser.add_argument('--static', type=bool,
|
||||
help='If True, then use only first video frame for inference', default=False)
|
||||
|
||||
# NB: I'm adjusting the FPS to match that of the end bump I made. The old default FPS was "25." (float 25)
|
||||
parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 29.97)',
|
||||
default=29.97, required=False)
|
||||
|
||||
parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0],
|
||||
help='Padding (top, bottom, left, right). Please adjust to include chin at least')
|
||||
|
||||
parser.add_argument('--face_det_batch_size', type=int,
|
||||
help='Batch size for face detection', default=16)
|
||||
parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
|
||||
|
||||
parser.add_argument('--resize_factor', default=1, type=int,
|
||||
help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
|
||||
|
||||
parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1],
|
||||
help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. '
|
||||
'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
|
||||
|
||||
parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1],
|
||||
help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
|
||||
'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
|
||||
|
||||
parser.add_argument('--rotate', default=False, action='store_true',
|
||||
help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
|
||||
'Use if you get a flipped result, despite feeding a normal looking video')
|
||||
|
||||
parser.add_argument('--nosmooth', default=False, action='store_true',
|
||||
help='Prevent smoothing face detections over a short temporal window')
|
||||
|
||||
parser.add_argument('--audio_start_pad_millis', default=0, type=int,
|
||||
help='Seconds to pad the start of the audio')
|
||||
|
||||
parser.add_argument('--audio_end_pad_millis', default=0, type=int,
|
||||
help='Seconds to pad the end of the audio')
|
||||
|
||||
# Purely for debugging on the host machine:
|
||||
parser.add_argument('--preserve_tempdir', default=False, action='store_true',
|
||||
help='Keep the tempdir arround for debugging')
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = 96
|
||||
|
||||
def datagen(frames, face_det_results, mels):
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
for i, m in enumerate(mels):
|
||||
idx = 0 if args.static else i%len(frames)
|
||||
frame_to_save = frames[idx].copy()
|
||||
face, coords = face_det_results[idx].copy()
|
||||
|
||||
face = cv2.resize(face, (args.img_size, args.img_size))
|
||||
|
||||
img_batch.append(face)
|
||||
mel_batch.append(m)
|
||||
frame_batch.append(frame_to_save)
|
||||
coords_batch.append(coords)
|
||||
|
||||
if len(img_batch) >= args.wav2lip_batch_size:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if len(img_batch) > 0:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
|
||||
mel_step_size = 16
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
print('Using {} for inference.'.format(device))
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if device == 'cuda':
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_model(path):
|
||||
model = Wav2Lip()
|
||||
print("Load checkpoint from: {}".format(path), flush=True)
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
|
||||
model = model.to(device)
|
||||
return model.eval()
|
||||
|
||||
def maybe_pad_audio_file(tempdir, args):
|
||||
if not args.audio_start_pad_millis and not args.audio_end_pad_millis:
|
||||
return
|
||||
|
||||
print('Padding audio with {} millis at start, {} millis at end.'.format(
|
||||
args.audio_start_pad_millis, args.audio_end_pad_millis), flush=True)
|
||||
|
||||
padded_wav_filename = os.path.join(tempdir, 'padded_audio.wav')
|
||||
# ffmpeg padding: https://superuser.com/a/579110
|
||||
# Structure of filters: https://stackoverflow.com/a/55463101
|
||||
command = ' '.join([
|
||||
'ffmpeg',
|
||||
'-i {}'.format(args.audio_filename),
|
||||
'-filter_complex "[0] adelay={}ms:all=true [a] ; [a] apad=pad_dur={}ms"'.format(args.audio_start_pad_millis, args.audio_end_pad_millis),
|
||||
'{}'.format(padded_wav_filename)
|
||||
])
|
||||
print('command:', command, flush=True)
|
||||
subprocess.call(command, shell=True)
|
||||
args.audio_filename = padded_wav_filename
|
||||
|
||||
|
||||
def maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h):
|
||||
if not args.end_bump_file:
|
||||
print('No end bump file.', flush=True)
|
||||
return
|
||||
|
||||
print('Adding end bump file: {}'.format(args.end_bump_file), flush=True)
|
||||
|
||||
bumpfile = args.end_bump_file
|
||||
|
||||
# May need to resize our bump file...
|
||||
# After many attempts to get this into the same ffmpeg pipeline,
|
||||
# I'm giving up and just adding another step.
|
||||
#
|
||||
# NOTE: The FPS of the bump needs to match the video. Especially
|
||||
# important for videos generated from still photos, because that defaults
|
||||
# to the odd choice of "25 fps".
|
||||
if frame_w is not 1920 or frame_h is not 1080:
|
||||
print('Need to resize bump.', flush=True)
|
||||
bumpfile = os.path.join(tempdir, 'end_bump_resized.mp4')
|
||||
|
||||
# Scale / aspect ratio / sar / resizing / padding:
|
||||
# https://stackoverflow.com/a/51946719
|
||||
command = ' '.join([
|
||||
'ffmpeg',
|
||||
'-i {}'.format(args.end_bump_file),
|
||||
'-vf "scale={}:{}:force_original_aspect_ratio=decrease,pad={}:{}:(ow-iw)/2:(oh-ih)/2,setsar=1"'.format(frame_w, frame_h, frame_w, frame_h),
|
||||
'{}'.format(bumpfile)
|
||||
])
|
||||
print('command:', command, flush=True)
|
||||
failure_code = subprocess.call(command, shell=True)
|
||||
|
||||
if failure_code:
|
||||
print('Failed to resize end bump')
|
||||
return
|
||||
|
||||
padded_wav_filename = os.path.join(tempdir, 'output_with_end_bump.mp4')
|
||||
# ffmpeg video concatenation: https://stackoverflow.com/a/11175851
|
||||
# Structure of filters: https://stackoverflow.com/a/55463101
|
||||
# More info on filters: https://stackoverflow.com/a/22958746
|
||||
command = ' '.join([
|
||||
'ffmpeg',
|
||||
'-i {}'.format(args.output_video_filename),
|
||||
'-i {}'.format(bumpfile),
|
||||
'-filter_complex "concat=n=2:v=1:a=1"',
|
||||
'{}'.format(padded_wav_filename)
|
||||
])
|
||||
print('command:', command, flush=True)
|
||||
failure_code = subprocess.call(command, shell=True)
|
||||
|
||||
if failure_code:
|
||||
print('Failed to concatenate end bump')
|
||||
else:
|
||||
print('Done concatenating end bump; renaming file...', flush=True)
|
||||
os.rename(padded_wav_filename, args.output_video_filename)
|
||||
|
||||
|
||||
def main(tempdir):
|
||||
frame_w = 0
|
||||
frame_h = 0
|
||||
fps = 0
|
||||
|
||||
if not os.path.isfile(args.image_or_video_filename):
|
||||
raise ValueError('image_or_video_filename is not a file')
|
||||
|
||||
if not os.path.isfile(args.cached_faces_filename):
|
||||
raise ValueError('cached_faces_filename is not a file')
|
||||
|
||||
if not os.path.isfile(args.audio_filename):
|
||||
raise ValueError('audio_filename is not a file')
|
||||
|
||||
if os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg'] \
|
||||
or args.is_image:
|
||||
full_frames = [cv2.imread(args.image_or_video_filename)]
|
||||
fps = args.fps
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
|
||||
else:
|
||||
video_stream = cv2.VideoCapture(args.image_or_video_filename)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
|
||||
print('Reading video frames...')
|
||||
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
if args.resize_factor > 1:
|
||||
frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
|
||||
|
||||
if args.rotate:
|
||||
frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
|
||||
|
||||
y1, y2, x1, x2 = args.crop
|
||||
if x2 == -1: x2 = frame.shape[1]
|
||||
if y2 == -1: y2 = frame.shape[0]
|
||||
|
||||
frame = frame[y1:y2, x1:x2]
|
||||
|
||||
full_frames.append(frame)
|
||||
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
|
||||
print("Number of frames available for inference: "+str(len(full_frames)), flush=True)
|
||||
print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True)
|
||||
|
||||
# TODO(bt): This could be more efficient
|
||||
if not args.audio_filename.endswith('.wav'):
|
||||
print('Extracting raw audio...')
|
||||
temp_wav_filename = os.path.join(tempdir, 'temp.wav') # previously 'temp/temp.wav'
|
||||
command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio_filename, temp_wav_filename)
|
||||
|
||||
print('command:', command, flush=True)
|
||||
subprocess.call(command, shell=True)
|
||||
args.audio_filename = temp_wav_filename
|
||||
|
||||
maybe_pad_audio_file(tempdir, args)
|
||||
|
||||
# TODO(bt): Save the spectrogram?
|
||||
wav = audio.load_wav(args.audio_filename, 16000)
|
||||
mel = audio.melspectrogram(wav)
|
||||
print(mel.shape)
|
||||
|
||||
if np.isnan(mel.reshape(-1)).sum() > 0:
|
||||
raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
|
||||
|
||||
mel_chunks = []
|
||||
mel_idx_multiplier = 80./fps
|
||||
i = 0
|
||||
while 1:
|
||||
start_idx = int(i * mel_idx_multiplier)
|
||||
if start_idx + mel_step_size > len(mel[0]):
|
||||
mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
|
||||
break
|
||||
mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
|
||||
i += 1
|
||||
|
||||
print("Length of mel chunks: {}".format(len(mel_chunks)), flush=True)
|
||||
|
||||
if len(mel_chunks) > 2000:
|
||||
# TODO(bt): Catch that the file is too long earlier on.
|
||||
raise Exception("Too many mel chunks: {}".format(len(mel_chunks)))
|
||||
|
||||
full_frames = full_frames[:len(mel_chunks)]
|
||||
|
||||
# Load face detection results from pickle file
|
||||
video_faces_pickle_file = args.cached_faces_filename
|
||||
with open(video_faces_pickle_file, 'rb') as f:
|
||||
face_det_results = pickle.load(f)
|
||||
|
||||
face_det_results = face_det_results[:len(mel_chunks)]
|
||||
|
||||
batch_size = args.wav2lip_batch_size
|
||||
gen = datagen(full_frames.copy(), face_det_results.copy(), mel_chunks)
|
||||
|
||||
output_video_filename = os.path.join(tempdir, 'result.avi') # previously 'temp/result.avi'
|
||||
|
||||
for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen,
|
||||
total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
|
||||
if i == 0:
|
||||
model = load_model(args.checkpoint_path)
|
||||
print ("Model loaded")
|
||||
|
||||
#frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
out = cv2.VideoWriter(output_video_filename,
|
||||
cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
|
||||
|
||||
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
|
||||
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(mel_batch, img_batch)
|
||||
|
||||
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
|
||||
|
||||
for p, f, c in zip(pred, frames, coords):
|
||||
y1, y2, x1, x2 = c
|
||||
p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
|
||||
|
||||
f[y1:y2, x1:x2] = p
|
||||
out.write(f)
|
||||
|
||||
out.release()
|
||||
|
||||
command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio_filename, output_video_filename, args.output_video_filename)
|
||||
print('command:', command, flush=True)
|
||||
subprocess.call(command, shell=True)
|
||||
|
||||
maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h)
|
||||
|
||||
# ==== METADATA (1) ====
|
||||
|
||||
command = [
|
||||
"ffprobe",
|
||||
"-loglevel", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format",
|
||||
"-show_streams",
|
||||
args.output_video_filename
|
||||
]
|
||||
|
||||
pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
||||
out, err = pipe.communicate()
|
||||
ffmpeg_metadata = json.loads(out)
|
||||
|
||||
# ==== METADATA (2) ====
|
||||
mime_type = magic.from_file(args.output_video_filename, mime=True)
|
||||
file_size_bytes = os.path.getsize(args.output_video_filename)
|
||||
|
||||
print(ffmpeg_metadata)
|
||||
|
||||
# We can't really get frame count anymore. The original source is invalid.
|
||||
# ffmpeg reports fps, but it's a ratio, eg. 25/1
|
||||
#frame_count = len(full_frames)
|
||||
#duration_seconds = frame_count / fps
|
||||
|
||||
duration_millis = int(float(ffmpeg_metadata['format']['duration']) * 1000)
|
||||
|
||||
metadata = {
|
||||
'is_video': True,
|
||||
'width': frame_w,
|
||||
'height': frame_h,
|
||||
#'num_frames': frame_count,
|
||||
#'fps': fps,
|
||||
'duration_millis': duration_millis,
|
||||
'mimetype': mime_type,
|
||||
'file_size_bytes': file_size_bytes,
|
||||
}
|
||||
|
||||
with open(args.output_metadata_filename, 'w') as json_file:
|
||||
json.dump(metadata, json_file)
|
||||
|
||||
if __name__ == '__main__':
|
||||
tempdir = tempfile.mkdtemp()
|
||||
print('Using Python tempdir: {}'.format(tempdir), flush=True)
|
||||
try:
|
||||
main(tempdir)
|
||||
finally:
|
||||
if not args.preserve_tempdir:
|
||||
shutil.rmtree(tempdir)
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
from os import listdir, path
|
||||
import numpy as np
|
||||
import scipy, cv2, os, sys, argparse, audio
|
||||
import json, subprocess, random, string
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch, face_detection
|
||||
from models import Wav2Lip
|
||||
import tempfile
|
||||
import shutil
|
||||
import datetime
|
||||
import pickle
|
||||
import magic
|
||||
from PIL import Image
|
||||
|
||||
"""
|
||||
Current worker filesystem architecture:
|
||||
|
||||
/tmp
|
||||
/tmp/models
|
||||
/tmp/models/wav2lip_gan.pth
|
||||
/tmp/.tmppd1Kvq
|
||||
/tmp/.tmppd1Kvq/input_audio_file
|
||||
/tmp/templates_images
|
||||
/tmp/templates_images/elon-musk.jpg
|
||||
/tmp/templates_videos
|
||||
/tmp/templates_videos/ben-shapiro-baby.mp4
|
||||
/tmp/templates_videos/ben-shapiro-baby.mp4.faces
|
||||
/tmp/templates_videos/dr-phil-bubble.mp4
|
||||
/tmp/tmp7um0yxjo
|
||||
/tmp/tmp7um0yxjo/padded_audio.wav
|
||||
/tmp/tmp7um0yxjo/temp.wav
|
||||
/tmp/assets
|
||||
/tmp/assets/vocodes-short-end-bump.mp4
|
||||
|
||||
"""
|
||||
|
||||
if 'PRINT_ENV' in os.environ and os.environ['PRINT_ENV']:
|
||||
for k, v in os.environ.items():
|
||||
print('{}={}'.format(k, v))
|
||||
|
||||
print('========================================')
|
||||
print('Python interpreter', sys.executable)
|
||||
print('PyTorch version', torch.__version__)
|
||||
print('CUDA Available?', torch.cuda.is_available())
|
||||
print('CUDA Device count', torch.cuda.device_count())
|
||||
print('========================================', flush=True)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
|
||||
|
||||
parser.add_argument('--checkpoint_path', type=str,
|
||||
help='Name of saved checkpoint to load weights from', required=True)
|
||||
|
||||
parser.add_argument('--image_or_video_filename', type=str,
|
||||
help='Filepath of video/image that contains faces to use', required=True)
|
||||
|
||||
parser.add_argument('--output_cached_faces_filename', type=str,
|
||||
help='Output filename for the cached faces file', required=True)
|
||||
|
||||
parser.add_argument('--output_metadata_filename', type=str,
|
||||
help='Output filename for the JSON containing width, height, etc.', required=True)
|
||||
|
||||
# NB: Not needed for processing the upload file for faces:
|
||||
#
|
||||
# parser.add_argument('--audio', type=str,
|
||||
# help='Filepath of video/audio file to use as raw audio source', required=True)
|
||||
# parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.',
|
||||
# default='results/result_voice.mp4')
|
||||
#
|
||||
|
||||
parser.add_argument('--static', type=bool,
|
||||
help='If True, then use only first video frame for inference', default=False)
|
||||
|
||||
# NB: I'm adjusting the FPS to match that of the end bump I made. The old default FPS was "25." (float 25)
|
||||
parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 29.97)',
|
||||
default=29.97, required=False)
|
||||
|
||||
parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0],
|
||||
help='Padding (top, bottom, left, right). Please adjust to include chin at least')
|
||||
|
||||
parser.add_argument('--face_det_batch_size', type=int,
|
||||
help='Batch size for face detection', default=16)
|
||||
parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
|
||||
|
||||
parser.add_argument('--resize_factor', default=1, type=int,
|
||||
help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
|
||||
|
||||
parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1],
|
||||
help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. '
|
||||
'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
|
||||
|
||||
parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1],
|
||||
help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
|
||||
'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
|
||||
|
||||
parser.add_argument('--rotate', default=False, action='store_true',
|
||||
help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
|
||||
'Use if you get a flipped result, despite feeding a normal looking video')
|
||||
|
||||
parser.add_argument('--nosmooth', default=False, action='store_true',
|
||||
help='Prevent smoothing face detections over a short temporal window')
|
||||
|
||||
parser.add_argument('--audio_start_pad_millis', default=0, type=int,
|
||||
help='Seconds to pad the start of the audio')
|
||||
|
||||
parser.add_argument('--audio_end_pad_millis', default=0, type=int,
|
||||
help='Seconds to pad the end of the audio')
|
||||
|
||||
parser.add_argument('--end_bump_file', type=str,
|
||||
help='Video file to concatenate at the end')
|
||||
|
||||
# This is a little more deliberate than "static" and causes us to pick an FPS, etc.
|
||||
# This is useful when the filename doesn't have an extension.
|
||||
parser.add_argument('--is_image', default=False, action='store_true',
|
||||
help='Denote that the input is a single-frame image, not a video')
|
||||
|
||||
# Purely for debugging on the host machine:
|
||||
parser.add_argument('--preserve_tempdir', default=False, action='store_true',
|
||||
help='Keep the tempdir arround for debugging')
|
||||
|
||||
args = parser.parse_args()
|
||||
args.img_size = 96
|
||||
|
||||
if os.path.isfile(args.image_or_video_filename) \
|
||||
and os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg']:
|
||||
args.static = True
|
||||
elif args.is_image:
|
||||
args.static = True
|
||||
|
||||
|
||||
class FacelessFrameError(Exception):
|
||||
"""Throw this when there are frames without faces."""
|
||||
pass
|
||||
|
||||
def get_smoothened_boxes(boxes, T):
|
||||
for i in range(len(boxes)):
|
||||
if i + T > len(boxes):
|
||||
window = boxes[len(boxes) - T:]
|
||||
else:
|
||||
window = boxes[i : i + T]
|
||||
boxes[i] = np.mean(window, axis=0)
|
||||
return boxes
|
||||
|
||||
def face_detect(images):
|
||||
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
||||
flip_input=False, device=device)
|
||||
|
||||
batch_size = args.face_det_batch_size
|
||||
|
||||
while 1:
|
||||
predictions = []
|
||||
try:
|
||||
for i in tqdm(range(0, len(images), batch_size)):
|
||||
predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
|
||||
except RuntimeError:
|
||||
if batch_size == 1:
|
||||
raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
|
||||
batch_size //= 2
|
||||
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
||||
continue
|
||||
break
|
||||
|
||||
results = []
|
||||
pady1, pady2, padx1, padx2 = args.pads
|
||||
for rect, image in zip(predictions, images):
|
||||
if rect is None:
|
||||
cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
|
||||
raise FacelessFrameError('Face not detected! Ensure the video contains a face in all the frames.')
|
||||
|
||||
y1 = max(0, rect[1] - pady1)
|
||||
y2 = min(image.shape[0], rect[3] + pady2)
|
||||
x1 = max(0, rect[0] - padx1)
|
||||
x2 = min(image.shape[1], rect[2] + padx2)
|
||||
|
||||
results.append([x1, y1, x2, y2])
|
||||
|
||||
boxes = np.array(results)
|
||||
if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
|
||||
results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
|
||||
|
||||
del detector
|
||||
return results
|
||||
|
||||
def detect_faces_in_frames(frames, video_faces_pickle_file):
|
||||
t0 = datetime.datetime.now()
|
||||
if args.box[0] == -1:
|
||||
if not args.static:
|
||||
face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
|
||||
else:
|
||||
face_det_results = face_detect([frames[0]])
|
||||
else:
|
||||
print('Using the specified bounding box instead of face detection...')
|
||||
y1, y2, x1, x2 = args.box
|
||||
face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
|
||||
t1 = datetime.datetime.now()
|
||||
|
||||
d = t1 - t0
|
||||
print('Total seconds to detect faces in {} frames: {}'.format(len(frames), d.total_seconds()))
|
||||
|
||||
with open(video_faces_pickle_file, 'wb') as f:
|
||||
print('Saving pickle file: {}'.format(video_faces_pickle_file))
|
||||
pickle.dump(face_det_results, f)
|
||||
|
||||
return face_det_results
|
||||
|
||||
#def datagen(frames, mels):
|
||||
def datagen(frames, face_det_results, mels):
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
#face_det_results = detect_faces_in_frames(frames)
|
||||
|
||||
for i, m in enumerate(mels):
|
||||
idx = 0 if args.static else i%len(frames)
|
||||
frame_to_save = frames[idx].copy()
|
||||
face, coords = face_det_results[idx].copy()
|
||||
|
||||
face = cv2.resize(face, (args.img_size, args.img_size))
|
||||
|
||||
img_batch.append(face)
|
||||
mel_batch.append(m)
|
||||
frame_batch.append(frame_to_save)
|
||||
coords_batch.append(coords)
|
||||
|
||||
if len(img_batch) >= args.wav2lip_batch_size:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
|
||||
|
||||
if len(img_batch) > 0:
|
||||
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
||||
|
||||
img_masked = img_batch.copy()
|
||||
img_masked[:, args.img_size//2:] = 0
|
||||
|
||||
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
||||
mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
||||
|
||||
yield img_batch, mel_batch, frame_batch, coords_batch
|
||||
|
||||
mel_step_size = 16
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
print('Using {} for inference.'.format(device))
|
||||
|
||||
#def _load(checkpoint_path):
|
||||
# if device == 'cuda':
|
||||
# checkpoint = torch.load(checkpoint_path)
|
||||
# else:
|
||||
# checkpoint = torch.load(checkpoint_path,
|
||||
# map_location=lambda storage, loc: storage)
|
||||
# return checkpoint
|
||||
|
||||
#def load_model(path):
|
||||
# model = Wav2Lip()
|
||||
# print("Load checkpoint from: {}".format(path), flush=True)
|
||||
# checkpoint = _load(path)
|
||||
# s = checkpoint["state_dict"]
|
||||
# new_s = {}
|
||||
# for k, v in s.items():
|
||||
# new_s[k.replace('module.', '')] = v
|
||||
# model.load_state_dict(new_s)
|
||||
#
|
||||
# model = model.to(device)
|
||||
# return model.eval()
|
||||
|
||||
|
||||
def main(tempdir):
|
||||
video_faces_pickle_file = args.output_cached_faces_filename
|
||||
print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True)
|
||||
|
||||
is_video = False
|
||||
frame_w = 0
|
||||
frame_h = 0
|
||||
frame_count = 0
|
||||
fps = 0.0
|
||||
duration_millis = 0.0
|
||||
|
||||
if not os.path.isfile(args.image_or_video_filename):
|
||||
raise ValueError('--image_or_video_filename argument must be a valid path to video/image file')
|
||||
|
||||
elif os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg'] \
|
||||
or args.is_image:
|
||||
full_frames = [cv2.imread(args.image_or_video_filename)]
|
||||
fps = args.fps
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
is_video = False
|
||||
frame_count = len(full_frames)
|
||||
|
||||
else:
|
||||
video_stream = cv2.VideoCapture(args.image_or_video_filename)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
|
||||
print('Reading video frames...')
|
||||
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
if args.resize_factor > 1:
|
||||
frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
|
||||
|
||||
if args.rotate:
|
||||
frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
|
||||
|
||||
y1, y2, x1, x2 = args.crop
|
||||
if x2 == -1: x2 = frame.shape[1]
|
||||
if y2 == -1: y2 = frame.shape[0]
|
||||
|
||||
frame = frame[y1:y2, x1:x2]
|
||||
|
||||
full_frames.append(frame)
|
||||
|
||||
frame_h, frame_w = full_frames[0].shape[:-1]
|
||||
|
||||
if len(full_frames) < 2:
|
||||
is_video = False
|
||||
frame_count = len(full_frames)
|
||||
else:
|
||||
is_video = True
|
||||
frame_count = len(full_frames)
|
||||
duration_seconds = frame_count / fps
|
||||
duration_millis = int(duration_seconds * 1000)
|
||||
|
||||
print("Number of frames available for inference: "+str(len(full_frames)), flush=True)
|
||||
print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True)
|
||||
|
||||
# NB: Instead of truncating, let's detect the entire file.
|
||||
# full_frames = full_frames[:len(mel_chunks)]
|
||||
|
||||
print('Face detection time...', flush=True)
|
||||
print('Number of frames to detect faces in: {}'.format(len(full_frames)), flush=True)
|
||||
|
||||
print('Detecting faces...', flush=True)
|
||||
try:
|
||||
_face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
|
||||
except FacelessFrameError as e:
|
||||
print('One or more frames without faces were encountered.')
|
||||
print(e)
|
||||
sys.exit(5)
|
||||
|
||||
print('Done detecting faces!', flush=True)
|
||||
|
||||
mime_type = magic.from_file(args.image_or_video_filename, mime=True)
|
||||
file_size_bytes = os.path.getsize(args.image_or_video_filename)
|
||||
|
||||
metadata = {
|
||||
'is_video': is_video,
|
||||
'width': frame_w,
|
||||
'height': frame_h,
|
||||
'num_frames': len(full_frames),
|
||||
'mimetype': mime_type,
|
||||
'file_size_bytes': file_size_bytes,
|
||||
}
|
||||
|
||||
if is_video:
|
||||
metadata['fps'] = fps
|
||||
metadata['duration_millis'] = duration_millis
|
||||
|
||||
with open(args.output_metadata_filename, 'w') as json_file:
|
||||
json.dump(metadata, json_file)
|
||||
|
||||
if __name__ == '__main__':
|
||||
tempdir = tempfile.mkdtemp()
|
||||
print('Using Python tempdir: {}'.format(tempdir), flush=True)
|
||||
try:
|
||||
main(tempdir)
|
||||
finally:
|
||||
if not args.preserve_tempdir:
|
||||
shutil.rmtree(tempdir)
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
from os.path import dirname, join, basename, isfile
|
||||
from tqdm import tqdm
|
||||
|
||||
from models import SyncNet_color as SyncNet
|
||||
from models import Wav2Lip as Wav2Lip
|
||||
import audio
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch import optim
|
||||
import torch.backends.cudnn as cudnn
|
||||
from torch.utils import data as data_utils
|
||||
import numpy as np
|
||||
|
||||
from glob import glob
|
||||
|
||||
import os, random, cv2, argparse
|
||||
from hparams import hparams, get_image_list
|
||||
|
||||
parser = argparse.ArgumentParser(description='Code to train the Wav2Lip model without the visual quality discriminator')
|
||||
|
||||
parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True, type=str)
|
||||
|
||||
parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str)
|
||||
parser.add_argument('--syncnet_checkpoint_path', help='Load the pre-trained Expert discriminator', required=True, type=str)
|
||||
|
||||
parser.add_argument('--checkpoint_path', help='Resume from this checkpoint', default=None, type=str)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
global_step = 0
|
||||
global_epoch = 0
|
||||
use_cuda = torch.cuda.is_available()
|
||||
print('use_cuda: {}'.format(use_cuda))
|
||||
|
||||
syncnet_T = 5
|
||||
syncnet_mel_step_size = 16
|
||||
|
||||
class Dataset(object):
|
||||
def __init__(self, split):
|
||||
self.all_videos = get_image_list(args.data_root, split)
|
||||
|
||||
def get_frame_id(self, frame):
|
||||
return int(basename(frame).split('.')[0])
|
||||
|
||||
def get_window(self, start_frame):
|
||||
start_id = self.get_frame_id(start_frame)
|
||||
vidname = dirname(start_frame)
|
||||
|
||||
window_fnames = []
|
||||
for frame_id in range(start_id, start_id + syncnet_T):
|
||||
frame = join(vidname, '{}.jpg'.format(frame_id))
|
||||
if not isfile(frame):
|
||||
return None
|
||||
window_fnames.append(frame)
|
||||
return window_fnames
|
||||
|
||||
def read_window(self, window_fnames):
|
||||
if window_fnames is None: return None
|
||||
window = []
|
||||
for fname in window_fnames:
|
||||
img = cv2.imread(fname)
|
||||
if img is None:
|
||||
return None
|
||||
try:
|
||||
img = cv2.resize(img, (hparams.img_size, hparams.img_size))
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
window.append(img)
|
||||
|
||||
return window
|
||||
|
||||
def crop_audio_window(self, spec, start_frame):
|
||||
if type(start_frame) == int:
|
||||
start_frame_num = start_frame
|
||||
else:
|
||||
start_frame_num = self.get_frame_id(start_frame) # 0-indexing ---> 1-indexing
|
||||
start_idx = int(80. * (start_frame_num / float(hparams.fps)))
|
||||
|
||||
end_idx = start_idx + syncnet_mel_step_size
|
||||
|
||||
return spec[start_idx : end_idx, :]
|
||||
|
||||
def get_segmented_mels(self, spec, start_frame):
|
||||
mels = []
|
||||
assert syncnet_T == 5
|
||||
start_frame_num = self.get_frame_id(start_frame) + 1 # 0-indexing ---> 1-indexing
|
||||
if start_frame_num - 2 < 0: return None
|
||||
for i in range(start_frame_num, start_frame_num + syncnet_T):
|
||||
m = self.crop_audio_window(spec, i - 2)
|
||||
if m.shape[0] != syncnet_mel_step_size:
|
||||
return None
|
||||
mels.append(m.T)
|
||||
|
||||
mels = np.asarray(mels)
|
||||
|
||||
return mels
|
||||
|
||||
def prepare_window(self, window):
|
||||
# 3 x T x H x W
|
||||
x = np.asarray(window) / 255.
|
||||
x = np.transpose(x, (3, 0, 1, 2))
|
||||
|
||||
return x
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_videos)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
while 1:
|
||||
idx = random.randint(0, len(self.all_videos) - 1)
|
||||
vidname = self.all_videos[idx]
|
||||
img_names = list(glob(join(vidname, '*.jpg')))
|
||||
if len(img_names) <= 3 * syncnet_T:
|
||||
continue
|
||||
|
||||
img_name = random.choice(img_names)
|
||||
wrong_img_name = random.choice(img_names)
|
||||
while wrong_img_name == img_name:
|
||||
wrong_img_name = random.choice(img_names)
|
||||
|
||||
window_fnames = self.get_window(img_name)
|
||||
wrong_window_fnames = self.get_window(wrong_img_name)
|
||||
if window_fnames is None or wrong_window_fnames is None:
|
||||
continue
|
||||
|
||||
window = self.read_window(window_fnames)
|
||||
if window is None:
|
||||
continue
|
||||
|
||||
wrong_window = self.read_window(wrong_window_fnames)
|
||||
if wrong_window is None:
|
||||
continue
|
||||
|
||||
try:
|
||||
wavpath = join(vidname, "audio.wav")
|
||||
wav = audio.load_wav(wavpath, hparams.sample_rate)
|
||||
|
||||
orig_mel = audio.melspectrogram(wav).T
|
||||
except Exception as e:
|
||||
continue
|
||||
|
||||
mel = self.crop_audio_window(orig_mel.copy(), img_name)
|
||||
|
||||
if (mel.shape[0] != syncnet_mel_step_size):
|
||||
continue
|
||||
|
||||
indiv_mels = self.get_segmented_mels(orig_mel.copy(), img_name)
|
||||
if indiv_mels is None: continue
|
||||
|
||||
window = self.prepare_window(window)
|
||||
y = window.copy()
|
||||
window[:, :, window.shape[2]//2:] = 0.
|
||||
|
||||
wrong_window = self.prepare_window(wrong_window)
|
||||
x = np.concatenate([window, wrong_window], axis=0)
|
||||
|
||||
x = torch.FloatTensor(x)
|
||||
mel = torch.FloatTensor(mel.T).unsqueeze(0)
|
||||
indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1)
|
||||
y = torch.FloatTensor(y)
|
||||
return x, indiv_mels, mel, y
|
||||
|
||||
def save_sample_images(x, g, gt, global_step, checkpoint_dir):
|
||||
x = (x.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
|
||||
g = (g.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
|
||||
gt = (gt.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
|
||||
|
||||
refs, inps = x[..., 3:], x[..., :3]
|
||||
folder = join(checkpoint_dir, "samples_step{:09d}".format(global_step))
|
||||
if not os.path.exists(folder): os.mkdir(folder)
|
||||
collage = np.concatenate((refs, inps, g, gt), axis=-2)
|
||||
for batch_idx, c in enumerate(collage):
|
||||
for t in range(len(c)):
|
||||
cv2.imwrite('{}/{}_{}.jpg'.format(folder, batch_idx, t), c[t])
|
||||
|
||||
logloss = nn.BCELoss()
|
||||
def cosine_loss(a, v, y):
|
||||
d = nn.functional.cosine_similarity(a, v)
|
||||
loss = logloss(d.unsqueeze(1), y)
|
||||
|
||||
return loss
|
||||
|
||||
device = torch.device("cuda" if use_cuda else "cpu")
|
||||
syncnet = SyncNet().to(device)
|
||||
for p in syncnet.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
recon_loss = nn.L1Loss()
|
||||
def get_sync_loss(mel, g):
|
||||
g = g[:, :, :, g.size(3)//2:]
|
||||
g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1)
|
||||
# B, 3 * T, H//2, W
|
||||
a, v = syncnet(mel, g)
|
||||
y = torch.ones(g.size(0), 1).float().to(device)
|
||||
return cosine_loss(a, v, y)
|
||||
|
||||
def train(device, model, train_data_loader, test_data_loader, optimizer,
|
||||
checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
|
||||
|
||||
global global_step, global_epoch
|
||||
resumed_step = global_step
|
||||
|
||||
while global_epoch < nepochs:
|
||||
print('Starting Epoch: {}'.format(global_epoch))
|
||||
running_sync_loss, running_l1_loss = 0., 0.
|
||||
prog_bar = tqdm(enumerate(train_data_loader))
|
||||
for step, (x, indiv_mels, mel, gt) in prog_bar:
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Move data to CUDA device
|
||||
x = x.to(device)
|
||||
mel = mel.to(device)
|
||||
indiv_mels = indiv_mels.to(device)
|
||||
gt = gt.to(device)
|
||||
|
||||
g = model(indiv_mels, x)
|
||||
|
||||
if hparams.syncnet_wt > 0.:
|
||||
sync_loss = get_sync_loss(mel, g)
|
||||
else:
|
||||
sync_loss = 0.
|
||||
|
||||
l1loss = recon_loss(g, gt)
|
||||
|
||||
loss = hparams.syncnet_wt * sync_loss + (1 - hparams.syncnet_wt) * l1loss
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if global_step % checkpoint_interval == 0:
|
||||
save_sample_images(x, g, gt, global_step, checkpoint_dir)
|
||||
|
||||
global_step += 1
|
||||
cur_session_steps = global_step - resumed_step
|
||||
|
||||
running_l1_loss += l1loss.item()
|
||||
if hparams.syncnet_wt > 0.:
|
||||
running_sync_loss += sync_loss.item()
|
||||
else:
|
||||
running_sync_loss += 0.
|
||||
|
||||
if global_step == 1 or global_step % checkpoint_interval == 0:
|
||||
save_checkpoint(
|
||||
model, optimizer, global_step, checkpoint_dir, global_epoch)
|
||||
|
||||
if global_step == 1 or global_step % hparams.eval_interval == 0:
|
||||
with torch.no_grad():
|
||||
average_sync_loss = eval_model(test_data_loader, global_step, device, model, checkpoint_dir)
|
||||
|
||||
if average_sync_loss < .75:
|
||||
hparams.set_hparam('syncnet_wt', 0.01) # without image GAN a lesser weight is sufficient
|
||||
|
||||
prog_bar.set_description('L1: {}, Sync Loss: {}'.format(running_l1_loss / (step + 1),
|
||||
running_sync_loss / (step + 1)))
|
||||
|
||||
global_epoch += 1
|
||||
|
||||
|
||||
def eval_model(test_data_loader, global_step, device, model, checkpoint_dir):
|
||||
eval_steps = 700
|
||||
print('Evaluating for {} steps'.format(eval_steps))
|
||||
sync_losses, recon_losses = [], []
|
||||
step = 0
|
||||
while 1:
|
||||
for x, indiv_mels, mel, gt in test_data_loader:
|
||||
step += 1
|
||||
model.eval()
|
||||
|
||||
# Move data to CUDA device
|
||||
x = x.to(device)
|
||||
gt = gt.to(device)
|
||||
indiv_mels = indiv_mels.to(device)
|
||||
mel = mel.to(device)
|
||||
|
||||
g = model(indiv_mels, x)
|
||||
|
||||
sync_loss = get_sync_loss(mel, g)
|
||||
l1loss = recon_loss(g, gt)
|
||||
|
||||
sync_losses.append(sync_loss.item())
|
||||
recon_losses.append(l1loss.item())
|
||||
|
||||
if step > eval_steps:
|
||||
averaged_sync_loss = sum(sync_losses) / len(sync_losses)
|
||||
averaged_recon_loss = sum(recon_losses) / len(recon_losses)
|
||||
|
||||
print('L1: {}, Sync loss: {}'.format(averaged_recon_loss, averaged_sync_loss))
|
||||
|
||||
return averaged_sync_loss
|
||||
|
||||
def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch):
|
||||
|
||||
checkpoint_path = join(
|
||||
checkpoint_dir, "checkpoint_step{:09d}.pth".format(global_step))
|
||||
optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None
|
||||
torch.save({
|
||||
"state_dict": model.state_dict(),
|
||||
"optimizer": optimizer_state,
|
||||
"global_step": step,
|
||||
"global_epoch": epoch,
|
||||
}, checkpoint_path)
|
||||
print("Saved checkpoint:", checkpoint_path)
|
||||
|
||||
|
||||
def _load(checkpoint_path):
|
||||
if use_cuda:
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path,
|
||||
map_location=lambda storage, loc: storage)
|
||||
return checkpoint
|
||||
|
||||
def load_checkpoint(path, model, optimizer, reset_optimizer=False, overwrite_global_states=True):
|
||||
global global_step
|
||||
global global_epoch
|
||||
|
||||
print("Load checkpoint from: {}".format(path))
|
||||
checkpoint = _load(path)
|
||||
s = checkpoint["state_dict"]
|
||||
new_s = {}
|
||||
for k, v in s.items():
|
||||
new_s[k.replace('module.', '')] = v
|
||||
model.load_state_dict(new_s)
|
||||
if not reset_optimizer:
|
||||
optimizer_state = checkpoint["optimizer"]
|
||||
if optimizer_state is not None:
|
||||
print("Load optimizer state from {}".format(path))
|
||||
optimizer.load_state_dict(checkpoint["optimizer"])
|
||||
if overwrite_global_states:
|
||||
global_step = checkpoint["global_step"]
|
||||
global_epoch = checkpoint["global_epoch"]
|
||||
|
||||
return model
|
||||
|
||||
if __name__ == "__main__":
|
||||
checkpoint_dir = args.checkpoint_dir
|
||||
|
||||
# Dataset and Dataloader setup
|
||||
train_dataset = Dataset('train')
|
||||
test_dataset = Dataset('val')
|
||||
|
||||
train_data_loader = data_utils.DataLoader(
|
||||
train_dataset, batch_size=hparams.batch_size, shuffle=True,
|
||||
num_workers=hparams.num_workers)
|
||||
|
||||
test_data_loader = data_utils.DataLoader(
|
||||
test_dataset, batch_size=hparams.batch_size,
|
||||
num_workers=4)
|
||||
|
||||
device = torch.device("cuda" if use_cuda else "cpu")
|
||||
|
||||
# Model
|
||||
model = Wav2Lip().to(device)
|
||||
print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad)))
|
||||
|
||||
optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
|
||||
lr=hparams.initial_learning_rate)
|
||||
|
||||
if args.checkpoint_path is not None:
|
||||
load_checkpoint(args.checkpoint_path, model, optimizer, reset_optimizer=False)
|
||||
|
||||
load_checkpoint(args.syncnet_checkpoint_path, syncnet, None, reset_optimizer=True, overwrite_global_states=False)
|
||||
|
||||
if not os.path.exists(checkpoint_dir):
|
||||
os.mkdir(checkpoint_dir)
|
||||
|
||||
# Train!
|
||||
train(device, model, train_data_loader, test_data_loader, optimizer,
|
||||
checkpoint_dir=checkpoint_dir,
|
||||
checkpoint_interval=hparams.checkpoint_interval,
|
||||
nepochs=hparams.nepochs)
|
||||
@@ -0,0 +1,29 @@
|
||||
FROM nvidia/cuda:11.0.3-base-ubuntu20.04
|
||||
ENV TZ=Etc/GMT
|
||||
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
|
||||
RUN apt-get update && apt-get install -y git ffmpeg python3.8 python3-pip curl
|
||||
RUN git clone -q https://github.com/SortAnon/ControllableTalkNet /talknet
|
||||
RUN git clone -q https://github.com/SortAnon/hifi-gan /talknet/hifi-gan
|
||||
RUN python3.8 -m pip --no-cache-dir install -r "/talknet/requirements.txt" -f https://download.pytorch.org/whl/torch_stable.html
|
||||
RUN python3.8 -m pip --no-cache-dir install git+https://github.com/SortAnon/NeMo.git
|
||||
RUN python3.8 -m pip uninstall -y pesq
|
||||
RUN python3.8 -m pip install pesq==0.0.2
|
||||
RUN touch /talknet/is_docker
|
||||
WORKDIR /talknet
|
||||
EXPOSE 8050
|
||||
|
||||
RUN printf "#!/bin/bash \
|
||||
\necho Updating TalkNet... \
|
||||
\ngit -C /talknet reset --hard origin/main -q \
|
||||
\ngit -C /talknet pull origin main -q \
|
||||
\necho Updating HiFi-GAN... \
|
||||
\ngit -C /talknet reset --hard origin/main -q \
|
||||
\ngit -C /talknet pull origin main -q \
|
||||
\necho Updating Python dependencies... \
|
||||
\npython3.8 -m pip --quiet --no-cache-dir install -r /talknet/requirements.txt -f https://download.pytorch.org/whl/torch_stable.html \
|
||||
\npython3.8 -m pip --quiet --no-cache-dir install git+https://github.com/SortAnon/NeMo.git \
|
||||
\necho Launching TalkNet... \
|
||||
\npython3.8 talknet_offline.py\n" > /talknet/docker_launch.sh
|
||||
|
||||
RUN chmod +x /talknet/docker_launch.sh
|
||||
CMD ["./docker_launch.sh"]
|
||||
@@ -0,0 +1,121 @@
|
||||
Creative Commons Legal Code
|
||||
|
||||
CC0 1.0 Universal
|
||||
|
||||
CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE
|
||||
LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN
|
||||
ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS
|
||||
INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES
|
||||
REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS
|
||||
PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM
|
||||
THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED
|
||||
HEREUNDER.
|
||||
|
||||
Statement of Purpose
|
||||
|
||||
The laws of most jurisdictions throughout the world automatically confer
|
||||
exclusive Copyright and Related Rights (defined below) upon the creator
|
||||
and subsequent owner(s) (each and all, an "owner") of an original work of
|
||||
authorship and/or a database (each, a "Work").
|
||||
|
||||
Certain owners wish to permanently relinquish those rights to a Work for
|
||||
the purpose of contributing to a commons of creative, cultural and
|
||||
scientific works ("Commons") that the public can reliably and without fear
|
||||
of later claims of infringement build upon, modify, incorporate in other
|
||||
works, reuse and redistribute as freely as possible in any form whatsoever
|
||||
and for any purposes, including without limitation commercial purposes.
|
||||
These owners may contribute to the Commons to promote the ideal of a free
|
||||
culture and the further production of creative, cultural and scientific
|
||||
works, or to gain reputation or greater distribution for their Work in
|
||||
part through the use and efforts of others.
|
||||
|
||||
For these and/or other purposes and motivations, and without any
|
||||
expectation of additional consideration or compensation, the person
|
||||
associating CC0 with a Work (the "Affirmer"), to the extent that he or she
|
||||
is an owner of Copyright and Related Rights in the Work, voluntarily
|
||||
elects to apply CC0 to the Work and publicly distribute the Work under its
|
||||
terms, with knowledge of his or her Copyright and Related Rights in the
|
||||
Work and the meaning and intended legal effect of CC0 on those rights.
|
||||
|
||||
1. Copyright and Related Rights. A Work made available under CC0 may be
|
||||
protected by copyright and related or neighboring rights ("Copyright and
|
||||
Related Rights"). Copyright and Related Rights include, but are not
|
||||
limited to, the following:
|
||||
|
||||
i. the right to reproduce, adapt, distribute, perform, display,
|
||||
communicate, and translate a Work;
|
||||
ii. moral rights retained by the original author(s) and/or performer(s);
|
||||
iii. publicity and privacy rights pertaining to a person's image or
|
||||
likeness depicted in a Work;
|
||||
iv. rights protecting against unfair competition in regards to a Work,
|
||||
subject to the limitations in paragraph 4(a), below;
|
||||
v. rights protecting the extraction, dissemination, use and reuse of data
|
||||
in a Work;
|
||||
vi. database rights (such as those arising under Directive 96/9/EC of the
|
||||
European Parliament and of the Council of 11 March 1996 on the legal
|
||||
protection of databases, and under any national implementation
|
||||
thereof, including any amended or successor version of such
|
||||
directive); and
|
||||
vii. other similar, equivalent or corresponding rights throughout the
|
||||
world based on applicable law or treaty, and any national
|
||||
implementations thereof.
|
||||
|
||||
2. Waiver. To the greatest extent permitted by, but not in contravention
|
||||
of, applicable law, Affirmer hereby overtly, fully, permanently,
|
||||
irrevocably and unconditionally waives, abandons, and surrenders all of
|
||||
Affirmer's Copyright and Related Rights and associated claims and causes
|
||||
of action, whether now known or unknown (including existing as well as
|
||||
future claims and causes of action), in the Work (i) in all territories
|
||||
worldwide, (ii) for the maximum duration provided by applicable law or
|
||||
treaty (including future time extensions), (iii) in any current or future
|
||||
medium and for any number of copies, and (iv) for any purpose whatsoever,
|
||||
including without limitation commercial, advertising or promotional
|
||||
purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each
|
||||
member of the public at large and to the detriment of Affirmer's heirs and
|
||||
successors, fully intending that such Waiver shall not be subject to
|
||||
revocation, rescission, cancellation, termination, or any other legal or
|
||||
equitable action to disrupt the quiet enjoyment of the Work by the public
|
||||
as contemplated by Affirmer's express Statement of Purpose.
|
||||
|
||||
3. Public License Fallback. Should any part of the Waiver for any reason
|
||||
be judged legally invalid or ineffective under applicable law, then the
|
||||
Waiver shall be preserved to the maximum extent permitted taking into
|
||||
account Affirmer's express Statement of Purpose. In addition, to the
|
||||
extent the Waiver is so judged Affirmer hereby grants to each affected
|
||||
person a royalty-free, non transferable, non sublicensable, non exclusive,
|
||||
irrevocable and unconditional license to exercise Affirmer's Copyright and
|
||||
Related Rights in the Work (i) in all territories worldwide, (ii) for the
|
||||
maximum duration provided by applicable law or treaty (including future
|
||||
time extensions), (iii) in any current or future medium and for any number
|
||||
of copies, and (iv) for any purpose whatsoever, including without
|
||||
limitation commercial, advertising or promotional purposes (the
|
||||
"License"). The License shall be deemed effective as of the date CC0 was
|
||||
applied by Affirmer to the Work. Should any part of the License for any
|
||||
reason be judged legally invalid or ineffective under applicable law, such
|
||||
partial invalidity or ineffectiveness shall not invalidate the remainder
|
||||
of the License, and in such case Affirmer hereby affirms that he or she
|
||||
will not (i) exercise any of his or her remaining Copyright and Related
|
||||
Rights in the Work or (ii) assert any associated claims and causes of
|
||||
action with respect to the Work, in either case contrary to Affirmer's
|
||||
express Statement of Purpose.
|
||||
|
||||
4. Limitations and Disclaimers.
|
||||
|
||||
a. No trademark or patent rights held by Affirmer are waived, abandoned,
|
||||
surrendered, licensed or otherwise affected by this document.
|
||||
b. Affirmer offers the Work as-is and makes no representations or
|
||||
warranties of any kind concerning the Work, express, implied,
|
||||
statutory or otherwise, including without limitation warranties of
|
||||
title, merchantability, fitness for a particular purpose, non
|
||||
infringement, or the absence of latent or other defects, accuracy, or
|
||||
the present or absence of errors, whether or not discoverable, all to
|
||||
the greatest extent permissible under applicable law.
|
||||
c. Affirmer disclaims responsibility for clearing rights of other persons
|
||||
that may apply to the Work or any use thereof, including without
|
||||
limitation any person's Copyright and Related Rights in the Work.
|
||||
Further, Affirmer disclaims responsibility for obtaining any necessary
|
||||
consents, permissions or other rights required for any use of the
|
||||
Work.
|
||||
d. Affirmer understands and acknowledges that Creative Commons is not a
|
||||
party to this document and has no duty or obligation with respect to
|
||||
this CC0 or use of the Work.
|
||||
@@ -0,0 +1,26 @@
|
||||
# Controllable TalkNet
|
||||
Controllable TalkNet is a web application that lets you synthesize speech,
|
||||
which mimics the pitch and pacing of an existing audio clip. It's based on [NVIDIA's implementation](https://github.com/NVIDIA/NeMo)
|
||||
of TalkNet 2, with some changes to support singing synthesis and higher audio quality.
|
||||
|
||||
## Requirements
|
||||
* A Google account to run Colab, or...
|
||||
* An NVIDIA GPU with 4+ GB of VRAM
|
||||
* 10 GB of free space
|
||||
|
||||
## How to run
|
||||
### Google Colab
|
||||
* [Go to the Colab notebook](https://colab.research.google.com/drive/1aj6Jk8cpRw7SsN3JSYCv57CrR6s0gYPB) and follow the instructions.
|
||||
|
||||
### TalkNet Offline (Windows)
|
||||
* [Download the setup script](https://github.com/SortAnon/ControllableTalkNet/releases/latest/download/TalkNetOffline.zip)
|
||||
and extract it to a folder.
|
||||
* Run setup.bat. The initial setup will take about 20 minutes.
|
||||
* When it's done, run talknet.bat to start TalkNet on http://127.0.0.1:8050/. To download updates, run update.bat.
|
||||
|
||||
### Docker (Linux)
|
||||
* Install Docker and NVIDIA Container Toolkit.
|
||||
* [Download the Dockerfile.](https://raw.githubusercontent.com/SortAnon/ControllableTalkNet/main/Dockerfile)
|
||||
Open a terminal, and navigate to the directory where you saved it.
|
||||
* Run ```docker build -t talknet-offline .``` to build the image. Add ```sudo``` if you're not using rootless Docker.
|
||||
* Run ```docker run -it --gpus all -p 8050:8050 talknet-offline``` to start TalkNet on http://127.0.0.1:8050/.
|
||||
@@ -0,0 +1,925 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# TalkNet Training (offline)\n",
|
||||
"\n",
|
||||
"To train a 22KHz TalkNet, run the cells below and follow the instructions.\n",
|
||||
"\n",
|
||||
"The notebook will automatically resume training any models from the last saved checkpoint. If you're resuming from a new session, always re-run step 1 first."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Gss5Ox_RNiba"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 1:** Configure training data paths. Create the following and change the paths below:\n",
|
||||
"* A dataset of .wav files, packaged as a .zip or .tar file\n",
|
||||
"* Training and validation filelists, in LJSpeech format with relative paths (note: ARPABET transcripts are not supported)\n",
|
||||
"* An output path for checkpoints"
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"### CHANGE THESE ###\n",
|
||||
"\n",
|
||||
"dataset = \"/absolute/path_to_dataset.zip\"\n",
|
||||
"train_filelist = \"/absolute/train_filelist.txt\"\n",
|
||||
"val_filelist = \"/absolute/val_filelist.txt\"\n",
|
||||
"output_dir = \"/absolute/name_of_character\"\n",
|
||||
"\n",
|
||||
"### ------------ ###\n",
|
||||
"\n",
|
||||
"import os\n",
|
||||
"import gdown\n",
|
||||
"import platform\n",
|
||||
"\n",
|
||||
"assert os.path.exists(dataset), \"Cannot find dataset\"\n",
|
||||
"assert os.path.exists(train_filelist), \"Cannot find training filelist\"\n",
|
||||
"assert os.path.exists(val_filelist), \"Cannot find validation filelist\"\n",
|
||||
"if not os.path.exists(output_dir):\n",
|
||||
" os.makedirs(output_dir)\n",
|
||||
"\n",
|
||||
"if os.getcwd() != output_dir:\n",
|
||||
" cwd = os.getcwd()\n",
|
||||
"else:\n",
|
||||
" os.chdir(cwd)\n",
|
||||
" \n",
|
||||
"# Download pre-trained models\n",
|
||||
"if not os.path.exists(\"talknet_spect.nemo\"):\n",
|
||||
" print(\"Downloading pre-trained models...\")\n",
|
||||
" zip_path = \"TalkNet_Pretrained.zip\"\n",
|
||||
" if not os.path.exists(zip_path) or os.stat(zip_path).st_size < 100:\n",
|
||||
" d = 'https://drive.google.com/uc?id='\n",
|
||||
" gdown.download(d+\"19wSym9mNEnmzLS9XdPlfNAW9_u-mP1hR\", zip_path, quiet=False)\n",
|
||||
" print(platform.system())\n",
|
||||
" if \"Windows\" in platform.system():\n",
|
||||
" !tar -xf {zip_path}\n",
|
||||
" else:\n",
|
||||
" !unzip -qo {zip_path}\n",
|
||||
" os.remove(zip_path)\n",
|
||||
"\n",
|
||||
"print(\"OK\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "nfSawDUD5tqv"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 2:** Dataset processing, part 1.\n",
|
||||
"\n",
|
||||
"If this step fails, try the following:\n",
|
||||
"* Make sure your filelists are correct. They should have relative \n",
|
||||
"paths that match the contents of the archive."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import shutil\n",
|
||||
"import sys\n",
|
||||
"import json\n",
|
||||
"import nemo\n",
|
||||
"import torch\n",
|
||||
"import torchaudio\n",
|
||||
"import numpy as np\n",
|
||||
"#from pysptk import sptk\n",
|
||||
"from pathlib import Path\n",
|
||||
"from tqdm.notebook import tqdm\n",
|
||||
"import ffmpeg\n",
|
||||
"\n",
|
||||
"def fix_transcripts(inpath):\n",
|
||||
" found_arpabet = False\n",
|
||||
" found_grapheme = False\n",
|
||||
" with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
|
||||
" lines = f.readlines()\n",
|
||||
" with open(inpath, \"w\", encoding=\"utf8\") as f:\n",
|
||||
" for l in lines:\n",
|
||||
" if l.strip() == \"\":\n",
|
||||
" continue\n",
|
||||
" if \"{\" in l:\n",
|
||||
" if not found_arpabet:\n",
|
||||
" print(\"Warning: Skipping ARPABET lines (not supported).\")\n",
|
||||
" found_arpabet = True\n",
|
||||
" else:\n",
|
||||
" f.write(l)\n",
|
||||
" found_grapheme = True\n",
|
||||
" assert found_grapheme, \"No non-ARPABET lines found in \" + inpath\n",
|
||||
"\n",
|
||||
"def generate_json(inpath, outpath):\n",
|
||||
" output = \"\"\n",
|
||||
" sample_rate = 22050\n",
|
||||
" with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
|
||||
" for l in f.readlines():\n",
|
||||
" lpath = l.split(\"|\")[0].strip()\n",
|
||||
" if lpath[:5] != \"wavs/\":\n",
|
||||
" lpath = \"wavs/\" + lpath\n",
|
||||
" lpath = os.path.join(output_dir, \"_training_data\", lpath)\n",
|
||||
" size = os.stat(lpath).st_size\n",
|
||||
" x = {\n",
|
||||
" \"audio_filepath\": lpath,\n",
|
||||
" \"duration\": size / (sample_rate * 2),\n",
|
||||
" \"text\": l.split(\"|\")[1].strip(),\n",
|
||||
" }\n",
|
||||
" output += json.dumps(x) + \"\\n\"\n",
|
||||
" with open(outpath, \"w\", encoding=\"utf8\") as w:\n",
|
||||
" w.write(output)\n",
|
||||
"\n",
|
||||
"def convert_to_22k(inpath):\n",
|
||||
" if inpath.strip()[-4:].lower() != \".wav\":\n",
|
||||
" print(\"Warning: \" + inpath.strip() + \" is not a .wav file!\")\n",
|
||||
" return\n",
|
||||
" ffmpeg.input(inpath).output(\n",
|
||||
" inpath + \"_22k.wav\",\n",
|
||||
" ar=\"22050\",\n",
|
||||
" ac=\"1\",\n",
|
||||
" acodec=\"pcm_s16le\",\n",
|
||||
" map_metadata=\"-1\",\n",
|
||||
" fflags=\"+bitexact\",\n",
|
||||
" ).overwrite_output().run(quiet=True)\n",
|
||||
" os.remove(inpath)\n",
|
||||
" os.rename(inpath + \"_22k.wav\", inpath)\n",
|
||||
"\n",
|
||||
"# Extract dataset\n",
|
||||
"os.chdir(output_dir)\n",
|
||||
"\n",
|
||||
"if os.path.exists(\"_training_data\"):\n",
|
||||
" shutil.rmtree(\"_training_data\")\n",
|
||||
"os.mkdir(\"_training_data\")\n",
|
||||
"os.chdir(\"_training_data\")\n",
|
||||
"os.mkdir(\"wavs\")\n",
|
||||
"os.chdir(\"wavs\")\n",
|
||||
"if dataset[-4:] == \".zip\":\n",
|
||||
" if \"Windows\" in platform.system():\n",
|
||||
" !tar -xf \"{dataset}\"\n",
|
||||
" else:\n",
|
||||
" !unzip -qo \"{dataset}\"\n",
|
||||
"elif dataset[-4:] == \".tar\":\n",
|
||||
" !tar -xf \"{dataset}\"\n",
|
||||
"else:\n",
|
||||
" raise Exception(\"Unknown extension for dataset\")\n",
|
||||
"if os.path.exists(os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\")):\n",
|
||||
" shutil.move(\n",
|
||||
" os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\"), \n",
|
||||
" os.path.join(output_dir, \"_training_data\", \"tempwavs\")\n",
|
||||
" )\n",
|
||||
" shutil.rmtree(os.path.join(output_dir, \"_training_data\", \"wavs\"))\n",
|
||||
" shutil.move(\n",
|
||||
" os.path.join(output_dir, \"_training_data\", \"tempwavs\"), \n",
|
||||
" os.path.join(output_dir, \"_training_data\", \"wavs\")\n",
|
||||
" )\n",
|
||||
"# Filelist for preprocessing\n",
|
||||
"os.chdir(output_dir)\n",
|
||||
"shutil.copy(train_filelist, \"trainfiles.txt\")\n",
|
||||
"shutil.copy(val_filelist, \"valfiles.txt\")\n",
|
||||
"fix_transcripts(\"trainfiles.txt\")\n",
|
||||
"fix_transcripts(\"valfiles.txt\")\n",
|
||||
"seen_files = []\n",
|
||||
"with open(\"trainfiles.txt\") as f:\n",
|
||||
" t = f.read().split(\"\\n\")\n",
|
||||
"with open(\"valfiles.txt\") as f:\n",
|
||||
" v = f.read().split(\"\\n\")\n",
|
||||
" all_filelist = t[:] + v[:]\n",
|
||||
"with open(\"allfiles.txt\", \"w\") as f:\n",
|
||||
" for x in all_filelist:\n",
|
||||
" if x.strip() == \"\":\n",
|
||||
" continue\n",
|
||||
" if x.split(\"|\")[0] not in seen_files:\n",
|
||||
" seen_files.append(x.split(\"|\")[0])\n",
|
||||
" f.write(x.strip() + \"\\n\")\n",
|
||||
"\n",
|
||||
"# Ensure audio is 22k\n",
|
||||
"print(\"Converting audio...\")\n",
|
||||
"for r, _, f in os.walk(os.path.join(\"_training_data\", \"wavs\")):\n",
|
||||
" for name in tqdm(f):\n",
|
||||
" convert_to_22k(os.path.join(r, name))\n",
|
||||
"\n",
|
||||
"# Convert to JSON\n",
|
||||
"generate_json(\"trainfiles.txt\", \"trainfiles.json\")\n",
|
||||
"generate_json(\"valfiles.txt\", \"valfiles.json\")\n",
|
||||
"generate_json(\"allfiles.txt\", \"allfiles.json\")\n",
|
||||
"\n",
|
||||
"print(\"OK\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "bxFr3Fdi_kOC"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 3:** Dataset processing, part 2. This takes a while, but\n",
|
||||
"you only have to run this once per dataset (results are saved to Drive).\n",
|
||||
"\n",
|
||||
"If this step fails, try the following:\n",
|
||||
"* Make sure your dataset only contains WAV files."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"# Extract phoneme duration\n",
|
||||
"\n",
|
||||
"import json\n",
|
||||
"from nemo.collections.asr.models import EncDecCTCModel\n",
|
||||
"asr_model = EncDecCTCModel.from_pretrained(model_name=\"asr_talknet_aligner\").cpu().eval()\n",
|
||||
"\n",
|
||||
"def forward_extractor(tokens, log_probs, blank):\n",
|
||||
" \"\"\"Computes states f and p.\"\"\"\n",
|
||||
" n, m = len(tokens), log_probs.shape[0]\n",
|
||||
" # `f[s, t]` -- max sum of log probs for `s` first codes\n",
|
||||
" # with `t` first timesteps with ending in `tokens[s]`.\n",
|
||||
" f = np.empty((n + 1, m + 1), dtype=float)\n",
|
||||
" f.fill(-(10 ** 9))\n",
|
||||
" p = np.empty((n + 1, m + 1), dtype=int)\n",
|
||||
" f[0, 0] = 0.0 # Start\n",
|
||||
" for s in range(1, n + 1):\n",
|
||||
" c = tokens[s - 1]\n",
|
||||
" for t in range((s + 1) // 2, m + 1):\n",
|
||||
" f[s, t] = log_probs[t - 1, c]\n",
|
||||
" # Option #1: prev char is equal to current one.\n",
|
||||
" if s == 1 or c == blank or c == tokens[s - 3]:\n",
|
||||
" options = f[s : (s - 2 if s > 1 else None) : -1, t - 1]\n",
|
||||
" else: # Is not equal to current one.\n",
|
||||
" options = f[s : (s - 3 if s > 2 else None) : -1, t - 1]\n",
|
||||
" f[s, t] += np.max(options)\n",
|
||||
" p[s, t] = np.argmax(options)\n",
|
||||
" return f, p\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def backward_extractor(f, p):\n",
|
||||
" \"\"\"Computes durs from f and p.\"\"\"\n",
|
||||
" n, m = f.shape\n",
|
||||
" n -= 1\n",
|
||||
" m -= 1\n",
|
||||
" durs = np.zeros(n, dtype=int)\n",
|
||||
" if f[-1, -1] >= f[-2, -1]:\n",
|
||||
" s, t = n, m\n",
|
||||
" else:\n",
|
||||
" s, t = n - 1, m\n",
|
||||
" while s > 0:\n",
|
||||
" durs[s - 1] += 1\n",
|
||||
" s -= p[s, t]\n",
|
||||
" t -= 1\n",
|
||||
" assert durs.shape[0] == n\n",
|
||||
" assert np.sum(durs) == m\n",
|
||||
" assert np.all(durs[1::2] > 0)\n",
|
||||
" return durs\n",
|
||||
"\n",
|
||||
"def preprocess_tokens(tokens, blank):\n",
|
||||
" new_tokens = [blank]\n",
|
||||
" for c in tokens:\n",
|
||||
" new_tokens.extend([c, blank])\n",
|
||||
" tokens = new_tokens\n",
|
||||
" return tokens\n",
|
||||
"\n",
|
||||
"data_config = {\n",
|
||||
" 'manifest_filepath': \"allfiles.json\",\n",
|
||||
" 'sample_rate': 22050,\n",
|
||||
" 'labels': asr_model.decoder.vocabulary,\n",
|
||||
" 'batch_size': 1,\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"parser = nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset.make_vocab(\n",
|
||||
" notation='phonemes', punct=True, spaces=True, stresses=False, add_blank_at=\"last\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"dataset = nemo.collections.asr.data.audio_to_text._AudioTextDataset(\n",
|
||||
" manifest_filepath=data_config['manifest_filepath'], sample_rate=data_config['sample_rate'], parser=parser,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"dl = torch.utils.data.DataLoader(\n",
|
||||
" dataset=dataset, batch_size=data_config['batch_size'], collate_fn=dataset.collate_fn, shuffle=False,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"blank_id = asr_model.decoder.num_classes_with_blank - 1\n",
|
||||
"\n",
|
||||
"if os.path.exists(os.path.join(output_dir, \"durations.pt\")):\n",
|
||||
" print(\"durations.pt already exists; skipping\")\n",
|
||||
"else:\n",
|
||||
" dur_data = {}\n",
|
||||
" for sample_idx, test_sample in tqdm(enumerate(dl), total=len(dl)):\n",
|
||||
" log_probs, _, greedy_predictions = asr_model(\n",
|
||||
" input_signal=test_sample[0], input_signal_length=test_sample[1]\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" log_probs = log_probs[0].cpu().detach().numpy()\n",
|
||||
" seq_ids = test_sample[2][0].cpu().detach().numpy()\n",
|
||||
"\n",
|
||||
" target_tokens = preprocess_tokens(seq_ids, blank_id)\n",
|
||||
"\n",
|
||||
" f, p = forward_extractor(target_tokens, log_probs, blank_id)\n",
|
||||
" durs = backward_extractor(f, p)\n",
|
||||
"\n",
|
||||
" dur_key = Path(dl.dataset.collection[sample_idx].audio_file).stem\n",
|
||||
" dur_data[dur_key] = {\n",
|
||||
" 'blanks': torch.tensor(durs[::2], dtype=torch.long).cpu().detach(), \n",
|
||||
" 'tokens': torch.tensor(durs[1::2], dtype=torch.long).cpu().detach()\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" del test_sample\n",
|
||||
"\n",
|
||||
" torch.save(dur_data, os.path.join(output_dir, \"durations.pt\"))\n",
|
||||
"\n",
|
||||
"#Extract F0 (pitch)\n",
|
||||
"import crepe\n",
|
||||
"from scipy.io import wavfile\n",
|
||||
"\n",
|
||||
"def crepe_f0(audio_file, hop_length=256):\n",
|
||||
" sr, audio = wavfile.read(audio_file)\n",
|
||||
" audio_x = np.arange(0, len(audio)) / 22050.0\n",
|
||||
" time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True)\n",
|
||||
"\n",
|
||||
" x = np.arange(0, len(audio), hop_length) / 22050.0\n",
|
||||
" freq_interp = np.interp(x, time, frequency)\n",
|
||||
" conf_interp = np.interp(x, time, confidence)\n",
|
||||
" audio_interp = np.interp(x, audio_x, np.absolute(audio)) / 32768.0\n",
|
||||
" weights = [0.5, 0.25, 0.25]\n",
|
||||
" audio_smooth = np.convolve(audio_interp, np.array(weights)[::-1], \"same\")\n",
|
||||
"\n",
|
||||
" conf_threshold = 0.25\n",
|
||||
" audio_threshold = 0.0005\n",
|
||||
" for i in range(len(freq_interp)):\n",
|
||||
" if conf_interp[i] < conf_threshold:\n",
|
||||
" freq_interp[i] = 0.0\n",
|
||||
" if audio_smooth[i] < audio_threshold:\n",
|
||||
" freq_interp[i] = 0.0\n",
|
||||
"\n",
|
||||
" # Hack to make f0 and mel lengths equal\n",
|
||||
" if len(audio) % hop_length == 0:\n",
|
||||
" freq_interp = np.pad(freq_interp, pad_width=[0, 1])\n",
|
||||
" return torch.from_numpy(freq_interp.astype(np.float32))\n",
|
||||
"\n",
|
||||
"if os.path.exists(os.path.join(output_dir, \"f0s.pt\")):\n",
|
||||
" print(\"f0s.pt already exists; skipping\")\n",
|
||||
"else:\n",
|
||||
" f0_data = {}\n",
|
||||
" with open(\"allfiles.json\") as f:\n",
|
||||
" for i, l in enumerate(f.readlines()):\n",
|
||||
" print(str(i))\n",
|
||||
" audio_path = json.loads(l)[\"audio_filepath\"]\n",
|
||||
" f0_data[Path(audio_path).stem] = crepe_f0(audio_path)\n",
|
||||
"\n",
|
||||
" # calculate f0 stats (mean & std) only for train set\n",
|
||||
" with open(\"trainfiles.json\") as f:\n",
|
||||
" train_ids = {Path(json.loads(l)[\"audio_filepath\"]).stem for l in f}\n",
|
||||
" all_f0 = torch.cat([f0[f0 >= 1e-5] for f0_id, f0 in f0_data.items() if f0_id in train_ids])\n",
|
||||
"\n",
|
||||
" F0_MEAN, F0_STD = all_f0.mean().item(), all_f0.std().item() \n",
|
||||
" print(\"F0_MEAN: \" + str(F0_MEAN) + \", F0_STD: \" + str(F0_STD))\n",
|
||||
" torch.save(f0_data, os.path.join(output_dir, \"f0s.pt\"))\n",
|
||||
" with open(os.path.join(output_dir, \"f0_info.json\"), \"w\") as f:\n",
|
||||
" f.write(json.dumps({\"FO_MEAN\": F0_MEAN, \"F0_STD\": F0_STD}))\n",
|
||||
"\n",
|
||||
"print(\"OK\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "sos9vsxPkIN7"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 4:** Train duration predictor.\n",
|
||||
"\n",
|
||||
"If CUDA runs out of memory, try the following:\n",
|
||||
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
||||
"* If that doesn't help, reduce the batch size (default 64)."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"batch_size = 64\n",
|
||||
"\n",
|
||||
"import os\n",
|
||||
"from hydra.experimental import compose, initialize\n",
|
||||
"from hydra.core.global_hydra import GlobalHydra\n",
|
||||
"from omegaconf import OmegaConf\n",
|
||||
"import pytorch_lightning as pl\n",
|
||||
"from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
|
||||
"from nemo.collections.tts.models import TalkNetDursModel\n",
|
||||
"from nemo.core.config import hydra_runner\n",
|
||||
"from nemo.utils.exp_manager import exp_manager\n",
|
||||
"\n",
|
||||
"epochs = 20\n",
|
||||
"learning_rate = 1e-3\n",
|
||||
"min_learning_rate = 3e-6\n",
|
||||
"load_checkpoints = True\n",
|
||||
"\n",
|
||||
"def train(cfg):\n",
|
||||
" cfg.sample_rate = 22050\n",
|
||||
" cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
|
||||
" cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
|
||||
" cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
|
||||
" cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
|
||||
" cfg.trainer.accelerator = \"dp\"\n",
|
||||
" cfg.trainer.max_epochs = epochs\n",
|
||||
" cfg.trainer.check_val_every_n_epoch = 5\n",
|
||||
" cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
|
||||
" cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
|
||||
" cfg.model.optim.lr = learning_rate\n",
|
||||
" cfg.model.optim.sched.min_lr = min_learning_rate\n",
|
||||
" cfg.exp_manager.exp_dir = output_dir\n",
|
||||
"\n",
|
||||
" # Find checkpoints\n",
|
||||
" ckpt_path = \"\"\n",
|
||||
" if load_checkpoints:\n",
|
||||
" path0 = os.path.join(output_dir, \"TalkNetDurs\")\n",
|
||||
" if os.path.exists(path0):\n",
|
||||
" path1 = sorted(os.listdir(path0))\n",
|
||||
" for i in range(len(path1)):\n",
|
||||
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
||||
" if os.path.exists(path2):\n",
|
||||
" match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
|
||||
" if len(match) > 0:\n",
|
||||
" ckpt_path = os.path.join(path2, match[0])\n",
|
||||
" print(\"Resuming training from \" + match[0])\n",
|
||||
" break\n",
|
||||
" \n",
|
||||
" if ckpt_path != \"\":\n",
|
||||
" trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
|
||||
" model = TalkNetDursModel(cfg=cfg.model, trainer=trainer)\n",
|
||||
" else:\n",
|
||||
" warmstart_path = os.path.join(cwd, \"talknet_durs.nemo\")\n",
|
||||
" trainer = pl.Trainer(**cfg.trainer)\n",
|
||||
" model = TalkNetDursModel.restore_from(warmstart_path, override_config_path=cfg)\n",
|
||||
" model.set_trainer(trainer)\n",
|
||||
" model.setup_training_data(cfg.model.train_ds)\n",
|
||||
" model.setup_validation_data(cfg.model.validation_ds)\n",
|
||||
" model.setup_optimization(cfg.model.optim)\n",
|
||||
" print(\"Warm-starting from \" + warmstart_path)\n",
|
||||
" exp_manager(trainer, cfg.get('exp_manager', None))\n",
|
||||
" trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n",
|
||||
" trainer.fit(model)\n",
|
||||
"\n",
|
||||
"os.chdir(cwd)\n",
|
||||
"GlobalHydra().clear()\n",
|
||||
"initialize(config_path=\"conf\")\n",
|
||||
"cfg = compose(config_name=\"talknet-durs\")\n",
|
||||
"train(cfg)\n"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "nM7-bMpKO7U2"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 5:** Train pitch predictor.\n",
|
||||
"\n",
|
||||
"If CUDA runs out of memory, try the following:\n",
|
||||
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
||||
"* If that doesn't help, reduce the batch size (default 64)."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"batch_size = 64\n",
|
||||
"epochs = 50\n",
|
||||
"\n",
|
||||
"import json\n",
|
||||
"\n",
|
||||
"with open(os.path.join(output_dir, \"f0_info.json\"), \"r\") as f:\n",
|
||||
" f0_info = json.load(f)\n",
|
||||
" f0_mean = f0_info[\"FO_MEAN\"]\n",
|
||||
" f0_std = f0_info[\"F0_STD\"]\n",
|
||||
"\n",
|
||||
"learning_rate = 1e-3\n",
|
||||
"min_learning_rate = 3e-6\n",
|
||||
"load_checkpoints = True\n",
|
||||
"\n",
|
||||
"import os\n",
|
||||
"from hydra.experimental import compose, initialize\n",
|
||||
"from hydra.core.global_hydra import GlobalHydra\n",
|
||||
"from omegaconf import OmegaConf\n",
|
||||
"import pytorch_lightning as pl\n",
|
||||
"from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
|
||||
"from nemo.collections.tts.models import TalkNetPitchModel\n",
|
||||
"from nemo.core.config import hydra_runner\n",
|
||||
"from nemo.utils.exp_manager import exp_manager\n",
|
||||
"\n",
|
||||
"def train(cfg):\n",
|
||||
" cfg.sample_rate = 22050\n",
|
||||
" cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
|
||||
" cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
|
||||
" cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
|
||||
" cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
|
||||
" cfg.trainer.accelerator = \"dp\"\n",
|
||||
" cfg.trainer.max_epochs = epochs\n",
|
||||
" cfg.trainer.check_val_every_n_epoch = 5\n",
|
||||
" cfg.model.f0_mean=f0_mean\n",
|
||||
" cfg.model.f0_std=f0_std\n",
|
||||
" cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
|
||||
" cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
|
||||
" cfg.model.optim.lr = learning_rate\n",
|
||||
" cfg.model.optim.sched.min_lr = min_learning_rate\n",
|
||||
" cfg.exp_manager.exp_dir = output_dir\n",
|
||||
"\n",
|
||||
" # Find checkpoints\n",
|
||||
" ckpt_path = \"\"\n",
|
||||
" if load_checkpoints:\n",
|
||||
" path0 = os.path.join(output_dir, \"TalkNetPitch\")\n",
|
||||
" if os.path.exists(path0):\n",
|
||||
" path1 = sorted(os.listdir(path0))\n",
|
||||
" for i in range(len(path1)):\n",
|
||||
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
||||
" if os.path.exists(path2):\n",
|
||||
" match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
|
||||
" if len(match) > 0:\n",
|
||||
" ckpt_path = os.path.join(path2, match[0])\n",
|
||||
" print(\"Resuming training from \" + match[0])\n",
|
||||
" break\n",
|
||||
" \n",
|
||||
" if ckpt_path != \"\":\n",
|
||||
" trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
|
||||
" model = TalkNetPitchModel(cfg=cfg.model, trainer=trainer)\n",
|
||||
" else:\n",
|
||||
" warmstart_path = os.path.join(cwd, \"talknet_pitch.nemo\")\n",
|
||||
" trainer = pl.Trainer(**cfg.trainer)\n",
|
||||
" model = TalkNetPitchModel.restore_from(warmstart_path, override_config_path=cfg)\n",
|
||||
" model.set_trainer(trainer)\n",
|
||||
" model.setup_training_data(cfg.model.train_ds)\n",
|
||||
" model.setup_validation_data(cfg.model.validation_ds)\n",
|
||||
" model.setup_optimization(cfg.model.optim)\n",
|
||||
" print(\"Warm-starting from \" + warmstart_path)\n",
|
||||
" exp_manager(trainer, cfg.get('exp_manager', None))\n",
|
||||
" trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n",
|
||||
" trainer.fit(model)\n",
|
||||
"\n",
|
||||
"os.chdir(cwd)\n",
|
||||
"GlobalHydra().clear()\n",
|
||||
"initialize(config_path=\"conf\")\n",
|
||||
"cfg = compose(config_name=\"talknet-pitch\")\n",
|
||||
"train(cfg)"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "JLfm00NuJfon"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 6:** Train spectrogram generator. 200+ epochs are recommended. \n",
|
||||
"\n",
|
||||
"This is the slowest of the three models to train, and the hardest to\n",
|
||||
"get good results from. If your character sounds noisy or robotic,\n",
|
||||
"try improving the dataset, or adjusting the epochs and learning rate.\n",
|
||||
"\n",
|
||||
"If CUDA runs out of memory, try the following:\n",
|
||||
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
||||
"* If that doesn't help, reduce the batch size (default 32)."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"epochs = 200\n",
|
||||
"batch_size = 32\n",
|
||||
"\n",
|
||||
"# Advanced settings. You can probably leave these at their defaults (1e-3, 3e-6, empty, checked).\n",
|
||||
"learning_rate = 1e-3\n",
|
||||
"min_learning_rate = 3e-6\n",
|
||||
"pretrained_path = \"\"\n",
|
||||
"load_checkpoints = True\n",
|
||||
"\n",
|
||||
"import os\n",
|
||||
"from hydra.experimental import compose, initialize\n",
|
||||
"from hydra.core.global_hydra import GlobalHydra\n",
|
||||
"from omegaconf import OmegaConf\n",
|
||||
"import pytorch_lightning as pl\n",
|
||||
"from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
|
||||
"from nemo.collections.tts.models import TalkNetSpectModel\n",
|
||||
"from nemo.core.config import hydra_runner\n",
|
||||
"from nemo.utils.exp_manager import exp_manager\n",
|
||||
"\n",
|
||||
"def train(cfg):\n",
|
||||
" cfg.sample_rate = 22050\n",
|
||||
" cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
|
||||
" cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
|
||||
" cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
|
||||
" cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
|
||||
" cfg.trainer.accelerator = \"dp\"\n",
|
||||
" cfg.trainer.max_epochs = epochs\n",
|
||||
" cfg.trainer.check_val_every_n_epoch = 5\n",
|
||||
" cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
|
||||
" cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
|
||||
" cfg.model.optim.lr = learning_rate\n",
|
||||
" cfg.model.optim.sched.min_lr = min_learning_rate\n",
|
||||
" cfg.exp_manager.exp_dir = output_dir\n",
|
||||
"\n",
|
||||
" # Find checkpoints\n",
|
||||
" ckpt_path = \"\"\n",
|
||||
" if load_checkpoints:\n",
|
||||
" path0 = os.path.join(output_dir, \"TalkNetSpect\")\n",
|
||||
" if os.path.exists(path0):\n",
|
||||
" path1 = sorted(os.listdir(path0))\n",
|
||||
" for i in range(len(path1)):\n",
|
||||
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
||||
" if os.path.exists(path2):\n",
|
||||
" match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
|
||||
" if len(match) > 0:\n",
|
||||
" ckpt_path = os.path.join(path2, match[0])\n",
|
||||
" print(\"Resuming training from \" + match[0])\n",
|
||||
" break\n",
|
||||
" \n",
|
||||
" if ckpt_path != \"\":\n",
|
||||
" trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
|
||||
" model = TalkNetSpectModel(cfg=cfg.model, trainer=trainer)\n",
|
||||
" else:\n",
|
||||
" if pretrained_path != \"\":\n",
|
||||
" warmstart_path = pretrained_path\n",
|
||||
" else:\n",
|
||||
" warmstart_path = os.path.join(cwd, \"talknet_spect.nemo\")\n",
|
||||
" trainer = pl.Trainer(**cfg.trainer)\n",
|
||||
" model = TalkNetSpectModel.restore_from(warmstart_path, override_config_path=cfg)\n",
|
||||
" model.set_trainer(trainer)\n",
|
||||
" model.setup_training_data(cfg.model.train_ds)\n",
|
||||
" model.setup_validation_data(cfg.model.validation_ds)\n",
|
||||
" model.setup_optimization(cfg.model.optim)\n",
|
||||
" print(\"Warm-starting from \" + warmstart_path)\n",
|
||||
" exp_manager(trainer, cfg.get('exp_manager', None))\n",
|
||||
" trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n",
|
||||
" trainer.fit(model)\n",
|
||||
"\n",
|
||||
"os.chdir(cwd)\n",
|
||||
"GlobalHydra().clear()\n",
|
||||
"initialize(config_path=\"conf\")\n",
|
||||
"cfg = compose(config_name=\"talknet-spect\")\n",
|
||||
"train(cfg)"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "N9hh4WPHbCcn"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 7:** Generate GTA spectrograms. This will help HiFi-GAN learn what your TalkNet model sounds like.\n",
|
||||
"\n",
|
||||
"If this step fails, make sure you've finished training the spectrogram generator."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"import sys\n",
|
||||
"import os\n",
|
||||
"import torch\n",
|
||||
"import numpy as np\n",
|
||||
"from tqdm import tqdm\n",
|
||||
"from nemo.collections.tts.models import TalkNetSpectModel\n",
|
||||
"import shutil\n",
|
||||
"\n",
|
||||
"def fix_paths(inpath):\n",
|
||||
" output = \"\"\n",
|
||||
" with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
|
||||
" for l in f.readlines():\n",
|
||||
" if l[:5].lower() != \"wavs/\":\n",
|
||||
" output += \"wavs/\" + l\n",
|
||||
" else:\n",
|
||||
" output += l\n",
|
||||
" with open(inpath, \"w\", encoding=\"utf8\") as w:\n",
|
||||
" w.write(output)\n",
|
||||
"\n",
|
||||
"shutil.copyfile(train_filelist, os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n",
|
||||
"shutil.copyfile(val_filelist, os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n",
|
||||
"fix_paths(os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n",
|
||||
"fix_paths(os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n",
|
||||
"fix_paths(os.path.join(output_dir, \"allfiles.txt\"))\n",
|
||||
"\n",
|
||||
"os.chdir(cwd)\n",
|
||||
"outdir = os.path.join(output_dir, \"_training_data\", \"wavs\")\n",
|
||||
"if not os.path.exists(outdir):\n",
|
||||
" os.mkdir(outdir)\n",
|
||||
"\n",
|
||||
"model_path = \"\"\n",
|
||||
"path0 = os.path.join(output_dir, \"TalkNetSpect\")\n",
|
||||
"if os.path.exists(path0):\n",
|
||||
" path1 = sorted(os.listdir(path0))\n",
|
||||
" for i in range(len(path1)):\n",
|
||||
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
||||
" if os.path.exists(path2):\n",
|
||||
" match = [x for x in os.listdir(path2) if \"TalkNetSpect.nemo\" in x]\n",
|
||||
" if len(match) > 0:\n",
|
||||
" model_path = os.path.join(path2, match[0])\n",
|
||||
" break\n",
|
||||
"assert model_path != \"\", \"TalkNetSpect.nemo not found\"\n",
|
||||
"\n",
|
||||
"dur_path = os.path.join(output_dir, \"durations.pt\")\n",
|
||||
"f0_path = os.path.join(output_dir, \"f0s.pt\")\n",
|
||||
"\n",
|
||||
"model = TalkNetSpectModel.restore_from(model_path)\n",
|
||||
"model.eval()\n",
|
||||
"with open(os.path.join(output_dir, \"allfiles.txt\"), \"r\", encoding=\"utf-8\") as f:\n",
|
||||
" dataset = f.readlines()\n",
|
||||
"durs = torch.load(dur_path)\n",
|
||||
"f0s = torch.load(f0_path)\n",
|
||||
"\n",
|
||||
"for x in tqdm(dataset):\n",
|
||||
" x_name = os.path.splitext(os.path.basename(x.split(\"|\")[0].strip()))[0]\n",
|
||||
" x_tokens = model.parse(text=x.split(\"|\")[1].strip())\n",
|
||||
" x_durs = (\n",
|
||||
" torch.stack(\n",
|
||||
" (\n",
|
||||
" durs[x_name][\"blanks\"],\n",
|
||||
" torch.cat((durs[x_name][\"tokens\"], torch.zeros(1).int())),\n",
|
||||
" ),\n",
|
||||
" dim=1,\n",
|
||||
" )\n",
|
||||
" .view(-1)[:-1]\n",
|
||||
" .view(1, -1)\n",
|
||||
" .to(\"cuda:0\")\n",
|
||||
" )\n",
|
||||
" x_f0s = f0s[x_name].view(1, -1).to(\"cuda:0\")\n",
|
||||
" x_spect = model.force_spectrogram(tokens=x_tokens, durs=x_durs, f0=x_f0s)\n",
|
||||
" rel_path = os.path.splitext(x.split(\"|\")[0].strip())[0][5:]\n",
|
||||
" abs_dir = os.path.join(outdir, os.path.dirname(rel_path))\n",
|
||||
" if abs_dir != \"\" and not os.path.exists(abs_dir):\n",
|
||||
" os.makedirs(abs_dir, exist_ok=True)\n",
|
||||
" np.save(os.path.join(outdir, rel_path + \".npy\"), x_spect.detach().cpu().numpy())\n"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "Ajfyfz2p9Ior"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 8:** Train HiFi-GAN. 2,000+ steps are recommended.\n",
|
||||
"Stop this cell to finish training the model.\n",
|
||||
"\n",
|
||||
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
||||
"If this step still fails to start, make sure step 7 finished successfully.\n",
|
||||
"\n",
|
||||
"Note: If the training process starts at step 2500000, delete the HiFiGAN folder and try again."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"import gdown\n",
|
||||
"d = 'https://drive.google.com/uc?id='\n",
|
||||
"\n",
|
||||
"os.chdir(os.path.join(cwd, \"hifi-gan\"))\n",
|
||||
"\n",
|
||||
"if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\")):\n",
|
||||
" os.makedirs(os.path.join(output_dir, \"HiFiGAN\"))\n",
|
||||
"if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\")):\n",
|
||||
" print(\"Downloading universal model...\")\n",
|
||||
" gdown.download(d+\"1qpgI41wNXFcH-iKq1Y42JlBC9j0je8PW\", os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), quiet=False)\n",
|
||||
" gdown.download(d+\"1O63eHZR9t1haCdRHQcEgMfMNxiOciSru\", os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\"), quiet=False)\n",
|
||||
" start_from_universal = \"--warm_start True \"\n",
|
||||
"else:\n",
|
||||
" start_from_universal = \"\"\n",
|
||||
"\n",
|
||||
"hifi_train = os.path.join(cwd, \"hifi-gan\", \"training.txt\")\n",
|
||||
"hifi_val = os.path.join(cwd, \"hifi-gan\", \"validation.txt\")\n",
|
||||
"hifi_wavs = os.path.join(output_dir, \"_training_data\")\n",
|
||||
"\n",
|
||||
"!python train.py --fine_tuning True --config config_v1b.json \\\n",
|
||||
"{start_from_universal} \\\n",
|
||||
"--checkpoint_interval 250 --checkpoint_path \"{os.path.join(output_dir, 'HiFiGAN')}\" \\\n",
|
||||
"--input_training_file \"{hifi_train}\" \\\n",
|
||||
"--input_validation_file \"{hifi_val}\" \\\n",
|
||||
"--input_wavs_dir \"{hifi_wavs}\"\n"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "yVBjGhRB9hUJ"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"**Step 9:** Package the models. They'll be saved to the output directory as [character_name]_TalkNet.zip.\n",
|
||||
"\n",
|
||||
"When done, upload it to Google Drive, with permissions set to \"Anyone with the link\". \n",
|
||||
"You can then use it with TalkNet by selecting \"Custom model\" as your character. \n",
|
||||
"\n",
|
||||
"This cell will also delete the training checkpoints and logs.\n",
|
||||
"That should free up roughly 2 GB of space.\n",
|
||||
"If you wish to keep them, set delete_checkpoints to False."
|
||||
],
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"source": [
|
||||
"character_name = \"Character\"\n",
|
||||
"delete_checkpoints = True\n",
|
||||
"\n",
|
||||
"import shutil\n",
|
||||
"from zipfile import ZipFile\n",
|
||||
"\n",
|
||||
"def find_talknet(model_dir):\n",
|
||||
" ckpt_path = \"\"\n",
|
||||
" path0 = os.path.join(output_dir, model_dir)\n",
|
||||
" if os.path.exists(path0):\n",
|
||||
" path1 = sorted(os.listdir(path0))\n",
|
||||
" for i in range(len(path1)):\n",
|
||||
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
||||
" if os.path.exists(path2):\n",
|
||||
" match = [x for x in os.listdir(path2) if \".nemo\" in x]\n",
|
||||
" if len(match) > 0:\n",
|
||||
" ckpt_path = os.path.join(path2, match[0])\n",
|
||||
" break\n",
|
||||
" assert ckpt_path != \"\", \"Couldn't find \" + model_dir\n",
|
||||
" return ckpt_path\n",
|
||||
"\n",
|
||||
"durs_path = find_talknet(\"TalkNetDurs\")\n",
|
||||
"pitch_path = find_talknet(\"TalkNetPitch\")\n",
|
||||
"spect_path = find_talknet(\"TalkNetSpect\")\n",
|
||||
"assert os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\")), \"Couldn't find HiFi-GAN\"\n",
|
||||
"\n",
|
||||
"zip = ZipFile(os.path.join(output_dir, character_name + \"_TalkNet.zip\"), 'w')\n",
|
||||
"zip.write(durs_path, \"TalkNetDurs.nemo\")\n",
|
||||
"zip.write(pitch_path, \"TalkNetPitch.nemo\")\n",
|
||||
"zip.write(spect_path, \"TalkNetSpect.nemo\")\n",
|
||||
"zip.write(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), \"hifiganmodel\")\n",
|
||||
"zip.write(os.path.join(output_dir, \"HiFiGAN\", \"config.json\"), \"config.json\")\n",
|
||||
"zip.write(os.path.join(output_dir, \"f0_info.json\"), \"f0_info.json\")\n",
|
||||
"zip.close()\n",
|
||||
"print(\"Archived model to \" + os.path.join(output_dir, character_name + \"_TalkNet.zip\"))\n",
|
||||
"\n",
|
||||
"if delete_checkpoints:\n",
|
||||
" shutil.rmtree((os.path.join(output_dir, \"TalkNetDurs\")))\n",
|
||||
" shutil.rmtree((os.path.join(output_dir, \"TalkNetPitch\")))\n",
|
||||
" shutil.rmtree((os.path.join(output_dir, \"TalkNetSpect\")))\n",
|
||||
" shutil.rmtree((os.path.join(output_dir, \"HiFiGAN\")))\n"
|
||||
],
|
||||
"outputs": [],
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "5OtwfrTT-blU"
|
||||
}
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "TalkNet_Training_Offline.ipynb",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.8.10"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Binary file not shown.
@@ -0,0 +1,420 @@
|
||||
@font-face {
|
||||
font-family: 'EquestriaWebfont';
|
||||
src: url('Equestria.woff') format('woff');
|
||||
}
|
||||
|
||||
/* Table of contents
|
||||
––––––––––––––––––––––––––––––––––––––––––––––––––
|
||||
- Plotly.js
|
||||
- Grid
|
||||
- Base Styles
|
||||
- Typography
|
||||
- Links
|
||||
- Buttons
|
||||
- Forms
|
||||
- Lists
|
||||
- Code
|
||||
- Tables
|
||||
- Spacing
|
||||
- Utilities
|
||||
- Clearing
|
||||
- Media Queries
|
||||
*/
|
||||
|
||||
/* PLotly.js
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
/* plotly.js's modebar's z-index is 1001 by default
|
||||
* https://github.com/plotly/plotly.js/blob/7e4d8ab164258f6bd48be56589dacd9bdd7fded2/src/css/_modebar.scss#L5
|
||||
* In case a dropdown is above the graph, the dropdown's options
|
||||
* will be rendered below the modebar
|
||||
* Increase the select option's z-index
|
||||
*/
|
||||
|
||||
/* This was actually not quite right -
|
||||
dropdowns were overlapping each other (edited October 26)
|
||||
|
||||
.Select {
|
||||
z-index: 1002;
|
||||
}*/
|
||||
|
||||
|
||||
/* Grid
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
.container {
|
||||
position: relative;
|
||||
width: 100%;
|
||||
max-width: 960px;
|
||||
margin: 0 auto;
|
||||
padding: 0 20px;
|
||||
box-sizing: border-box; }
|
||||
.column,
|
||||
.columns {
|
||||
width: 100%;
|
||||
float: left;
|
||||
box-sizing: border-box; }
|
||||
|
||||
/* For devices larger than 400px */
|
||||
@media (min-width: 400px) {
|
||||
.container {
|
||||
width: 85%;
|
||||
padding: 0; }
|
||||
}
|
||||
|
||||
/* For devices larger than 550px */
|
||||
@media (min-width: 550px) {
|
||||
.container {
|
||||
width: 80%; }
|
||||
.column,
|
||||
.columns {
|
||||
margin-left: 4%; }
|
||||
.column:first-child,
|
||||
.columns:first-child {
|
||||
margin-left: 0; }
|
||||
|
||||
.one.column,
|
||||
.one.columns { width: 4.66666666667%; }
|
||||
.two.columns { width: 13.3333333333%; }
|
||||
.three.columns { width: 22%; }
|
||||
.four.columns { width: 30.6666666667%; }
|
||||
.five.columns { width: 39.3333333333%; }
|
||||
.six.columns { width: 48%; }
|
||||
.seven.columns { width: 56.6666666667%; }
|
||||
.eight.columns { width: 65.3333333333%; }
|
||||
.nine.columns { width: 74.0%; }
|
||||
.ten.columns { width: 82.6666666667%; }
|
||||
.eleven.columns { width: 91.3333333333%; }
|
||||
.twelve.columns { width: 100%; margin-left: 0; }
|
||||
|
||||
.one-third.column { width: 30.6666666667%; }
|
||||
.two-thirds.column { width: 65.3333333333%; }
|
||||
|
||||
.one-half.column { width: 48%; }
|
||||
|
||||
/* Offsets */
|
||||
.offset-by-one.column,
|
||||
.offset-by-one.columns { margin-left: 8.66666666667%; }
|
||||
.offset-by-two.column,
|
||||
.offset-by-two.columns { margin-left: 17.3333333333%; }
|
||||
.offset-by-three.column,
|
||||
.offset-by-three.columns { margin-left: 26%; }
|
||||
.offset-by-four.column,
|
||||
.offset-by-four.columns { margin-left: 34.6666666667%; }
|
||||
.offset-by-five.column,
|
||||
.offset-by-five.columns { margin-left: 43.3333333333%; }
|
||||
.offset-by-six.column,
|
||||
.offset-by-six.columns { margin-left: 52%; }
|
||||
.offset-by-seven.column,
|
||||
.offset-by-seven.columns { margin-left: 60.6666666667%; }
|
||||
.offset-by-eight.column,
|
||||
.offset-by-eight.columns { margin-left: 69.3333333333%; }
|
||||
.offset-by-nine.column,
|
||||
.offset-by-nine.columns { margin-left: 78.0%; }
|
||||
.offset-by-ten.column,
|
||||
.offset-by-ten.columns { margin-left: 86.6666666667%; }
|
||||
.offset-by-eleven.column,
|
||||
.offset-by-eleven.columns { margin-left: 95.3333333333%; }
|
||||
|
||||
.offset-by-one-third.column,
|
||||
.offset-by-one-third.columns { margin-left: 34.6666666667%; }
|
||||
.offset-by-two-thirds.column,
|
||||
.offset-by-two-thirds.columns { margin-left: 69.3333333333%; }
|
||||
|
||||
.offset-by-one-half.column,
|
||||
.offset-by-one-half.columns { margin-left: 52%; }
|
||||
|
||||
}
|
||||
|
||||
|
||||
/* Base Styles
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
/* NOTE
|
||||
html is set to 62.5% so that all the REM measurements throughout Skeleton
|
||||
are based on 10px sizing. So basically 1.5rem = 15px :) */
|
||||
html {
|
||||
font-size: 62.5%; }
|
||||
body {
|
||||
font-size: 1.5em; /* currently ems cause chrome bug misinterpreting rems on body element */
|
||||
line-height: 1.6;
|
||||
font-weight: 400;
|
||||
font-family: "Open Sans", "HelveticaNeue", "Helvetica Neue", Helvetica, Arial, sans-serif;
|
||||
color: rgb(50, 50, 50); }
|
||||
|
||||
|
||||
/* Typography
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
h1, h2, h3, h4, h5, h6 {
|
||||
margin-top: 0;
|
||||
margin-bottom: 0;
|
||||
font-weight: 300; }
|
||||
h1 { font-size: 4.5rem; line-height: 1.2; letter-spacing: -.1rem; margin-bottom: 2rem; }
|
||||
h2 { font-size: 3.6rem; line-height: 1.25; letter-spacing: -.1rem; margin-bottom: 1.8rem; margin-top: 1.8rem;}
|
||||
h3 { font-size: 3.0rem; line-height: 1.3; letter-spacing: -.1rem; margin-bottom: 1.5rem; margin-top: 1.5rem;}
|
||||
h4 { font-size: 2.6rem; line-height: 1.35; letter-spacing: -.08rem; margin-bottom: 1.2rem; margin-top: 1.2rem;}
|
||||
h5 { font-size: 2.2rem; line-height: 1.5; letter-spacing: -.05rem; margin-bottom: 0.6rem; margin-top: 0.6rem;}
|
||||
h6 { font-size: 2.0rem; line-height: 1.6; letter-spacing: 0; margin-bottom: 0.75rem; margin-top: 0.75rem;}
|
||||
|
||||
p {
|
||||
margin-top: 0; }
|
||||
|
||||
|
||||
/* Blockquotes
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
blockquote {
|
||||
border-left: 4px lightgrey solid;
|
||||
padding-left: 1rem;
|
||||
margin-top: 2rem;
|
||||
margin-bottom: 2rem;
|
||||
margin-left: 0rem;
|
||||
}
|
||||
|
||||
|
||||
/* Links
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
a {
|
||||
color: #1EAEDB;
|
||||
text-decoration: underline;
|
||||
cursor: pointer;}
|
||||
a:hover {
|
||||
color: #0FA0CE; }
|
||||
|
||||
|
||||
/* Buttons
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
.button,
|
||||
button,
|
||||
input[type="submit"],
|
||||
input[type="reset"],
|
||||
input[type="button"] {
|
||||
display: inline-block;
|
||||
height: 38px;
|
||||
padding: 0 30px;
|
||||
color: #555;
|
||||
text-align: center;
|
||||
font-size: 11px;
|
||||
font-weight: 600;
|
||||
line-height: 38px;
|
||||
letter-spacing: .1rem;
|
||||
text-transform: uppercase;
|
||||
text-decoration: none;
|
||||
white-space: nowrap;
|
||||
background-color: transparent;
|
||||
border-radius: 4px;
|
||||
border: 1px solid #bbb;
|
||||
cursor: pointer;
|
||||
box-sizing: border-box; }
|
||||
.button:hover,
|
||||
button:hover,
|
||||
input[type="submit"]:hover,
|
||||
input[type="reset"]:hover,
|
||||
input[type="button"]:hover,
|
||||
.button:focus,
|
||||
button:focus,
|
||||
input[type="submit"]:focus,
|
||||
input[type="reset"]:focus,
|
||||
input[type="button"]:focus {
|
||||
color: #333;
|
||||
border-color: #888;
|
||||
outline: 0; }
|
||||
.button.button-primary,
|
||||
button.button-primary,
|
||||
input[type="submit"].button-primary,
|
||||
input[type="reset"].button-primary,
|
||||
input[type="button"].button-primary {
|
||||
color: #FFF;
|
||||
background-color: #33C3F0;
|
||||
border-color: #33C3F0; }
|
||||
.button.button-primary:hover,
|
||||
button.button-primary:hover,
|
||||
input[type="submit"].button-primary:hover,
|
||||
input[type="reset"].button-primary:hover,
|
||||
input[type="button"].button-primary:hover,
|
||||
.button.button-primary:focus,
|
||||
button.button-primary:focus,
|
||||
input[type="submit"].button-primary:focus,
|
||||
input[type="reset"].button-primary:focus,
|
||||
input[type="button"].button-primary:focus {
|
||||
color: #FFF;
|
||||
background-color: #1EAEDB;
|
||||
border-color: #1EAEDB; }
|
||||
|
||||
|
||||
/* Forms
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
input[type="email"],
|
||||
input[type="number"],
|
||||
input[type="search"],
|
||||
input[type="text"],
|
||||
input[type="tel"],
|
||||
input[type="url"],
|
||||
input[type="password"],
|
||||
textarea,
|
||||
select {
|
||||
height: 38px;
|
||||
padding: 6px 10px; /* The 6px vertically centers text on FF, ignored by Webkit */
|
||||
background-color: #fff;
|
||||
border: 1px solid #D1D1D1;
|
||||
border-radius: 4px;
|
||||
box-shadow: none;
|
||||
box-sizing: border-box;
|
||||
font-family: inherit;
|
||||
font-size: inherit; /*https://stackoverflow.com/questions/6080413/why-doesnt-input-inherit-the-font-from-body*/}
|
||||
/* Removes awkward default styles on some inputs for iOS */
|
||||
input[type="email"],
|
||||
input[type="number"],
|
||||
input[type="search"],
|
||||
input[type="text"],
|
||||
input[type="tel"],
|
||||
input[type="url"],
|
||||
input[type="password"],
|
||||
textarea {
|
||||
-webkit-appearance: none;
|
||||
-moz-appearance: none;
|
||||
appearance: none; }
|
||||
textarea {
|
||||
min-height: 65px;
|
||||
padding-top: 6px;
|
||||
padding-bottom: 6px; }
|
||||
input[type="email"]:focus,
|
||||
input[type="number"]:focus,
|
||||
input[type="search"]:focus,
|
||||
input[type="text"]:focus,
|
||||
input[type="tel"]:focus,
|
||||
input[type="url"]:focus,
|
||||
input[type="password"]:focus,
|
||||
textarea:focus,
|
||||
select:focus {
|
||||
border: 1px solid #33C3F0;
|
||||
outline: 0; }
|
||||
label,
|
||||
legend {
|
||||
display: block;
|
||||
margin-bottom: 0px; }
|
||||
fieldset {
|
||||
padding: 0;
|
||||
border-width: 0; }
|
||||
input[type="checkbox"],
|
||||
input[type="radio"] {
|
||||
display: inline; }
|
||||
label > .label-body {
|
||||
display: inline-block;
|
||||
margin-left: .5rem;
|
||||
font-weight: normal; }
|
||||
|
||||
|
||||
/* Lists
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
ul {
|
||||
list-style: circle inside; }
|
||||
ol {
|
||||
list-style: decimal inside; }
|
||||
ol, ul {
|
||||
padding-left: 0;
|
||||
margin-top: 0; }
|
||||
ul ul,
|
||||
ul ol,
|
||||
ol ol,
|
||||
ol ul {
|
||||
margin: 1.5rem 0 1.5rem 3rem;
|
||||
font-size: 90%; }
|
||||
li {
|
||||
margin-bottom: 1rem; }
|
||||
|
||||
|
||||
/* Tables
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
table {
|
||||
border-collapse: collapse;
|
||||
}
|
||||
th:not(.CalendarDay),
|
||||
td:not(.CalendarDay) {
|
||||
padding: 12px 15px;
|
||||
text-align: left;
|
||||
border-bottom: 1px solid #E1E1E1; }
|
||||
th:first-child:not(.CalendarDay),
|
||||
td:first-child:not(.CalendarDay) {
|
||||
padding-left: 0; }
|
||||
th:last-child:not(.CalendarDay),
|
||||
td:last-child:not(.CalendarDay) {
|
||||
padding-right: 0; }
|
||||
|
||||
|
||||
/* Spacing
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
button,
|
||||
.button {
|
||||
margin-bottom: 0rem; }
|
||||
input,
|
||||
textarea,
|
||||
select,
|
||||
fieldset {
|
||||
margin-bottom: 0rem; }
|
||||
pre,
|
||||
dl,
|
||||
figure,
|
||||
table,
|
||||
form {
|
||||
margin-bottom: 0rem; }
|
||||
p,
|
||||
ul,
|
||||
ol {
|
||||
margin-bottom: 0.75rem; }
|
||||
|
||||
/* Utilities
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
.u-full-width {
|
||||
width: 100%;
|
||||
box-sizing: border-box; }
|
||||
.u-max-full-width {
|
||||
max-width: 100%;
|
||||
box-sizing: border-box; }
|
||||
.u-pull-right {
|
||||
float: right; }
|
||||
.u-pull-left {
|
||||
float: left; }
|
||||
|
||||
|
||||
/* Misc
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
hr {
|
||||
margin-top: 3rem;
|
||||
margin-bottom: 3.5rem;
|
||||
border-width: 0;
|
||||
border-top: 1px solid #E1E1E1; }
|
||||
|
||||
|
||||
/* Clearing
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
|
||||
/* Self Clearing Goodness */
|
||||
.container:after,
|
||||
.row:after,
|
||||
.u-cf {
|
||||
content: "";
|
||||
display: table;
|
||||
clear: both; }
|
||||
|
||||
|
||||
/* Media Queries
|
||||
–––––––––––––––––––––––––––––––––––––––––––––––––– */
|
||||
/*
|
||||
Note: The best way to structure the use of media queries is to create the queries
|
||||
near the relevant code. For example, if you wanted to change the styles for buttons
|
||||
on small devices, paste the mobile query code up in the buttons section and style it
|
||||
there.
|
||||
*/
|
||||
|
||||
|
||||
/* Larger than mobile */
|
||||
@media (min-width: 400px) {}
|
||||
|
||||
/* Larger than phablet (also point when grid becomes active) */
|
||||
@media (min-width: 550px) {}
|
||||
|
||||
/* Larger than tablet */
|
||||
@media (min-width: 750px) {}
|
||||
|
||||
/* Larger than desktop */
|
||||
@media (min-width: 1000px) {}
|
||||
|
||||
/* Larger than Desktop HD */
|
||||
@media (min-width: 1200px) {}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,368 @@
|
||||
import sys
|
||||
import os
|
||||
import io
|
||||
import base64
|
||||
import dash
|
||||
from jupyter_dash import JupyterDash
|
||||
import dash_core_components as dcc
|
||||
import dash_html_components as html
|
||||
from dash.exceptions import PreventUpdate
|
||||
import gdown
|
||||
import traceback
|
||||
from scipy.io import wavfile
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
sys.path.append("DiffSVC_inference_only")
|
||||
from end2end import load_e2e_diffsvc, endtoend_from_path, write_to_file
|
||||
|
||||
app = JupyterDash(__name__)
|
||||
UPLOAD_DIRECTORY = "/content"
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
app.layout = html.Div(
|
||||
children=[
|
||||
html.H1(
|
||||
children="DiffSVC",
|
||||
style={
|
||||
"font-family": "EquestriaWebfont",
|
||||
"color": "#280e5f",
|
||||
"font-size": "4em",
|
||||
"text-align": "center",
|
||||
"margin-top": "0em",
|
||||
"margin-bottom": "0em",
|
||||
},
|
||||
),
|
||||
html.Label("Character selection", htmlFor="speaker-dropdown"),
|
||||
dcc.Dropdown(
|
||||
id="speaker-dropdown",
|
||||
options=[
|
||||
{
|
||||
"label": "Twilight Sparkle",
|
||||
"value": "Twilight",
|
||||
},
|
||||
{
|
||||
"label": "Discord",
|
||||
"value": "Discord",
|
||||
},
|
||||
{
|
||||
"label": "Pinkie Pie",
|
||||
"value": "Pinkie",
|
||||
},
|
||||
{"label": "Nancy", "value": "Nancy"},
|
||||
],
|
||||
value=None,
|
||||
style={
|
||||
"max-width": "90vw",
|
||||
"width": "20em",
|
||||
"margin-bottom": "0.7em",
|
||||
},
|
||||
),
|
||||
html.Label(
|
||||
"Upload reference audio to " + UPLOAD_DIRECTORY,
|
||||
htmlFor="reference-dropdown",
|
||||
),
|
||||
dcc.Store(id="current-f0s"),
|
||||
dcc.Store(id="current-filename"),
|
||||
dcc.Loading(
|
||||
id="audio-loading",
|
||||
children=[
|
||||
html.Div(
|
||||
[
|
||||
html.Button(
|
||||
"Update file list",
|
||||
id="update-button",
|
||||
style={
|
||||
"margin-right": "10px",
|
||||
},
|
||||
),
|
||||
dcc.Dropdown(
|
||||
id="reference-dropdown",
|
||||
options=[],
|
||||
value=None,
|
||||
style={
|
||||
"max-width": "80vw",
|
||||
"width": "30em",
|
||||
},
|
||||
),
|
||||
],
|
||||
style={
|
||||
"width": "100%",
|
||||
"display": "flex",
|
||||
"align-items": "center",
|
||||
"justify-content": "center",
|
||||
"flex-direction": "row",
|
||||
"margin-left": "50px",
|
||||
"vertical-align": "middle",
|
||||
},
|
||||
),
|
||||
html.Div(
|
||||
id="audio-loading-output",
|
||||
style={
|
||||
"font-style": "italic",
|
||||
"margin-bottom": "0.7em",
|
||||
"text-align": "center",
|
||||
},
|
||||
),
|
||||
],
|
||||
type="default",
|
||||
),
|
||||
html.Div(
|
||||
[
|
||||
dcc.Checklist(
|
||||
id="svc-options",
|
||||
options=[
|
||||
{"label": "Singing mode", "value": "no_pc"},
|
||||
],
|
||||
value=[],
|
||||
)
|
||||
],
|
||||
style={
|
||||
"width": "100%",
|
||||
"display": "flex",
|
||||
"align-items": "center",
|
||||
"justify-content": "center",
|
||||
"flex-direction": "row",
|
||||
"margin-left": "50px",
|
||||
"margin-bottom": "0.7em",
|
||||
},
|
||||
),
|
||||
html.Div(
|
||||
[
|
||||
html.Label("Steps", htmlFor="linsteps-input"),
|
||||
dcc.Input(
|
||||
id="linsteps-input",
|
||||
type="number",
|
||||
value="1000",
|
||||
style={"width": "7em", "margin-left": "10px"},
|
||||
min=1,
|
||||
max=1000,
|
||||
step=1,
|
||||
),
|
||||
html.Label("Step strength", htmlFor="linend-input"),
|
||||
dcc.Input(
|
||||
id="linend-input",
|
||||
type="number",
|
||||
value="0.06",
|
||||
style={"width": "7em", "margin-left": "10px"},
|
||||
min=0.01,
|
||||
max=0.3,
|
||||
step=0.01,
|
||||
),
|
||||
],
|
||||
style={
|
||||
"display": "grid",
|
||||
"grid-template-columns": "auto 100px",
|
||||
"margin-bottom": "0.7em",
|
||||
},
|
||||
),
|
||||
dcc.Loading(
|
||||
html.Div(
|
||||
[
|
||||
html.Button(
|
||||
"Generate",
|
||||
id="gen-button",
|
||||
),
|
||||
html.Audio(
|
||||
id="audio-out",
|
||||
controls=True,
|
||||
style={
|
||||
"display": "none",
|
||||
},
|
||||
),
|
||||
html.Div(
|
||||
id="generated-info",
|
||||
style={
|
||||
"font-style": "italic",
|
||||
},
|
||||
),
|
||||
],
|
||||
style={
|
||||
"width": "100%",
|
||||
"display": "flex",
|
||||
"align-items": "center",
|
||||
"justify-content": "center",
|
||||
"flex-direction": "column",
|
||||
},
|
||||
)
|
||||
),
|
||||
html.Footer(
|
||||
children="""
|
||||
Presented by the Pony Preservation Project. Models by Cookie.
|
||||
""",
|
||||
style={"margin-top": "2em", "font-size": "0.7em"},
|
||||
),
|
||||
],
|
||||
style={
|
||||
"width": "100%",
|
||||
"display": "flex",
|
||||
"align-items": "center",
|
||||
"justify-content": "center",
|
||||
"flex-direction": "column",
|
||||
"background-color": "#FFF",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
playback_style = {
|
||||
"margin-top": "0.3em",
|
||||
"margin-bottom": "0.3em",
|
||||
"display": "block",
|
||||
"width": "600px",
|
||||
"max-width": "90vw",
|
||||
}
|
||||
|
||||
playback_hide = {
|
||||
"display": "none",
|
||||
}
|
||||
|
||||
|
||||
@app.callback(
|
||||
dash.dependencies.Output("reference-dropdown", "options"),
|
||||
[
|
||||
dash.dependencies.Input("update-button", "n_clicks"),
|
||||
],
|
||||
)
|
||||
def update_filelist(n_clicks):
|
||||
filelist = []
|
||||
supported_formats = [".wav", ".ogg", ".mp3", "flac"]
|
||||
for x in os.listdir(UPLOAD_DIRECTORY):
|
||||
if x[-4:].lower() in supported_formats:
|
||||
filelist.append({"label": x, "value": x})
|
||||
return filelist
|
||||
|
||||
|
||||
@app.callback(
|
||||
[
|
||||
dash.dependencies.Output("audio-loading-output", "children"),
|
||||
dash.dependencies.Output("current-filename", "data"),
|
||||
],
|
||||
[
|
||||
dash.dependencies.Input("reference-dropdown", "value"),
|
||||
],
|
||||
)
|
||||
def select_file(dropdown_value):
|
||||
if dropdown_value is not None:
|
||||
return [
|
||||
"Selected " + dropdown_value,
|
||||
dropdown_value,
|
||||
]
|
||||
else:
|
||||
return ["No audio selected", None]
|
||||
|
||||
|
||||
def download_model(drive_id, outname):
|
||||
d = "https://drive.google.com/uc?id="
|
||||
model_dir = os.path.join(UPLOAD_DIRECTORY, "models")
|
||||
out_fullpath = os.path.join(model_dir, outname)
|
||||
out_dir = os.path.dirname(out_fullpath)
|
||||
|
||||
if not os.path.exists(model_dir):
|
||||
os.mkdir(model_dir)
|
||||
if not os.path.exists(out_fullpath):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
gdown.download(
|
||||
d + drive_id,
|
||||
out_fullpath,
|
||||
quiet=False,
|
||||
)
|
||||
return out_fullpath
|
||||
|
||||
|
||||
svc_model, svc_name, last_end, last_step = None, None, None, None
|
||||
|
||||
|
||||
@app.callback(
|
||||
[
|
||||
dash.dependencies.Output("audio-out", "src"),
|
||||
dash.dependencies.Output("generated-info", "children"),
|
||||
dash.dependencies.Output("audio-out", "style"),
|
||||
],
|
||||
[dash.dependencies.Input("gen-button", "n_clicks")],
|
||||
[
|
||||
dash.dependencies.State("speaker-dropdown", "value"),
|
||||
dash.dependencies.State("svc-options", "value"),
|
||||
dash.dependencies.State("current-filename", "data"),
|
||||
dash.dependencies.State("linend-input", "value"),
|
||||
dash.dependencies.State("linsteps-input", "value"),
|
||||
],
|
||||
)
|
||||
def generate_audio(n_clicks, speaker, svc_options, wav_name, lin_end, lin_n_steps):
|
||||
global svc_model, svc_name, last_end, last_step
|
||||
|
||||
if n_clicks is None:
|
||||
raise PreventUpdate
|
||||
if speaker is None:
|
||||
return [None, "No character selected", playback_hide]
|
||||
if wav_name is None:
|
||||
return [None, "No reference audio selected", playback_hide]
|
||||
|
||||
try:
|
||||
diffsvc_id = "1Uh17L1JtynFDgX9X6jKLPtD_t4IcTEp9"
|
||||
asr_id = "1qt0pGhCbH0TltFFwSP2Oy3NClthcpr7N"
|
||||
hifigan_id = "1QQT0HjMhGgDuPhyesZYbvck2lEY7GHAB"
|
||||
hifigan_config_id = "1mi_O54zi6nW2eU6VuAO_Y4gv_Df-H9Yv"
|
||||
|
||||
diffsvc_path = download_model(
|
||||
diffsvc_id,
|
||||
os.path.join(diffsvc_id, "diffsvc_model"),
|
||||
)
|
||||
asr_path = download_model(
|
||||
asr_id,
|
||||
os.path.join(asr_id, "asr_model"),
|
||||
)
|
||||
hifigan_path = download_model(
|
||||
hifigan_id,
|
||||
os.path.join(hifigan_id, "hifigan_model"),
|
||||
)
|
||||
download_model(
|
||||
hifigan_config_id,
|
||||
os.path.join(hifigan_id, "config.json"),
|
||||
)
|
||||
if svc_name != diffsvc_id:
|
||||
svc_model = load_e2e_diffsvc(
|
||||
diffsvc_path=diffsvc_path,
|
||||
dilated_asr_path=asr_path,
|
||||
hifigan_path=hifigan_path,
|
||||
device=DEVICE,
|
||||
)
|
||||
svc_name = diffsvc_id
|
||||
if last_end != lin_end or last_step != lin_n_steps:
|
||||
svc_model[0].generator.diffusion.set_noise_schedule(
|
||||
1e-4, float(lin_end), int(lin_n_steps), device=DEVICE
|
||||
)
|
||||
last_end = lin_end
|
||||
last_step = lin_n_steps
|
||||
|
||||
pred_audio = endtoend_from_path(
|
||||
*svc_model,
|
||||
os.path.join(UPLOAD_DIRECTORY, wav_name),
|
||||
speaker,
|
||||
"no_pc" not in svc_options,
|
||||
"no_pc" not in svc_options,
|
||||
"no_pc" not in svc_options,
|
||||
t_max_step=int(lin_n_steps),
|
||||
)
|
||||
pred_audio *= 32768.0
|
||||
|
||||
buffer = io.BytesIO()
|
||||
wavfile.write(
|
||||
buffer,
|
||||
svc_model[4].sampling_rate,
|
||||
pred_audio.squeeze().cpu().numpy().astype(np.int16),
|
||||
)
|
||||
b64 = base64.b64encode(buffer.getvalue())
|
||||
sound = "data:audio/x-wav;base64," + b64.decode("ascii")
|
||||
return [sound, "Conversion complete", playback_style]
|
||||
except Exception:
|
||||
return [None, str(traceback.format_exc()), playback_hide]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run_server(
|
||||
mode="external",
|
||||
debug=True,
|
||||
dev_tools_ui=True,
|
||||
dev_tools_hot_reload=True,
|
||||
threaded=True,
|
||||
)
|
||||
@@ -0,0 +1,165 @@
|
||||
[
|
||||
{"source": "MLP",
|
||||
"characters": [
|
||||
{
|
||||
"name": "Applejack",
|
||||
"drive_id": "1kpEjZ3YqMN3chKSXODOqayEm581rxj4r",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Celestia",
|
||||
"drive_id": "1whXXcnXu9XPcI60xIkTEofpaDYDOw5yB",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Chrysalis",
|
||||
"drive_id": "1bb5jKAcQcEQbx1feVwT1UmAEgocINh-E",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Discord",
|
||||
"drive_id": "1Cg9Oc_K9UDe5WgVDAcaCSbbBoo-Npj1E",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Fluttershy",
|
||||
"drive_id": "1KgVnjrnxZTXgjnI56ilkq5G4UJCbbwZZ",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Luna",
|
||||
"drive_id": "1_ztAbe5YArCMwyyQ_G9lUiz74ym5xJKC",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Maud Pie",
|
||||
"drive_id": "132G6oD0HHPPn4t1H6IkYv18_F0UVLWgi",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Pinkie Pie",
|
||||
"drive_id": "1CdYZ2r52mtgJsFs88U0ZViMSnzpQ_HRp",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Rainbow Dash",
|
||||
"drive_id": "1k3EMXxLC0fLvfxzGbeP6B6plgu9hqCSx",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Rarity",
|
||||
"drive_id": "1QWBvQSso4guc1LRUD40WRJ8DY2CfqHGK",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Spike",
|
||||
"drive_id": "1TKFdmFLttjjzByj2fZW8J70ZHjR-RTwc",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Starlight Glimmer",
|
||||
"drive_id": "1M1AMBq_xjwGTNzRUCXtSLIDJHbcSs3zR",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Sunset Shimmer",
|
||||
"drive_id": "1x1aJt06lBvzUWRlxJ9CEKcFHxQxZPpST",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Trixie",
|
||||
"drive_id": "1a3CYt0-oTTSFjxtZvAVMpClTmQteYua5",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Twilight Sparkle",
|
||||
"drive_id": "1QnOliOAmerMUNuo2wXoH-YoainoSjZen",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Twilight Sparkle (whispering)",
|
||||
"drive_id": "14_TUQVirITdyBh9etfNV8KFFhi_PUs30",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Zecora",
|
||||
"drive_id": "1gL0hqqB7952Q1S185moQd_DRCFfIa3_g",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Twilight Sparkle (singing)",
|
||||
"drive_id": "10CENYWV5ugTXZbnsldN6OKR7wkDEe7V7",
|
||||
"is_singing": true
|
||||
},
|
||||
{
|
||||
"name": "Pinkie Pie (singing)",
|
||||
"drive_id": "1lo-8wI8iB3GArrzIDlaNTA1If8WX0OEf",
|
||||
"is_singing": true
|
||||
},
|
||||
{
|
||||
"name": "Applejack (singing)",
|
||||
"drive_id": "1kegd0J2G4nogaX4IOB2O_vd8msaqMx8M",
|
||||
"is_singing": true
|
||||
},
|
||||
{
|
||||
"name": "Fluttershy (singing)",
|
||||
"drive_id": "1BBdTHis91MwnHTt7tD_xtZ-nQ9SgvqD6",
|
||||
"is_singing": true
|
||||
},
|
||||
{
|
||||
"name": "Rarity (singing)",
|
||||
"drive_id": "1EI7npfQeQYuP6b_Gy8iwA5zAg9Yj5MJA",
|
||||
"is_singing": true
|
||||
},
|
||||
{
|
||||
"name": "Rainbow Dash (singing)",
|
||||
"drive_id": "1q4RgCEVJCSaR_1Dh0rxs6JVRJt30hQRK",
|
||||
"is_singing": true
|
||||
},
|
||||
{
|
||||
"name": "Apple Bloom",
|
||||
"drive_id": "1Wm-7gqws0B3j8mtSQa77Tdk3RZwHEwT5",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Scootaloo",
|
||||
"drive_id": "1YkV1VtP1w5XOx3jYYarrCKSzXCB_FLCy",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Sweetie Belle",
|
||||
"drive_id": "1jLX0Py6j8uY93Fjf2l0HOZQYXiShfWUO",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Cozy Glow",
|
||||
"drive_id": "1cUQQ_4KXTXbmH3MqwwyZu77SpMr20HAL",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Big McIntosh",
|
||||
"drive_id": "1UzL5vo_ykmRdbi-5JOlsQzNKkoV4tJBP",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Cadance",
|
||||
"drive_id": "1w98tMJPfGhWvl797gImMfErPXDz_hA-6",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Granny Smith",
|
||||
"drive_id": "1FJuxLB6MhcHLtmEnOIxhN-yxJ85kmacu",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Shining Armor",
|
||||
"drive_id": "1PUUS71w2ik0uuycNB30nXFze8C7O8OzY",
|
||||
"is_singing": false
|
||||
},
|
||||
{
|
||||
"name": "Tirek",
|
||||
"drive_id": "1rcPDqgDeCIHGDdvfOo-fxfA1XeM4g3CB",
|
||||
"is_singing": false
|
||||
}
|
||||
]}
|
||||
]
|
||||
@@ -0,0 +1,114 @@
|
||||
import os
|
||||
from flask import Flask, request, send_file
|
||||
import sys
|
||||
import torch
|
||||
import numpy as np
|
||||
from scipy.io import wavfile
|
||||
import io
|
||||
from nemo.collections.tts.models import TalkNetSpectModel
|
||||
from nemo.collections.tts.models import TalkNetPitchModel
|
||||
from nemo.collections.tts.models import TalkNetDursModel
|
||||
import json
|
||||
|
||||
sys.path.append("hifi-gan")
|
||||
from env import AttrDict
|
||||
from meldataset import MAX_WAV_VALUE
|
||||
from models import Generator
|
||||
from denoiser import Denoiser
|
||||
|
||||
|
||||
app = Flask(__name__)
|
||||
RUN_PATH = os.path.dirname(os.path.realpath(__file__))
|
||||
DEVICE = "cuda:0"
|
||||
|
||||
|
||||
def load_hifigan(model_name, conf_name):
|
||||
# Load HiFi-GAN
|
||||
conf = os.path.join("hifi-gan", conf_name + ".json")
|
||||
with open(conf) as f:
|
||||
json_config = json.loads(f.read())
|
||||
h = AttrDict(json_config)
|
||||
torch.manual_seed(h.seed)
|
||||
hifigan = Generator(h).to(torch.device(DEVICE))
|
||||
state_dict_g = torch.load(model_name, map_location=torch.device(DEVICE))
|
||||
hifigan.load_state_dict(state_dict_g["generator"])
|
||||
hifigan.eval()
|
||||
hifigan.remove_weight_norm()
|
||||
denoiser = Denoiser(hifigan, mode="normal")
|
||||
return hifigan, h, denoiser
|
||||
|
||||
|
||||
def generate_json(input, outpath):
|
||||
output = ""
|
||||
sample_rate = 22050
|
||||
lpath = input.split("|")[0].strip()
|
||||
size = os.stat(lpath).st_size
|
||||
x = {
|
||||
"audio_filepath": lpath,
|
||||
"duration": size / (sample_rate * 2),
|
||||
"text": input.split("|")[1].strip(),
|
||||
}
|
||||
output += json.dumps(x) + "\n"
|
||||
with open(outpath, "w", encoding="utf8") as w:
|
||||
w.write(output)
|
||||
|
||||
|
||||
def load_talknet(talknet_path):
|
||||
with torch.no_grad():
|
||||
tnmodel = TalkNetSpectModel.restore_from(talknet_path)
|
||||
durs_path = os.path.join(os.path.dirname(talknet_path), "TalkNetDurs.nemo")
|
||||
tndurs = TalkNetDursModel.restore_from(durs_path)
|
||||
tnmodel.add_module("_durs_model", tndurs)
|
||||
pitch_path = os.path.join(os.path.dirname(talknet_path), "TalkNetPitch.nemo")
|
||||
tnpitch = TalkNetPitchModel.restore_from(pitch_path)
|
||||
tnmodel.add_module("_pitch_model", tnpitch)
|
||||
tnmodel.eval()
|
||||
return tnmodel
|
||||
|
||||
|
||||
def generate_audio(transcript, tnmodel, hifigan, denoiser):
|
||||
with torch.no_grad():
|
||||
tokens = tnmodel.parse(text=transcript.strip())
|
||||
spect = tnmodel.generate_spectrogram(tokens=tokens)
|
||||
|
||||
y_g_hat = hifigan(spect.float())
|
||||
audio = y_g_hat.squeeze()
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio_denoised = denoiser(audio.view(1, -1), strength=35)[:, 0]
|
||||
audio_np = audio_denoised.detach().cpu().numpy().reshape(-1).astype(np.int16)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
wavfile.write(buffer, 22050, audio_np)
|
||||
return buffer
|
||||
|
||||
|
||||
@app.route("/", methods=["GET"])
|
||||
def get_check():
|
||||
return "TalkNet server online"
|
||||
|
||||
|
||||
@app.route("/api/tts", methods=["GET"])
|
||||
def get_tts():
|
||||
if "text" not in request.args:
|
||||
return ""
|
||||
transcript = request.args.get("text")
|
||||
return send_file(
|
||||
generate_audio(transcript, tnmodel, hifigan, denoiser),
|
||||
attachment_filename="audio.wav",
|
||||
mimetype="audio/x-wav",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
hifigan, h, denoiser = load_hifigan(
|
||||
os.path.join(
|
||||
RUN_PATH, "models", "1QnOliOAmerMUNuo2wXoH-YoainoSjZen", "hifiganmodel"
|
||||
),
|
||||
"config_v1",
|
||||
)
|
||||
tnmodel = load_talknet(
|
||||
os.path.join(
|
||||
RUN_PATH, "models", "1QnOliOAmerMUNuo2wXoH-YoainoSjZen", "TalkNetSpect.nemo"
|
||||
)
|
||||
)
|
||||
app.run(debug=False)
|
||||
@@ -0,0 +1,28 @@
|
||||
tensorflow==2.4.1
|
||||
tensorboard==2.4.1
|
||||
dash==1.21.0
|
||||
dash-bootstrap-components==0.13.0
|
||||
jupyter-dash==0.4.0
|
||||
psola==0.0.1
|
||||
wget==3.2
|
||||
unidecode==1.2.0
|
||||
pysptk==0.1.18
|
||||
frozendict==2.0.3
|
||||
torch==1.8.1+cu111
|
||||
torchvision==0.9.1+cu111
|
||||
torchaudio==0.8.1
|
||||
torchtext==0.9.1
|
||||
torch_stft==0.1.4
|
||||
kaldiio==2.17.2
|
||||
pydub==0.25.1
|
||||
pyannote.audio==1.1.2
|
||||
g2p_en==2.1.0
|
||||
pesq==0.0.2
|
||||
pystoi==0.3.3
|
||||
crepe==0.0.12
|
||||
resampy==0.2.2
|
||||
ffmpeg-python==0.2.0
|
||||
tqdm
|
||||
gdown==3.13.0
|
||||
editdistance==0.5.3
|
||||
ipywidgets==7.6.3
|
||||
@@ -0,0 +1,17 @@
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
from controllable_talknet import *
|
||||
|
||||
if os.path.exists("/talknet/is_docker"):
|
||||
app.run_server(
|
||||
host="0.0.0.0",
|
||||
mode="external",
|
||||
debug=False,
|
||||
threaded=True,
|
||||
)
|
||||
else:
|
||||
app.run_server(
|
||||
mode="external",
|
||||
debug=False,
|
||||
threaded=True,
|
||||
)
|
||||
@@ -0,0 +1,225 @@
|
||||
from collections import OrderedDict
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from hydra.utils import instantiate
|
||||
from omegaconf import DictConfig
|
||||
from pytorch_lightning import Trainer
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
import numpy as np
|
||||
|
||||
from nemo.collections.asr.data.audio_to_text import AudioToCharWithDursF0Dataset
|
||||
from nemo.collections.tts.helpers.helpers import get_mask_from_lengths
|
||||
from nemo.collections.tts.models.base import SpectrogramGenerator
|
||||
from nemo.collections.tts.modules.talknet import (
|
||||
GaussianEmbedding,
|
||||
MaskedInstanceNorm1d,
|
||||
StyleResidual,
|
||||
)
|
||||
from nemo.core import Exportable
|
||||
from nemo.core.classes import ModelPT, PretrainedModelInfo, typecheck
|
||||
from nemo.core.neural_types import MelSpectrogramType, NeuralType
|
||||
|
||||
|
||||
class TalkNetSingerModel(SpectrogramGenerator, Exportable):
|
||||
"""TalkNet's mel spectrogram prediction pipeline."""
|
||||
|
||||
@property
|
||||
def output_types(self):
|
||||
return OrderedDict(
|
||||
{"mel-spectrogram": NeuralType(("B", "D", "T"), MelSpectrogramType())}
|
||||
)
|
||||
|
||||
def __init__(self, cfg: DictConfig, trainer: "Trainer" = None):
|
||||
super().__init__(cfg=cfg, trainer=trainer)
|
||||
typecheck.set_typecheck_enabled(enabled=False)
|
||||
|
||||
cfg = self._cfg
|
||||
self.vocab = AudioToCharWithDursF0Dataset.make_vocab(
|
||||
**cfg.train_ds.dataset.vocab
|
||||
)
|
||||
self.blanking = cfg.train_ds.dataset.blanking
|
||||
self.preprocessor = instantiate(cfg.preprocessor)
|
||||
self.embed = GaussianEmbedding(self.vocab, cfg.d_char)
|
||||
self.norm_f0 = MaskedInstanceNorm1d(1)
|
||||
self.res_f0 = StyleResidual(cfg.d_char, 1, kernel_size=3)
|
||||
self.model = instantiate(cfg.model)
|
||||
d_out = cfg.model.jasper[-1].filters
|
||||
self.proj = nn.Conv1d(d_out, cfg.n_mels, kernel_size=1)
|
||||
|
||||
def forward(self, text, text_len, durs, f0):
|
||||
x, x_len = self.embed(text, durs).transpose(1, 2), durs.sum(-1)
|
||||
f0_sin = np.nan_to_num(np.sin(9.06472 * np.log(0.0363636 * f0.cpu())))
|
||||
f0_cos = np.nan_to_num(np.cos(9.06472 * np.log(0.0363636 * f0.cpu())))
|
||||
f0_note = torch.stack(
|
||||
(torch.from_numpy(f0_sin), torch.from_numpy(f0_cos)), 1
|
||||
).to(self.device)
|
||||
f0, f0_mask = f0.clone(), f0 > 0.0
|
||||
f0 = self.norm_f0(f0.unsqueeze(1), f0_mask)
|
||||
f0[~f0_mask.unsqueeze(1)] = 0.0
|
||||
x = self.res_f0(x, f0)
|
||||
x = torch.cat((x, f0_note), 1)
|
||||
y, _ = self.model(x, x_len)
|
||||
mel = self.proj(y)
|
||||
return mel
|
||||
|
||||
@staticmethod
|
||||
def _metrics(true_mel, true_mel_len, pred_mel):
|
||||
loss = F.mse_loss(pred_mel, true_mel, reduction="none").mean(dim=-2)
|
||||
mask = get_mask_from_lengths(true_mel_len)
|
||||
loss *= mask.float()
|
||||
loss = loss.sum() / mask.sum()
|
||||
return loss
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
audio, audio_len, text, text_len, durs, f0, f0_mask = batch
|
||||
mel, mel_len = self.preprocessor(audio, audio_len)
|
||||
pred_mel = self(text=text, text_len=text_len, durs=durs, f0=f0)
|
||||
loss = self._metrics(true_mel=mel, true_mel_len=mel_len, pred_mel=pred_mel)
|
||||
train_log = {"train_loss": loss}
|
||||
return {"loss": loss, "progress_bar": train_log, "log": train_log}
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
audio, audio_len, text, text_len, durs, f0, f0_mask = batch
|
||||
mel, mel_len = self.preprocessor(audio, audio_len)
|
||||
pred_mel = self(text=text, text_len=text_len, durs=durs, f0=f0)
|
||||
loss = self._metrics(true_mel=mel, true_mel_len=mel_len, pred_mel=pred_mel)
|
||||
val_log = {"val_loss": loss}
|
||||
self.log_dict(
|
||||
val_log, prog_bar=False, on_epoch=True, logger=True, sync_dist=True
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _loader(cfg):
|
||||
dataset = instantiate(cfg.dataset)
|
||||
return torch.utils.data.DataLoader( # noqa
|
||||
dataset=dataset,
|
||||
collate_fn=dataset.collate_fn,
|
||||
**cfg.dataloader_params,
|
||||
)
|
||||
|
||||
def setup_training_data(self, cfg):
|
||||
self._train_dl = self._loader(cfg)
|
||||
|
||||
def setup_validation_data(self, cfg):
|
||||
self._validation_dl = self._loader(cfg)
|
||||
|
||||
def setup_test_data(self, cfg):
|
||||
"""Omitted."""
|
||||
pass
|
||||
|
||||
def parse(self, text: str, **kwargs) -> torch.Tensor:
|
||||
return torch.tensor(self.vocab.encode(text)).long().unsqueeze(0).to(self.device)
|
||||
|
||||
def generate_spectrogram(self, tokens: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||
assert hasattr(self, "_durs_model") and hasattr(self, "_pitch_model")
|
||||
|
||||
if self.blanking:
|
||||
tokens = [
|
||||
AudioToCharWithDursF0Dataset.interleave(
|
||||
x=torch.empty(len(t) + 1, dtype=torch.long, device=t.device).fill_(
|
||||
self.vocab.blank
|
||||
),
|
||||
y=t,
|
||||
)
|
||||
for t in tokens
|
||||
]
|
||||
tokens = AudioToCharWithDursF0Dataset.merge(
|
||||
tokens, value=self.vocab.pad, dtype=torch.long
|
||||
)
|
||||
|
||||
text_len = torch.tensor(tokens.shape[-1], dtype=torch.long).unsqueeze(0)
|
||||
durs = self._durs_model(tokens, text_len)
|
||||
durs = durs.exp() - 1
|
||||
durs[durs < 0.0] = 0.0
|
||||
durs = durs.round().long()
|
||||
|
||||
# Pitch
|
||||
f0_sil, f0_body = self._pitch_model(tokens, text_len, durs)
|
||||
sil_mask = f0_sil.sigmoid() > 0.5
|
||||
f0 = f0_body * self._pitch_model.f0_std + self._pitch_model.f0_mean
|
||||
f0 = (~sil_mask * f0).float()
|
||||
|
||||
# Spect
|
||||
mel = self(tokens, text_len, durs, f0)
|
||||
|
||||
return mel
|
||||
|
||||
def forward_for_export(self, tokens: torch.Tensor, text_len: torch.Tensor):
|
||||
durs = self._durs_model(tokens, text_len)
|
||||
durs = durs.exp() - 1
|
||||
durs[durs < 0.0] = 0.0
|
||||
durs = durs.round().long()
|
||||
|
||||
# Pitch
|
||||
f0_sil, f0_body = self._pitch_model(tokens, text_len, durs)
|
||||
sil_mask = f0_sil.sigmoid() > 0.5
|
||||
f0 = f0_body * self._pitch_model.f0_std + self._pitch_model.f0_mean
|
||||
f0 = (~sil_mask * f0).float()
|
||||
|
||||
# Spect
|
||||
x, x_len = self.embed(tokens, durs).transpose(1, 2), durs.sum(-1)
|
||||
f0_sin = np.nan_to_num(np.sin(9.06472 * np.log(0.0363636 * f0.cpu())))
|
||||
f0_cos = np.nan_to_num(np.cos(9.06472 * np.log(0.0363636 * f0.cpu())))
|
||||
f0_note = torch.stack(
|
||||
(torch.from_numpy(f0_sin), torch.from_numpy(f0_cos)), 1
|
||||
).to(self.device)
|
||||
f0, f0_mask = f0.clone(), f0 > 0.0
|
||||
f0 = self.norm_f0(f0.unsqueeze(1), f0_mask)
|
||||
f0[~f0_mask.unsqueeze(1)] = 0.0
|
||||
x = self.res_f0(x, f0)
|
||||
x = torch.cat((x, f0_note), 1)
|
||||
y, _ = self.model(x, x_len)
|
||||
mel = self.proj(y)
|
||||
|
||||
return mel
|
||||
|
||||
def force_spectrogram(
|
||||
self, tokens: torch.Tensor, durs: torch.Tensor, f0: torch.Tensor, **kwargs
|
||||
) -> torch.Tensor:
|
||||
if self.blanking:
|
||||
tokens = [
|
||||
AudioToCharWithDursF0Dataset.interleave(
|
||||
x=torch.empty(len(t) + 1, dtype=torch.long, device=t.device).fill_(
|
||||
self.vocab.blank
|
||||
),
|
||||
y=t,
|
||||
)
|
||||
for t in tokens
|
||||
]
|
||||
tokens = AudioToCharWithDursF0Dataset.merge(
|
||||
tokens, value=self.vocab.pad, dtype=torch.long
|
||||
)
|
||||
|
||||
text_len = torch.tensor(tokens.shape[-1], dtype=torch.long).unsqueeze(0)
|
||||
durs_len = torch.tensor(durs.shape[-1], dtype=torch.long).unsqueeze(0)
|
||||
assert text_len == durs_len
|
||||
|
||||
# Spect
|
||||
mel = self(tokens, text_len, durs, f0)
|
||||
return mel
|
||||
|
||||
@classmethod
|
||||
def list_available_models(cls) -> "List[PretrainedModelInfo]":
|
||||
"""
|
||||
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
|
||||
Returns:
|
||||
List of available pre-trained models.
|
||||
"""
|
||||
list_of_models = []
|
||||
model = PretrainedModelInfo(
|
||||
pretrained_model_name="tts_en_talknet",
|
||||
location=(
|
||||
"https://api.ngc.nvidia.com/v2/models/nvidia/nemo/tts_en_talknet/versions/1.0.0rc1/files"
|
||||
"/talknet_spect.nemo"
|
||||
),
|
||||
description=(
|
||||
"This model is trained on LJSpeech sampled at 22050Hz, and can be used to generate female "
|
||||
"English voices with an American accent."
|
||||
),
|
||||
class_=cls, # noqa
|
||||
aliases=["TalkNet-22050Hz"],
|
||||
)
|
||||
list_of_models.append(model)
|
||||
return list_of_models
|
||||
@@ -0,0 +1,19 @@
|
||||
__pycache__
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.pyd
|
||||
.Python
|
||||
env
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
.tox
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*,cover
|
||||
*.log
|
||||
.git
|
||||
**/*.nemo
|
||||
**/*.ckpt
|
||||
@@ -0,0 +1,42 @@
|
||||
---
|
||||
name: Bug report
|
||||
about: Create a report to help us improve
|
||||
title: ''
|
||||
labels: bug
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
**Describe the bug**
|
||||
|
||||
A clear and concise description of what the bug is.
|
||||
|
||||
**Steps/Code to reproduce bug**
|
||||
|
||||
Please list *minimal* steps or code snippet for us to be able to reproduce the bug.
|
||||
|
||||
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
|
||||
|
||||
|
||||
**Expected behavior**
|
||||
|
||||
A clear and concise description of what you expected to happen.
|
||||
|
||||
**Environment overview (please complete the following information)**
|
||||
|
||||
- Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
|
||||
- Method of NeMo install: [pip install or from source]. Please specify exact commands you used to install.
|
||||
- If method of install is [Docker], provide `docker pull` & `docker run` commands used
|
||||
|
||||
**Environment details**
|
||||
|
||||
If NVIDIA docker image is used you don't need to specify these.
|
||||
Otherwise, please provide:
|
||||
- OS version
|
||||
- PyTorch version
|
||||
- Python version
|
||||
|
||||
**Additional context**
|
||||
|
||||
Add any other context about the problem here.
|
||||
Example: GPU model
|
||||
@@ -0,0 +1,25 @@
|
||||
---
|
||||
name: Feature request
|
||||
about: Suggest an idea for this project
|
||||
title: ''
|
||||
labels: feature request
|
||||
assignees: okuchaiev
|
||||
|
||||
---
|
||||
|
||||
**Is your feature request related to a problem? Please describe.**
|
||||
|
||||
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
|
||||
|
||||
**Describe the solution you'd like**
|
||||
|
||||
A clear and concise description of what you want to happen.
|
||||
Provide a code snippet on how new APIs/changes would be used by others.
|
||||
|
||||
**Describe alternatives you've considered**
|
||||
|
||||
A clear and concise description of any alternative solutions or features you've considered.
|
||||
|
||||
**Additional context**
|
||||
|
||||
Add any other context or screenshots about the feature request here.
|
||||
@@ -0,0 +1,176 @@
|
||||
# log and data files
|
||||
*.model
|
||||
*.pkl
|
||||
#*.ipynb
|
||||
output
|
||||
result
|
||||
*.pt
|
||||
tests/data/asr
|
||||
.DS_Store
|
||||
bert.pt.json
|
||||
work
|
||||
runs
|
||||
fastspeech_output
|
||||
.hydra
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
**.pyc
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.idea
|
||||
.Python
|
||||
wandb
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
#parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
pip-wheel-metadata/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/build
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don’t work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# celery beat schedule file
|
||||
celerybeat-schedule
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# VSCode project settins
|
||||
.vscode/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
/docs/html
|
||||
/docs/docs_zh/zh
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# Emacs backup files
|
||||
*~
|
||||
|
||||
cifar-10-batches-py
|
||||
*.tar.gz
|
||||
|
||||
# Test data.
|
||||
tests/.data
|
||||
tests/data
|
||||
|
||||
# outputs folder
|
||||
examples/*/outputs
|
||||
examples/*/NeMo_experiments
|
||||
examples/*/nemo_experiments
|
||||
examples/*/.hydra
|
||||
examples/*/wandb
|
||||
examples/*/data
|
||||
wandb
|
||||
dump.py
|
||||
|
||||
docs/sources/source/test_build/
|
||||
|
||||
# Checkpoints, config files and temporary files created in tutorials.
|
||||
examples/neural_graphs/*.chkpt
|
||||
examples/neural_graphs/*.yml
|
||||
|
||||
.hydra/
|
||||
nemo_experiments/
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
# =============================================================================
|
||||
# Copyright (c) 2020 NVIDIA. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# =============================================================================
|
||||
|
||||
# Read the Docs configuration file
|
||||
# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
|
||||
|
||||
# Required field.
|
||||
version: 2
|
||||
|
||||
# Build documentation in the docs/ directory with Sphinx.
|
||||
sphinx:
|
||||
configuration: docs/source/conf.py
|
||||
|
||||
# Set the version of Python and requirements required to build your docs
|
||||
python:
|
||||
version: 3.7
|
||||
install:
|
||||
- requirements: requirements/requirements_docs.txt
|
||||
@@ -0,0 +1,68 @@
|
||||
# Contributions are welcome!
|
||||
|
||||
We do all of NeMo's development in the open. Contributions from NeMo community are welcome.
|
||||
|
||||
|
||||
# Pull Requests (PR) Guidelines
|
||||
|
||||
**Send your PRs to the `main` branch**
|
||||
|
||||
1) Make sure your PR does one thing. Have a clear answer to "What does this PR do?".
|
||||
2) Read General Principles and style guide below
|
||||
3) Make sure you sign your commits. E.g. use ``git commit -s`` when before your commit
|
||||
4) Make sure all unittests finish successfully before sending PR ``pytest`` or (if yor dev box does not have GPU) ``pytest --cpu`` from NeMo's root folder
|
||||
5) Send your PR and request a review
|
||||
|
||||
## Unit tests
|
||||
Quick tests (locally, while developing)
|
||||
```
|
||||
pytest
|
||||
# If you don't have NVIDIA GPU do:
|
||||
# pytest --cpu
|
||||
```
|
||||
Full tests, including pre-trained model downloads
|
||||
```
|
||||
pytest --with_downloads
|
||||
```
|
||||
|
||||
## Whom should you ask for review:
|
||||
1. For changes to NeMo's core: @ericharper, @titu1994, @blisc, or @okuchaiev
|
||||
1. For changes to NeMo's ASR collection: @titu1994, @redoctopus, @jbalam-nv, or @okuchaiev
|
||||
1. For changes to NeMo's NLP collection: @MaximumEntropy, @ericharper, @ekmb, @yzhang123, @VahidooX, @vladgets, or @okuchaiev
|
||||
1. For changes to NeMo's TTS collection: @blisc or @stasbel, or @okuchaiev
|
||||
|
||||
Note that some people may self-assign to review your PR - in which case, please wait for them to add a review.
|
||||
|
||||
Your pull requests must pass all checks and peer-review before they can be merged.
|
||||
|
||||
# General principles
|
||||
1. **User-oriented**: make it easy for end users, even at the cost of writing more code in the background
|
||||
1. **Robust**: make it hard for users to make mistakes.
|
||||
1. **Well-tested**: please add simple, fast unittests. Consider adding CI tests for end-to-end functionality.
|
||||
1. **Reusable**: for every piece of code, think about how it can be reused in the future and make it easy to be reused.
|
||||
1. **Readable**: code should be easier to read.
|
||||
1. **Legal**: if you copy even one line of code from the Internet, make sure that the code allows the license that NeMo supports. Give credit and link back to the code.
|
||||
1. **Sensible**: code should make sense. If you think a piece of code might be confusing, write comments.
|
||||
|
||||
## Python style
|
||||
We use ``black`` as our style guide. To check whether your code will pass style check (from the NeMo's repo folder) run:
|
||||
``python setup.py style`` and if it does not pass run ``python setup.py style --fix``.
|
||||
|
||||
1. Include docstrings for every class and method exposed to the user.
|
||||
1. Use Python 3 type hints for every class and method exposed to the user.
|
||||
1. Avoid wild import: ``from X import *`` unless in ``X.py``, ``__all__`` is defined.
|
||||
1. Minimize the use of ``**kwargs``.
|
||||
1. ``RaiseError`` is preferred to ``assert``. Write: ```if X: raise Error``` instead of ```assert X```.
|
||||
1. Classes are preferred to standalone methods.
|
||||
1. Methods should be atomic. A method shouldn't be longer than 75 lines, e.g. can be fit into the computer screen without scrolling.
|
||||
1. If a method has arguments that don't fit into one line, each argument should be in its own line for readability.
|
||||
1. Add ``__init__.py`` for every folder.
|
||||
1. F-strings are prefered to formatted strings.
|
||||
1. Loggers are preferred to print. In NeMo, you can use logger from ``from nemo.utils import logging``
|
||||
1. Private functions (functions start with ``_``) shouldn't be called outside its host file.
|
||||
1. If a comment lasts multiple lines, use ``'''`` instead of ``#``.
|
||||
|
||||
# Collections
|
||||
Collection is a logical grouping of related Neural Modules. It is a grouping of modules that share a domain area or semantics.
|
||||
When contributing module to a collection, please make sure it belongs to that category.
|
||||
If you would like to start a new one and contribute back to the platform, you are very welcome to do so.
|
||||
@@ -0,0 +1,104 @@
|
||||
# syntax=docker/dockerfile:experimental
|
||||
|
||||
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:21.05-py3
|
||||
|
||||
|
||||
# build an image that includes only the nemo dependencies, ensures that dependencies
|
||||
# are included first for optimal caching, and useful for building a development
|
||||
# image (by specifying build target as `nemo-deps`)
|
||||
FROM ${BASE_IMAGE} as nemo-deps
|
||||
|
||||
# Ensure apt-get won't prompt for selecting options
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
RUN apt-get update && \
|
||||
apt-get install -y \
|
||||
libsndfile1 sox \
|
||||
libfreetype6 \
|
||||
python-setuptools swig \
|
||||
python-dev ffmpeg && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# uninstall stuff from base container
|
||||
RUN pip uninstall -y sacrebleu torchtext
|
||||
|
||||
# build torchaudio (change latest release version to match pytorch)
|
||||
WORKDIR /tmp/torchaudio_build
|
||||
RUN git clone --depth 1 --branch release/0.7 https://github.com/pytorch/audio.git && \
|
||||
cd audio && \
|
||||
BUILD_SOX=1 python setup.py install && \
|
||||
cd .. && rm -r audio
|
||||
|
||||
# TODO: remove when 21.04 container is released
|
||||
# build torchtext
|
||||
WORKDIR /tmp/torchtext_build
|
||||
RUN git clone --branch v0.8.1 https://github.com/pytorch/text.git && \
|
||||
cd text && \
|
||||
git submodule update --init --recursive && \
|
||||
python setup.py clean install && \
|
||||
cd .. && rm -r text
|
||||
|
||||
#install TRT tools: PT quantization support and ONNX graph optimizer
|
||||
WORKDIR /tmp/trt_build
|
||||
RUN git clone https://github.com/NVIDIA/TensorRT.git && \
|
||||
cd TensorRT/tools/onnx-graphsurgeon && python setup.py install && \
|
||||
cd ../pytorch-quantization && \
|
||||
python setup.py install && \
|
||||
rm -fr /tmp/trt_build
|
||||
|
||||
# install nemo dependencies
|
||||
WORKDIR /tmp/nemo
|
||||
COPY requirements .
|
||||
RUN for f in $(ls requirements*.txt); do pip install --disable-pip-version-check --no-cache-dir -r $f; done
|
||||
|
||||
# install nemo_text_processing dependencies
|
||||
COPY nemo_text_processing /tmp/nemo/nemo_text_processing/
|
||||
RUN /bin/bash /tmp/nemo/nemo_text_processing/setup.sh
|
||||
|
||||
# copy nemo source into a scratch image
|
||||
FROM scratch as nemo-src
|
||||
COPY . .
|
||||
|
||||
# start building the final container
|
||||
FROM nemo-deps as nemo
|
||||
ARG NEMO_VERSION=1.0.2
|
||||
|
||||
# Check that NEMO_VERSION is set. Build will fail without this. Expose NEMO and base container
|
||||
# version information as runtime environment variable for introspection purposes
|
||||
RUN /usr/bin/test -n "$NEMO_VERSION" && \
|
||||
/bin/echo "export NEMO_VERSION=${NEMO_VERSION}" >> /root/.bashrc && \
|
||||
/bin/echo "export BASE_IMAGE=${BASE_IMAGE}" >> /root/.bashrc
|
||||
RUN --mount=from=nemo-src,target=/tmp/nemo cd /tmp/nemo && pip install ".[all]" && \
|
||||
python -c "import nemo.collections.asr as nemo_asr" && \
|
||||
python -c "import nemo.collections.nlp as nemo_nlp" && \
|
||||
python -c "import nemo.collections.tts as nemo_tts" && \
|
||||
python -c "import nemo_text_processing.text_normalization as text_normalization"
|
||||
|
||||
# TODO: Remove once 21.04 container is base container
|
||||
# install latest numba version
|
||||
RUN conda update -c numba numba -y
|
||||
|
||||
# copy scripts/examples/tests into container for end user
|
||||
WORKDIR /workspace/nemo
|
||||
COPY scripts /workspace/nemo/scripts
|
||||
COPY examples /workspace/nemo/examples
|
||||
COPY tests /workspace/nemo/tests
|
||||
COPY tutorials /workspace/nemo/tutorials
|
||||
# COPY README.rst LICENSE /workspace/nemo/
|
||||
|
||||
RUN printf "#!/bin/bash\njupyter lab --no-browser --allow-root --ip=0.0.0.0" >> start-jupyter.sh && \
|
||||
chmod +x start-jupyter.sh
|
||||
|
||||
+1451
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,174 @@
|
||||
|
||||
|status| |documentation| |license| |lgtm_grade| |lgtm_alerts| |black|
|
||||
|
||||
.. |status| image:: http://www.repostatus.org/badges/latest/active.svg
|
||||
:target: http://www.repostatus.org/#active
|
||||
:alt: Project Status: Active – The project has reached a stable, usable state and is being actively developed.
|
||||
|
||||
.. |documentation| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=main
|
||||
:alt: Documentation
|
||||
:target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/
|
||||
|
||||
.. |license| image:: https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg
|
||||
:target: https://github.com/NVIDIA/NeMo/blob/master/LICENSE
|
||||
:alt: NeMo core license and license for collections in this repo
|
||||
|
||||
.. |lgtm_grade| image:: https://img.shields.io/lgtm/grade/python/g/NVIDIA/NeMo.svg?logo=lgtm&logoWidth=18
|
||||
:target: https://lgtm.com/projects/g/NVIDIA/NeMo/context:python
|
||||
:alt: Language grade: Python
|
||||
|
||||
.. |lgtm_alerts| image:: https://img.shields.io/lgtm/alerts/g/NVIDIA/NeMo.svg?logo=lgtm&logoWidth=18
|
||||
:target: https://lgtm.com/projects/g/NVIDIA/NeMo/alerts/
|
||||
:alt: Total alerts
|
||||
|
||||
.. |black| image:: https://img.shields.io/badge/code%20style-black-000000.svg
|
||||
:target: https://github.com/psf/black
|
||||
:alt: Code style: black
|
||||
|
||||
.. _main-readme:
|
||||
**NVIDIA NeMo**
|
||||
===============
|
||||
|
||||
Introduction
|
||||
------------
|
||||
|
||||
NVIDIA NeMo is a conversational AI toolkit built for researchers working on automatic speech recognition (ASR), natural language processing (NLP), and text-to-speech synthesis (TTS).
|
||||
The primary objective of NeMo is to help researchers from industry and academia to reuse prior work (code and pretrained models and make it easier to create new `conversational AI models <https://developer.nvidia.com/conversational-ai#started>`_.
|
||||
|
||||
|
||||
`Introductory video. <https://www.youtube.com/embed/wBgpMf_KQVw>`_
|
||||
|
||||
Key Features
|
||||
------------
|
||||
|
||||
* Speech processing
|
||||
* `Automatic Speech recognition (ASR) <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/intro.html>`_: Jasper, QuartzNet, CitriNet, Conformer
|
||||
* `Speech Classification <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speech_classification/intro.html>`_: MatchboxNet (command recognition), MarbleNet (voice activity detection)
|
||||
* `Speaker Recognition <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speaker_recognition/intro.html>`_: SpeakerNet
|
||||
* `Speaker Diarization <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speaker_diarization/intro.html>`_: MarbleNet + SpeakerNet
|
||||
* `NGC collection of pre-trained speech processing models. <https://ngc.nvidia.com/catalog/collections/nvidia:nemo_asr>`_
|
||||
* Natural Language Processing
|
||||
* `Compatible with Hugging Face Transformers and NVIDIA Megatron <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/megatron_finetuning.html>`_
|
||||
* `Neural Machine Translation (NMT) <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/machine_translation.html>`_
|
||||
* `Punctuation and Capitalization <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/punctuation_and_capitalization.html>`_
|
||||
* `Token classification (named entity recognition) <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/token_classification.html>`_
|
||||
* `Text classification <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/text_classification.html>`_
|
||||
* `Joint Intent and Slot Classification <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/joint_intent_slot.html>`_
|
||||
* `BERT pre-training <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/bert_pretraining.html>`_
|
||||
* `Question answering <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/question_answering.html>`_
|
||||
* `GLUE benchmark <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/glue_benchmark.html>`_
|
||||
* `Information retrieval <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/information_retrieval.html>`_
|
||||
* `NGC collection of pre-trained NLP models. <https://ngc.nvidia.com/catalog/collections/nvidia:nemo_nlp>`_
|
||||
* `Speech synthesis (TTS) <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/tts/intro.html#>`_
|
||||
* Spectrogram generation: Tacotron2, GlowTTS, FastSpeech2, FastPitch,
|
||||
* Vocoders: WaveGlow, SqueezeWave, UniGlow, MelGAN, HiFiGAN
|
||||
* End-to-end speech generation: FastPitch_HifiGan_E2E, FastSpeech2_HifiGan_E2E
|
||||
* `NGC collection of pre-trained TTS models. <https://ngc.nvidia.com/catalog/collections/nvidia:nemo_tts>`_
|
||||
|
||||
|
||||
Built for speed, NeMo can utilize NVIDIA's Tensor Cores and scale out training to multiple GPUs and multiple nodes.
|
||||
|
||||
Requirements
|
||||
------------
|
||||
|
||||
1) Python 3.6, 3.7 or 3.8
|
||||
2) Pytorch 1.8.1 or above
|
||||
3) NVIDIA GPU for training
|
||||
|
||||
Documentation
|
||||
-------------
|
||||
|
||||
.. |main| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=main
|
||||
:alt: Documentation Status
|
||||
:scale: 100%
|
||||
:target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/
|
||||
|
||||
.. |v1.0.2| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=v1.0.2
|
||||
:alt: Documentation Status
|
||||
:scale: 100%
|
||||
:target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/v1.0.2/
|
||||
|
||||
.. |stable| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=stable
|
||||
:alt: Documentation Status
|
||||
:scale: 100%
|
||||
:target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/
|
||||
|
||||
+---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+
|
||||
| Version | Status | Description |
|
||||
+=========+=============+==================================================================================================================================+
|
||||
| Latest | |main| | `Documentation of the latest (i.e. main) branch. <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/>`_ |
|
||||
+---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+
|
||||
| Next | |v1.0.2| | `Documentation of the most recent release: v1.0.2 <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/v1.0.2/>`_ |
|
||||
+---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+
|
||||
| Stable | |stable| | `Documentation of the stable (i.e. stable) branch. <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/>`_ |
|
||||
+---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+
|
||||
|
||||
Tutorials
|
||||
---------
|
||||
A great way to start with NeMo is by checking `one of our tutorials <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/v1.0.2/starthere/tutorials.html>`_.
|
||||
|
||||
Getting help with NeMo
|
||||
----------------------
|
||||
FAQ can be found on NeMo's `Discussions board <https://github.com/NVIDIA/NeMo/discussions>`_. You are welcome to ask questions or start discussions there.
|
||||
|
||||
|
||||
Installation
|
||||
------------
|
||||
|
||||
Pip
|
||||
~~~
|
||||
Use this installation mode if you want the latest released version.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
apt-get update && apt-get install -y libsndfile1 ffmpeg
|
||||
pip install Cython
|
||||
pip install nemo_toolkit['all']
|
||||
|
||||
Pip from source
|
||||
~~~~~~~~~~~~~~~
|
||||
Use this installation mode if you want the a version from particular GitHub branch (e.g main).
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
apt-get update && apt-get install -y libsndfile1 ffmpeg
|
||||
pip install Cython
|
||||
python -m pip install git+https://github.com/NVIDIA/NeMo.git@{BRANCH}#egg=nemo_toolkit[all]
|
||||
|
||||
|
||||
From source
|
||||
~~~~~~~~~~~
|
||||
Use this installation mode if you are contributing to NeMo.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
apt-get update && apt-get install -y libsndfile1 ffmpeg
|
||||
git clone https://github.com/NVIDIA/NeMo
|
||||
cd NeMo
|
||||
./reinstall.sh
|
||||
|
||||
Docker containers:
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
If you chose to work with main branch, we recommend using NVIDIA's PyTorch container version 21.05-py3 and then installing from GitHub.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
docker run --gpus all -it --rm -v <nemo_github_folder>:/NeMo --shm-size=8g \
|
||||
-p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit \
|
||||
stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:21.05-py3
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
Many example can be found under `"Examples" <https://github.com/NVIDIA/NeMo/tree/main/examples>`_ folder.
|
||||
|
||||
|
||||
Contributing
|
||||
------------
|
||||
|
||||
We welcome community contributions! Please refer to the `CONTRIBUTING.md <https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md>`_ CONTRIBUTING.md for the process.
|
||||
|
||||
License
|
||||
-------
|
||||
NeMo is under `Apache 2.0 license <https://github.com/NVIDIA/NeMo/blob/main/LICENSE>`_.
|
||||
@@ -0,0 +1,216 @@
|
||||
# Makefile for Sphinx documentation
|
||||
#
|
||||
|
||||
# You can set these variables from the command line.
|
||||
SPHINXOPTS =
|
||||
SPHINXBUILD = sphinx-build
|
||||
PAPER =
|
||||
BUILDDIR = build
|
||||
|
||||
# User-friendly check for sphinx-build
|
||||
ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1)
|
||||
$(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/)
|
||||
endif
|
||||
|
||||
# Internal variables.
|
||||
PAPEROPT_a4 = -D latex_paper_size=a4
|
||||
PAPEROPT_letter = -D latex_paper_size=letter
|
||||
ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) source
|
||||
# the i18n builder cannot share the environment and doctrees with the others
|
||||
I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) source
|
||||
|
||||
.PHONY: help
|
||||
help:
|
||||
@echo "Please use \`make <target>' where <target> is one of"
|
||||
@echo " html to make standalone HTML files"
|
||||
@echo " dirhtml to make HTML files named index.html in directories"
|
||||
@echo " singlehtml to make a single large HTML file"
|
||||
@echo " pickle to make pickle files"
|
||||
@echo " json to make JSON files"
|
||||
@echo " htmlhelp to make HTML files and a HTML help project"
|
||||
@echo " qthelp to make HTML files and a qthelp project"
|
||||
@echo " applehelp to make an Apple Help Book"
|
||||
@echo " devhelp to make HTML files and a Devhelp project"
|
||||
@echo " epub to make an epub"
|
||||
@echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter"
|
||||
@echo " latexpdf to make LaTeX files and run them through pdflatex"
|
||||
@echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx"
|
||||
@echo " text to make text files"
|
||||
@echo " man to make manual pages"
|
||||
@echo " texinfo to make Texinfo files"
|
||||
@echo " info to make Texinfo files and run them through makeinfo"
|
||||
@echo " gettext to make PO message catalogs"
|
||||
@echo " changes to make an overview of all changed/added/deprecated items"
|
||||
@echo " xml to make Docutils-native XML files"
|
||||
@echo " pseudoxml to make pseudoxml-XML files for display purposes"
|
||||
@echo " linkcheck to check all external links for integrity"
|
||||
@echo " doctest to run all doctests embedded in the documentation (if enabled)"
|
||||
@echo " coverage to run coverage check of the documentation (if enabled)"
|
||||
|
||||
.PHONY: clean
|
||||
clean:
|
||||
rm -rf $(BUILDDIR)/*
|
||||
|
||||
.PHONY: html
|
||||
html:
|
||||
$(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html
|
||||
@echo
|
||||
@echo "Build finished. The HTML pages are in $(BUILDDIR)/html."
|
||||
|
||||
.PHONY: dirhtml
|
||||
dirhtml:
|
||||
$(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml
|
||||
@echo
|
||||
@echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml."
|
||||
|
||||
.PHONY: singlehtml
|
||||
singlehtml:
|
||||
$(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml
|
||||
@echo
|
||||
@echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml."
|
||||
|
||||
.PHONY: pickle
|
||||
pickle:
|
||||
$(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle
|
||||
@echo
|
||||
@echo "Build finished; now you can process the pickle files."
|
||||
|
||||
.PHONY: json
|
||||
json:
|
||||
$(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json
|
||||
@echo
|
||||
@echo "Build finished; now you can process the JSON files."
|
||||
|
||||
.PHONY: htmlhelp
|
||||
htmlhelp:
|
||||
$(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp
|
||||
@echo
|
||||
@echo "Build finished; now you can run HTML Help Workshop with the" \
|
||||
".hhp project file in $(BUILDDIR)/htmlhelp."
|
||||
|
||||
.PHONY: qthelp
|
||||
qthelp:
|
||||
$(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp
|
||||
@echo
|
||||
@echo "Build finished; now you can run "qcollectiongenerator" with the" \
|
||||
".qhcp project file in $(BUILDDIR)/qthelp, like this:"
|
||||
@echo "# qcollectiongenerator $(BUILDDIR)/qthelp/OpenSeq2Seq.qhcp"
|
||||
@echo "To view the help file:"
|
||||
@echo "# assistant -collectionFile $(BUILDDIR)/qthelp/OpenSeq2Seq.qhc"
|
||||
|
||||
.PHONY: applehelp
|
||||
applehelp:
|
||||
$(SPHINXBUILD) -b applehelp $(ALLSPHINXOPTS) $(BUILDDIR)/applehelp
|
||||
@echo
|
||||
@echo "Build finished. The help book is in $(BUILDDIR)/applehelp."
|
||||
@echo "N.B. You won't be able to view it unless you put it in" \
|
||||
"~/Library/Documentation/Help or install it in your application" \
|
||||
"bundle."
|
||||
|
||||
.PHONY: devhelp
|
||||
devhelp:
|
||||
$(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp
|
||||
@echo
|
||||
@echo "Build finished."
|
||||
@echo "To view the help file:"
|
||||
@echo "# mkdir -p $$HOME/.local/share/devhelp/OpenSeq2Seq"
|
||||
@echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/OpenSeq2Seq"
|
||||
@echo "# devhelp"
|
||||
|
||||
.PHONY: epub
|
||||
epub:
|
||||
$(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub
|
||||
@echo
|
||||
@echo "Build finished. The epub file is in $(BUILDDIR)/epub."
|
||||
|
||||
.PHONY: latex
|
||||
latex:
|
||||
$(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
|
||||
@echo
|
||||
@echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex."
|
||||
@echo "Run \`make' in that directory to run these through (pdf)latex" \
|
||||
"(use \`make latexpdf' here to do that automatically)."
|
||||
|
||||
.PHONY: latexpdf
|
||||
latexpdf:
|
||||
$(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
|
||||
@echo "Running LaTeX files through pdflatex..."
|
||||
$(MAKE) -C $(BUILDDIR)/latex all-pdf
|
||||
@echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
|
||||
|
||||
.PHONY: latexpdfja
|
||||
latexpdfja:
|
||||
$(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
|
||||
@echo "Running LaTeX files through platex and dvipdfmx..."
|
||||
$(MAKE) -C $(BUILDDIR)/latex all-pdf-ja
|
||||
@echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
|
||||
|
||||
.PHONY: text
|
||||
text:
|
||||
$(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text
|
||||
@echo
|
||||
@echo "Build finished. The text files are in $(BUILDDIR)/text."
|
||||
|
||||
.PHONY: man
|
||||
man:
|
||||
$(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man
|
||||
@echo
|
||||
@echo "Build finished. The manual pages are in $(BUILDDIR)/man."
|
||||
|
||||
.PHONY: texinfo
|
||||
texinfo:
|
||||
$(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo
|
||||
@echo
|
||||
@echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo."
|
||||
@echo "Run \`make' in that directory to run these through makeinfo" \
|
||||
"(use \`make info' here to do that automatically)."
|
||||
|
||||
.PHONY: info
|
||||
info:
|
||||
$(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo
|
||||
@echo "Running Texinfo files through makeinfo..."
|
||||
make -C $(BUILDDIR)/texinfo info
|
||||
@echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo."
|
||||
|
||||
.PHONY: gettext
|
||||
gettext:
|
||||
$(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale
|
||||
@echo
|
||||
@echo "Build finished. The message catalogs are in $(BUILDDIR)/locale."
|
||||
|
||||
.PHONY: changes
|
||||
changes:
|
||||
$(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes
|
||||
@echo
|
||||
@echo "The overview file is in $(BUILDDIR)/changes."
|
||||
|
||||
.PHONY: linkcheck
|
||||
linkcheck:
|
||||
$(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck
|
||||
@echo
|
||||
@echo "Link check complete; look for any errors in the above output " \
|
||||
"or in $(BUILDDIR)/linkcheck/output.txt."
|
||||
|
||||
.PHONY: doctest
|
||||
doctest:
|
||||
$(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest
|
||||
@echo "Testing of doctests in the sources finished, look at the " \
|
||||
"results in $(BUILDDIR)/doctest/output.txt."
|
||||
|
||||
.PHONY: coverage
|
||||
coverage:
|
||||
$(SPHINXBUILD) -b coverage $(ALLSPHINXOPTS) $(BUILDDIR)/coverage
|
||||
@echo "Testing of coverage in the sources finished, look at the " \
|
||||
"results in $(BUILDDIR)/coverage/python.txt."
|
||||
|
||||
.PHONY: xml
|
||||
xml:
|
||||
$(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(BUILDDIR)/xml
|
||||
@echo
|
||||
@echo "Build finished. The XML files are in $(BUILDDIR)/xml."
|
||||
|
||||
.PHONY: pseudoxml
|
||||
pseudoxml:
|
||||
$(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml
|
||||
@echo
|
||||
@echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml."
|
||||
@@ -0,0 +1,199 @@
|
||||
|
||||
/* Width of template */
|
||||
|
||||
.wy-nav-content {
|
||||
max-width: 1200px !important;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/* Standard Text Formatting */
|
||||
|
||||
h1
|
||||
{
|
||||
color: #76b900;
|
||||
text-align: center;
|
||||
background-color: #333333;
|
||||
}
|
||||
|
||||
h2
|
||||
{
|
||||
color: #ffffff;
|
||||
background-color: #76b900;
|
||||
Padding: 5px;
|
||||
}
|
||||
|
||||
h3
|
||||
{
|
||||
padding-top: 0px;
|
||||
border-top: solid 3px #76b900;
|
||||
border-bottom: solid 3px #76b900;
|
||||
}
|
||||
|
||||
p
|
||||
{
|
||||
margin-bottom: 24px;
|
||||
}
|
||||
/* Link Colors */
|
||||
a
|
||||
{
|
||||
color: #76b900;
|
||||
}
|
||||
|
||||
a:visited
|
||||
{
|
||||
color: #218219;
|
||||
}
|
||||
|
||||
|
||||
/* ----------------------------------------------TABLES--------------------------------------- */
|
||||
|
||||
/* Table head Color */
|
||||
thead td
|
||||
{
|
||||
background-color: #333333 !important;
|
||||
}
|
||||
|
||||
.row-odd p
|
||||
{
|
||||
/*padding-bottom: 0px;*/
|
||||
/*margin-bottom: 0px;*/
|
||||
}
|
||||
/* even rows*/
|
||||
|
||||
.row-even tr
|
||||
{
|
||||
background-color: #e5f1e6 !important;
|
||||
}
|
||||
|
||||
/* odd rows*/
|
||||
|
||||
|
||||
.wy-table-responsive table tr {
|
||||
background-color: #ffffff !important;
|
||||
}
|
||||
|
||||
|
||||
|
||||
.wy-table-responsive table td {
|
||||
white-space: normal;
|
||||
}
|
||||
|
||||
|
||||
/* Removes bottom margin in tables*/
|
||||
|
||||
.rst-content .line-block {
|
||||
margin-bottom: 0px;
|
||||
}
|
||||
|
||||
.wy-table-responsive {
|
||||
overflow: visible !important;
|
||||
}
|
||||
|
||||
/* reduces the size of text in multiline table columns. */
|
||||
|
||||
.rst-content table.docutils td
|
||||
{
|
||||
font-size: 80%;
|
||||
}
|
||||
|
||||
.rst-content dl:not(.docutils) dt
|
||||
{
|
||||
|
||||
background-color: inherit;
|
||||
color: #000000;
|
||||
border-top: solid 0px #000000;
|
||||
|
||||
}
|
||||
|
||||
.rst-content dl:not(.docutils) dt:before {
|
||||
color: #333333;
|
||||
}
|
||||
|
||||
.rst-content .line-block {
|
||||
margin-bottom: 0px;
|
||||
}
|
||||
|
||||
.wy-side-nav-search, .wy-nav-top
|
||||
{
|
||||
background-color: #000000;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
.wy-side-nav-search img
|
||||
{
|
||||
padding: 0px;
|
||||
padding: 0px 0px;
|
||||
margin-bottom: 0;
|
||||
}
|
||||
|
||||
.wy-side-nav-search input[type=text]
|
||||
{
|
||||
border-radius: 0px;
|
||||
}
|
||||
|
||||
|
||||
.wy-menu-vertical p.caption
|
||||
{
|
||||
color: #76b900;
|
||||
}
|
||||
|
||||
|
||||
.wy-side-nav-search>a img.logo, .wy-side-nav-search .wy-dropdown>a img.logo
|
||||
{
|
||||
margin: 0px 0px 0px 0px;
|
||||
}
|
||||
|
||||
.wy-nav-content
|
||||
{
|
||||
margin: 0;
|
||||
min-height: 100%;
|
||||
height: 100%;
|
||||
background: #ffffff;
|
||||
}
|
||||
|
||||
/* List (numbered, bulleted) padding Fix */
|
||||
|
||||
|
||||
.wy-plain-list-decimal li
|
||||
{
|
||||
margin-top: -6px;
|
||||
margin-bottom: -6px;
|
||||
}
|
||||
|
||||
.rst-content .section ol.loweralpha
|
||||
{
|
||||
margin-top: -6px;
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
.wy-plain-list-disc, .rst-content .toctree-wrapper ul, article ul
|
||||
{
|
||||
margin-top: 0px !important;
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
/* Alert Boxes */
|
||||
/* Background color of Alert Box Title */
|
||||
|
||||
.rst-content .section ul
|
||||
{
|
||||
margin-top: -12px;
|
||||
margin-bottom: 16px;
|
||||
}
|
||||
|
||||
.wy-alert.wy-alert-info .wy-alert-title, .rst-content .note .wy-alert-title, .rst-content .wy-alert-info.attention .wy-alert-title, .rst-content .wy-alert-info.caution .wy-alert-title, .rst-content .wy-alert-info.danger .wy-alert-title, .rst-content .wy-alert-info.error .wy-alert-title, .rst-content .wy-alert-info.hint .wy-alert-title, .rst-content .wy-alert-info.important .wy-alert-title, .rst-content .wy-alert-info.tip .wy-alert-title, .rst-content .wy-alert-info.warning .wy-alert-title, .rst-content .seealso .wy-alert-title, .rst-content .wy-alert-info.admonition-todo .wy-alert-title, .rst-content .wy-alert-info.admonition .wy-alert-title, .wy-alert.wy-alert-info .rst-content .admonition-title, .rst-content .wy-alert.wy-alert-info .admonition-title, .rst-content .note .admonition-title, .rst-content .wy-alert-info.attention .admonition-title, .rst-content .wy-alert-info.caution .admonition-title, .rst-content .wy-alert-info.danger .admonition-title, .rst-content .wy-alert-info.error .admonition-title, .rst-content .wy-alert-info.hint .admonition-title, .rst-content .wy-alert-info.important .admonition-title, .rst-content .wy-alert-info.tip .admonition-title, .rst-content .wy-alert-info.warning .admonition-title, .rst-content .seealso .admonition-title, .rst-content .wy-alert-info.admonition-todo .admonition-title, .rst-content .wy-alert-info.admonition .admonition-title
|
||||
{
|
||||
background: #76b900;
|
||||
}
|
||||
|
||||
/* Background and Font Color of Alert Box Main Body*/
|
||||
.wy-alert.wy-alert-info, .rst-content .note, .rst-content .wy-alert-info.attention, .rst-content .wy-alert-info.caution, .rst-content .wy-alert-info.danger, .rst-content .wy-alert-info.error, .rst-content .wy-alert-info.hint, .rst-content .wy-alert-info.important, .rst-content .wy-alert-info.tip, .rst-content .wy-alert-info.warning, .rst-content .seealso, .rst-content .wy-alert-info.admonition-todo, .rst-content .wy-alert-info.admonition {
|
||||
background: #333333;
|
||||
color: #999999;
|
||||
}
|
||||
|
||||
.section
|
||||
{
|
||||
margin-top: 50px;
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
document.addEventListener("DOMContentLoaded", function () {
|
||||
var params = window.location.search.substring(1).split("&").reduce(function (params, param) {
|
||||
if (!param) {
|
||||
return params;
|
||||
}
|
||||
|
||||
var values = param.split("=");
|
||||
var name = values[0];
|
||||
var value = values[1];
|
||||
params[name] = value;
|
||||
return params;
|
||||
}, {});
|
||||
|
||||
var form = document.getElementById("feedback-form");
|
||||
for (var name in params) {
|
||||
var input = form.querySelector("[name=" + name + "]");
|
||||
input.value = params[name];
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,203 @@
|
||||
NeMo ASR collection API
|
||||
=======================
|
||||
|
||||
|
||||
Model Classes
|
||||
-------------
|
||||
|
||||
.. autoclass:: nemo.collections.asr.models.EncDecCTCModel
|
||||
:show-inheritance:
|
||||
:members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
|
||||
|
||||
|
||||
.. autoclass:: nemo.collections.asr.models.EncDecCTCModelBPE
|
||||
:show-inheritance:
|
||||
:members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
|
||||
|
||||
|
||||
.. autoclass:: nemo.collections.asr.models.EncDecRNNTModel
|
||||
:show-inheritance:
|
||||
:members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
|
||||
|
||||
|
||||
.. autoclass:: nemo.collections.asr.models.EncDecRNNTBPEModel
|
||||
:show-inheritance:
|
||||
:members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
|
||||
|
||||
|
||||
.. autoclass:: nemo.collections.asr.models.EncDecClassificationModel
|
||||
:show-inheritance:
|
||||
:members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
|
||||
|
||||
|
||||
.. autoclass:: nemo.collections.asr.models.EncDecSpeakerLabelModel
|
||||
:show-inheritance:
|
||||
:members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
|
||||
|
||||
|
||||
Modules
|
||||
-------
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.ConvASREncoder
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.ConvASRDecoder
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.ConvASRDecoderClassification
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.SpeakerDecoder
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.ConformerEncoder
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
|
||||
Parts
|
||||
-----
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.submodules.jasper.JasperBlock
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
|
||||
Mixins
|
||||
------
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.mixins.mixins.ASRBPEMixin
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.mixins.mixins.ASRModuleMixin
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
Datasets
|
||||
--------
|
||||
|
||||
Character Encoding Datasets
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. autoclass:: nemo.collections.asr.data.audio_to_text.AudioToCharDataset
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
Subword Encoding Datasets
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. autoclass:: nemo.collections.asr.data.audio_to_text.AudioToBPEDataset
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
Audio Preprocessors
|
||||
-------------------
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.AudioToMFCCPreprocessor
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
Audio Augmentors
|
||||
----------------
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.SpectrogramAugmentation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.SpeedPerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.TimeStretchPerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.GainPerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.ImpulsePerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.ShiftPerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.NoisePerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.WhiteNoisePerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.RirAndNoisePerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.TranscodePerturbation
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
Miscellaneous Classes
|
||||
---------------------
|
||||
|
||||
RNNT Decoding
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
.. autoclass:: nemo.collections.asr.metrics.rnnt_wer.RNNTDecoding
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.metrics.rnnt_wer_bpe.RNNTBPEDecoding
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.submodules.rnnt_greedy_decoding.GreedyRNNTInfer
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.submodules.rnnt_greedy_decoding.GreedyBatchedRNNTInfer
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.submodules.rnnt_beam_decoding.BeamRNNTInfer
|
||||
:show-inheritance:
|
||||
:members:
|
||||
|
||||
Hypotheses
|
||||
~~~~~~~~~~
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.utils.rnnt_utils.Hypothesis
|
||||
:show-inheritance:
|
||||
:no-members:
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.utils.rnnt_utils.NBestHypotheses
|
||||
:show-inheritance:
|
||||
:no-members:
|
||||
@@ -0,0 +1,980 @@
|
||||
@article{matchboxnet,
|
||||
title={{MatchboxNet}: 1D Time-Channel Separable Convolutional Neural Network Architecture for Speech Commands Recognition},
|
||||
author={Majumdar, Somshubra and Ginsburg, Boris},
|
||||
journal={Proc. Interspeech 2020},
|
||||
year={2020}
|
||||
}
|
||||
|
||||
@article{marblenet,
|
||||
title={MarbleNet: Deep 1D Time-Channel Separable Convolutional Neural Network for Voice Activity Detection},
|
||||
author={Jia, Fei and Majumdar, Somshubra and Ginsburg, Boris},
|
||||
journal={arXiv preprint arXiv:2010.13886},
|
||||
year={2020}
|
||||
}
|
||||
|
||||
@inproceedings{panayotov2015librispeech,
|
||||
title={Librispeech: an ASR corpus based on public domain audio books},
|
||||
author={Panayotov, Vassil and Chen, Guoguo and Povey, Daniel and Khudanpur, Sanjeev},
|
||||
booktitle={Acoustics, Speech and Signal Processing (ICASSP), 2015 IEEE International Conference on},
|
||||
pages={5206--5210},
|
||||
year={2015},
|
||||
organization={IEEE}
|
||||
}
|
||||
|
||||
@article{luong17,
|
||||
author = {Minh{-}Thang Luong and Eugene Brevdo and Rui Zhao},
|
||||
title = {Neural Machine Translation (seq2seq) Tutorial},
|
||||
journal = {https://github.com/tensorflow/nmt},
|
||||
year = {2017},
|
||||
}
|
||||
|
||||
@INPROCEEDINGS{LaurentSeqWiseBN,
|
||||
author={C. {Laurent} and G. {Pereyra} and P. {Brakel} and Y. {Zhang} and Y. {Bengio}},
|
||||
booktitle={2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
|
||||
title={Batch normalized recurrent neural networks},
|
||||
year={2016},
|
||||
volume={},
|
||||
number={},
|
||||
pages={2657-2661},
|
||||
keywords={feedforward neural nets;learning (artificial intelligence);recurrent neural nets;speech recognition;batch normalized recurrent neural networks;RNN;sequential data;long-term dependency learning;convergence rate improvement;intermediate representation normalization;feedforward neural networks;speech recognition task;language modeling;training criterion;Training;Recurrent neural networks;Convergence;Speech recognition;Computer architecture;Speech;batch normalization;RNN;LSTM;optimization},
|
||||
doi={10.1109/ICASSP.2016.7472159},
|
||||
ISSN={2379-190X},
|
||||
month={March},}
|
||||
|
||||
@article{graves2005,
|
||||
author = {Alex Graves and Jurgen Schmidhuber},
|
||||
title = {Framewise phoneme classification with bidirectional LSTM and other neural network architectures},
|
||||
journal = {Neural Networks, vol. 18},
|
||||
pages={602–-610},
|
||||
year = {2005},
|
||||
}
|
||||
|
||||
@inproceedings{graves2006,
|
||||
title={Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks},
|
||||
author={Graves, Alex and Fern{\'a}ndez, Santiago and Gomez, Faustino and Schmidhuber, J{\"u}rgen},
|
||||
booktitle={Proceedings of the 23rd international conference on Machine learning},
|
||||
pages={369--376},
|
||||
year={2006},
|
||||
organization={ACM}
|
||||
}
|
||||
|
||||
@article{li2019jasper,
|
||||
title={Jasper: An End-to-End Convolutional Neural Acoustic Model},
|
||||
author={Li, Jason and Lavrukhin, Vitaly and Ginsburg, Boris and Leary, Ryan and Kuchaiev, Oleksii and Cohen, Jonathan M and Nguyen, Huyen and Gadde, Ravi Teja},
|
||||
journal={arXiv preprint arXiv:1904.03288},
|
||||
year={2019}
|
||||
}
|
||||
|
||||
@misc{ardila2019common,
|
||||
title={Common Voice: A Massively-Multilingual Speech Corpus},
|
||||
author={Rosana Ardila and Megan Branson and Kelly Davis and Michael Henretty and Michael Kohler and Josh Meyer and Reuben Morais and Lindsay Saunders and Francis M. Tyers and Gregor Weber},
|
||||
year={2019},
|
||||
eprint={1912.06670},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
|
||||
@article{graves2012,
|
||||
title={Sequence Transduction with Recurrent Neural Networks},
|
||||
author={Graves, Alex},
|
||||
journal={arXiv preprint arXiv:1211.3711},
|
||||
year={2012}
|
||||
}
|
||||
|
||||
|
||||
@article{graves2013,
|
||||
title={Generating sequences with recurrent neural networks},
|
||||
author={Graves, Alex},
|
||||
journal={arXiv preprint arXiv:1308.0850},
|
||||
year={2013}
|
||||
}
|
||||
|
||||
@article{sergeev2018horovod,
|
||||
title={Horovod: fast and easy distributed deep learning in TensorFlow},
|
||||
author={Sergeev, Alexander and Del Balso, Mike},
|
||||
journal={arXiv preprint arXiv:1802.05799},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@misc{NVVolta,
|
||||
title = {NVIDIA TESLA V100 GPU ARCHITECTURE},
|
||||
howpublished = {\url{http://images.nvidia.com/content/volta-architecture/pdf/volta-architecture-whitepaper.pdf}},
|
||||
note = {Accessed: 2018-10-09}
|
||||
}
|
||||
|
||||
@article{NVTuring,
|
||||
title = {NVIDIA TURING GPU ARCHITECTURE},
|
||||
howpublished = {\url{https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/technologies/turing-architecture/NVIDIA-Turing-Architecture-Whitepaper.pdf}},
|
||||
author = {NVIDIA},
|
||||
year = {2018},
|
||||
note = {Accessed: 2018-10-09}
|
||||
}
|
||||
|
||||
@misc{Rygaard2015,
|
||||
title = {Using Synthesized Speech to Improve Speech Recognition for Low-Resource Languages},
|
||||
author = {Luise Valentin Rygaard},
|
||||
howpublished = {\url{https://parasol.tamu.edu/dreu2015/Rygaard/report.pdf}},
|
||||
year = {2015},
|
||||
}
|
||||
|
||||
@misc{OpenSeq2Seq,
|
||||
title = {OpenSeq2Seq: extensible toolkit for distributed and mixed precision training of sequence-to-sequence models},
|
||||
author = {Kuchaiev, Oleksii and Ginsburg, Boris and Gitman, Igor and Lavrukhin,Vitaly and Case, Carl and Micikevicius, Paulius},
|
||||
howpublished = {\url{https://arxiv.org/abs/1805.10387}},
|
||||
year = {2018},
|
||||
}
|
||||
|
||||
@misc{MPGuide,
|
||||
title = {Training with Mixed Precision},
|
||||
howpublished = {\url{http://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/}},
|
||||
note = {Accessed: 2018-04-06},
|
||||
}
|
||||
|
||||
@misc{Mozilla,
|
||||
title = {Mozilla: A Journey to less than 10\% Word Error Rate},
|
||||
howpublished = {\url{https://hacks.mozilla.org/2017/11/a-journey-to-10-word-error-rate/}},
|
||||
note = {Accessed: 2018-04-06},
|
||||
}
|
||||
|
||||
@article{Waibel1989,
|
||||
title={A time-delay neural network architecture for isolated word recognition},
|
||||
author={Waibel, Alexander, and Hanazawa, Toshiyki and Hinton,Geoffrey and Shirano, Kiyohiro and Lang, Kevin },
|
||||
journal={IEEE Trans. on Acoustics, Speech and Signal Processing},
|
||||
year={1989}
|
||||
}
|
||||
|
||||
@article{Lang1990,
|
||||
title={A time-delay neural network architecture for isolated word recognition},
|
||||
author={Lang, Kevin and Waibel, Alexander, and Hinton,Geoffrey },
|
||||
journal={Neural Networks},
|
||||
year={1990}
|
||||
}
|
||||
|
||||
@book{Bengio1996,
|
||||
Author = {Bengio, Y.},
|
||||
Publisher = {International Thomson Computer Press},
|
||||
Title = {Neural Networks for Speech and Sequence Recognition},
|
||||
Year = {1996}
|
||||
}
|
||||
|
||||
@article{Bengio1992,
|
||||
title={Global optimization of a neural network-hidden Markov model hybrid},
|
||||
author={Bengio, Y., and De Mori, R., and Flammia, G., and Kompe, R. },
|
||||
journal={IEEE Transactions on Neural Networks, 3(2), 252–259},
|
||||
year={1992}
|
||||
}
|
||||
|
||||
@article{Bourlard1994,
|
||||
title={Connectionist speech recognition: a hybrid approach},
|
||||
author={Bourlard, H. A. and Morgan, N.},
|
||||
journal={volume 247 Springer },
|
||||
year={1994}
|
||||
}
|
||||
|
||||
@article{srivastava14a,
|
||||
author = {Nitish Srivastava, and Geoffrey Hinton, and Alex Krizhevsky, and Ilya Sutskever, and Ruslan Salakhutdinov},
|
||||
title = {Dropout: A Simple Way to Prevent Neural Networks from Overfitting},
|
||||
journal = {Journal of Machine Learning Research},
|
||||
year = {2014},
|
||||
volume = {15},
|
||||
pages = {1929-1958},
|
||||
url = {http://jmlr.org/papers/v15/srivastava14a.html}
|
||||
}
|
||||
|
||||
|
||||
@article{Hinton2012,
|
||||
title={Deep Neural Networks for Acoustic Modeling in Speech Recognition},
|
||||
author={ Hinton,Geoffrey and Deng, Li and Yu, Dong and Dahl,George and Mohamed,Abdel-rahman and Jaitly, Navdeep and Senior, Andrew and Vanhoucke, Vincent and Nguyen, Patrick and Kingsbury, Brian and Sainath, Tara},
|
||||
journal={IEEE Signal Processing Magazine},
|
||||
year={2012}
|
||||
}
|
||||
|
||||
@article{Graves2014,
|
||||
title={Towards End-to-End Speech Recognition with Recurrent Neural Networks},
|
||||
author={Graves, Alex and Jaitly, Navdeep},
|
||||
journal={International Conference on Machine Learning},
|
||||
year={2014}
|
||||
}
|
||||
|
||||
@article{Chorowski2014,
|
||||
title={End-to-end Continuous Speech Recognition using Attention-based Recurrent NN: First Results},
|
||||
author={ Chorowski, Jan, and Bahdanau, Dzmitry , and Cho, Kyunghyun , and Bengio, Yoshua },
|
||||
journal={Neural Information Processing Systems: Workshop Deep Learning and Representation Learning Workshop },
|
||||
year={2014}
|
||||
}
|
||||
|
||||
@article{Sak2014,
|
||||
title={Long short-term memory recurrent neural network architectures for large scale acoustic modeling},
|
||||
author={Sak, Hasim and Senior, Andrew and Beaufays, Francoise },
|
||||
journal={Interspeech 2014},
|
||||
year={2014}
|
||||
}
|
||||
|
||||
@article{Ko2015,
|
||||
title={Audio Augmentation for Speech Recognition},
|
||||
author={Tom, Ko and Vijayaditya, Peddinti and Daniel, Povey
|
||||
and Sanjeev, Khudanpur },
|
||||
journal={Interspeech 2015},
|
||||
year={2015}
|
||||
}
|
||||
|
||||
@article{Tjandra2017,
|
||||
title={Listening while Speaking: Speech Chain by Deep Learning},
|
||||
author={Andros, Tjandra and Sakriani, Sakti and Satoshi, Nakamura },
|
||||
journal={ASRU 2017},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@article{Tjandra2018,
|
||||
title={Machine Speech Chain with One-shot Speaker Adaptation},
|
||||
author={Andros, Tjandra and Sakriani, Sakti and Satoshi, Nakamura },
|
||||
journal={Interspeech 2018},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@article{bahdanau2014neural,
|
||||
title={Neural machine translation by jointly learning to align and translate},
|
||||
author={Bahdanau, Dzmitry and Cho, Kyunghyun and Bengio, Yoshua},
|
||||
journal={arXiv preprint arXiv:1409.0473},
|
||||
year={2014}
|
||||
}
|
||||
|
||||
@article{cho2014learning,
|
||||
title={Learning phrase representations using RNN encoder-decoder for statistical machine translation},
|
||||
author={Cho, Kyunghyun and Van Merri{\"e}nboer, Bart and Gulcehre, Caglar and Bahdanau, Dzmitry and Bougares, Fethi and Schwenk, Holger and Bengio, Yoshua},
|
||||
journal={arXiv preprint arXiv:1406.1078},
|
||||
year={2014}
|
||||
}
|
||||
|
||||
@article{rush2015neural,
|
||||
title={A neural attention model for abstractive sentence summarization},
|
||||
author={Rush, Alexander M and Chopra, Sumit and Weston, Jason},
|
||||
journal={arXiv preprint arXiv:1509.00685},
|
||||
year={2015}
|
||||
}
|
||||
|
||||
@article{micikevicius2017mixed,
|
||||
title={Mixed precision training},
|
||||
author={Micikevicius, Paulius and Narang, Sharan and Alben, Jonah and Diamos, Gregory and Elsen, Erich and Garcia, David and Ginsburg, Boris and Houston, Michael and Kuchaev, Oleksii and Venkatesh, Ganesh and others},
|
||||
journal={arXiv preprint arXiv:1710.03740},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@ARTICLE{Britz:2017,
|
||||
author = {{Britz}, Denny and {Goldie}, Anna and {Luong}, Thang and {Le}, Quoc},
|
||||
title = "{Massive Exploration of Neural Machine Translation Architectures}",
|
||||
journal = {ArXiv e-prints arXiv:1703.03906},
|
||||
archivePrefix = "arXiv",
|
||||
eprinttype = {arxiv},
|
||||
eprint = {1703.03906},
|
||||
primaryClass = "cs.CL",
|
||||
keywords = {Computer Science - Computation and Language},
|
||||
year = 2017,
|
||||
month = mar,
|
||||
}
|
||||
|
||||
@inproceedings{vaswani2017attention,
|
||||
title={Attention is all you need},
|
||||
author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, {\L}ukasz and Polosukhin, Illia},
|
||||
booktitle={Advances in Neural Information Processing Systems},
|
||||
pages={6000--6010},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@inproceedings{abadi2016tensorflow,
|
||||
title={TensorFlow: A System for Large-Scale Machine Learning.},
|
||||
author={Abadi, Mart{\'\i}n and Barham, Paul and Chen, Jianmin and Chen, Zhifeng and Davis, Andy and Dean, Jeffrey and Devin, Matthieu and Ghemawat, Sanjay and Irving, Geoffrey and Isard, Michael and others},
|
||||
booktitle={OSDI},
|
||||
volume={16},
|
||||
pages={265--283},
|
||||
year={2016}
|
||||
}
|
||||
|
||||
@article{tensor2tensor,
|
||||
author = {Ashish Vaswani and Samy Bengio and Eugene Brevdo and Francois Chollet and Aidan N. Gomez and Stephan Gouws and Llion Jones and \L{}ukasz Kaiser and Nal Kalchbrenner and Niki Parmar and Ryan Sepassi and
|
||||
Noam Shazeer and Jakob Uszkoreit},
|
||||
title = {Tensor2Tensor for Neural Machine Translation},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1803.07416},
|
||||
year = {2018},
|
||||
url = {http://arxiv.org/abs/1803.07416},
|
||||
}
|
||||
|
||||
@article{gehring2017convs2s,
|
||||
author = {Gehring, Jonas, and Auli, Michael and Grangier, David and Yarats, Denis and Dauphin, Yann N},
|
||||
title = "{Convolutional Sequence to Sequence Learning}",
|
||||
journal = {ArXiv e-prints arXiv:1705.03122},
|
||||
archivePrefix = "arXiv",
|
||||
eprinttype = {arxiv},
|
||||
eprint = {1705.03122},
|
||||
primaryClass = "cs.CL",
|
||||
keywords = {Computer Science - Computation and Language},
|
||||
year = 2017,
|
||||
month = May,
|
||||
}
|
||||
|
||||
@inproceedings{chan2015,
|
||||
title={Listen, attend and spell},
|
||||
author={Chan, William and Jaitly, Navdeep and Le, Quoc V and Vinyals, Oriol},
|
||||
booktitle={Acoustics, Speech and Signal Processing (ICASSP), 2016 IEEE International Conference on},
|
||||
pages={5206--5210},
|
||||
year={2016},
|
||||
organization={IEEE}
|
||||
}
|
||||
|
||||
@inproceedings{xu2015show,
|
||||
title={Show, attend and tell: Neural image caption generation with visual attention},
|
||||
author={Xu, Kelvin and Ba, Jimmy and Kiros, Ryan and Cho, Kyunghyun and Courville, Aaron and Salakhudinov, Ruslan and Zemel, Rich and Bengio, Yoshua},
|
||||
booktitle={International Conference on Machine Learning},
|
||||
pages={2048--2057},
|
||||
year={2015}
|
||||
}
|
||||
|
||||
@incollection{Sutskever2014,
|
||||
title = {Sequence to Sequence Learning with Neural Networks},
|
||||
author = {Sutskever, Ilya and Vinyals, Oriol and Le, Quoc V},
|
||||
booktitle = {Advances in Neural Information Processing Systems 27},
|
||||
editor = {Z. Ghahramani and M. Welling and C. Cortes and N. D. Lawrence and K. Q. Weinberger},
|
||||
pages = {3104--3112},
|
||||
year = {2014},
|
||||
publisher = {Curran Associates, Inc.},
|
||||
url = {http://papers.nips.cc/paper/5346-sequence-to-sequence-learning-with-neural-networks.pdf}
|
||||
}
|
||||
|
||||
@article{DeepSpeech2014,
|
||||
title = {Deep Speech: Scaling up end-to-end speech recognition},
|
||||
author = {Awni Y. Hannun and Carl Case and Jared Casper and Bryan Catanzaro and Greg Diamos and Erich Elsen and Ryan Prenger and Sanjeev Satheesh and Shubho Sengupta and Adam Coates and Andrew Y. Ng},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1412.5567},
|
||||
year = {2014},
|
||||
url = {http://arxiv.org/abs/1412.5567},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1412.5567},
|
||||
timestamp = {Mon, 13 Aug 2018 16:48:07 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/HannunCCCDEPSSCN14},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@inproceedings{DeepSpeech2,
|
||||
author = {Amodei, Dario and Ananthanarayanan, Sundaram and Anubhai, Rishita and Bai, Jingliang and Battenberg, Eric and Case, Carl and Casper, Jared and Catanzaro, Bryan and Cheng, Qiang and Chen, Guoliang and Chen, Jie and Chen, Jingdong and Chen, Zhijie and Chrzanowski, Mike and Coates, Adam and Diamos, Greg and Ding, Ke and Du, Niandong and Elsen, Erich and Engel, Jesse and Fang, Weiwei and Fan, Linxi and Fougner, Christopher and Gao, Liang and Gong, Caixia and Hannun, Awni and Han, Tony and Johannes, Lappi Vaino and Jiang, Bing and Ju, Cai and Jun, Billy and LeGresley, Patrick and Lin, Libby and Liu, Junjie and Liu, Yang and Li, Weigao and Li, Xiangang and Ma, Dongpeng and Narang, Sharan and Ng, Andrew and Ozair, Sherjil and Peng, Yiping and Prenger, Ryan and Qian, Sheng and Quan, Zongfeng and Raiman, Jonathan and Rao, Vinay and Satheesh, Sanjeev and Seetapun, David and Sengupta, Shubho and Srinet, Kavya and Sriram, Anuroop and Tang, Haiyuan and Tang, Liliang and Wang, Chong and Wang, Jidong and Wang, Kaifu and Wang, Yi and Wang, Zhijian and Wang, Zhiqian and Wu, Shuang and Wei, Likai and Xiao, Bo and Xie, Wen and Xie, Yan and Yogatama, Dani and Yuan, Bin and Zhan, Jun and Zhu, Zhenyao},
|
||||
title = {Deep Speech 2: End-to-end Speech Recognition in English and Mandarin},
|
||||
booktitle = {Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48},
|
||||
series = {ICML'16},
|
||||
year = {2016},
|
||||
location = {New York, NY, USA},
|
||||
pages = {173--182},
|
||||
numpages = {10},
|
||||
url = {http://dl.acm.org/citation.cfm?id=3045390.3045410},
|
||||
acmid = {3045410},
|
||||
publisher = {JMLR.org},
|
||||
}
|
||||
|
||||
@inproceedings{prabhavalkar2017comparison,
|
||||
title={A comparison of sequence-to-sequence models for speech recognition},
|
||||
author={Prabhavalkar, Rohit and Rao, Kanishka and Sainath, Tara N and Li, Bo and Johnson, Leif and Jaitly, Navdeep},
|
||||
booktitle={Proc. Interspeech},
|
||||
pages={939--943},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@article{chiu2017state,
|
||||
title={State-of-the-art speech recognition with sequence-to-sequence models},
|
||||
author={Chiu, Chung-Cheng and Sainath, Tara N and Wu, Yonghui and Prabhavalkar, Rohit and Nguyen, Patrick and Chen, Zhifeng and Kannan, Anjuli and Weiss, Ron J and Rao, Kanishka and Gonina, Katya and others},
|
||||
journal={arXiv preprint arXiv:1712.01769},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@misc{NVMixed,
|
||||
title = {{NVIDA's Mixed-Precision Training - TensorFlow example}},
|
||||
howpublished = {\url{https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/#example_tensorflow}},
|
||||
author={NVIDIA},
|
||||
note = {Accessed: 2018-10-09},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@article{gehring2017,
|
||||
title={Convolutional sequence to sequence learning},
|
||||
author={Gehring, Jonas and Auli, Michael and Grangier, David and Yarats, Denis and Dauphin, Yann N},
|
||||
journal={arXiv preprint arXiv:1705.03122},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@article{collobert2016,
|
||||
title={Wav2letter: an end-to-end convnet-based speech recognition system},
|
||||
author={Collobert, Ronan and Puhrsch, Christian and Synnaeve, Gabriel},
|
||||
journal={arXiv preprint arXiv:1609.03193},
|
||||
year={2016}
|
||||
}
|
||||
|
||||
@inproceedings{Zhang2016,
|
||||
author={Ying Zhang and Mohammad Pezeshki and Philémon Brakel and Saizheng Zhang and César Laurent and Yoshua Bengio and Aaron Courville},
|
||||
title={Towards End-to-End Speech Recognition with Deep Convolutional Neural Networks},
|
||||
year=2016,
|
||||
booktitle={Interspeech 2016},
|
||||
doi={10.21437/Interspeech.2016-1446},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2016-1446},
|
||||
pages={410--414}
|
||||
}
|
||||
|
||||
@inproceedings{Zhang2017,
|
||||
title={Very deep convolutional networks for end-to-end speech recognition},
|
||||
author={Zhang, Yu, and Chan, William, and Jaitly, Navdeep},
|
||||
booktitle={Acoustics, Speech and Signal Processing (ICASSP), 2017 IEEE International Conference on},
|
||||
year={2017},
|
||||
organization={IEEE}
|
||||
}
|
||||
|
||||
|
||||
@article{Wang2017,
|
||||
title={Tacotron: Towards End-to-End Speech Synthesis},
|
||||
author={ Wang, Yuxuan, and Skerry-Ryan, RJ, and Stanton, Daisy and Wu, Yonghui and Weiss, Ron, and Jaitly, Navdeep and Yang, Zongheng and Xiao, Ying and Chen,Zhifeng and Bengio, Samy and Le, Quoc and Agiomyrgiannakis, Yannis and Clark,Rob and Saurous, Rif A.},
|
||||
journal={arXiv preprint arXiv:1703.10135},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@inproceedings{shen2018natural,
|
||||
title={Natural TTS synthesis by conditioning wavenet on mel spectrogram predictions},
|
||||
author={Shen, Jonathan and Pang, Ruoming and Weiss, Ron J and Schuster, Mike and Jaitly, Navdeep and Yang, Zongheng and Chen, Zhifeng and Zhang, Yu and Wang, Yuxuan and Skerrv-Ryan, Rj and others},
|
||||
booktitle={2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
|
||||
pages={4779--4783},
|
||||
year={2018},
|
||||
organization={IEEE}
|
||||
}
|
||||
|
||||
@article{griffin1984signal,
|
||||
title={Signal estimation from modified short-time Fourier transform},
|
||||
author={Griffin, Daniel and Lim, Jae},
|
||||
journal={IEEE Transactions on Acoustics, Speech, and Signal Processing},
|
||||
volume={32},
|
||||
number={2},
|
||||
pages={236--243},
|
||||
year={1984},
|
||||
publisher={IEEE}
|
||||
}
|
||||
|
||||
@misc{ito2017lj,
|
||||
title={The LJ speech dataset},
|
||||
author={Ito, Keith and others},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@misc{mailabs,
|
||||
title = {{The M-AILABS Speech Dataset}},
|
||||
howpublished = {\url{http://www.m-ailabs.bayern/en/the-mailabs-speech-dataset/}},
|
||||
author={M-AILABS},
|
||||
note = {Accessed: 2018-10-09},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@article{merity2016pointer,
|
||||
title={Pointer sentinel mixture models},
|
||||
author={Merity, Stephen and Xiong, Caiming and Bradbury, James and Socher, Richard},
|
||||
journal={arXiv preprint arXiv:1609.07843},
|
||||
year={2016}
|
||||
}
|
||||
|
||||
@inproceedings{socher2013recursive,
|
||||
title={Recursive deep models for semantic compositionality over a sentiment treebank},
|
||||
author={Socher, Richard and Perelygin, Alex and Wu, Jean and Chuang, Jason and Manning, Christopher D and Ng, Andrew and Potts, Christopher},
|
||||
booktitle={Proceedings of the 2013 conference on empirical methods in natural language processing},
|
||||
pages={1631--1642},
|
||||
year={2013}
|
||||
}
|
||||
|
||||
@InProceedings{maas-EtAl:2011:ACL-HLT2011,
|
||||
author = {Maas, Andrew L. and Daly, Raymond E. and Pham, Peter T. and Huang, Dan and Ng, Andrew Y. and Potts, Christopher},
|
||||
title = {Learning Word Vectors for Sentiment Analysis},
|
||||
booktitle = {Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies},
|
||||
month = {June},
|
||||
year = {2011},
|
||||
address = {Portland, Oregon, USA},
|
||||
publisher = {Association for Computational Linguistics},
|
||||
pages = {142--150},
|
||||
url = {http://www.aclweb.org/anthology/P11-1015}
|
||||
}
|
||||
|
||||
@inproceedings{Povey2018SemiOrthogonalLM,
|
||||
title={Semi-Orthogonal Low-Rank Matrix Factorization for Deep Neural Networks},
|
||||
author={Daniel Povey and Gaofeng Cheng and Yiming Wang and Ke Li and Hainan Xu and Mahsa Yarmohammadi and Sanjeev Khudanpur},
|
||||
booktitle={Interspeech},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@article{CAPIO2017,
|
||||
author = {Kyu J. Han and Akshay Chandrashekaran and Jungsuk Kim and Ian R. Lane},
|
||||
title = {The {CAPIO} 2017 Conversational Speech Recognition System},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1801.00059},
|
||||
year = {2018},
|
||||
url = {http://arxiv.org/abs/1801.00059},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1801.00059},
|
||||
timestamp = {Mon, 13 Aug 2018 16:49:10 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/abs-1801-00059},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{WaveNet,
|
||||
author = {A{\"{a}}ron van den Oord and Sander Dieleman and Heiga Zen and Karen Simonyan and Oriol Vinyals and Alex Graves and Nal Kalchbrenner and Andrew W. Senior and Koray Kavukcuoglu},
|
||||
title = {WaveNet: {A} Generative Model for Raw Audio},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1609.03499},
|
||||
year = {2016},
|
||||
url = {http://arxiv.org/abs/1609.03499},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1609.03499},
|
||||
timestamp = {Mon, 13 Aug 2018 16:49:15 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/OordDZSVGKSK16},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{FacebookGERENGBackTranslation,
|
||||
author = {Rico Sennrich and Barry Haddow and Alexandra Birch},
|
||||
title = {Improving Neural Machine Translation Models with Monolingual Data},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1511.06709},
|
||||
year = {2015},
|
||||
url = {http://arxiv.org/abs/1511.06709},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1511.06709},
|
||||
timestamp = {Mon, 13 Aug 2018 16:47:05 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/SennrichHB15a},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{GlobalStyleTokens,
|
||||
author = {Yuxuan Wang and Daisy Stanton and Yu Zhang and R. J. Skerry{-}Ryan and Eric Battenberg and Joel Shor and Ying Xiao and Fei Ren and Ye Jia and Rif A. Saurous},
|
||||
title = {Style Tokens: Unsupervised Style Modeling, Control and Transfer in End-to-End Speech Synthesis},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1803.09017},
|
||||
year = {2018},
|
||||
url = {http://arxiv.org/abs/1803.09017},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1803.09017},
|
||||
timestamp = {Mon, 13 Aug 2018 16:46:53 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/abs-1803-09017},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{IoffeS15BatchNorm,
|
||||
author = {Sergey Ioffe and Christian Szegedy},
|
||||
title = {Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1502.03167},
|
||||
year = {2015},
|
||||
url = {http://arxiv.org/abs/1502.03167},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1502.03167},
|
||||
timestamp = {Mon, 13 Aug 2018 16:47:06 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/IoffeS15},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{kingma,
|
||||
author = {Diederik P. Kingma and
|
||||
Jimmy Ba},
|
||||
title = {Adam: {A} Method for Stochastic Optimization},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1412.6980},
|
||||
year = {2014},
|
||||
url = {http://arxiv.org/abs/1412.6980},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1412.6980},
|
||||
timestamp = {Mon, 13 Aug 2018 01:00:00 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/KingmaB14},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@incollection{Salimans2016WeightNorm,
|
||||
title = {Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks},
|
||||
author = {Salimans, Tim and Kingma, Durk P},
|
||||
booktitle = {Advances in Neural Information Processing Systems 29},
|
||||
editor = {D. D. Lee and M. Sugiyama and U. V. Luxburg and I. Guyon and R. Garnett},
|
||||
pages = {901--909},
|
||||
year = {2016},
|
||||
publisher = {Curran Associates, Inc.},
|
||||
url = {http://papers.nips.cc/paper/6114-weight-normalization-a-simple-reparameterization-to-accelerate-training-of-deep-neural-networks.pdf}
|
||||
}
|
||||
|
||||
@article{wu2016google,
|
||||
title={Google's neural machine translation system: Bridging the gap between human and machine translation},
|
||||
author={Wu, Yonghui and Schuster, Mike and Chen, Zhifeng and Le, Quoc V and Norouzi, Mohammad and Macherey, Zolfgang and Krikun, Maxim and Cao, Yuan and Gao, Qin and Macherey, Klaus and others},
|
||||
journal={arXiv preprint arXiv:1609.08144},
|
||||
year={2016}
|
||||
}
|
||||
|
||||
@inproceedings{opennmt,
|
||||
author = {Guillaume Klein and Yoon Kim and Yuntian Deng and Jean Senellart and Alexander M. Rush},
|
||||
title = {OpenNMT: Open-Source Toolkit for Neural Machine Translation},
|
||||
booktitle = {Proc. ACL},
|
||||
year = {2017},
|
||||
url = {https://doi.org/10.18653/v1/P17-4012},
|
||||
doi = {10.18653/v1/P17-4012}
|
||||
}
|
||||
|
||||
@article{paszke2017automatic,
|
||||
title={Automatic differentiation in PyTorch},
|
||||
author={Paszke, Adam and Gross, Sam and Chintala, Soumith and Chanan, Gregory and Yang, Edward and DeVito, Zachary and Lin, Zeming and Desmaison, Alban and Antiga, Luca and Lerer, Adam},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@article{yu2014introduction,
|
||||
title={An introduction to computational networks and the computational network toolkit},
|
||||
author={Yu, Dong and Eversole, Adam and Seltzer, Mike and Yao, Kaisheng and Huang, Zhiheng and Guenter, Brian and Kuchaiev, Oleksii and Zhang, Yu and Seide, Frank and Wang, Huaming and others},
|
||||
journal={Microsoft Technical Report MSR-TR-2014--112},
|
||||
year={2014}
|
||||
}
|
||||
|
||||
@article{nvidia2017v100,
|
||||
title={V100 GPU architecture. The world’s most advanced data center GPU. Version WP-08608-001\_v1. 1},
|
||||
author={NVIDIA, Tesla},
|
||||
journal={NVIDIA. Aug},
|
||||
pages={108},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@article{post2018call,
|
||||
title={A call for clarity in reporting bleu scores},
|
||||
author={Post, Matt},
|
||||
journal={arXiv preprint arXiv:1804.08771},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@article{Ba2016LayerNorm,
|
||||
author = {Jimmy Lei Ba and Jamie Ryan Kiros and Geoffrey E Hinton},
|
||||
title = {Layer normalization},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1607.06450},
|
||||
year = {2016},
|
||||
url = {http://arxiv.org/abs/1607.06450},
|
||||
archivePrefix = {arXiv},
|
||||
}
|
||||
|
||||
@inproceedings{Dauphin2017GLU,
|
||||
author = {Dauphin, Yann N. and Fan, Angela and Auli, Michael and Grangier, David},
|
||||
title = {Language Modeling with Gated Convolutional Networks},
|
||||
booktitle = {Proceedings of the 34th International Conference on Machine Learning - Volume 70},
|
||||
series = {ICML'17},
|
||||
year = {2017},
|
||||
location = {Sydney, NSW, Australia},
|
||||
pages = {933--941},
|
||||
numpages = {9},
|
||||
url = {http://dl.acm.org/citation.cfm?id=3305381.3305478},
|
||||
acmid = {3305478},
|
||||
publisher = {JMLR.org},
|
||||
}
|
||||
|
||||
@incollection{Oord2016PixelCNN,
|
||||
title = {Conditional Image Generation with PixelCNN Decoders},
|
||||
author = {van den Oord, Aaron and Kalchbrenner, Nal and Espeholt, Lasse and kavukcuoglu, koray and Vinyals, Oriol and Graves, Alex},
|
||||
booktitle = {Advances in Neural Information Processing Systems 29},
|
||||
editor = {D. D. Lee and M. Sugiyama and U. V. Luxburg and I. Guyon and R. Garnett},
|
||||
pages = {4790--4798},
|
||||
year = {2016},
|
||||
publisher = {Curran Associates, Inc.},
|
||||
url = {http://papers.nips.cc/paper/6527-conditional-image-generation-with-pixelcnn-decoders.pdf}
|
||||
}
|
||||
|
||||
@article{he2015,
|
||||
title={Deep residual learning for image recognition},
|
||||
author={K. He, and X. Zhang, and S. Ren, and J. Sun},
|
||||
journal={arXiv preprint arXiv:1512.03385},
|
||||
year={2015}
|
||||
}
|
||||
|
||||
@article{huang2016,
|
||||
title={Densely Connected Convolutional Networks},
|
||||
author={Gao Huang, and Zhuang Liu, and Laurens van der Maaten, and Kilian Q. Weinberger},
|
||||
journal={arXiv preprint arXiv:1608.06993},
|
||||
year={2016}
|
||||
}
|
||||
|
||||
@inproceedings{heafield2011kenlm,
|
||||
title={KenLM: Faster and smaller language model queries},
|
||||
author={Heafield, Kenneth},
|
||||
booktitle={Proceedings of the sixth workshop on statistical machine translation},
|
||||
pages={187--197},
|
||||
year={2011},
|
||||
organization={Association for Computational Linguistics}
|
||||
}
|
||||
|
||||
@article{dai2018transformer,
|
||||
title={Transformer-XL: Language Modeling with Longer-Term Dependency},
|
||||
author={Dai, Zihang and Yang, Zhilin and Yang, Yiming and Cohen, William W and Carbonell, Jaime and Le, Quoc V and Salakhutdinov, Ruslan},
|
||||
year={2018},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1901.02860},
|
||||
url = {http://arxiv.org/abs/1901.02860},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1901.02860},
|
||||
timestamp = {Fri, 01 Feb 2019 13:39:59 +0100},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/abs-1901-02860},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@inproceedings{Saon+2016,
|
||||
author={George Saon and Tom Sercu and Steven Rennie and Hong-Kwang J. Kuo},
|
||||
title={The IBM 2016 English Conversational Telephone Speech Recognition System},
|
||||
year=2016,
|
||||
booktitle={Interspeech 2016},
|
||||
doi={10.21437/Interspeech.2016-1460},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2016-1460},
|
||||
pages={7--11}
|
||||
}
|
||||
|
||||
@INPROCEEDINGS{Sercu-2016,
|
||||
author={T. {Sercu} and C. {Puhrsch} and B. {Kingsbury} and Y. {LeCun}},
|
||||
booktitle={2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
|
||||
title={Very deep multilingual convolutional neural networks for LVCSR},
|
||||
year={2016},
|
||||
volume={},
|
||||
number={},
|
||||
pages={4955-4959},
|
||||
keywords={natural language processing;neural nets;speech recognition;very deep multilingual convolutional neural networks;LVCSR;CNN;large vocabulary continuous speech recognition systems;word error rate;Training;Context;Hidden Markov models;Neural networks;Computer architecture;Kernel;Training data;Convolutional Networks;Multilingual;Acoustic Modeling;Speech Recognition;Neural Networks},
|
||||
doi={10.1109/ICASSP.2016.7472620},
|
||||
ISSN={2379-190X},
|
||||
month={March},}
|
||||
|
||||
|
||||
@inproceedings{Sercu+2016,
|
||||
author={Tom Sercu and Vaibhava Goel},
|
||||
title={Advances in Very Deep Convolutional Neural Networks for LVCSR},
|
||||
year=2016,
|
||||
booktitle={Interspeech 2016},
|
||||
doi={10.21437/Interspeech.2016-1033},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2016-1033},
|
||||
pages={3429--3433}
|
||||
}
|
||||
|
||||
@INPROCEEDINGS{Xiong-2018,
|
||||
author={W. {Xiong} and L. {Wu} and F. {Alleva} and J. {Droppo} and X. {Huang} and A. {Stolcke}},
|
||||
booktitle={2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
|
||||
title={The Microsoft 2017 Conversational Speech Recognition System},
|
||||
year={2018},
|
||||
volume={},
|
||||
number={},
|
||||
pages={5934-5938},
|
||||
keywords={convolution;feedforward neural nets;natural language processing;speaker recognition;speech processing;language model rescoring step;senone level;switchboard domains;character-based LSTM language models;NIST 2000 switchboard test set;frame level;word-level voting;acoustic model posteriors;dialog session aware LSTM language models;CNN-BLSTM acoustic model;Microsoft 2017 conversational speech recognition system;Acoustics;Error analysis;Training;Speech recognition;Switches;Computational modeling;Context modeling;Conversational speech recognition;CNN;LACE;BLSTM;LSTM-LM;system combination;human parity},
|
||||
doi={10.1109/ICASSP.2018.8461870},
|
||||
ISSN={2379-190X},
|
||||
month={April},}
|
||||
|
||||
@inproceedings{zeyer2018improved,
|
||||
author={Albert Zeyer and Kazuki Irie and Ralf Schlüter and Hermann Ney},
|
||||
title={Improved Training of End-to-end Attention Models for Speech Recognition},
|
||||
year=2018,
|
||||
booktitle={Proc. Interspeech 2018},
|
||||
pages={7--11},
|
||||
doi={10.21437/Interspeech.2018-1616},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2018-1616}
|
||||
}
|
||||
|
||||
@article{Wav2LetterV2,
|
||||
author = {Vitaliy Liptchinsky and
|
||||
Gabriel Synnaeve and
|
||||
Ronan Collobert},
|
||||
title = {Letter-Based Speech Recognition with Gated ConvNets},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1712.09444},
|
||||
year = {2017},
|
||||
url = {http://arxiv.org/abs/1712.09444},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1712.09444},
|
||||
timestamp = {Mon, 13 Aug 2018 16:46:33 +0200},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/abs-1712-09444},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{zeghidour2018,
|
||||
author = {Neil Zeghidour and
|
||||
Qiantong Xu and
|
||||
Vitaliy Liptchinsky and
|
||||
Nicolas Usunier and
|
||||
Gabriel Synnaeve and
|
||||
Ronan Collobert},
|
||||
title = {Fully Convolutional Speech Recognition},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1812.06864},
|
||||
year = {2018},
|
||||
url = {http://arxiv.org/abs/1812.06864},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1812.06864},
|
||||
timestamp = {Tue, 01 Jan 2019 15:01:25 +0100},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/abs-1812-06864},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@inproceedings{Hadian2018,
|
||||
author={Hossein Hadian and Hossein Sameti and Daniel Povey and Sanjeev Khudanpur},
|
||||
title={End-to-end Speech Recognition Using Lattice-free MMI},
|
||||
year=2018,
|
||||
booktitle={Proc. Interspeech 2018},
|
||||
pages={12--16},
|
||||
doi={10.21437/Interspeech.2018-1423},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2018-1423}
|
||||
}
|
||||
|
||||
@inproceedings{Tang2018,
|
||||
author={Jian Tang and Yan Song and Lirong Dai and Ian McLoughlin},
|
||||
title={Acoustic Modeling with Densely Connected Residual Network for Multichannel Speech Recognition},
|
||||
year=2018,
|
||||
booktitle={Proc. Interspeech 2018},
|
||||
pages={1783--1787},
|
||||
doi={10.21437/Interspeech.2018-1089},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2018-1089}
|
||||
}
|
||||
|
||||
@article{Kurata2017LanguageMW,
|
||||
title={Language modeling with highway LSTM},
|
||||
author={Gakuto Kurata and Bhuvana Ramabhadran and George Saon and Abhinav Sethy},
|
||||
journal={2017 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)},
|
||||
year={2017},
|
||||
pages={244-251}
|
||||
}
|
||||
|
||||
@inproceedings{Saon2017,
|
||||
author={George Saon and Gakuto Kurata and Tom Sercu and Kartik Audhkhasi and Samuel Thomas and Dimitrios Dimitriadis and Xiaodong Cui and Bhuvana Ramabhadran and Michael Picheny and Lynn-Li Lim and Bergul Roomi and Phil Hall},
|
||||
title={English Conversational Telephone Speech Recognition by Humans and Machines},
|
||||
year=2017,
|
||||
booktitle={Proc. Interspeech 2017},
|
||||
pages={132--136},
|
||||
doi={10.21437/Interspeech.2017-405},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2017-405}
|
||||
}
|
||||
|
||||
@inproceedings{Povey+2016,
|
||||
author={Daniel Povey and Vijayaditya Peddinti and Daniel Galvez and Pegah Ghahremani and Vimal Manohar and Xingyu Na and Yiming Wang and Sanjeev Khudanpur},
|
||||
title={Purely Sequence-Trained Neural Networks for ASR Based on Lattice-Free MMI},
|
||||
year=2016,
|
||||
booktitle={Interspeech 2016},
|
||||
doi={10.21437/Interspeech.2016-595},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2016-595},
|
||||
pages={2751--2755}
|
||||
}
|
||||
|
||||
@article{Yang2018,
|
||||
author = {Xuerui Yang and
|
||||
Jiwei Li and
|
||||
Xi Zhou},
|
||||
title = {A novel pyramidal-FSMN architecture with lattice-free {MMI} for speech
|
||||
recognition},
|
||||
journal = {CoRR},
|
||||
volume = {abs/1810.11352},
|
||||
year = {2018},
|
||||
url = {http://arxiv.org/abs/1810.11352},
|
||||
archivePrefix = {arXiv},
|
||||
eprint = {1810.11352},
|
||||
timestamp = {Wed, 31 Oct 2018 14:24:29 +0100},
|
||||
biburl = {https://dblp.org/rec/bib/journals/corr/abs-1810-11352},
|
||||
bibsource = {dblp computer science bibliography, https://dblp.org}
|
||||
}
|
||||
|
||||
@article{liptchinsky2017based,
|
||||
title={Letter-Based Speech Recognition with Gated ConvNets},
|
||||
author={Liptchinsky, Vitaliy and Synnaeve, Gabriel and Collobert, Ronan},
|
||||
journal={arXiv preprint arXiv:1712.09444},
|
||||
year={2017}
|
||||
}
|
||||
|
||||
@inproceedings{Weng2018,
|
||||
author={Chao Weng and Jia Cui and Guangsen Wang and Jun Wang and Chengzhu Yu and Dan Su and Dong Yu},
|
||||
title={Improving Attention Based Sequence-to-Sequence Models for End-to-End English Conversational Speech Recognition},
|
||||
year=2018,
|
||||
booktitle={Proc. Interspeech 2018},
|
||||
pages={761--765},
|
||||
doi={10.21437/Interspeech.2018-1030},
|
||||
url={http://dx.doi.org/10.21437/Interspeech.2018-1030}
|
||||
}
|
||||
|
||||
@INPROCEEDINGS{Battenberg2017,
|
||||
author={E. {Battenberg} and J. {Chen} and R. {Child} and A. {Coates} and Y. G. Y. {Li} and H. {Liu} and S. {Satheesh} and A. {Sriram} and Z. {Zhu}},
|
||||
booktitle={2017 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)},
|
||||
title={Exploring neural transducers for end-to-end speech recognition},
|
||||
year={2017},
|
||||
volume={},
|
||||
number={},
|
||||
pages={206-213},
|
||||
keywords={recurrent neural nets;speech recognition;Hub500 benchmark;CTC models;speech recognition pipeline;RNN-Transducer models;language model;Seq2Seq models;end-to-end speech recognition;neural transducers;Decoding;Hidden Markov models;Transducers;Task analysis;Speech;Mathematical model;Neural networks},
|
||||
doi={10.1109/ASRU.2017.8268937},
|
||||
ISSN={},
|
||||
month={Dec},
|
||||
}
|
||||
|
||||
@inproceedings{
|
||||
loshchilov2018,
|
||||
title={Decoupled Weight Decay Regularization},
|
||||
author={Ilya Loshchilov and Frank Hutter},
|
||||
booktitle={International Conference on Learning Representations},
|
||||
year={2019},
|
||||
url={https://openreview.net/forum?id=Bkg6RiCqY7},
|
||||
}
|
||||
|
||||
@article{zhang2017ndadam,
|
||||
author = {Zijun Zhang and Lin Ma and Zongpeng Li and Chuan Wu},
|
||||
title = {Normalized Direction-preserving Adam},
|
||||
journal = {arXiv e-prints arXiv:1709.04546},
|
||||
year = {2017},
|
||||
}
|
||||
|
||||
@article{park2019,
|
||||
author = {{Park}, Daniel S. and {Chan}, William and {Zhang}, Yu and
|
||||
{Chiu}, Chung-Cheng and {Zoph}, Barret and {Cubuk}, Ekin D. and
|
||||
{Le}, Quoc V.},
|
||||
title = "{SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition}",
|
||||
journal = {arXiv e-prints},
|
||||
year = "2019",
|
||||
eid = {arXiv:1904.08779},
|
||||
eprint = {1904.08779},
|
||||
}
|
||||
|
||||
@article{novograd2019,
|
||||
author = {{Ginsburg}, Boris and {Castonguay}, Patrice and {Hrinchuk}, Oleksii and
|
||||
{Kuchaiev}, Oleksii and {Lavrukhin}, Vitaly and {Leary}, Ryan and
|
||||
{Li}, Jason and {Nguyen}, Huyen and {Cohen}, Jonathan M.},
|
||||
title = "{Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks}",
|
||||
journal = {arXiv e-prints},
|
||||
year = "2019",
|
||||
eid = {arXiv:1905.11286},
|
||||
eprint = {1905.11286},
|
||||
}
|
||||
|
||||
@article{kriman2019quartznet,
|
||||
title={Quartznet: {Deep} automatic speech recognition with 1d time-channel separable convolutions},
|
||||
author={Kriman, Samuel and Beliaev, Stanislav and Ginsburg, Boris and Huang, Jocelyn and Kuchaiev, Oleksii and Lavrukhin, Vitaly and Leary, Ryan and Li, Jason and Zhang, Yang},
|
||||
journal={arXiv preprint arXiv:1910.10261},
|
||||
year={2019}
|
||||
}
|
||||
|
||||
@misc{itu1988g711,
|
||||
title={{ITU-T} {G.711} - {Pulse} code modulation ({PCM}) of voice frequencies},
|
||||
author={ITU-T Geneva Switzerland},
|
||||
year={1988},
|
||||
}
|
||||
|
||||
@article{han2020contextnet,
|
||||
title={ContextNet: Improving convolutional neural networks for automatic speech recognition with global context},
|
||||
author={Han, Wei and Zhang, Zhengdong and Zhang, Yu and Yu, Jiahui and Chiu, Chung-Cheng and Qin, James and Gulati, Anmol and Pang, Ruoming and Wu, Yonghui},
|
||||
journal={arXiv:2005.03191},
|
||||
year={2020}
|
||||
}
|
||||
|
||||
@inproceedings{hu2018squeeze,
|
||||
title={Squeeze-and-excitation networks},
|
||||
author={Hu, Jie and Shen, Li and Sun, Gang},
|
||||
booktitle={ICVPR},
|
||||
year={2018}
|
||||
}
|
||||
|
||||
@article{koluguri2020speakernet,
|
||||
title={SpeakerNet: 1D Depth-wise Separable Convolutional Network for Text-Independent Speaker Recognition and Verification},
|
||||
author={Koluguri, Nithin Rao and Li, Jason and Lavrukhin, Vitaly and Ginsburg, Boris},
|
||||
journal={arXiv preprint arXiv:2010.12653},
|
||||
year={2020}
|
||||
}
|
||||
|
||||
@article{gulati2020conformer,
|
||||
title={Conformer: Convolution-augmented transformer for speech recognition},
|
||||
author={Gulati, Anmol and Qin, James and Chiu, Chung-Cheng and Parmar, Niki and Zhang, Yu and Yu, Jiahui and Han, Wei and Wang, Shibo and Zhang, Zhengdong and Wu, Yonghui and others},
|
||||
journal={arXiv preprint arXiv:2005.08100},
|
||||
year={2020}
|
||||
}
|
||||
@@ -0,0 +1,292 @@
|
||||
#####################
|
||||
ASR Language Modeling
|
||||
#####################
|
||||
|
||||
Language models have shown to help the accuracy of ASR models. NeMo supports the following two approaches to incorporate language models into the ASR models:
|
||||
|
||||
* :ref:`ngram_modeling`
|
||||
* :ref:`neural_rescoring`
|
||||
|
||||
It is possible to use both approaches on the same ASR model.
|
||||
|
||||
|
||||
.. _ngram_modeling:
|
||||
|
||||
************************
|
||||
N-gram Language Modeling
|
||||
************************
|
||||
|
||||
In this approach, an N-gram LM is trained on text data, then it is used in fusion with beam search decoding to find the
|
||||
best candidates. The beam search decoders in NeMo support language models trained with KenLM library (
|
||||
`https://github.com/kpu/kenlm <https://github.com/kpu/kenlm>`__).
|
||||
The beam search decoders and KenLM library are not installed by default in NeMo, and you need to install them to be
|
||||
able to use beam search decoding and N-gram LM.
|
||||
Please refer to `scripts/ngram_lm/install_beamsearch_decoders.sh` on how to install them.
|
||||
|
||||
NeMo supports both character-based and BPE-based models for N-gram LMs. An N-gram LM can be used with beam search
|
||||
decoders on top of the ASR models to produce more accurate candidates. The beam search decoder would incorporate
|
||||
the scores produced by the N-gram LM into its score calculations as the following:
|
||||
|
||||
.. code::
|
||||
|
||||
final_score = acoustic_score + beam_alpha*lm_score + beam_beta*seq_length
|
||||
|
||||
where acoustic_score is the score predicted by the acoustic encoder and lm_score is the one estimated by the LM.
|
||||
Parameter 'beam_alpha' specifies amount of importance to place on the N-gram language model, and 'beam_beta' is a
|
||||
penalty term to consider the sequence length in the scores. Larger alpha means more importance on the LM and less
|
||||
importance on the acoustic model. Negative values for beta will give penalty to longer sequences and make the decoder
|
||||
to prefer shorter predictions, while positive values would result in longer candidates.
|
||||
|
||||
|
||||
Train N-gram LM
|
||||
===============
|
||||
|
||||
The script to train an N-gram language model with KenLM can be found at
|
||||
`scripts/asr_language_modeling/ngram_lm/train_kenlm.py <https://github.com/NVIDIA/NeMo/blob/main/scripts/asr_language_modeling/ngram_lm/train_kenlm.py>`__.
|
||||
|
||||
This script would train an N-gram language model with KenLM library which can be used with the beam search decoders
|
||||
on top of the ASR models. This script supports both character level and BPE level encodings and models which is
|
||||
detected automatically from the type of the model.
|
||||
|
||||
|
||||
You may train the N-gram model as the following:
|
||||
|
||||
.. code::
|
||||
|
||||
python train_kenlm.py --nemo_model_file <path to the .nemo file of the model> \
|
||||
--train_file <path to the training text or JSON manifest file \
|
||||
--kenlm_bin_path <path to the bin folder of KenLM library> \
|
||||
--kenlm_model_file <path to store the binary KenLM model> \
|
||||
--ngram_length <order of N-gram model>
|
||||
|
||||
The train file specified by `--train_file` can be a text file or JSON manifest. If the file's extension is anything
|
||||
other than `.json`, it assumes that data format is plain text. For plain text format, each line should contain one
|
||||
sample. For JSON manifest file, the file need to contain json formatted samples per each line like this:
|
||||
|
||||
.. code::
|
||||
|
||||
{"audio_filepath": "/data_path/file1.wav", "text": "The transcript of the audio file."}
|
||||
|
||||
It just extracts the `text` field from each line to create the training text file. After the N-gram model is trained,
|
||||
it is stored at the path specified by `--kenlm_model_file`.
|
||||
|
||||
The following is the list of the arguments for the training script:
|
||||
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| **Argument** | **Type** | **Default** | **Description** |
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| nemo_model_file | str | Required | The path of the `.nemo` file of the ASR model. It is needed to extract the tokenizer. |
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| train_file | str | Required | Path to the training file, it can be a text file or JSON manifest. |
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| kenlm_model_file | str | Required | The path to store the KenLM binary model file. |
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| kenlm_bin_path | str | Required | The path to the bin folder of KenLM. It is a folder named `bin` under where KenLM is installed.|
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| ngram_length** | int | Required | Specifies order of N-gram LM. |
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
| do_lower_case | bool | ``False`` | Whether to make the training text all lower case. |
|
||||
+------------------+----------+-------------+------------------------------------------------------------------------------------------------+
|
||||
|
||||
** Note: Recommend to use 6 as the order of the N-gram model for BPE-based models. Higher orders may need the re-compilation of KenLM to support it.
|
||||
|
||||
Evaluate by Beam Search Decoding and N-gram LM
|
||||
==============================================
|
||||
|
||||
NeMo's beam search decoders are capable of using the KenLM's N-gram models to find the best candidates.
|
||||
The script to evaluate an ASR model with beam search decoding and N-gram models can be found at
|
||||
`scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py <https://github.com/NVIDIA/NeMo/blob/main/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py>`__.
|
||||
|
||||
You may evaluate an ASR model as the following:
|
||||
|
||||
.. code::
|
||||
|
||||
python eval_beamsearch_ngram.py --nemo_model_file <path to the .nemo file of the model> \
|
||||
--input_manifest <path to the evaluation JSON manifest file \
|
||||
--kenlm_model_file <path to the binary KenLM model> \
|
||||
--acoustic_batch_size <batch size for calculating log probabilities> \
|
||||
--beam_width <list of the beam widths> \
|
||||
--beam_alpha <list of the beam alphas> \
|
||||
--beam_width <list of the beam betas> \
|
||||
--preds_output_folder <optional folder to store the predictions> \
|
||||
--decoding_mode beam_search_ngram
|
||||
|
||||
It can evaluate a model in the three following modes by setting the argument `--decoding_mode`:
|
||||
|
||||
* greedy: Just greedy decoding is done, and no beam search decoding is performed.
|
||||
* beamsearch: The beam search decoding is done but without using the N-gram language model, final results would be equivalent to setting the weight of LM (beam_beta) to zero.
|
||||
* beamsearch_ngram: The beam search decoding is done with N-gram LM.
|
||||
|
||||
The `beamsearch` mode would evaluate by beam search decoding without any language model.
|
||||
It would report the performances in terms of Word Error Rate (WER) and Character Error Rate (CER). Moreover,
|
||||
the WER/CER of the model when the best candidate is selected among the candidates is also reported as the best WER/CER.
|
||||
It can be an indicator of how good the predicted candidates are.
|
||||
|
||||
The script would initially load the ASR model and predict the outputs of the model's encoder as log probabilities.
|
||||
This part would be computed in batches on a device selected by `--device`, which can be CPU (`--device=cpu`) or a
|
||||
single GPU (`--device=cuda:0`). The batch size of this part can get specified by `--acoustic_batch_size`. You may use
|
||||
the largest batch size feasible to speed up the step of calculating the log probabilities. You may also use `--use_amp`
|
||||
to speed up the calculation of log probabilities and make it possible to use larger sizes for `--acoustic_batch_size`.
|
||||
Currently multi-GPU is not supported for calculating the log probabilities, but using `--probs_cache_file` can help.
|
||||
It stores the log probabilities produced from the model's encoder into a pickle file so that next time the first step
|
||||
can get skipped.
|
||||
|
||||
The following is the list of the arguments for the evaluation script:
|
||||
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| **Argument** |**Type**| **Default** | **Description** |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| nemo_model_file | str | Required | The path of the `.nemo` file of the ASR model to extract the tokenizer. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| input_manifest | str | Required | Path to the training file, it can be a text file or JSON manifest. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| kenlm_model_file | str | Required | The path to store the KenLM binary model file. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| preds_output_folder | str | None | The path to an optional folder to store the predictions. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| probs_cache_file | str | None | The cache file for storing the outputs of the model. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| acoustic_batch_size | int | 16 | The batch size to calculate log probabilities. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| use_amp | bool | ``False`` | Whether to use AMP if available to calculate log probabilities. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| device | str | cuda | The device to load the model onto to calculate log probabilities. |
|
||||
| | | | It can `cpu`, `cuda`, `cuda:0`, `cuda:1`, ... |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| decoding_mode | str | beamsearch_ngram | The decoding scheme to be used for evaluation. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beam_width | float | Required | The width or list of the widths of the beam search decoding. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beam_alpha | float | Required | The alpha parameter or list of the alphas for the beam search decoding. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beam_beta | float | Required | The beta parameter or list of the betas for the beam search decoding. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beam_batch_size | int | 128 | The batch size to be used for beam search decoding. |
|
||||
| | | | Larger batch size can be a little faster, but uses larger memory. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
|
||||
Width of the beam search (`--beam_width`) specifies the number of top candidates/predictions the beam search decoder
|
||||
would search for. Larger beams result in more accurate but slower predictions.
|
||||
|
||||
There is also a tutorial to learn more about evaluating the ASR models with N-gram LM here:
|
||||
`Offline ASR Inference with Beam Search and External Language Model Rescoring <https://colab.research.google.com/github/NVIDIA/NeMo/blob/v1.0.2/tutorials/asr/Offline_ASR.ipynb>`_
|
||||
|
||||
Hyperparameter Grid Search
|
||||
--------------------------
|
||||
|
||||
Beam search decoding with N-gram LM has three main hyperparameters: `beam_width`, `beam_alpha`, and `beam_beta`.
|
||||
The accuracy of the model is dependent to the values of these parameters, specially beam_alpha and beam_beta.
|
||||
You may specify a single or list of values for each of these parameters to perform grid search. It would perform the
|
||||
beam search decoding on all the combinations of the these three hyperparameters.
|
||||
For instance, the following set of parameters would results in 2*1*2=4 beam search decodings:
|
||||
|
||||
.. code::
|
||||
|
||||
python eval_beamsearch_ngram.py ... \
|
||||
--beam_width 64 128 \
|
||||
--beam_alpha 1.0 \
|
||||
--beam_beta 1.0 0.5
|
||||
|
||||
|
||||
.. _neural_rescoring:
|
||||
|
||||
****************
|
||||
Neural Rescoring
|
||||
****************
|
||||
|
||||
In this approach a neural network is used which can gives scores to a candidate. A candidate is the text transcript predicted by the decoder of the ASR model.
|
||||
The top K candidates produced by the beam search decoding (beam width of K) are given to a neural language model to rank them.
|
||||
Ranking can be done by a language model which gives a score to each candidate.
|
||||
This score is usually combined with the scores from the beam search decoding to produce the final scores and rankings.
|
||||
|
||||
Train Neural Rescorer
|
||||
=====================
|
||||
|
||||
An example script to train such a language model with Transformer can be found at `examples/nlp/language_modeling/transformer_lm.py <https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/transformer_lm.py>`__.
|
||||
It trains a TransformerLMModel which can be used as a neural rescorer for an ASR system.
|
||||
|
||||
|
||||
Evaluation
|
||||
==========
|
||||
|
||||
Given a trained TransformerLMModel `.nemo` file, the script available at
|
||||
`scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py <https://github.com/NVIDIA/NeMo/blob/main/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py>`__
|
||||
can be used to re-score beams obtained with ASR model. You need the `.tsv` file containing the candidates produced
|
||||
by the acoustic model and the beam search decoding to use this script. The candidates can be the result of just the beam
|
||||
search decoding or the result of fusion with an N-gram LM. You may generate this file by specifying `--preds_output_folder' for
|
||||
`scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py <https://github.com/NVIDIA/NeMo/blob/main/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py>`__.
|
||||
|
||||
The neural rescorer would rescore the beams/candidates by using two parameters of `rescorer_alpha` and `rescorer_beta` as the following:
|
||||
|
||||
.. code::
|
||||
|
||||
final_score = beam_search_score + rescorer_alpha*neural_rescorer_score + rescorer_beta*seq_length
|
||||
|
||||
Parameter `rescorer_alpha` specifies amount of importance to place on the neural rescorer model, and `rescorer_beta` is
|
||||
a penalty term to consider the sequence length in the scores. They have similar effects like the parameters
|
||||
`beam_alpha` and `beam_beta` of beam search decoder and N-gram LM.
|
||||
|
||||
You may follow the following steps to evaluate a neural LM:
|
||||
|
||||
#. Obtain `.tsv` file with beams and their corresponding scores. Scores can be from a regular beam search decoder or
|
||||
in fusion with an N-gram LM scores. For a given beam size `beam_size` and a number of examples
|
||||
for evaluation `num_eval_examples`, it should contain (`num_eval_examples` x `beam_size`) lines of
|
||||
form `beam_candidate_text \t score`. This file can be generated by `scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py <https://github.com/NVIDIA/NeMo/blob/main/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py>`__
|
||||
|
||||
#. Rescore the candidates by `scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py <https://github.com/NVIDIA/NeMo/blob/main/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py>`__.
|
||||
|
||||
.. code::
|
||||
python eval_neural_rescorer.py
|
||||
--lm_model=[path to .nemo file of the LM]
|
||||
--beams_file=[path to beams .tsv file]
|
||||
--beam_size=[size of the beams]
|
||||
--eval_manifest=[path to eval manifest .json file]
|
||||
--batch_size=[batch size used for inference on the LM model]
|
||||
--alpha=[the value for the parameter rescorer_alpha]
|
||||
--beta=[the value for the parameter rescorer_beta]
|
||||
--scores_output_file=[the optional path to store the rescored candidates]
|
||||
|
||||
The candidates along with their new scores would be stored at the file specified by `--scores_output_file`.
|
||||
|
||||
The following is the list of the arguments for the evaluation script:
|
||||
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| **Argument** |**Type**| **Default** | **Description** |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| lm_model | str | Required | The path of the '.nemo' file of the ASR model to extract the tokenizer. |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| eval_manifest | str | Required | Path to the evaluation manifest file (.json manifest file) |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beams_file | str | Required | path to beams file (.tsv) containing the candidates and their scores |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beam_size | int | Required | The width of the beams (number of candidates) generated by the decoder |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| alpha | float | None | The value for parameter rescorer_alpha |
|
||||
| | | | Not passing value would enable linear search for rescorer_alpha |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| beta | float | None | The value for parameter rescorer_beta |
|
||||
| | | | Not passing value would enable linear search for rescorer_beta |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| batch_size | int | 16 | The batch size used to calculate the scores |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| scores_output_file | str | None | The optional file to store the rescored beams |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| use_amp | bool | ``False`` | Whether to use AMP if available calculate the scores |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
| device | str | cuda | The device to load LM model onto to calculate the scores |
|
||||
| | | | It can be 'cpu', 'cuda', 'cuda:0', 'cuda:1', ... |
|
||||
+---------------------+--------+------------------+-------------------------------------------------------------------------+
|
||||
|
||||
|
||||
Hyperparameter Linear Search
|
||||
----------------------------
|
||||
|
||||
This script also supports linear search for parameters `rescorer_alpha` and `rescorer_beta`. If any of the two is not
|
||||
provided, a linear search is performed to find the best value for that parameter. When linear search is used, initially
|
||||
`rescorer_beta` is set to zero and the best value for `rescorer_alpha` is found, then `rescorer_alpha` is fixed with
|
||||
that value and another linear search is done to find the best value for `rescorer_beta`.
|
||||
If any of the of these two parameters is already specified, then search for that one is skipped. After each search for a
|
||||
parameter, the plot of WER% for different values of the parameter is also shown.
|
||||
|
||||
It is recommended to first use the linear search for both parameters on a validation set by not providing any values for `--alpha` and `--beta`.
|
||||
Then check the WER curves and decide on the best values for each parameter. Finally, evaluate the best values on the test set.
|
||||
@@ -0,0 +1,528 @@
|
||||
NeMo ASR Configuration Files
|
||||
============================
|
||||
|
||||
This section describes the NeMo configuration file setup that is specific to models in the ASR collection. For general information
|
||||
about how to set up and run experiments that is common to all NeMo models (e.g. Experiment Manager and PyTorch Lightning trainer
|
||||
parameters), see the :doc:`../core/core` section.
|
||||
|
||||
The model section of the NeMo ASR configuration files generally requires information about the dataset(s) being used, the preprocessor
|
||||
for audio files, parameters for any augmentation being performed, as well as the model architecture specification. The sections on
|
||||
this page cover each of these in more detail.
|
||||
|
||||
Example configuration files for all of the NeMo ASR scripts can be found in the
|
||||
`config directory of the examples <https://github.com/NVIDIA/NeMo/tree/v1.0.2/examples/asr/conf>`_.
|
||||
|
||||
|
||||
Dataset Configuration
|
||||
---------------------
|
||||
|
||||
Training, validation, and test parameters are specified using the ``train_ds``, ``validation_ds``, and
|
||||
``test_ds`` sections in the configuration file, respectively. Depending on the task, there may be arguments specifying the sample rate
|
||||
of the audio files, the vocabulary of the dataset (for character prediction), whether or not to shuffle the dataset, and so on. You may
|
||||
also decide to leave fields such as the ``manifest_filepath`` blank, to be specified via the command-line at runtime.
|
||||
|
||||
Any initialization parameter that is accepted for the Dataset class used in the experiment can be set in the config file.
|
||||
Refer to the `Datasets <./api.html#Datasets>`__ section of the API for a list of Datasets and their respective parameters.
|
||||
|
||||
An example ASR train and validation configuration should look similar to the following:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
# Specified at the beginning of the config file
|
||||
labels: &labels [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
|
||||
"n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
|
||||
|
||||
model:
|
||||
train_ds:
|
||||
manifest_filepath: ???
|
||||
sample_rate: 16000
|
||||
labels: *labels # Uses the labels above
|
||||
batch_size: 32
|
||||
trim_silence: True
|
||||
max_duration: 16.7
|
||||
shuffle: True
|
||||
is_tarred: False # If set to true, uses the tarred version of the Dataset
|
||||
tarred_audio_filepaths: null # Not used if is_tarred is false
|
||||
tarred_shard_strategy: "scatter" # Not used if is_tarred is false
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
validation_ds:
|
||||
manifest_filepath: ???
|
||||
sample_rate: 16000
|
||||
labels: *labels # Uses the labels above
|
||||
batch_size: 32
|
||||
shuffle: False # No need to shuffle the validation data
|
||||
num_workers: 8
|
||||
pin_memory: true
|
||||
|
||||
|
||||
Preprocessor Configuration
|
||||
--------------------------
|
||||
|
||||
If you are loading audio files for your experiment, you will likely want to use a preprocessor to convert from the
|
||||
raw audio signal to features (e.g. mel-spectrogram or MFCC). The ``preprocessor`` section of the config specifies the audio
|
||||
preprocessor to be used via the ``_target_`` field, as well as any initialization parameters for that preprocessor.
|
||||
|
||||
An example of specifying a preprocessor is as follows:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
preprocessor:
|
||||
# _target_ is the audio preprocessor module you want to use
|
||||
_target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
|
||||
normalize: "per_feature"
|
||||
window_size: 0.02
|
||||
...
|
||||
# Other parameters for the preprocessor
|
||||
|
||||
Refer to the `Audio Preprocessors <./api.html#Audio Preprocessors>`__ API section for the preprocessor options, expected arguments,
|
||||
and defaults.
|
||||
|
||||
Augmentation Configurations
|
||||
---------------------------
|
||||
|
||||
There are a few on-the-fly spectrogram augmentation options for NeMo ASR, which can be specified by the
|
||||
configuration file using a ``spec_augment`` section.
|
||||
|
||||
For example, there are options for `Cutout <https://arxiv.org/abs/1708.04552>`_ and
|
||||
`SpecAugment <https://arxiv.org/abs/1904.08779>`_ available via the ``SpectrogramAugmentation`` module.
|
||||
|
||||
The following example sets up both ``Cutout`` (via the ``rect_*`` parameters) and ``SpecAugment`` (via the ``freq_*``
|
||||
and ``time_*`` parameters).
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
spec_augment:
|
||||
_target_: nemo.collections.asr.modules.SpectrogramAugmentation
|
||||
# Cutout parameters
|
||||
rect_masks: 5 # Number of rectangles to cut from any given spectrogram
|
||||
rect_freq: 50 # Max cut of size 50 along the frequency dimension
|
||||
rect_time: 120 # Max cut of size 120 along the time dimension
|
||||
# SpecAugment parameters
|
||||
freq_masks: 2 # Cut two frequency bands
|
||||
freq_width: 15 # ... of width 15 at maximum
|
||||
time_masks: 5 # Cut out 10 time bands
|
||||
time_width: 25 # ... of width 25 at maximum
|
||||
|
||||
You can use any combination of ``Cutout``, frequency/time ``SpecAugment``, or neither of them.
|
||||
|
||||
With NeMo ASR, you can also add augmentation pipelines that can be used to simulate various kinds of noise
|
||||
added to audio in the channel. Augmentors in a pipeline are applied on the audio data read in the data layer. Online
|
||||
augmentors can be specified in the config file using an ``augmentor`` section in ``train_ds``. The following example
|
||||
adds an augmentation pipeline that first adds white noise to an audio sample with a probability of 0.5 and at a level
|
||||
randomly picked between -50 dB and -10 dB and then passes the resultant samples through a room impulse response randomly
|
||||
picked from the manifest file provided for ``impulse`` augmentation in the config file.
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
train_ds:
|
||||
...
|
||||
augmentor:
|
||||
white_noise:
|
||||
prob: 0.5
|
||||
min_level: -50
|
||||
max_level: -10
|
||||
impulse:
|
||||
prob: 0.3
|
||||
manifest_path: /path/to/impulse_manifest.json
|
||||
|
||||
Refer to the `Audio Augmentors <./api.html#Audio Augmentors>`__ API section for more details.
|
||||
|
||||
Tokenizer Configurations
|
||||
------------------------
|
||||
|
||||
Some models utilize sub-word encoding via an external tokenizer instead of explicitly defining their vocabulary.
|
||||
|
||||
For such models, a ``tokenizer`` section is added to the model config. ASR models currently support two types of
|
||||
custom tokenizers:
|
||||
|
||||
- Google Sentencepiece tokenizers (tokenizer type of ``bpe`` in the config)
|
||||
- HuggingFace WordPiece tokenizers (tokenizer type of ``wpe`` in the config)
|
||||
|
||||
In order to build custom tokenizers, refer to the ``ASR_with_Subword_Tokenization`` notebook available in the
|
||||
ASR tutorials directory.
|
||||
|
||||
The following example sets up a ``SentencePiece Tokenizer`` at a path specified by the user:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
tokenizer:
|
||||
dir: "<path to the directory that contains the custom tokenizer files>"
|
||||
type: "bpe" # can be "bpe" or "wpe"
|
||||
|
||||
For models which utilize sub-word tokenization, we share the decoder module (``ConvASRDecoder``) with character tokenization models.
|
||||
All parameters are shared, but for models which utilize sub-word encoding, there are minor differences when setting up the config. For
|
||||
such models, the tokenizer is utilized to fill in the missing information when the model is constructed automatically.
|
||||
|
||||
For example, a decoder config corresponding to a sub-word tokenization model should look similar to the following:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
decoder:
|
||||
_target_: nemo.collections.asr.modules.ConvASRDecoder
|
||||
feat_in: *enc_final
|
||||
num_classes: -1 # filled with vocabulary size from tokenizer at runtime
|
||||
vocabulary: [] # filled with vocabulary from tokenizer at runtime
|
||||
|
||||
|
||||
Model Architecture Configurations
|
||||
---------------------------------
|
||||
|
||||
Each configuration file should describe the model architecture being used for the experiment. Models in the NeMo ASR collection need
|
||||
an ``encoder`` section and a ``decoder`` section, with the ``_target_`` field specifying the module to use for each.
|
||||
|
||||
Here is the list of the parameters in the model section which are shared among most of the ASR models:
|
||||
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
| **Parameter** | **Datatype** | **Description** | **Supported Values** |
|
||||
+=========================+==================+===============================================================================================================+=================================+
|
||||
| :code:`log_prediction` | bool | Whether a random sample should be printed in the output at each step, along with its predicted transcript. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
| :code:`ctc_reduction` | string | Specifies the reduction type of CTC loss. Defaults to ``mean_batch`` which would take the average over the | :code:`none`, |
|
||||
| | | batch after taking the average over the length of each sample. | :code:`mean_batch` |
|
||||
| | | | :code:`mean`, :code:`sum` |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
|
||||
The following sections go into more detail about the specific configurations of each model architecture.
|
||||
|
||||
For more information about the ASR models, refer to the :doc:`Models <./models>` section.
|
||||
|
||||
Jasper and QuartzNet
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The `Jasper <./models.html#Jasper>`__ and `QuartzNet <./models.html#QuartzNet>`__ models are very similar, and as such the components in their
|
||||
configs are very similar as well.
|
||||
|
||||
Both architectures use the ``ConvASREncoder`` for the ``encoder``, with parameters detailed in the table below. The encoder parameters
|
||||
include details about the Jasper/QuartzNet ``[BxR]`` encoder architecture, including how many blocks to use (``B``), how many times
|
||||
to repeat each sub-block (``R``), and the convolution parameters for each block.
|
||||
|
||||
The number of blocks ``B`` is determined by the number of list elements under ``jasper`` minus the one prologue and two epilogue blocks.
|
||||
The number of sub-blocks ``R`` is determined by setting the ``repeat`` parameter.
|
||||
|
||||
To use QuartzNet (which uses more compact time-channel separable convolutions) instead of Jasper, add :code:`separable: true` to all
|
||||
but the last block in the architecture.
|
||||
|
||||
Change the parameter name ``jasper``.
|
||||
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+
|
||||
| **Parameter** | **Datatype** | **Description** | **Supported Values** |
|
||||
+=========================+==================+===============================================================================================================+=====================================+
|
||||
| :code:`feat_in` | int | The number of input features. Should be equal to :code:`features` in the preprocessor parameters. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+
|
||||
| :code:`activation` | string | Which activation function to use in the encoder. | :code:`hardtanh`, :code:`relu`, |
|
||||
| | | | :code:`selu`, :code:`swish` |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+
|
||||
| :code:`conv_mask` | bool | Whether to use masked convolutions in the encoder. Defaults to ``true``. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+
|
||||
| :code:`jasper` | | A list of blocks that specifies your encoder architecture. Each entry in this list represents one block in | |
|
||||
| | | the architecture and contains the parameters for that block, including convolution parameters, dropout, and | |
|
||||
| | | the number of times the block is repeated. Refer to the `Jasper <https://arxiv.org/pdf/1904.03288.pdf>`_ and | |
|
||||
| | | `QuartzNet <https://arxiv.org/pdf/1910.10261.pdf>`_ papers for details about specific model configurations. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+
|
||||
|
||||
A QuartzNet 15x5 (fifteen blocks, each sub-block repeated five times) encoder configuration should look similar to the following example:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
# Specified at the beginning of the file for convenience
|
||||
n_mels: &n_mels 64 # Used for both the preprocessor and encoder as number of input features
|
||||
repeat: &repeat 5 # R=5
|
||||
dropout: &dropout 0.0
|
||||
separable: &separable true # Set to true for QN. Set to false for Jasper.
|
||||
|
||||
model:
|
||||
...
|
||||
encoder:
|
||||
_target_: nemo.collections.asr.modules.ConvASREncoder
|
||||
feat_in: *n_mels # Should match "features" in the preprocessor.
|
||||
activation: relu
|
||||
conv_mask: true
|
||||
|
||||
jasper: # This field name should be "jasper" for both types of models.
|
||||
|
||||
# Prologue block
|
||||
- dilation: [1]
|
||||
dropout: *dropout
|
||||
filters: 256
|
||||
kernel: [33]
|
||||
repeat: 1 # Prologue block is not repeated.
|
||||
residual: false
|
||||
separable: *separable
|
||||
stride: [2]
|
||||
|
||||
# Block 1
|
||||
- dilation: [1]
|
||||
dropout: *dropout
|
||||
filters: 256
|
||||
kernel: [33]
|
||||
repeat: *repeat
|
||||
residual: true
|
||||
separable: *separable
|
||||
stride: [1]
|
||||
|
||||
... # Entries for blocks 2~14
|
||||
|
||||
# Block 15
|
||||
- dilation: [1]
|
||||
dropout: *dropout
|
||||
filters: 512
|
||||
kernel: [75]
|
||||
repeat: *repeat
|
||||
residual: true
|
||||
separable: *separable
|
||||
stride: [1]
|
||||
|
||||
# Two epilogue blocks
|
||||
- dilation: [2]
|
||||
dropout: *dropout
|
||||
filters: 512
|
||||
kernel: [87]
|
||||
repeat: 1 # Epilogue blocks are not repeated
|
||||
residual: false
|
||||
separable: *separable
|
||||
stride: [1]
|
||||
|
||||
- dilation: [1]
|
||||
dropout: *dropout
|
||||
filters: &enc_filters 1024
|
||||
kernel: [1]
|
||||
repeat: 1 # Epilogue blocks are not repeated
|
||||
residual: false
|
||||
stride: [1]
|
||||
|
||||
Both Jasper and QuartzNet use the ``ConvASRDecoder`` as the decoder. The decoder parameters are detailed in the following table.
|
||||
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
| **Parameter** | **Datatype** | **Description** | **Supported Values** |
|
||||
+=========================+==================+===============================================================================================================+=================================+
|
||||
| :code:`feat_in` | int | The number of input features to the decoder. Should be equal to the number of filters in the last block of | |
|
||||
| | | the encoder. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
| :code:`vocabulary` | list | A list of the valid output characters for your model. For example, for an English dataset, this could be a | |
|
||||
| | | list of all lowercase letters, space, and apostrophe. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
| :code:`num_classes` | int | Number of output classes, i.e. the length of :code:`vocabulary`. | |
|
||||
+-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+
|
||||
|
||||
For example, a decoder config corresponding to the encoder above should look similar to the following:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
decoder:
|
||||
_target_: nemo.collections.asr.modules.ConvASRDecoder
|
||||
feat_in: *enc_filters
|
||||
vocabulary: *labels
|
||||
num_classes: 28 # Length of the vocabulary list
|
||||
|
||||
Citrinet
|
||||
~~~~~~~~
|
||||
|
||||
The `Citrinet <./models.html#Citrinet>`__ and `QuartzNet <./models.html#QuartzNet>`__ models are very similar, and as such the
|
||||
components in their configs are very similar as well. Citrinet utilizes Squeeze and Excitation, as well as sub-word tokenization, in
|
||||
contrast to QuartzNet. Depending on the dataset, we utilize different tokenizers. For Librispeech, we utilize the HuggingFace WordPiece
|
||||
tokenizer, and for all other datasets we utilize the Google Sentencepiece tokenizer - usually the ``unigram`` tokenizer type.
|
||||
|
||||
Both architectures use the ``ConvASREncoder`` for the ``encoder``, with parameters detailed above. The encoder parameters include
|
||||
details about the Citrinet-C encoder architecture, including how many filters are used per channel (``C``). The Citrinet-C
|
||||
configuration is a shortform notation for Citrinet-21x5xC, such that ``B = 21`` and ``R = 5`` are the default and should generally
|
||||
not be changed.
|
||||
|
||||
To use Citrinet instead of QuartzNet, refer to the ``citrinet_512.yaml`` configuration found inside the ``examples/asr/conf/citrinet``
|
||||
directory. Citrinet is primarily comprised of the same :class:`~nemo.collections.asr.parts.submodules.jasper.JasperBlock` as ``Jasper`` or
|
||||
``QuartzNet`.
|
||||
|
||||
While the configs for Citrinet and QuartzNet are similar, we note the additional flags used for Citrinet below. Refer to the
|
||||
``JasperBlock`` documentation for the meaning of these arguments.
|
||||
|
||||
+---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+
|
||||
| **Parameter** | **Datatype** | **Description** | **Supported Values** |
|
||||
+===========================+==================+===========================================================================================================+===================================+
|
||||
| :code:`se` | bool | Whether to apply squeeze-and-excitation mechanism or not. | :code:`true` or :code:`false` |
|
||||
+---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+
|
||||
| :code:`se_context_size` | int | SE context size. -1 means global context. | :code:`-1` or :code:`+ve int` |
|
||||
+---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+
|
||||
| :code:`stride_last` | bool | Stride on the final repeated block or all repeated blocks. | :code:`true` or :code:`false` |
|
||||
+---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+
|
||||
| :code:`residual_mode` | str | Type of residual branch to construct. | :code:`"add"` or |
|
||||
| | | Can be pointwise residual addition or pointwise strided residual attention | :code:`"stride_add"` |
|
||||
+---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+
|
||||
|
||||
A Citrinet-512 config should look similar to the following:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
model:
|
||||
...
|
||||
# Specify some defaults across the entire model
|
||||
model_defaults:
|
||||
repeat: 5
|
||||
dropout: 0.1
|
||||
separable: true
|
||||
se: true
|
||||
se_context_size: -1
|
||||
...
|
||||
encoder:
|
||||
_target_: nemo.collections.asr.modules.ConvASREncoder
|
||||
feat_in: *n_mels # Should match "features" in the preprocessor.
|
||||
activation: relu
|
||||
conv_mask: true
|
||||
|
||||
jasper: # This field name should be "jasper" for the JasperBlock (which constructs Citrinet).
|
||||
|
||||
# Prologue block
|
||||
- filters: 512
|
||||
repeat: 1
|
||||
kernel: [5]
|
||||
stride: [1]
|
||||
dilation: [1]
|
||||
dropout: 0.0
|
||||
residual: false
|
||||
separable: ${model.model_defaults.separable}
|
||||
se: ${model.model_defaults.se}
|
||||
se_context_size: ${model.model_defaults.se_context_size}
|
||||
|
||||
# Block 1
|
||||
- filters: 512
|
||||
repeat: ${model.model_defaults.repeat}
|
||||
kernel: [11]
|
||||
stride: [2]
|
||||
dilation: [1]
|
||||
dropout: ${model.model_defaults.dropout}
|
||||
residual: true
|
||||
separable: ${model.model_defaults.separable}
|
||||
se: ${model.model_defaults.se}
|
||||
se_context_size: ${model.model_defaults.se_context_size}
|
||||
stride_last: true
|
||||
residual_mode: "stride_add"
|
||||
|
||||
... # Entries for blocks 2~21
|
||||
|
||||
# Block 22
|
||||
- filters: 512
|
||||
repeat: ${model.model_defaults.repeat}
|
||||
kernel: [39]
|
||||
stride: [1]
|
||||
dilation: [1]
|
||||
dropout: ${model.model_defaults.dropout}
|
||||
residual: true
|
||||
separable: ${model.model_defaults.separable}
|
||||
se: ${model.model_defaults.se}
|
||||
se_context_size: ${model.model_defaults.se_context_size}
|
||||
|
||||
# Epilogue block
|
||||
|
||||
- filters: &enc_final 640
|
||||
repeat: 1
|
||||
kernel: [41]
|
||||
stride: [1]
|
||||
dilation: [1]
|
||||
dropout: 0.0
|
||||
residual: false
|
||||
separable: ${model.model_defaults.separable}
|
||||
se: ${model.model_defaults.se}
|
||||
se_context_size: ${model.model_defaults.se_context_size}
|
||||
|
||||
As mentioned above, Citrinet uses the ``ConvASRDecoder`` as the decoder layer similar to QuartzNet. Only the configuration must be
|
||||
changed slightly as Citrinet utilizes sub-word tokenization.
|
||||
|
||||
.. note::
|
||||
The following information is relevant to any of the above models that implements its encoder as an :class:`~nemo.collections.asr.modules.conv_asr.ConvASREncoder`, and utilizes the ``SqueezeExcite`` mechanism.
|
||||
|
||||
The ``SqueezeExcite`` block within a :class:`~nemo.collections.asr.modules.conv_asr.ConvASREncoder` network can be modified to utilize a different context window after the model has been instantiated (even after the model has been trained) so as to evaluate the model with limited context. This can be achieved using the :meth:`~nemo.collections.asr.parts.mixins.mixins.ASRModuleMixin.change_conv_asr_se_context_window`
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Here, model can be any model that has a `ConvASREncoder` as its encoder, and utilized `SqueezeExcite` blocks
|
||||
# `context_window` : It is an integer representing the number of timeframes (each corresponding to some window stride).
|
||||
# `update_config` : Bool flag which determines whether the config of the model should be updated to reflect the new context window.
|
||||
|
||||
# Here, we specify that 128 timeframes of 0.01s stride should be the context window
|
||||
# This is equivalent to 128 * 0.01s context window for `SqueezeExcite`
|
||||
model.change_conv_asr_se_context_window(context_window=128, update_config=True)
|
||||
|
||||
Conformer-CTC
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
The config files for Conformer-CTC model contain character-based encoding and sub-word encoding at
|
||||
``<NeMo_git_root>/examples/asr/conf/conformer/conformer_ctc_char.yaml`` and ``<NeMo_git_root>/examples/asr/conf/conformer/conformer_ctc_bpe.yaml``
|
||||
respectively. Some components of the configs of `Conformer-CTC <./models.html#Conformer-CTC>`__ include the following datasets:
|
||||
|
||||
- ``train_ds``, ``validation_ds``, and ``test_ds``
|
||||
- opimizer (``optim``)
|
||||
- augmentation (``spec_augment``)
|
||||
- ``decoder``
|
||||
- ``trainer``
|
||||
``exp_manager``
|
||||
|
||||
These datasets are similar to other ASR models like `QuartzNet <./models.html#QuartzNet>`__. There should be a tokenizer section where you can
|
||||
specify the tokenizer if you want to use sub-word encoding instead of character-based encoding.
|
||||
|
||||
The encoder section includes the details about the Conformer-CTC encoder architecture. You may find more information in the
|
||||
config files and also :doc:`nemo.collections.asr.modules.ConformerEncoder<./api.html#nemo.collections.asr.modules.ConformerEncoder>`.
|
||||
|
||||
|
||||
Fine-tuning Configurations
|
||||
-------------------------
|
||||
|
||||
All ASR scripts support easy fine-tuning by partially/fully loading the pretrained weights from a checkpoint into the currently instantiated model. Pre-trained weights can be provided in multiple ways -
|
||||
|
||||
1) Providing a path to a NeMo model (via ``init_from_nemo_model``)
|
||||
2) Providing a name of a pretrained NeMo model (which will be downloaded via the cloud) (via ``init_from_pretrained_model``)
|
||||
3) Providing a path to a Pytorch Lightning checkpoint file (via ``init_from_ptl_ckpt``)
|
||||
|
||||
Fine-tuning via a NeMo model
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: sh
|
||||
|
||||
python examples/asr/script_to_<script_name>.py \
|
||||
--config-path=<path to dir of configs> \
|
||||
--config-name=<name of config without .yaml>) \
|
||||
model.train_ds.manifest_filepath="<path to manifest file>" \
|
||||
model.validation_ds.manifest_filepath="<path to manifest file>" \
|
||||
trainer.gpus=-1 \
|
||||
trainer.max_epochs=50 \
|
||||
+init_from_nemo_model="<path to .nemo model file>"
|
||||
|
||||
|
||||
Fine-tuning via a NeMo pretrained model name
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: sh
|
||||
|
||||
python examples/asr/script_to_<script_name>.py \
|
||||
--config-path=<path to dir of configs> \
|
||||
--config-name=<name of config without .yaml>) \
|
||||
model.train_ds.manifest_filepath="<path to manifest file>" \
|
||||
model.validation_ds.manifest_filepath="<path to manifest file>" \
|
||||
trainer.gpus=-1 \
|
||||
trainer.max_epochs=50 \
|
||||
+init_from_pretrained_model="<name of pretrained checkpoint>"
|
||||
|
||||
Fine-tuning via a Pytorch Lightning checkpoint
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: sh
|
||||
|
||||
python examples/asr/script_to_<script_name>.py \
|
||||
--config-path=<path to dir of configs> \
|
||||
--config-name=<name of config without .yaml>) \
|
||||
model.train_ds.manifest_filepath="<path to manifest file>" \
|
||||
model.validation_ds.manifest_filepath="<path to manifest file>" \
|
||||
trainer.gpus=-1 \
|
||||
trainer.max_epochs=50 \
|
||||
+init_from_ptl_ckpt="<name of pytorch lightning checkpoint>"
|
||||
@@ -0,0 +1,21 @@
|
||||
Model Name,Model Base Class,Model Card
|
||||
QuartzNet15x5Base-En,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels"
|
||||
stt_en_jasper10x5dr,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr"
|
||||
stt_en_citrinet_256,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256"
|
||||
stt_en_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512"
|
||||
stt_en_citrinet_1024,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024"
|
||||
stt_ca_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5"
|
||||
stt_it_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5"
|
||||
stt_fr_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_fr_quartznet15x5"
|
||||
stt_es_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5"
|
||||
stt_de_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5"
|
||||
stt_pl_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5"
|
||||
stt_ru_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5"
|
||||
stt_zh_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5"
|
||||
stt_zh_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512"
|
||||
stt_en_conformer_ctc_small,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small"
|
||||
stt_en_conformer_ctc_medium,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium"
|
||||
stt_en_conformer_ctc_large,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large"
|
||||
stt_en_conformer_ctc_small_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small_ls"
|
||||
stt_en_conformer_ctc_medium_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium_ls"
|
||||
stt_en_conformer_ctc_large_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls"
|
||||
|
@@ -0,0 +1,2 @@
|
||||
Model Name,Model Base Class,Model Card
|
||||
asrlm_en_transformer_large_ls,TransformerLMModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:asrlm_en_transformer_large_ls"
|
||||
|
@@ -0,0 +1,2 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_ca_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5"
|
||||
|
@@ -0,0 +1,3 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_de_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5"
|
||||
|
||||
|
@@ -0,0 +1,14 @@
|
||||
Model Name,Model Base Class,Model Card
|
||||
QuartzNet15x5Base-En,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels"
|
||||
stt_zh_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5"
|
||||
stt_en_jasper10x5dr,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr"
|
||||
stt_en_citrinet_256,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256"
|
||||
stt_en_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512"
|
||||
stt_en_citrinet_1024,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024"
|
||||
stt_zh_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512"
|
||||
stt_en_conformer_ctc_small,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small"
|
||||
stt_en_conformer_ctc_medium,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium"
|
||||
stt_en_conformer_ctc_large,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large"
|
||||
stt_en_conformer_ctc_small_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small_ls"
|
||||
stt_en_conformer_ctc_medium_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium_ls"
|
||||
stt_en_conformer_ctc_large_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls"
|
||||
|
@@ -0,0 +1,3 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_es_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5"
|
||||
|
||||
|
@@ -0,0 +1,3 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_fr_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_fr_quartznet15x5"
|
||||
|
||||
|
@@ -0,0 +1,3 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_it_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5"
|
||||
|
||||
|
@@ -0,0 +1,2 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_pl_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5"
|
||||
|
@@ -0,0 +1,3 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_ru_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5"
|
||||
|
||||
|
@@ -0,0 +1,3 @@
|
||||
Model,Model Base Class,Model Card
|
||||
stt_zh_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5"
|
||||
stt_zh_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512"
|
||||
|
@@ -0,0 +1,268 @@
|
||||
Datasets
|
||||
========
|
||||
|
||||
NeMo has scripts to convert several common ASR datasets into the format expected by the ``nemo_asr`` collection. You can get started
|
||||
with those datasets by following the instructions to run those scripts in the section appropriate to each dataset below.
|
||||
|
||||
If the user has their own data and want to preprocess it to use with NeMo ASR models, refer to the `Preparing Custom ASR Data`_ section.
|
||||
|
||||
If the user already has a dataset that you want to convert to a tarred format, refer to the `Tarred Datasets`_ section.
|
||||
|
||||
.. _LibriSpeech_dataset:
|
||||
|
||||
LibriSpeech
|
||||
-----------
|
||||
|
||||
Run the following scripts to download the LibriSpeech data and convert it into the format expected by `nemo_asr`. At least 250GB free
|
||||
space is required.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# install sox
|
||||
sudo apt-get install sox
|
||||
mkdir data
|
||||
python get_librispeech_data.py --data_root=data --data_set=ALL
|
||||
|
||||
After this, the ``data`` folder should contain wav files and ``.json`` manifests for NeMo ASR datalayer.
|
||||
|
||||
Each line is a training example. ``audio_filepath`` contains the path to the wav file, ``duration`` is the duration in seconds, and ``text`` is the transcript:
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{"audio_filepath": "<absolute_path_to>/1355-39947-0000.wav", "duration": 11.3, "text": "psychotherapy and the community both the physician and the patient find their place in the community the life interests of which are superior to the interests of the individual"}
|
||||
{"audio_filepath": "<absolute_path_to>/1355-39947-0001.wav", "duration": 15.905, "text": "it is an unavoidable question how far from the higher point of view of the social mind the psychotherapeutic efforts should be encouraged or suppressed are there any conditions which suggest suspicion of or direct opposition to such curative work"}
|
||||
|
||||
Fisher English Training Speech
|
||||
------------------------------
|
||||
|
||||
Run these scripts to convert the Fisher English Training Speech data into a format expected by the ``nemo_asr`` collection.
|
||||
|
||||
In brief, the following scripts convert the ``.sph`` files to ``.wav``, slices those files into smaller audio samples, matches the
|
||||
smaller slices with their corresponding transcripts, and splits the resulting audio segments into train, validation, and test sets
|
||||
(with one manifest each).
|
||||
|
||||
.. note::
|
||||
- 106 GB of space is required to run the ``.wav`` conversion
|
||||
- additional 105 GB is required for the slicing and matching
|
||||
- ``sph2pipe`` is required in order to run the ``.wav`` conversion
|
||||
|
||||
**Instructions**
|
||||
|
||||
The following scripts assume that you already have the Fisher dataset from the Linguistic Data Consortium, with a directory structure
|
||||
that looks similar to the following:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
FisherEnglishTrainingSpeech/
|
||||
├── LDC2004S13-Part1
|
||||
│ ├── fe_03_p1_transcripts
|
||||
│ ├── fisher_eng_tr_sp_d1
|
||||
│ ├── fisher_eng_tr_sp_d2
|
||||
│ ├── fisher_eng_tr_sp_d3
|
||||
│ └── ...
|
||||
└── LDC2005S13-Part2
|
||||
├── fe_03_p2_transcripts
|
||||
├── fe_03_p2_sph1
|
||||
├── fe_03_p2_sph2
|
||||
├── fe_03_p2_sph3
|
||||
└── ...
|
||||
|
||||
The transcripts that will be used are located in the ``fe_03_p<1,2>_transcripts/data/trans`` directory. The audio files (``.sph``)
|
||||
are located in the remaining directories in an ``audio`` subdirectory.
|
||||
|
||||
#. Convert the audio files from ``.sph`` to ``.wav`` by running:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
cd <nemo_root>/scripts
|
||||
python fisher_audio_to_wav.py \
|
||||
--data_root=<fisher_root> --dest_root=<conversion_target_dir>
|
||||
|
||||
This will place the unsliced ``.wav`` files in ``<conversion_target_dir>/LDC200[4,5]S13-Part[1,2]/audio-wav/``. It will take several
|
||||
minutes to run.
|
||||
|
||||
#. Process the transcripts and slice the audio data.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python process_fisher_data.py \
|
||||
--audio_root=<conversion_target_dir> --transcript_root=<fisher_root> \
|
||||
--dest_root=<processing_target_dir> \
|
||||
--remove_noises
|
||||
|
||||
This script splits the full dataset into train, validation, test sets, and places the audio slices in the corresponding folders
|
||||
in the destination directory. One manifest is written out per set, which includes each slice's transcript, duration, and path.
|
||||
|
||||
This will likely take around 20 minutes to run. Once finished, delete the 10 minute long ``.wav`` files.
|
||||
|
||||
2000 HUB5 English Evaluation Speech
|
||||
-----------------------------------
|
||||
|
||||
Run the following script to convert the HUB5 data into a format expected by the ``nemo_asr`` collection.
|
||||
|
||||
Similarly, to the Fisher dataset processing scripts, this script converts the ``.sph`` files to ``.wav``, slices the audio files and
|
||||
transcripts into utterances, and combines them into segments of some minimum length (default is 10 seconds). The resulting segments
|
||||
are all written out to an audio directory and the corresponding transcripts are written to a manifest JSON file.
|
||||
|
||||
.. note::
|
||||
- 5 GB of free space is required to run this script
|
||||
- ``sph2pipe`` is also required to be installed
|
||||
|
||||
This script assumes you already have the 2000 HUB5 dataset from the Linguistic Data Consortium.
|
||||
|
||||
Run the following command to process the 2000 HUB5 English Evaluation Speech samples:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python process_hub5_data.py \
|
||||
--data_root=<path_to_HUB5_data> \
|
||||
--dest_root=<target_dir>
|
||||
|
||||
You can optionally include ``--min_slice_duration=<num_seconds>`` if you would like to change the minimum audio segment duration.
|
||||
|
||||
AN4 Dataset
|
||||
-----------
|
||||
|
||||
This is a small dataset recorded and distributed by Carnegie Mellon University. It consists of recordings of people spelling out
|
||||
addresses, names, etc. Information about this dataset can be found on the `official CMU site <http://www.speech.cs.cmu.edu/databases/an4/>`_.
|
||||
|
||||
#. `Download and extract the dataset <http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz>`_ (which is labeled "NIST's Sphere audio (.sph) format (64M)".
|
||||
|
||||
#. Convert the ``.sph`` files to ``.wav`` using sox, and build one training and one test manifest.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python process_an4_data.py --data_root=<path_to_extracted_data>
|
||||
|
||||
After the script finishes, the ``train_manifest.json`` and ``test_manifest.json`` can be found in the ``<data_root>/an4/`` directory.
|
||||
|
||||
Aishell-1
|
||||
---------
|
||||
|
||||
To download the Aishell-1 data and convert it into a format expected by ``nemo_asr``, run:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# install sox
|
||||
sudo apt-get install sox
|
||||
mkdir data
|
||||
python get_aishell_data.py --data_root=data
|
||||
|
||||
After the script finishes, the ``data`` folder should contain a ``data_aishell`` folder which contains a wav file, a transcript folder, and related ``.json`` and ``vocab.txt`` files.
|
||||
|
||||
Aishell-2
|
||||
---------
|
||||
|
||||
To process the AIShell-2 dataset, in the command below, set the data folder of AIShell-2 using ``--audio_folder`` and where to push
|
||||
these files using ``--dest_folder``. In order to generate files in the supported format of ``nemo_asr``, run:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python process_aishell2_data.py --audio_folder=<data directory> --dest_folder=<destination directory>
|
||||
|
||||
After the script finishes, the ``train.json``, ``dev.json``, ``test.json``, and ``vocab.txt`` files can be found in the ``dest_folder`` directory.
|
||||
|
||||
Preparing Custom ASR Data
|
||||
-------------------------
|
||||
|
||||
The ``nemo_asr`` collection expects each dataset to consist of a set of utterances in individual audio files plus
|
||||
a manifest that describes the dataset, with information about one utterance per line (``.json``).
|
||||
The audio files can be of any format supported by `Pydub <https://github.com/jiaaro/pydub>`_, though we recommend
|
||||
WAV files as they are the default and have been most thoroughly tested.
|
||||
|
||||
There should be one manifest file per dataset that will be passed in, therefore, if the user wants separate training and validation
|
||||
datasets, they should also have separate manifests. Otherwise, thay will be loading validation data with their training data and vice
|
||||
versa.
|
||||
|
||||
Each line of the manifest should be in the following format:
|
||||
|
||||
.. code::
|
||||
|
||||
{"audio_filepath": "/path/to/audio.wav", "text": "the transcription of the utterance", "duration": 23.147}
|
||||
|
||||
The :code:`audio_filepath` field should provide an absolute path to the ``.wav`` file corresponding to the utterance.
|
||||
The :code:`text` field should contain the full transcript for the utterance, and the :code:`duration` field should
|
||||
reflect the duration of the utterance in seconds.
|
||||
|
||||
Each entry in the manifest (describing one audio file) should be bordered by '{' and '}' and must
|
||||
be contained on one line. The fields that describe the file should be separated by commas, and have the form :code:`"field_name": value`,
|
||||
as shown above. There should be no extra lines in the manifest, i.e. there should be exactly as many lines in the manifest as
|
||||
there are audio files in the dataset.
|
||||
|
||||
Since the manifest specifies the path for each utterance, the audio files do not have to be located
|
||||
in the same directory as the manifest, or even in any specific directory structure.
|
||||
|
||||
Once there is a manifest that describes each audio file in the dataset, use the dataset by passing
|
||||
in the manifest file path in the experiment config file, e.g. as ``training_ds.manifest_filepath=<path/to/manifest,json>``.
|
||||
|
||||
Tarred Datasets
|
||||
---------------
|
||||
|
||||
If experiments are run on a cluster with datasets stored on a distributed file system, the user will likely
|
||||
want to avoid constantly reading multiple small files and would prefer tarring their audio files.
|
||||
There are tarred versions of some NeMo ASR dataset classes for this case, such as the ``TarredAudioToCharDataset``
|
||||
(corresponding to the ``AudioToCharDataset``) and the ``TarredAudioToBPEDataset`` (corresponding to the
|
||||
``AudioToBPEDataset``). The tarred audio dataset classes in NeMo use `WebDataset <https://github.com/tmbdev/webdataset>`_.
|
||||
|
||||
To use an existing tarred dataset instead of a non-tarred dataset, set ``is_tarred: true`` in
|
||||
the experiment config file. Then, pass in the paths to all of the audio tarballs in ``tarred_audio_filepaths``, either as a list
|
||||
of filepaths, e.g. ``['/data/shard1.tar', '/data/shard2.tar']``, or in a single brace-expandable string, e.g.
|
||||
``'/data/shard_{1..64}.tar'`` or ``'/data/shard__OP_1..64_CL_'`` (recommended, see note below).
|
||||
|
||||
.. note::
|
||||
For brace expansion, there may be cases where ``{x..y}`` syntax cannot be used due to shell interference. This occurs most commonly
|
||||
inside SLURM scripts. Therefore, we provide a few equivalent replacements. Supported opening braces (equivalent to ``{``) are ``(``,
|
||||
``[``, ``<`` and the special tag ``_OP_``. Supported closing braces (equivalent to ``}``) are ``)``, ``]``, ``>`` and the special
|
||||
tag ``_CL_``. For SLURM based tasks, we suggest the use of the special tags for ease of use.
|
||||
|
||||
As with non-tarred datasets, the manifest file should be passed in ``manifest_filepath``. The dataloader assumes that the length
|
||||
of the manifest after filtering is the correct size of the dataset for reporting training progress.
|
||||
|
||||
The ``tarred_shard_strategy`` field of the config file can be set if you have multiple shards and are running an experiment with
|
||||
multiple workers. It defaults to ``scatter``, which preallocates a set of shards per worker which do not change during runtime.
|
||||
|
||||
For more information about the individual tarred datasets and the parameters available, including shuffling options,
|
||||
see the corresponding class APIs in the `Datasets <./api.html#Datasets>`__ section.
|
||||
|
||||
.. warning::
|
||||
If using multiple workers, the number of shards should be divisible by the world size to ensure an even
|
||||
split among workers. If it is not divisible, logging will give a warning but training will proceed, but likely hang at the last epoch.
|
||||
In addition, if using distributed processing, each shard must have the same number of entries after filtering is
|
||||
applied such that each worker ends up with the same number of files. We currently do not check for this in any dataloader, but the user's
|
||||
program may hang if the shards are uneven.
|
||||
|
||||
Conversion to Tarred Datasets
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
You can easily convert your existing NeMo-compatible ASR datasets using the
|
||||
`conversion script here <https://github.com/NVIDIA/NeMo/blob/v1.0.2/scripts/speech_recognition/convert_to_tarred_audio_dataset.py>`_.
|
||||
|
||||
.. code::
|
||||
|
||||
python convert_to_tarred_audio_dataset.py \
|
||||
--manifest_path=<path to the manifest file> \
|
||||
--target_dir=<path to output directory> \
|
||||
--num_shards=<number of tarfiles that will contain the audio>
|
||||
--max_duration=<float representing maximum duration of audio samples> \
|
||||
--min_duration=<float representing minimum duration of audio samples> \
|
||||
--shuffle --shuffle_seed=0
|
||||
|
||||
This script shuffles the entries in the given manifest (if ``--shuffle`` is set, which we recommend), filter
|
||||
audio files according to ``min_duration`` and ``max_duration``, and tar the remaining audio files to the directory
|
||||
``--target_dir`` in ``n`` shards, along with separate manifest and metadata files.
|
||||
|
||||
The files in the target directory should look similar to the following:
|
||||
|
||||
.. code::
|
||||
|
||||
target_dir/
|
||||
├── audio_1.tar
|
||||
├── audio_2.tar
|
||||
├── ...
|
||||
├── metadata.yaml
|
||||
└── tarred_audio_manifest.json
|
||||
|
||||
Note that file structures are flattened such that all audio files are at the top level in each tarball. This ensures that
|
||||
filenames are unique in the tarred dataset and the filepaths do not contain "-sub" and forward slashes in each ``audio_filepath`` are
|
||||
simply converted to underscores. For example, a manifest entry for ``/data/directory1/file.wav`` would be ``_data_directory1_file.wav``
|
||||
in the tarred dataset manifest, and ``/data/directory2/file.wav`` would be converted to ``_data_directory2_file.wav``.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 135 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 103 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 53 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 25 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 273 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 257 KiB |
@@ -0,0 +1,59 @@
|
||||
Automatic Speech Recognition (ASR)
|
||||
==================================
|
||||
|
||||
ASR, or Automatic Speech Recognition, refers to the problem of getting a program to automatically transcribe spoken language
|
||||
(speech-to-text). Our goal is usually to have a model that minimizes the Word Error Rate (WER) metric when transcribing speech input.
|
||||
In other words, given some audio file (e.g. a WAV file) containing speech, how do we transform this into the corresponding text with
|
||||
as few errors as possible?
|
||||
|
||||
Traditional speech recognition takes a generative approach, modeling the full pipeline of how speech sounds are produced in order to
|
||||
evaluate a speech sample. We would start from a language model that encapsulates the most likely orderings of words that are generated
|
||||
(e.g. an n-gram model), to a pronunciation model for each word in that ordering (e.g. a pronunciation table), to an acoustic model that
|
||||
translates those pronunciations to audio waveforms (e.g. a Gaussian Mixture Model).
|
||||
|
||||
Then, if we receive some spoken input, our goal would be to find the most likely sequence of text that would result in the given audio
|
||||
according to our generative pipeline of models. Overall, with traditional speech recognition, we try to model ``Pr(audio|transcript)*Pr(transcript)``,
|
||||
and take the argmax of this over possible transcripts.
|
||||
|
||||
Over time, neural nets advanced to the point where each component of the traditional speech recognition model could be replaced by a
|
||||
neural model that had better performance and that had a greater potential for generalization. For example, we could replace an n-gram
|
||||
model with a neural language model, and replace a pronunciation table with a neural pronunciation model, and so on. However, each of
|
||||
these neural models need to be trained individually on different tasks, and errors in any model in the pipeline could throw off the
|
||||
whole prediction.
|
||||
|
||||
Thus, we can see the appeal of end-to-end ASR architectures: discriminative models that simply take an audio input and give a textual
|
||||
output, and in which all components of the architecture are trained together towards the same goal. The model's encoder would be
|
||||
akin to an acoustic model for extracting speech features, which can then be directly piped to a decoder which outputs text. If desired,
|
||||
we could integrate a language model that would improve our predictions, as well.
|
||||
|
||||
And the entire end-to-end ASR model can be trained at once--a much easier pipeline to handle!
|
||||
|
||||
The full documentation tree is as follows:
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 8
|
||||
|
||||
models
|
||||
datasets
|
||||
asr_language_modeling
|
||||
results
|
||||
configs
|
||||
api
|
||||
|
||||
Resources and Documentation
|
||||
---------------------------
|
||||
|
||||
Hands-on speech recognition tutorial notebooks can be found under `the ASR tutorials folder <https://github.com/NVIDIA/NeMo/tree/v1.0.2/tutorials/asr/>`_.
|
||||
If you are a beginner to NeMo, consider trying out the `ASR with NeMo <https://github.com/NVIDIA/NeMo/tree/v1.0.2/tutorials/asr/01_ASR_with_NeMo.ipynb>`_ tutorial.
|
||||
This and most other tutorials can be run on Google Colab by specifying the link to the notebooks' GitHub pages on Colab.
|
||||
|
||||
If you are looking for information about a particular ASR model, or would like to find out more about the model
|
||||
architectures available in the `nemo_asr` collection, refer to the :doc:`Models <./models>` section.
|
||||
|
||||
NeMo includes preprocessing scripts for several common ASR datasets. The :doc:`Datasets <./datasets>` section contains instructions on
|
||||
running those scripts. It also includes guidance for creating your own NeMo-compatible dataset, if you have your own data.
|
||||
|
||||
Information about how to load model checkpoints (either local files or pretrained ones from NGC), as well as a list of the checkpoints
|
||||
available on NGC are located on the :doc:`Checkpoints <./results>` section.
|
||||
|
||||
Documentation regarding the configuration files specific to the ``nemo_asr`` models can be found on the :doc:`Configuration Files <./configs>` section.
|
||||
@@ -0,0 +1,97 @@
|
||||
Models
|
||||
======
|
||||
|
||||
This section gives a brief overview of the models that NeMo's ASR collection currently supports.
|
||||
|
||||
Each of these models can be used with the example ASR scripts (in the ``<NeMo_git_root>/examples/asr`` directory) by
|
||||
specifying the model architecture in the config file used. Examples of config files for each model can be found in
|
||||
the ``<NeMo_git_root>/examples/asr/conf`` directory.
|
||||
|
||||
For more information about the config files and how they should be structured, refer to the :doc:`./configs` section.
|
||||
|
||||
Pretrained checkpoints for all of these models, as well as instructions on how to load them, can be found in the :doc:`./results`
|
||||
section. You can use the available checkpoints for immediate inference, or fine-tune them on your own datasets. The checkpoints section
|
||||
also contains benchmark results for the available ASR models.
|
||||
|
||||
.. _Jasper_model:
|
||||
|
||||
Jasper
|
||||
------
|
||||
|
||||
Jasper ("Just Another Speech Recognizer") :cite:`asr-models-li2019jasper` is a deep time delay neural network (TDNN) comprising of
|
||||
blocks of 1D-convolutional layers. The Jasper family of models are denoted as ``Jasper_[BxR]`` where ``B`` is the number of blocks
|
||||
and ``R`` is the number of convolutional sub-blocks within a block. Each sub-block contains a 1-D convolution, batch normalization,
|
||||
ReLU, and dropout:
|
||||
|
||||
.. image:: images/jasper_vertical.png
|
||||
:align: center
|
||||
:alt: japer model
|
||||
:scale: 50%
|
||||
|
||||
Jasper models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModel` class.
|
||||
|
||||
QuartzNet
|
||||
---------
|
||||
|
||||
QuartzNet :cite:`asr-models-kriman2019quartznet` is a version of Jasper :cite:`asr-models-li2019jasper` model with separable
|
||||
convolutions and larger filters. It can achieve performance similar to Jasper but with an order of magnitude fewer parameters.
|
||||
Similarly to Jasper, the QuartzNet family of models are denoted as ``QuartzNet_[BxR]`` where ``B`` is the number of blocks and ``R``
|
||||
is the number of convolutional sub-blocks within a block. Each sub-block contains a 1-D *separable* convolution, batch normalization,
|
||||
ReLU, and dropout:
|
||||
|
||||
.. image:: images/quartz_vertical.png
|
||||
:align: center
|
||||
:alt: quartznet model
|
||||
:scale: 40%
|
||||
|
||||
QuartzNet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModel` class.
|
||||
|
||||
Citrinet
|
||||
--------
|
||||
|
||||
Citrinet is a version of QuartzNet :cite:`asr-models-kriman2019quartznet` that extends ContextNet :cite:`asr-models-han2020contextnet`,
|
||||
utilizing subword encoding (via Word Piece tokenization) and Squeeze-and-Excitation mechanism :cite:`asr-models-hu2018squeeze` to
|
||||
obtain highly accurate audio transcripts while utilizing a non-autoregressive CTC based decoding scheme for efficient inference.
|
||||
|
||||
.. image:: images/citrinet_vertical.png
|
||||
:align: center
|
||||
:alt: citrinet model
|
||||
:scale: 50%
|
||||
|
||||
Citrinet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModelBPE` class.
|
||||
|
||||
.. _Conformer-CTC_model:
|
||||
|
||||
Conformer-CTC
|
||||
-------------
|
||||
|
||||
Conformer-CTC is a CTC-based variant of the Conformer model introduced in :cite:`asr-models-gulati2020conformer`. Conformer-CTC has a
|
||||
similar encoder as the original Conformer but uses CTC loss and decoding instead of RNNT loss, which makes it a non-autoregressive model.
|
||||
We also drop the LSTM decoder and instead use a linear decoder on the top of the encoder. This model uses the combination of
|
||||
self-attention and convolution modules to achieve the best of the two approaches, the self-attention layers can learn the global
|
||||
interaction while the convolutions efficiently capture the local correlations. The self-attention modules support both regular
|
||||
self-attention with absolute positional encoding, and also Transformer-XL's self-attention with relative positional encodings.
|
||||
|
||||
Here is the overall architecture of the encoder of Conformer-CTC:
|
||||
|
||||
.. image:: images/conformer_ctc.png
|
||||
:align: center
|
||||
:alt: Conformer-CTC Model
|
||||
:scale: 50%
|
||||
|
||||
This model supports both the sub-word level and character level encodings. You can find more details on the config files for the
|
||||
Conformer-CTC models at `Conformer-CTC <./configs.html#conformer-ctc>`. The variant with sub-word encoding is a BPE-based model
|
||||
which can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModelBPE` class, while the
|
||||
character-based variant is based on :class:`~nemo.collections.asr.models.EncDecCTCModel`.
|
||||
|
||||
You may find the example config files of Conformer-CTC model with character-based encoding at
|
||||
``<NeMo_git_root>/examples/asr/conf/conformer/conformer_ctc_char.yaml`` and
|
||||
with sub-word encoding at ``<NeMo_git_root>/examples/asr/conf/conformer/conformer_ctc_bpe.yaml``.
|
||||
|
||||
References
|
||||
----------
|
||||
|
||||
.. bibliography:: asr_all.bib
|
||||
:style: plain
|
||||
:labelprefix: ASR-MODELS
|
||||
:keyprefix: asr-models-
|
||||
@@ -0,0 +1,181 @@
|
||||
Checkpoints
|
||||
===========
|
||||
|
||||
There are two main ways to load pretrained checkpoints in NeMo:
|
||||
|
||||
* Using the :code:`restore_from()` method to load a local checkpoint file (``.nemo``), or
|
||||
* Using the :code:`from_pretrained()` method to download and set up a checkpoint from NGC.
|
||||
|
||||
Refer to the following sections for instructions and examples for each.
|
||||
|
||||
Note that these instructions are for loading fully trained checkpoints for evaluation or fine-tuning. For resuming an unfinished
|
||||
training experiment, use the Experiment Manager to do so by setting the ``resume_if_exists`` flag to ``True``.
|
||||
|
||||
Loading Local Checkpoints
|
||||
-------------------------
|
||||
|
||||
NeMo automatically saves checkpoints of a model that is trained in a ``.nemo`` format. Alternatively, to manually save the model at any
|
||||
point, issue :code:`model.save_to(<checkpoint_path>.nemo)`.
|
||||
|
||||
If there is a local ``.nemo`` checkpoint that you'd like to load, use the :code:`restore_from()` method:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import nemo.collections.asr as nemo_asr
|
||||
model = nemo_asr.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>")
|
||||
|
||||
Where the model base class is the ASR model class of the original checkpoint, or the general ``ASRModel`` class.
|
||||
|
||||
NGC Pretrained Checkpoints
|
||||
--------------------------
|
||||
|
||||
The ASR collection has checkpoints of several models trained on various datasets for a variety of tasks. These checkpoints are
|
||||
obtainable via NGC `NeMo Automatic Speech Recognition collection <https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels>`_.
|
||||
The model cards on NGC contain more information about each of the checkpoints available.
|
||||
|
||||
The tables below list the ASR models available from NGC. The models can be accessed via the :code:`from_pretrained()` method inside
|
||||
the ASR Model class. In general, you can load any of these models with code in the following format:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import nemo.collections.asr as nemo_asr
|
||||
model = nemo_asr.models.ASRModel.from_pretrained(model_name="<MODEL_NAME>")
|
||||
|
||||
Where the model name is the value under "Model Name" entry in the tables below.
|
||||
|
||||
For example, to load the base English QuartzNet model for speech recognition, run:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model = nemo_asr.models.ASRModel.from_pretrained(model_name="QuartzNet15x5Base-En")
|
||||
|
||||
You can also call :code:`from_pretrained()` from the specific model class (such as :code:`EncDecCTCModel`
|
||||
for QuartzNet) if you need to access a specific model functionality.
|
||||
|
||||
If you would like to programatically list the models available for a particular base class, you can use the
|
||||
:code:`list_available_models()` method.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
nemo_asr.models.<MODEL_BASE_CLASS>.list_available_models()
|
||||
|
||||
Transcribing/Inference
|
||||
^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
To perform inference and transcribe a sample of speech after loading the model, use the ``transcribe()`` method:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model.transcribe(paths2audio_files=[list of audio files], batch_size=BATCH_SIZE, logprobs=False)
|
||||
|
||||
Setting the argument ``logprobs`` to ``True`` returns the log probabilities instead of transcriptions. For more information, see :doc:`./api.html#modules`.
|
||||
The audio files should be 16KHz monochannel wav files.
|
||||
|
||||
Fine-tuning on Different Datasets
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
There are multiple ASR tutorials provided in the :ref:`Tutorials <tutorials>` section. Most of these tutorials explain how to instantiate a pre-trained model, prepare the model for fine-tuning on some dataset (in the same language) as a demonstration.
|
||||
|
||||
|
||||
Automatic Speech Recognition Models
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
.. csv-table::
|
||||
:file: data/asr_results.csv
|
||||
:align: left
|
||||
:widths: 30, 30, 40
|
||||
:header-rows: 1
|
||||
|
||||
Language Models for ASR
|
||||
^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
.. csv-table::
|
||||
:file: data/asrlm_results.csv
|
||||
:align: left
|
||||
:widths: 30, 30, 40
|
||||
:header-rows: 1
|
||||
|
||||
|
||||
Speech Recognition (Languages)
|
||||
------------------------------
|
||||
|
||||
English
|
||||
^^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_en.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
-----------------------------
|
||||
|
||||
Mandarin
|
||||
^^^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_zh.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
-----------------------------
|
||||
|
||||
German
|
||||
^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_de.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
-----------------------------
|
||||
|
||||
Polish
|
||||
^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_pl.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
-----------------------------
|
||||
|
||||
Italian
|
||||
^^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_it.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
-----------------------------
|
||||
|
||||
Russian
|
||||
^^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_ru.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
-----------------------------
|
||||
|
||||
Spanish
|
||||
^^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_es.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
|
||||
-----------------------------
|
||||
|
||||
Catalan
|
||||
^^^^^^^
|
||||
.. csv-table::
|
||||
:file: data/benchmark_ca.csv
|
||||
:align: left
|
||||
:widths: 40, 10, 50
|
||||
:header-rows: 1
|
||||
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
NeMo Speaker Diarization API
|
||||
=============================
|
||||
|
||||
|
||||
Model Classes
|
||||
-------------
|
||||
.. autoclass:: nemo.collections.asr.models.ClusteringDiarizer
|
||||
:show-inheritance:
|
||||
:members: init_speaker_model, set_vad_model, init_vad_model, setup_vad_test_data, setup_spkr_test_data, run_vad, extract_embeddings, path2audio_files_to_manifest, diarize, make_nemo_file_from_folder, save_to, unpack_nemo_file, restore_from
|
||||
|
||||
|
||||
Mixins
|
||||
------
|
||||
|
||||
.. autoclass:: nemo.collections.asr.parts.mixins.mixins.DiarizationMixin
|
||||
:show-inheritance:
|
||||
:members:
|
||||
@@ -0,0 +1,65 @@
|
||||
NeMo Speaker Diarization Configuration Files
|
||||
============================================
|
||||
|
||||
Since speaker diarization model here is not a fully-trainble End-to-End model but an inference pipeline, we use **diarizer** instead of **model** which is used in other tasks.
|
||||
|
||||
The diarizer section will generally require information about the dataset(s) being
|
||||
used, models used in this pipline, as well as inference related parameters such as postprocessing of each models.
|
||||
The sections on this page cover each of these in more detail.
|
||||
|
||||
|
||||
Example configuration files for speaker diarization can be found in the
|
||||
``<NeMo_git_root>/examples/speaker_recognition/conf/speaker_diarization.yaml>``
|
||||
|
||||
.. note::
|
||||
For model details and deep understanding about configs, finetuning, tuning threshold, and evaluation,
|
||||
please refer to ``<NeMo_git_root>/tutorials/speaker_recognition/Speaker_Diarization_Inference.ipynb``;
|
||||
for other applications such as possible integration with ASR, have a look at ``<NeMo_git_root>/tutorials/speaker_recognition/ASR_with_SpeakerDiarization.ipynb``.
|
||||
|
||||
|
||||
Dataset Configuration
|
||||
-----------------------
|
||||
|
||||
In contrast to other ASR related tasks or models in NeMo, speaker diarization supported in NeMo is a modular inference pipeline.
|
||||
Datasets here denotes the data you would like to perform speaker diarization on.
|
||||
|
||||
An example Speaker Diarization dataset configuration could look like:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
diarizer:
|
||||
oracle_num_speakers: ??? # if number of speakers is known input it or pass null if not known. Accepts int or path to file containing uniq-id and num of speakers of that session
|
||||
max_num_speakers: 8 # max number of speakers for each recording. If oracle num speakers is passed this value is ignored
|
||||
out_dir: ???
|
||||
paths2audio_files: ??? # file containing paths to audio files for diarzation
|
||||
path2groundtruth_rttm_files: null # (Optional) file containing paths to rttm files (this can be passed if we need to calculate DER rate based on ground truth rttm files).
|
||||
...
|
||||
|
||||
.. note::
|
||||
We expect audio and the corresponding RTTM to have the same base name and the name should be unique.
|
||||
|
||||
|
||||
Diarizer Architecture Configurations
|
||||
-------------------------------------
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
diarizer:
|
||||
...
|
||||
vad:
|
||||
model_path: null #.nemo local model path or pretrained model name or none
|
||||
window_length_in_sec: 0.15
|
||||
shift_length_in_sec: 0.01
|
||||
threshold: 0.5 # tune threshold on dev set. Check <NeMo_git_root>/scripts/voice_activity_detection/vad_tune_threshold.py
|
||||
vad_decision_smoothing: True
|
||||
smoothing_params:
|
||||
method: "median"
|
||||
overlap: 0.875
|
||||
|
||||
speaker_embeddings:
|
||||
oracle_vad_manifest: null # leave it null if to perform diarization with above VAD model else path to manifest file genrerated as shown in Datasets section
|
||||
model_path: ??? #.nemo local model path or pretrained model name
|
||||
window_length_in_sec: 1.5 # for better performace try with 3 sec and 1.5 sec
|
||||
shift_length_in_sec: 0.75 # for better performace try with 1.5 sec and 0.75 sec
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user