|
from typing import Iterable, Iterator, List, Tuple |
|
|
|
import cv2 |
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
from omegaconf import DictConfig |
|
from tqdm import tqdm |
|
|
|
from config import hparams as hp |
|
from nota_wav2lip.models.util import count_params, load_model |
|
|
|
|
|
class Wav2LipInferenceImpl: |
|
def __init__(self, model_name: str, hp_inference_model: DictConfig, device='cpu'): |
|
self.model: nn.Module = load_model( |
|
model_name, |
|
device=device, |
|
**hp_inference_model |
|
) |
|
self.device = device |
|
self._params: str = self._format_param(count_params(self.model)) |
|
|
|
@property |
|
def params(self): |
|
return self._params |
|
|
|
@staticmethod |
|
def _format_param(num_params: int) -> str: |
|
params_in_million = num_params / 1e6 |
|
return f"{params_in_million:.1f}M" |
|
|
|
@staticmethod |
|
def _reset_batch() -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], List[List[int]]]: |
|
return [], [], [], [] |
|
|
|
def get_data_iterator( |
|
self, |
|
audio_iterable: Iterable[np.ndarray], |
|
video_iterable: List[Tuple[np.ndarray, List[int]]] |
|
) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray, List[int]]]: |
|
img_batch, mel_batch, frame_batch, coords_batch = self._reset_batch() |
|
|
|
for i, m in enumerate(audio_iterable): |
|
idx = i % len(video_iterable) |
|
_frame_to_save, coords = video_iterable[idx] |
|
frame_to_save = _frame_to_save.copy() |
|
face = frame_to_save[coords[0]:coords[1], coords[2]:coords[3]].copy() |
|
|
|
face: np.ndarray = cv2.resize(face, (hp.face.img_size, hp.face.img_size)) |
|
|
|
img_batch.append(face) |
|
mel_batch.append(m) |
|
frame_batch.append(frame_to_save) |
|
coords_batch.append(coords) |
|
|
|
if len(img_batch) >= hp.inference.batch_size: |
|
img_batch = np.asarray(img_batch) |
|
mel_batch = np.asarray(mel_batch) |
|
|
|
img_masked = img_batch.copy() |
|
img_masked[:, hp.face.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 = self._reset_batch() |
|
|
|
if len(img_batch) > 0: |
|
img_batch = np.asarray(img_batch) |
|
mel_batch = np.asarray(mel_batch) |
|
|
|
img_masked = img_batch.copy() |
|
img_masked[:, hp.face.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 |
|
|
|
@torch.no_grad() |
|
def inference_with_iterator( |
|
self, |
|
audio_iterable: Iterable[np.ndarray], |
|
video_iterable: List[Tuple[np.ndarray, List[int]]] |
|
) -> Iterator[np.ndarray]: |
|
data_iterator = self.get_data_iterator(audio_iterable, video_iterable) |
|
|
|
for (img_batch, mel_batch, frames, coords) in \ |
|
tqdm(data_iterator, total=int(np.ceil(float(len(audio_iterable)) / hp.inference.batch_size))): |
|
|
|
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(self.device) |
|
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(self.device) |
|
|
|
preds: torch.Tensor = self.forward(mel_batch, img_batch) |
|
|
|
preds = preds.cpu().numpy().transpose(0, 2, 3, 1) * 255. |
|
for pred, frame, coord in zip(preds, frames, coords): |
|
y1, y2, x1, x2 = coord |
|
pred = cv2.resize(pred.astype(np.uint8), (x2 - x1, y2 - y1)) |
|
|
|
frame[y1:y2, x1:x2] = pred |
|
yield frame |
|
|
|
@torch.no_grad() |
|
def forward(self, audio_sequences: torch.Tensor, face_sequences: torch.Tensor) -> torch.Tensor: |
|
return self.model(audio_sequences, face_sequences) |
|
|
|
def __call__(self, *args, **kwargs): |
|
return self.forward(*args, **kwargs) |
|
|
