116 lines
5.1 KiB
Python
116 lines
5.1 KiB
Python
import os.path
|
|
import torch
|
|
import os
|
|
import numpy as np
|
|
import time
|
|
from talkingface.run_utils import smooth_array, video_pts_process
|
|
from talkingface.run_utils import mouth_replace, prepare_video_data
|
|
from talkingface.utils import generate_face_mask, INDEX_LIPS_OUTER
|
|
from talkingface.data.few_shot_dataset import select_ref_index,get_ref_images_fromVideo,generate_input, generate_input_pixels
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
import pickle
|
|
import cv2
|
|
|
|
|
|
face_mask = generate_face_mask()
|
|
|
|
|
|
class RenderModel:
|
|
def __init__(self):
|
|
self.__net = None
|
|
|
|
self.__pts_driven = None
|
|
self.__mat_list = None
|
|
self.__pts_normalized_list = None
|
|
self.__face_mask_pts = None
|
|
self.__ref_img = None
|
|
self.__cap_input = None
|
|
self.frame_index = 0
|
|
self.__mouth_coords_array = None
|
|
|
|
def loadModel(self, ckpt_path):
|
|
from talkingface.models.DINet import DINet_five_Ref as DINet
|
|
n_ref = 5
|
|
source_channel = 6
|
|
ref_channel = n_ref * 6
|
|
self.__net = DINet(source_channel, ref_channel).cuda()
|
|
checkpoint = torch.load(ckpt_path)
|
|
self.__net.load_state_dict(checkpoint)
|
|
self.__net.eval()
|
|
|
|
def reset_charactor(self, video_path, Path_pkl, ref_img_index_list = None):
|
|
if self.__cap_input is not None:
|
|
self.__cap_input.release()
|
|
|
|
self.__pts_driven, self.__mat_list,self.__pts_normalized_list, self.__face_mask_pts, self.__ref_img, self.__cap_input = \
|
|
prepare_video_data(video_path, Path_pkl, ref_img_index_list)
|
|
|
|
ref_tensor = torch.from_numpy(self.__ref_img / 255.).float().permute(2, 0, 1).unsqueeze(0).cuda()
|
|
self.__net.ref_input(ref_tensor)
|
|
|
|
x_min, x_max = np.min(self.__pts_normalized_list[:, INDEX_LIPS_OUTER, 0]), np.max(self.__pts_normalized_list[:, INDEX_LIPS_OUTER, 0])
|
|
y_min, y_max = np.min(self.__pts_normalized_list[:, INDEX_LIPS_OUTER, 1]), np.max(self.__pts_normalized_list[:, INDEX_LIPS_OUTER, 1])
|
|
z_min, z_max = np.min(self.__pts_normalized_list[:, INDEX_LIPS_OUTER, 2]), np.max(self.__pts_normalized_list[:, INDEX_LIPS_OUTER, 2])
|
|
|
|
x_mid,y_mid,z_mid = (x_min + x_max)/2, (y_min + y_max)/2, (z_min + z_max)/2
|
|
x_len, y_len, z_len = (x_max - x_min)/2, (y_max - y_min)/2, (z_max - z_min)/2
|
|
x_min, x_max = x_mid - x_len*0.9, x_mid + x_len*0.9
|
|
y_min, y_max = y_mid - y_len*0.9, y_mid + y_len*0.9
|
|
z_min, z_max = z_mid - z_len*0.9, z_mid + z_len*0.9
|
|
|
|
# print(face_personal.shape, x_min, x_max, y_min, y_max, z_min, z_max)
|
|
coords_array = np.zeros([100, 150, 4])
|
|
for i in range(100):
|
|
for j in range(150):
|
|
coords_array[i, j, 0] = j/149
|
|
coords_array[i, j, 1] = i/100
|
|
# coords_array[i, j, 2] = int((-75 + abs(j - 75))*(2./3))
|
|
coords_array[i, j, 2] = ((j - 75)/ 75) ** 2
|
|
coords_array[i, j, 3] = 1
|
|
|
|
coords_array = coords_array*np.array([x_max - x_min, y_max - y_min, z_max - z_min, 1]) + np.array([x_min, y_min, z_min, 0])
|
|
self.__mouth_coords_array = coords_array.reshape(-1, 4).transpose(1, 0)
|
|
|
|
|
|
|
|
def interface(self, mouth_frame):
|
|
vid_frame_count = self.__cap_input.get(cv2.CAP_PROP_FRAME_COUNT)
|
|
if self.frame_index % vid_frame_count == 0:
|
|
self.__cap_input.set(cv2.CAP_PROP_POS_FRAMES, 0) # 设置要获取的帧号
|
|
ret, frame = self.__cap_input.read() # 按帧读取视频
|
|
|
|
epoch = self.frame_index // len(self.__mat_list)
|
|
if epoch % 2 == 0:
|
|
new_index = self.frame_index % len(self.__mat_list)
|
|
else:
|
|
new_index = -1 - self.frame_index % len(self.__mat_list)
|
|
|
|
# print(self.__face_mask_pts.shape, "ssssssss")
|
|
source_img, target_img, crop_coords = generate_input_pixels(frame, self.__pts_driven[new_index], self.__mat_list[new_index],
|
|
mouth_frame, self.__face_mask_pts[new_index],
|
|
self.__mouth_coords_array)
|
|
|
|
# tensor
|
|
source_tensor = torch.from_numpy(source_img / 255.).float().permute(2, 0, 1).unsqueeze(0).cuda()
|
|
target_tensor = torch.from_numpy(target_img / 255.).float().permute(2, 0, 1).unsqueeze(0).cuda()
|
|
|
|
source_tensor, source_prompt_tensor = source_tensor[:, :3], source_tensor[:, 3:]
|
|
fake_out = self.__net.interface(source_tensor, source_prompt_tensor)
|
|
|
|
image_numpy = fake_out.detach().squeeze(0).cpu().float().numpy()
|
|
image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
|
|
image_numpy = image_numpy.clip(0, 255)
|
|
image_numpy = image_numpy.astype(np.uint8)
|
|
|
|
image_numpy = target_img * face_mask + image_numpy * (1 - face_mask)
|
|
|
|
img_bg = frame
|
|
x_min, y_min, x_max, y_max = crop_coords
|
|
|
|
img_face = cv2.resize(image_numpy, (x_max - x_min, y_max - y_min))
|
|
img_bg[y_min:y_max, x_min:x_max] = img_face
|
|
self.frame_index += 1
|
|
return img_bg
|
|
|
|
def save(self, path):
|
|
torch.save(self.__net.state_dict(), path) |