import argparse import glob import json import os import pickle import shutil import cv2 import numpy as np import torch import torchvision.transforms as transforms from PIL import Image from diffusers import AutoencoderKL from face_alignment import NetworkSize from mmpose.apis import inference_topdown, init_model from mmpose.structures import merge_data_samples from tqdm import tqdm try: from utils.face_parsing import FaceParsing except ModuleNotFoundError: from musetalk.utils.face_parsing import FaceParsing def video2imgs(vid_path, save_path, ext='.png', cut_frame=10000000): cap = cv2.VideoCapture(vid_path) count = 0 while True: if count > cut_frame: break ret, frame = cap.read() if ret: cv2.imwrite(f"{save_path}/{count:08d}.png", frame) count += 1 else: break def read_imgs(img_list): frames = [] print('reading images...') for img_path in tqdm(img_list): frame = cv2.imread(img_path) frames.append(frame) return frames def get_landmark_and_bbox(img_list, upperbondrange=0): frames = read_imgs(img_list) batch_size_fa = 1 batches = [frames[i:i + batch_size_fa] for i in range(0, len(frames), batch_size_fa)] coords_list = [] landmarks = [] if upperbondrange != 0: print('get key_landmark and face bounding boxes with the bbox_shift:', upperbondrange) else: print('get key_landmark and face bounding boxes with the default value') average_range_minus = [] average_range_plus = [] coord_placeholder = (0.0, 0.0, 0.0, 0.0) for fb in tqdm(batches): results = inference_topdown(model, np.asarray(fb)[0]) results = merge_data_samples(results) keypoints = results.pred_instances.keypoints face_land_mark = keypoints[0][23:91] face_land_mark = face_land_mark.astype(np.int32) # get bounding boxes by face detetion bbox = fa.get_detections_for_batch(np.asarray(fb)) # adjust the bounding box refer to landmark # Add the bounding box to a tuple and append it to the coordinates list for j, f in enumerate(bbox): if f is None: # no face in the image coords_list += [coord_placeholder] continue half_face_coord = face_land_mark[29] # np.mean([face_land_mark[28], face_land_mark[29]], axis=0) range_minus = (face_land_mark[30] - face_land_mark[29])[1] range_plus = (face_land_mark[29] - face_land_mark[28])[1] average_range_minus.append(range_minus) average_range_plus.append(range_plus) if upperbondrange != 0: half_face_coord[1] = upperbondrange + half_face_coord[1] # 手动调整 + 向下（偏29） - 向上（偏28） half_face_dist = np.max(face_land_mark[:, 1]) - half_face_coord[1] upper_bond = half_face_coord[1] - half_face_dist f_landmark = ( np.min(face_land_mark[:, 0]), int(upper_bond), np.max(face_land_mark[:, 0]), np.max(face_land_mark[:, 1])) x1, y1, x2, y2 = f_landmark if y2 - y1 <= 0 or x2 - x1 <= 0 or x1 < 0: # if the landmark bbox is not suitable, reuse the bbox coords_list += [f] w, h = f[2] - f[0], f[3] - f[1] print("error bbox:", f) else: coords_list += [f_landmark] return coords_list, frames 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 # torch.backends.cuda.matmul.allow_tf32 = False # torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = False # torch.backends.cudnn.allow_tf32 = True print('cuda start') # 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 def get_mask_tensor(): """ Creates a mask tensor for image processing. :return: A mask tensor. """ mask_tensor = torch.zeros((256, 256)) mask_tensor[:256 // 2, :] = 1 mask_tensor[mask_tensor < 0.5] = 0 mask_tensor[mask_tensor >= 0.5] = 1 return mask_tensor def preprocess_img(img_name, half_mask=False): window = [] if isinstance(img_name, str): window_fnames = [img_name] for fname in window_fnames: img = cv2.imread(fname) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_LANCZOS4) window.append(img) else: img = cv2.cvtColor(img_name, cv2.COLOR_BGR2RGB) window.append(img) x = np.asarray(window) / 255. x = np.transpose(x, (3, 0, 1, 2)) x = torch.squeeze(torch.FloatTensor(x)) if half_mask: x = x * (get_mask_tensor() > 0.5) normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) x = normalize(x) x = x.unsqueeze(0) # [1, 3, 256, 256] torch tensor x = x.to(device) return x def encode_latents(image): with torch.no_grad(): init_latent_dist = vae.encode(image.to(vae.dtype)).latent_dist init_latents = vae.config.scaling_factor * init_latent_dist.sample() return init_latents def get_latents_for_unet(img): ref_image = preprocess_img(img, half_mask=True) # [1, 3, 256, 256] RGB, torch tensor masked_latents = encode_latents(ref_image) # [1, 4, 32, 32], torch tensor ref_image = preprocess_img(img, half_mask=False) # [1, 3, 256, 256] RGB, torch tensor ref_latents = encode_latents(ref_image) # [1, 4, 32, 32], torch tensor latent_model_input = torch.cat([masked_latents, ref_latents], dim=1) return latent_model_input def get_crop_box(box, expand): x, y, x1, y1 = box x_c, y_c = (x + x1) // 2, (y + y1) // 2 w, h = x1 - x, y1 - y s = int(max(w, h) // 2 * expand) crop_box = [x_c - s, y_c - s, x_c + s, y_c + s] return crop_box, s def face_seg(image): seg_image = fp(image) if seg_image is None: print("error, no person_segment") return None seg_image = seg_image.resize(image.size) return seg_image def get_image_prepare_material(image, face_box, upper_boundary_ratio=0.5, expand=1.2): body = Image.fromarray(image[:, :, ::-1]) x, y, x1, y1 = face_box # print(x1-x,y1-y) crop_box, s = get_crop_box(face_box, expand) x_s, y_s, x_e, y_e = crop_box face_large = body.crop(crop_box) ori_shape = face_large.size mask_image = face_seg(face_large) mask_small = mask_image.crop((x - x_s, y - y_s, x1 - x_s, y1 - y_s)) mask_image = Image.new('L', ori_shape, 0) mask_image.paste(mask_small, (x - x_s, y - y_s, x1 - x_s, y1 - y_s)) # keep upper_boundary_ratio of talking area width, height = mask_image.size top_boundary = int(height * upper_boundary_ratio) modified_mask_image = Image.new('L', ori_shape, 0) modified_mask_image.paste(mask_image.crop((0, top_boundary, width, height)), (0, top_boundary)) blur_kernel_size = int(0.1 * ori_shape[0] // 2 * 2) + 1 mask_array = cv2.GaussianBlur(np.array(modified_mask_image), (blur_kernel_size, blur_kernel_size), 0) return mask_array, crop_box ##todo 简单根据文件后缀判断 要更精确的可以自己修改 使用 magic def is_video_file(file_path): video_exts = ['.mp4', '.mkv', '.flv', '.avi', '.mov'] # 这里列出了一些常见的视频文件扩展名，可以根据需要添加更多 file_ext = os.path.splitext(file_path)[1].lower() # 获取文件扩展名并转换为小写 return file_ext in video_exts def create_dir(dir_path): if not os.path.exists(dir_path): os.makedirs(dir_path) current_dir = os.path.dirname(os.path.abspath(__file__)) def create_musetalk_human(file, avatar_id): # 保存文件设置 可以不动 save_path = os.path.join(current_dir, f'../data/avatars/avator_{avatar_id}') save_full_path = os.path.join(current_dir, f'../data/avatars/avator_{avatar_id}/full_imgs') create_dir(save_path) create_dir(save_full_path) mask_out_path = os.path.join(current_dir, f'../data/avatars/avator_{avatar_id}/mask') create_dir(mask_out_path) # 模型 mask_coords_path = os.path.join(current_dir, f'{save_path}/mask_coords.pkl') coords_path = os.path.join(current_dir, f'{save_path}/coords.pkl') latents_out_path = os.path.join(current_dir, f'{save_path}/latents.pt') with open(os.path.join(current_dir, f'{save_path}/avator_info.json'), "w") as f: json.dump({ "avatar_id": avatar_id, "video_path": file, "bbox_shift": 5 }, f) if os.path.isfile(file): if is_video_file(file): video2imgs(file, save_full_path, ext='png') else: shutil.copyfile(file, f"{save_full_path}/{os.path.basename(file)}") else: files = os.listdir(file) files.sort() files = [file for file in files if file.split(".")[-1] == "png"] for filename in files: shutil.copyfile(f"{file}/{filename}", f"{save_full_path}/{filename}") input_img_list = sorted(glob.glob(os.path.join(save_full_path, '*.[jpJP][pnPN]*[gG]'))) print("extracting landmarks...") coord_list, frame_list = get_landmark_and_bbox(input_img_list, 5) input_latent_list = [] idx = -1 # maker if the bbox is not sufficient coord_placeholder = (0.0, 0.0, 0.0, 0.0) for bbox, frame in zip(coord_list, frame_list): idx = idx + 1 if bbox == coord_placeholder: continue x1, y1, x2, y2 = bbox crop_frame = frame[y1:y2, x1:x2] resized_crop_frame = cv2.resize(crop_frame, (256, 256), interpolation=cv2.INTER_LANCZOS4) latents = get_latents_for_unet(resized_crop_frame) input_latent_list.append(latents) frame_list_cycle = frame_list + frame_list[::-1] coord_list_cycle = coord_list + coord_list[::-1] input_latent_list_cycle = input_latent_list + input_latent_list[::-1] mask_coords_list_cycle = [] mask_list_cycle = [] for i, frame in enumerate(tqdm(frame_list_cycle)): cv2.imwrite(f"{save_full_path}/{str(i).zfill(8)}.png", frame) face_box = coord_list_cycle[i] mask, crop_box = get_image_prepare_material(frame, face_box) cv2.imwrite(f"{mask_out_path}/{str(i).zfill(8)}.png", mask) mask_coords_list_cycle += [crop_box] mask_list_cycle.append(mask) with open(mask_coords_path, 'wb') as f: pickle.dump(mask_coords_list_cycle, f) with open(coords_path, 'wb') as f: pickle.dump(coord_list_cycle, f) torch.save(input_latent_list_cycle, os.path.join(latents_out_path)) # initialize the mmpose model device = "cuda" if torch.cuda.is_available() else "cpu" fa = FaceAlignment(1, flip_input=False, device=device) config_file = os.path.join(current_dir, 'utils/dwpose/rtmpose-l_8xb32-270e_coco-ubody-wholebody-384x288.py') checkpoint_file = os.path.abspath(os.path.join(current_dir, '../models/dwpose/dw-ll_ucoco_384.pth')) model = init_model(config_file, checkpoint_file, device=device) vae = AutoencoderKL.from_pretrained(os.path.abspath(os.path.join(current_dir, '../models/sd-vae-ft-mse'))) vae.to(device) fp = FaceParsing(os.path.abspath(os.path.join(current_dir, '../models/face-parse-bisent/resnet18-5c106cde.pth')), os.path.abspath(os.path.join(current_dir, '../models/face-parse-bisent/79999_iter.pth'))) if __name__ == '__main__': # 视频文件地址 parser = argparse.ArgumentParser() parser.add_argument("--file", type=str, default=r'D:\ok\00000000.png', ) parser.add_argument("--avatar_id", type=str, default='3', ) args = parser.parse_args() create_musetalk_human(args.file, args.avatar_id)