313 lines
9.9 KiB
Python
313 lines
9.9 KiB
Python
#encoding = utf8
|
|
import copy
|
|
import glob
|
|
import logging
|
|
import os
|
|
import pickle
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from tqdm import tqdm
|
|
from PIL import Image
|
|
|
|
import face_detection
|
|
from models import Wav2Lip, Wav2LipV2
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def mirror_index(size, index):
|
|
# size = len(self.coord_list_cycle)
|
|
turn = index // size
|
|
res = index % size
|
|
if turn % 2 == 0:
|
|
return res
|
|
else:
|
|
return size - res - 1
|
|
|
|
|
|
def read_image(path):
|
|
image = Image.open(path)
|
|
return image
|
|
|
|
|
|
def read_images(img_list):
|
|
frames = []
|
|
print('reading images...')
|
|
for img_path in tqdm(img_list):
|
|
# frame = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
|
|
frame = Image.open(img_path)
|
|
frame = np.array(frame)
|
|
frames.append(frame)
|
|
return frames
|
|
|
|
|
|
def read_files_path(path):
|
|
file_paths = []
|
|
files = os.listdir(path)
|
|
for file in files:
|
|
if not os.path.isdir(file) and file.endswith('.png') or file.endswith('.jpg'):
|
|
file_paths.append(os.path.join(path, file))
|
|
return file_paths
|
|
|
|
|
|
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 datagen_signal(frame, mel, face_det_results, img_size, wav2lip_batch_size=128):
|
|
img_batch, mel_batch, frame_batch, coord_batch = [], [], [], []
|
|
|
|
idx = 0
|
|
frame_to_save = frame.copy()
|
|
face, coord = face_det_results[idx].copy()
|
|
|
|
face = cv2.resize(face, (img_size, img_size))
|
|
|
|
for i, m in enumerate(mel):
|
|
img_batch.append(face)
|
|
mel_batch.append(m)
|
|
frame_batch.append(frame_to_save)
|
|
coord_batch.append(coord)
|
|
|
|
if len(img_batch) >= wav2lip_batch_size:
|
|
img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
|
|
|
|
img_masked = img_batch.copy()
|
|
img_masked[:, 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])
|
|
|
|
return img_batch, mel_batch, frame_batch, coord_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[:, 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])
|
|
|
|
return img_batch, mel_batch, frame_batch, coord_batch
|
|
|
|
|
|
def face_detect(images, device):
|
|
detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
|
|
flip_input=False, device=device)
|
|
batch_size = 16
|
|
|
|
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. Please use the --resize_factor argument')
|
|
batch_size //= 2
|
|
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
|
continue
|
|
break
|
|
|
|
results = []
|
|
pad_y1, pad_y2, pad_x1, pad_x2 = [0, 10, 0, 0]
|
|
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] - pad_y1)
|
|
y2 = min(image.shape[0], rect[3] + pad_y2)
|
|
x1 = max(0, rect[0] - pad_x1)
|
|
x2 = min(image.shape[1], rect[2] + pad_x2)
|
|
|
|
results.append([x1, y1, x2, y2])
|
|
|
|
boxes = np.array(results)
|
|
if not False:
|
|
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 get_device():
|
|
return 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
|
|
def _load(checkpoint_path):
|
|
device = get_device()
|
|
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 = Wav2LipV2()
|
|
print("Load checkpoint from: {}".format(path))
|
|
logging.info(f'Load checkpoint from {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)
|
|
device = get_device()
|
|
model = model.to(device)
|
|
return model.eval()
|
|
|
|
|
|
def load_avatar(path, img_size, device):
|
|
print(f'load avatar:{path}')
|
|
face_images_path = os.path.join(path, 'face')
|
|
face_images_path = read_files_path(face_images_path)
|
|
full_list_cycle = read_images(face_images_path)
|
|
|
|
face_det_results = face_detect(full_list_cycle, device)
|
|
|
|
face_frames = []
|
|
coord_frames = []
|
|
for face, coord in face_det_results:
|
|
resized_crop_frame = cv2.resize(face[:, :, :3], (img_size, img_size))
|
|
face_frames.append(resized_crop_frame)
|
|
coord_frames.append(coord)
|
|
|
|
return full_list_cycle, face_frames, coord_frames
|
|
|
|
|
|
def load_avatar_from_processed(base_path, avatar_name):
|
|
avatar_path = os.path.join(base_path, 'data', 'avatars', avatar_name)
|
|
print(f'load avatar from processed:{avatar_path}')
|
|
coord_path = os.path.join(avatar_path, 'coords.pkl')
|
|
print(f'load avatar_path from processed:{avatar_path}')
|
|
face_image_path = os.path.join(avatar_path, 'face_imgs')
|
|
print(f'load face_image_path from processed:{face_image_path}')
|
|
full_image_path = os.path.join(avatar_path, 'full_imgs')
|
|
print(f'load full_image_path from processed:{full_image_path}')
|
|
|
|
with open(coord_path, 'rb') as f:
|
|
coord_list_frames = pickle.load(f)
|
|
|
|
face_image_list = glob.glob(os.path.join(face_image_path, '*.[jpJP][pnPN]*[gG]'))
|
|
face_image_list = sorted(face_image_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
|
|
face_list_cycle = read_images(face_image_list)
|
|
|
|
full_image_list = glob.glob(os.path.join(full_image_path, '*.[jpJP][pnPN]*[gG]'))
|
|
full_image_list = sorted(full_image_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
|
|
frame_list_cycle = read_images(full_image_list)
|
|
|
|
return frame_list_cycle, face_list_cycle, coord_list_frames
|
|
|
|
|
|
def jpeg_to_png(image):
|
|
min_green = np.array([50, 100, 100])
|
|
max_green = np.array([70, 255, 255])
|
|
|
|
hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
|
|
mask = cv2.inRange(hsv, min_green, max_green)
|
|
mask_not = cv2.bitwise_not(mask)
|
|
green_not = cv2.bitwise_and(image, image, mask=mask_not)
|
|
b, g, r = cv2.split(green_not)
|
|
|
|
# todo 合成四通道
|
|
image = cv2.merge([b, g, r, mask_not])
|
|
return image
|
|
|
|
|
|
def load_avatar_from_256_processed(base_path, avatar_name, pkl):
|
|
avatar_path = os.path.join(base_path, 'data', 'avatars', avatar_name, pkl)
|
|
print(f'load avatar from processed:{avatar_path}')
|
|
|
|
with open(avatar_path, "rb") as f:
|
|
avatar_data = pickle.load(f)
|
|
|
|
face_list_cycle = []
|
|
frame_list_cycle = []
|
|
coord_list_frames = []
|
|
align_frames = []
|
|
m_frames = []
|
|
inv_m_frames = []
|
|
|
|
frame_info_list = avatar_data['frame_info_list']
|
|
for frame_info in tqdm(frame_info_list):
|
|
face_list_cycle.append(frame_info['img'])
|
|
frame_list_cycle.append(jpeg_to_png(frame_info['frame']))
|
|
coord_list_frames.append(frame_info['coords'])
|
|
align_frames.append(frame_info['align_frame'])
|
|
m_frames.append(frame_info['m'])
|
|
inv_m_frames.append(frame_info['inv_m'])
|
|
|
|
return frame_list_cycle, face_list_cycle, coord_list_frames, align_frames, m_frames, inv_m_frames
|
|
|
|
|
|
def config_logging(file_name: str, console_level: int = logging.INFO, file_level: int = logging.DEBUG):
|
|
file_handler = logging.FileHandler(file_name, mode='a', encoding="utf8")
|
|
file_handler.setFormatter(logging.Formatter(
|
|
'%(asctime)s [%(levelname)s] %(module)s.%(lineno)d %(name)s:\t%(message)s'
|
|
))
|
|
file_handler.setLevel(file_level)
|
|
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setFormatter(logging.Formatter(
|
|
'[%(asctime)s.%(msecs)03d %(levelname)s] %(message)s',
|
|
datefmt="%Y/%m/%d %H:%M:%S"
|
|
))
|
|
console_handler.setLevel(console_level)
|
|
|
|
logging.basicConfig(
|
|
level=min(console_level, file_level),
|
|
handlers=[file_handler, console_handler],
|
|
)
|
|
|
|
|
|
def object_stop(obj):
|
|
if obj is not None:
|
|
obj.stop()
|
|
|
|
|
|
def img_warp_back_inv_m(img, img_to, inv_m):
|
|
h_up, w_up, c = img_to.shape
|
|
mask = np.ones_like(img).astype(np.float32)
|
|
inv_mask = cv2.warpAffine(mask, inv_m, (w_up, h_up))
|
|
inv_img = cv2.warpAffine(img, inv_m, (w_up, h_up))
|
|
mask_indices = inv_mask == 1
|
|
if 4 == c:
|
|
img_to[:, :, :3][mask_indices] = inv_img[mask_indices]
|
|
else:
|
|
img_to[inv_mask == 1] = inv_img[inv_mask == 1]
|
|
return img_to
|
|
|
|
|
|
def render_image(context, frame):
|
|
res_frame, idx, type_ = frame
|
|
|
|
if type_ == 0:
|
|
combine_frame = context.frame_list_cycle[idx]
|
|
else:
|
|
bbox = context.coord_list_cycle[idx]
|
|
combine_frame = copy.deepcopy(context.frame_list_cycle[idx])
|
|
af = context.align_frames[idx]
|
|
inv_m = context.inv_m_frames[idx]
|
|
y1, y2, x1, x2 = bbox
|
|
try:
|
|
res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
|
|
af[y1:y2, x1:x2] = res_frame
|
|
combine_frame = img_warp_back_inv_m(af, combine_frame, inv_m)
|
|
except Exception as e:
|
|
logging.error(f'resize error:{e}')
|
|
return None
|
|
|
|
return combine_frame
|