115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
#encoding = utf8
|
|
import queue
|
|
import time
|
|
from queue import Queue
|
|
from threading import Event, Thread
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from .audio_handler import AudioHandler
|
|
from utils import load_model, mirror_index, get_device
|
|
|
|
|
|
class AudioInferenceHandler(AudioHandler):
|
|
def __init__(self, context, handler):
|
|
super().__init__(context, handler)
|
|
|
|
self._mal_queue = Queue()
|
|
self._audio_queue = Queue()
|
|
|
|
self._exit_event = Event()
|
|
self._run_thread = Thread(target=self.__on_run)
|
|
self._exit_event.set()
|
|
self._run_thread.start()
|
|
|
|
def on_handle(self, stream, type_):
|
|
if type_ == 1:
|
|
self._mal_queue.put(stream)
|
|
elif type_ == 0:
|
|
self._audio_queue.put(stream)
|
|
|
|
def __on_run(self):
|
|
model = load_model(r'.\checkpoints\wav2lip.pth')
|
|
print("Model loaded")
|
|
|
|
face_list_cycle = self._context.face_list_cycle
|
|
|
|
length = len(face_list_cycle)
|
|
index = 0
|
|
count = 0
|
|
count_time = 0
|
|
print('start inference')
|
|
|
|
device = get_device()
|
|
print(f'use device:{device}')
|
|
|
|
while True:
|
|
if self._exit_event.is_set():
|
|
start_time = time.perf_counter()
|
|
batch_size = self._context.batch_size
|
|
try:
|
|
mel_batch = self._mal_queue.get(block=True, timeout=0.1)
|
|
except queue.Empty:
|
|
continue
|
|
|
|
is_all_silence = True
|
|
audio_frames = []
|
|
for _ in range(batch_size * 2):
|
|
frame, type_ = self._audio_queue.get()
|
|
audio_frames.append((frame, type_))
|
|
if type_ == 0:
|
|
is_all_silence = False
|
|
if is_all_silence:
|
|
for i in range(batch_size):
|
|
self.on_next_handle((None, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
|
|
0)
|
|
index = index + 1
|
|
else:
|
|
print('infer=======')
|
|
t = time.perf_counter()
|
|
img_batch = []
|
|
for i in range(batch_size):
|
|
idx = mirror_index(length, index + i)
|
|
face = face_list_cycle[idx]
|
|
img_batch.append(face)
|
|
|
|
img_batch = np.asarray(img_batch)
|
|
mel_batch = np.asarray(mel_batch)
|
|
img_masked = img_batch.copy()
|
|
img_masked[:, face.shape[0] // 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])
|
|
|
|
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.
|
|
|
|
count_time += (time.perf_counter() - t)
|
|
count += batch_size
|
|
|
|
if count >= 100:
|
|
print(f"------actual avg infer fps:{count / count_time:.4f}")
|
|
count = 0
|
|
count_time = 0
|
|
|
|
image_index = 0
|
|
for i, res_frame in enumerate(pred):
|
|
self.on_next_handle(
|
|
(res_frame, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
|
|
0)
|
|
index = index + 1
|
|
image_index = image_index + 1
|
|
print('batch count', image_index)
|
|
print('total batch time:', time.perf_counter() - start_time)
|
|
else:
|
|
time.sleep(1)
|
|
break
|
|
print('musereal inference processor stop')
|