# !/usr/bin/env python # -*- coding: UTF-8 -*- import datetime import logging import os import torch from PIL import Image import numpy as np import cv2 import sys import pickle import gc from huggingface_hub import hf_hub_download try: from moviepy.editor import VideoFileClip, AudioFileClip except: try: from moviepy import VideoFileClip, AudioFileClip except: from moviepy import * import random from .src.utils.mp_utils import LMKExtractor from .src.utils.motion_utils import motion_sync from .src.utils.util import save_videos_grid, crop_and_pad,crop_and_pad_rectangle,center_crop from .echomimic_v2.src.utils.dwpose_util import draw_pose_select_v2 from comfy.utils import common_upscale,ProgressBar import folder_paths from tqdm import tqdm weight_dtype = torch.float16 cur_path = os.path.dirname(os.path.abspath(__file__)) device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" tensorrt_lite= os.path.join(folder_paths.get_input_directory(),"tensorrt_lite") def process_video_v2(ref_image_pil, uploaded_audio, width, height, length, seed, context_frames, context_overlap, cfg, steps, sample_rate, fps, pipe, save_video, pose_dir, audio_file_prefix,visualizer,video_images): origin_h = height origin_w = width # 处理输入图片的尺寸 panding_img=img_padding(height, width, ref_image_pil) # 不管输出图片是何种尺寸，为保证图片质量，将输入图片转为为正方形，横裁切，竖填充，长宽为输出尺寸最大 infer_image_pil=Image.fromarray(cv2.cvtColor(panding_img,cv2.COLOR_BGR2RGB)) #将高宽改成最大图幅，方便裁切 height = max(height,width) width = max(height, width) if visualizer and isinstance(video_images,torch.Tensor): logging.info("***** start infer video to npy files for drive pose ! ***** ") video_len, _, _, _ = video_images.size() if video_len < 50: raise "input video has not much frames for driver,change your input video!" else: tensor_list = list(torch.chunk(video_images, chunks=video_len)) input_frames_cv2 = [img_padding(height,width,i) for i in tensor_list] pose_dir = os.path.join(tensorrt_lite, audio_file_prefix) if not os.path.exists(pose_dir): os.makedirs(pose_dir) else: timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") pose_dir = os.path.join(tensorrt_lite, f"{audio_file_prefix}_{timestamp}") os.makedirs(pose_dir) visualizer.move_to_cuda() # 首帧手势对齐输入图片，keypoint数据左眼，肩膀及手肘， _, first_key, first_box_xy = visualizer(np.asarray(input_frames_cv2[0]),None) _, input_key, input_box_xy = visualizer(np.asarray(panding_img),None) first_length,first_left_eye_y=estimate_ratio(first_key,first_box_xy) input_length,input_left_eye_y = estimate_ratio(input_key,input_box_xy) #print(first_length,first_left_eye_y,input_length,input_left_eye_y) #160.0 [236.0] 158.0 [151.0] 眼睛高度为绝对值 if first_length and first_left_eye_y and input_length and input_left_eye_y: if abs(input_length / height - first_length / height) > 0.005: # 比例不同须基于输入图片对齐 logging.info( "Starting the first frame gesture alignment based on the input image *** 基于输入图片，开始首帧手势缩放对齐 !") input_left_eye_y_=input_left_eye_y[0] first_left_eye_y_=first_left_eye_y[0] input_frames_cv2=[align_img(input_length, first_length, height, width, i, input_left_eye_y_, first_left_eye_y_) for i in input_frames_cv2] else:# 人体比例接近，但是高度不对，也需要对齐 logging.info( "Starting the first frame shift based on the input image *** 基于输入图片，开始首帧手势平移对齐 !") if abs(input_left_eye_y[0] / height - first_left_eye_y[ 0] / height) > 0.005: input_frames_cv2 = [affine_img(input_left_eye_y, first_left_eye_y, i) for i in input_frames_cv2] xxx=input_frames_cv2 for i,xx in enumerate(xxx): cv2.imwrite(f"{i}.png",xx) if i >2: break index = 0 for i in tqdm(input_frames_cv2): pose_img,_,_=visualizer(np.asarray(i),[5]) np.save(os.path.join(pose_dir, f"{index}"), pose_img) index+=1 USE_Default = False visualizer.enable_model_cpu_offload() gc.collect() torch.cuda.empty_cache() else: if pose_dir == "none": logging.info("when use echo v2,need choice a pose dir,using default pose for testing !") pose_dir = os.path.join(cur_path, "echomimic_v2/assets/halfbody_demo/pose/01") USE_Default = True else: logging.info( "Use NPY files for custom videos, which must be located in directory comfyui/input/tensorrt_lite") pose_dir = os.path.join(tensorrt_lite, pose_dir) USE_Default = False if seed is not None and seed > -1: generator = torch.manual_seed(seed) else: generator = torch.manual_seed(random.randint(100, 1000000)) #final_fps = fps start_idx = 0 audio_clip = AudioFileClip(uploaded_audio) L = min(int(audio_clip.duration * fps),length,len(os.listdir(pose_dir))) # if above will cause error #L=min(length,L) #length is definitely print(f"***** infer length is {L}") pose_list = [] for index in range(start_idx, start_idx + L): tgt_musk_path = os.path.join(pose_dir, "{}.npy".format(index)) if USE_Default: detected_pose = np.load(tgt_musk_path, allow_pickle=True).tolist() imh_new, imw_new, rb, re, cb, ce = detected_pose['draw_pose_params'] #print(imh_new, imw_new, rb, re, cb, ce) 官方示例蒙版的尺寸是768*768 im = draw_pose_select_v2(detected_pose, imh_new, imw_new, ref_w=800) #缩放比例为1，im也是768 ref_w！=768 im = np.transpose(np.array(im), (1, 2, 0)) tgt_musk = np.zeros((imw_new, imh_new, 3)).astype('uint8') tgt_musk[rb:re, cb:ce, :] = im else: tgt_musk = np.load(tgt_musk_path, allow_pickle=True) tgt_musk = center_resize_pad(tgt_musk, width, height) # 缩放裁剪遮罩，防止遮罩非正方形 tgt_musk_pil = Image.fromarray(np.array(tgt_musk)).convert('RGB') pose_list.append( torch.Tensor(np.array(tgt_musk_pil)).to(dtype=weight_dtype, device=device).permute(2, 0, 1) / 255.0) poses_tensor = torch.stack(pose_list, dim=1).unsqueeze(0) print(f"poses_tensor:{poses_tensor.shape}") # audio_clip = AudioFileClip(audio_clip) # # audio_clip = audio_clip.set_duration(L / fps) video = pipe( infer_image_pil, uploaded_audio, poses_tensor[:, :, :L, ...], width, height, L, steps, cfg, generator=generator, audio_sample_rate=sample_rate, context_frames=context_frames, fps=fps, context_overlap=context_overlap, start_idx=start_idx, ).videos final_length = min(video.shape[2],poses_tensor.shape[2], L) video_sig = video[:, :, :final_length, :, :] output_file = os.path.join(folder_paths.output_directory, f"{audio_file_prefix}_echo.mp4") print(f"**** final_length is : {final_length} ****") if origin_h!=origin_w: ouput_list = save_videos_grid(video_sig, output_file, n_rows=1, fps=fps, save_video=save_video,size=(origin_w,origin_h),ref_image_pil=ref_image_pil) else: ouput_list = save_videos_grid(video_sig, output_file, n_rows=1, fps=fps, save_video=save_video) if save_video: output_video_path = os.path.join(folder_paths.output_directory, f"{audio_file_prefix}_audio.mp4") video_clip = VideoFileClip(output_file) audio_clip = AudioFileClip(uploaded_audio) final_clip = video_clip.set_audio(audio_clip) final_clip.write_videofile( output_video_path, codec="libx264", audio_codec="aac") print(f"**** saving{output_file} at {output_video_path} ****") video_clip.reader.close() audio_clip.close() final_clip.reader.close() return ouput_list def process_video(face_img, uploaded_audio, width, height, length, seed, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, pipe, face_detector, save_video, pose_dir, video_images, audio_file_prefix, visualizer=None,): if seed is not None and seed > -1: generator = torch.manual_seed(seed) else: generator = torch.manual_seed(random.randint(100, 1000000)) #### face musk prepare face_mask = np.zeros((face_img.shape[0], face_img.shape[1])).astype('uint8') det_bboxes, probs = face_detector.detect(face_img) select_bbox = select_face(det_bboxes, probs) if select_bbox is None: face_mask[:, :] = 255 else: xyxy = select_bbox[:4].astype(float) # 面部处理出来是浮点数，无法实现整形 xyxy = np.round(xyxy).astype("int") rb, re, cb, ce = xyxy[1], xyxy[3], xyxy[0], xyxy[2] #56 227 132 268 r_pad = int((re - rb) * facemask_dilation_ratio) # ratio：0.1 遮罩膨胀系数 17*2 c_pad = int((ce - cb) * facemask_dilation_ratio) # ratio：0.1 遮罩膨胀系数 14*2 face_mask[rb - r_pad: re + r_pad, cb - c_pad: ce + c_pad] = 255 #### face crop #### if facecrop_dilation_ratio<1.0: if facecrop_dilation_ratio==0: facecrop_dilation_ratio=1 r_pad_crop = int((re - rb) * facecrop_dilation_ratio) # ratio 0.5 r_pad_crop：85,c_pad_crop:68 c_pad_crop = int((ce - cb) * facecrop_dilation_ratio) # ratio 1.0 r_pad_crop：171,c_pad_crop:136 crop_rect = [max(0, cb - c_pad_crop), max(0, rb - r_pad_crop), min(ce + c_pad_crop, face_img.shape[1]), min(re + r_pad_crop, face_img.shape[0])] if width == height: # 输出图片指定尺寸，如果是非方形，则会变形 face_img_i, ori_face_rect_i = crop_and_pad(face_img, crop_rect) face_mask_m, ori_mask_rect_m = crop_and_pad(face_mask, crop_rect) # (0, 7, 384, 391) face_img = cv2.resize(face_img_i, (width, height)) face_mask = cv2.resize(face_mask_m, (width, height)) else: face_img,face_mask=crop_and_pad_rectangle(face_img,face_mask,crop_rect) face_img= cv2tensor(face_img).permute(0, 2, 3, 1)#[1, 3, 357, 245] =>[[1,357, 245,3]] face_mask = cv2tensor(face_mask).permute(0, 2, 3, 1) face_img=tensor_upscale(face_img, width, height) face_img=tensor2cv(face_img) face_mask = tensor_upscale(face_mask, width, height) face_mask=cv2.cvtColor(tensor2cv(face_mask), cv2.COLOR_BGR2GRAY)#二值化 ret, face_mask = cv2.threshold(face_mask, 0, 255, cv2.THRESH_BINARY) else: #when ratio=1 no crop print("when facecrop_ratio=1.0,The maximum image size will be obtained, but there may be edge deformation.** 选择最大裁切为1.0时，边缘可能会出现形变！") ref_image_pil = Image.fromarray(face_img[:, :, [2, 1, 0]]) if visualizer: if pose_dir == "none": # motion sync if isinstance(video_images,torch.Tensor): print("**** Use video pose drive video! ****") pose_dir_path,video_len = motion_sync_main(visualizer, width, height, video_images, face_img,facecrop_dilation_ratio, audio_file_prefix) else: raise ("**** You need link video_images for drive video ****") #pose_dir = os.path.join(cur_path, "assets", "test_pose_demo_pose") # default else: print("**** Use pkl drive video! ****") pose_dir_path = os.path.join(tensorrt_lite, pose_dir) files_and_directories = os.listdir(pose_dir_path) # 过滤出文件，排除子目录 files = [f for f in files_and_directories if os.path.isfile(os.path.join(pose_dir_path, f))] video_len=len(files) if length>video_len: print(f"**** video length {video_len} is less than length,use {video_len} as {length} ****") length=video_len pose_list = [] for index in range(len(os.listdir(pose_dir_path))): tgt_musk_path = os.path.join(pose_dir_path, f"{index}.pkl") with open(tgt_musk_path, "rb") as f: tgt_kpts = pickle.load(f) tgt_musk = visualizer.draw_landmarks((width, height), tgt_kpts,facecrop_dilation_ratio) tgt_musk_pil = Image.fromarray(np.array(tgt_musk).astype(np.uint8)).convert('RGB') pose_list.append( torch.Tensor(np.array(tgt_musk_pil)).to(dtype=weight_dtype, device="cuda").permute(2, 0, 1) / 255.0) face_mask_tensor = torch.stack(pose_list, dim=1).unsqueeze(0) else: print("**** Use audio drive video! ****") face_mask_tensor = torch.Tensor(face_mask).to(dtype=weight_dtype, device="cuda").unsqueeze(0).unsqueeze( 0).unsqueeze(0) / 255.0 video = pipe( ref_image_pil, uploaded_audio, face_mask_tensor, width, height, length, steps, cfg, generator=generator, audio_sample_rate=sample_rate, context_frames=context_frames, fps=fps, context_overlap=context_overlap ).videos final_length = min(video.shape[2], face_mask_tensor.shape[2], length) output_file = os.path.join(folder_paths.output_directory, f"{audio_file_prefix}_echo.mp4") print(f"**** final_length is : {final_length} ****") ouput_list = save_videos_grid(video, output_file, n_rows=1, fps=fps, save_video=save_video) if save_video: output_video_path = os.path.join(folder_paths.output_directory, f"{audio_file_prefix}_audio.mp4") video_clip = VideoFileClip(output_file) audio_clip = AudioFileClip(uploaded_audio) final_clip = video_clip.set_audio(audio_clip) final_clip.write_videofile( output_video_path, codec="libx264", audio_codec="aac") print(f"**** saving{output_file} at {output_video_path} ****") video_clip.reader.close() audio_clip.close() final_clip.reader.close() return ouput_list def motion_sync_main(vis, width, height, video_images, face_img,facecrop_dilation_ratio, audio_file_prefix): lmk_extractor = LMKExtractor() ref_det = lmk_extractor(face_img) #driver_video = os.path.join(folder_paths.input_directory, driver_video) # if audio_form_video: # audio_path = os.path.join(folder_paths.input_directory, f"{audio_file_prefix}_audio.wav") # video_clip = VideoFileClip(driver_video) # audio_clip = video_clip.audio # audio_clip.write_audiofile(audio_path) # video_clip.close() # audio_clip.close() video_len,_,_,_=video_images.size() if video_len<25: raise "input video has not much frames for driver,change your input video!" else: tensor_list = list(torch.chunk(video_images, chunks=video_len)) input_frames_cv2=[tensor2cv(tensor_upscale(i, width, height)) for i in tensor_list] # print(ref_det) sequence_driver_det = [] if input_frames_cv2: try: print("**** Starting process video ****") for frame in input_frames_cv2: result = lmk_extractor(frame) assert result is not None, "bad video, face not detected" sequence_driver_det.append(result) except: print("face detection failed") else: raise "input video error,change your input video!" print("**** motion sync lenght " f"{len(sequence_driver_det)} ****") if vis: if facecrop_dilation_ratio==0: facecrop_dilation_ratio=1 pose_frames_driver = [vis.draw_landmarks((width, height), i["lmks"],facecrop_dilation_ratio, normed=True) for i in sequence_driver_det] poses_add_driver = [(i * 0.5 + j * 0.5).clip(0, 255).astype(np.uint8) for i, j in zip(input_frames_cv2, pose_frames_driver)] #print(f"**** poses_add_driver is done in len : {len(poses_add_driver)} ****") save_dir = os.path.join(tensorrt_lite, audio_file_prefix) if not os.path.exists(save_dir): os.makedirs(save_dir) sequence_det_ms = motion_sync(sequence_driver_det, ref_det) for i in range(len(sequence_det_ms)): with open('{}/{}.pkl'.format(save_dir, i), 'wb') as file: pickle.dump(sequence_det_ms[i], file) print(f"**** motion_sync {save_dir} is done ****") else: #即便有文件夹，还是重新生成，避免出错 timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") save_dir=os.path.join(tensorrt_lite,f"{audio_file_prefix}_{timestamp}") os.makedirs(save_dir) sequence_det_ms = motion_sync(sequence_driver_det, ref_det) for i in range(len(sequence_det_ms)): with open('{}/{}.pkl'.format(save_dir, i), 'wb') as file: pickle.dump(sequence_det_ms[i], file) print(f"**** motion_sync {save_dir} is done ****") return save_dir,video_len def select_face(det_bboxes, probs): ## max face from faces that the prob is above 0.8 ## box: xyxy if det_bboxes is None or probs is None: return None filtered_bboxes = [] for bbox_i in range(len(det_bboxes)): if probs[bbox_i] > 0.8: filtered_bboxes.append(det_bboxes[bbox_i]) if len(filtered_bboxes) == 0: return None sorted_bboxes = sorted(filtered_bboxes, key=lambda x: (x[3] - x[1]) * (x[2] - x[0]), reverse=True) return sorted_bboxes[0] def find_directories(base_path): directories = [] for root, dirs, files in os.walk(base_path): for name in dirs: directories.append(name) return directories def download_weights(file_dir,repo_id,subfolder="",pt_name=""): if subfolder: file_path = os.path.join(file_dir,subfolder, pt_name) sub_dir=os.path.join(file_dir,subfolder) if not os.path.exists(sub_dir): os.makedirs(sub_dir) if not os.path.exists(file_path): pt_path = hf_hub_download( repo_id=repo_id, subfolder=subfolder, filename=pt_name, local_dir = file_dir, ) else: pt_path=get_instance_path(file_path) return pt_path else: file_path = os.path.join(file_dir, pt_name) if not os.path.exists(file_dir): os.makedirs(file_dir) if not os.path.exists(file_path): pt_path = hf_hub_download( repo_id=repo_id, filename=pt_name, local_dir=file_dir, ) else: pt_path=get_instance_path(file_path) return pt_path def pil2narry(img): img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img def narry_list(list_in): for i in range(len(list_in)): value = list_in[i] modified_value = pil2narry(value) list_in[i] = modified_value return list_in def get_video_img(tensor): if tensor == None: return None outputs = [] for x in tensor: x = tensor_to_pil(x) outputs.append(x) yield outputs def instance_path(path, repo): if repo == "": if path == "none": repo = "none" else: model_path = get_local_path(folder_paths.base_path, path) repo = get_instance_path(model_path) return repo def gen_img_form_video(tensor): pil = [] for x in tensor: pil[x] = tensor_to_pil(x) yield pil def phi_list(list_in): for i in range(len(list_in)): value = list_in[i] list_in[i] = value return list_in def tensor_to_pil(tensor): image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def nomarl_upscale(img_tensor, width, height): samples = img_tensor.movedim(-1, 1) img = common_upscale(samples, width, height, "nearest-exact", "center") samples = img.movedim(1, -1) img_pil = tensor_to_pil(samples) return img_pil def tensor_upscale(img_tensor, width, height): samples = img_tensor.movedim(-1, 1) img = common_upscale(samples, width, height, "nearest-exact", "center") samples = img.movedim(1, -1) return samples def get_local_path(comfy_file_path, model_path): path = os.path.join(comfy_file_path, "models", "diffusers", model_path) model_path = os.path.normpath(path) if sys.platform == 'win32': model_path = model_path.replace('\\', "/") return model_path def get_instance_path(path): instance_path = os.path.normpath(path) if sys.platform == 'win32': instance_path = instance_path.replace('\\', "/") return instance_path def tensor2cv(tensor_image): if len(tensor_image.shape)==4:# b hwc to hwc tensor_image=tensor_image.squeeze(0) if tensor_image.is_cuda: tensor_image = tensor_image.cpu() tensor_image=tensor_image.numpy() #反归一化 maxValue=tensor_image.max() tensor_image=tensor_image*255/maxValue img_cv2=np.uint8(tensor_image)#32 to uint8 img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR) return img_cv2 def cvargb2tensor(img): assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img)) img = torch.from_numpy(img.transpose((2, 0, 1))) return img.float().div(255).unsqueeze(0) # 255也可以改为256 def cv2tensor(img): assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = torch.from_numpy(img.transpose((2, 0, 1))) return img.float().div(255).unsqueeze(0) # 255也可以改为256 def images_generator(img_list: list,): #get img size sizes = {} for image_ in img_list: if isinstance(image_,Image.Image): count = sizes.get(image_.size, 0) sizes[image_.size] = count + 1 elif isinstance(image_,np.ndarray): count = sizes.get(image_.shape[:2][::-1], 0) sizes[image_.shape[:2][::-1]] = count + 1 else: raise "unsupport image list,must be pil or cv2!!!" size = max(sizes.items(), key=lambda x: x[1])[0] yield size[0], size[1] # any to tensor def load_image(img_in): if isinstance(img_in, Image.Image): img_in=img_in.convert("RGB") i = np.array(img_in, dtype=np.float32) i = torch.from_numpy(i).div_(255) if i.shape[0] != size[1] or i.shape[1] != size[0]: i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0) i = common_upscale(i, size[0], size[1], "lanczos", "center") i = i.squeeze(0).movedim(0, -1).numpy() return i elif isinstance(img_in,np.ndarray): i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32) i = torch.from_numpy(i).div_(255) #print(i.shape) return i else: raise "unsupport image list,must be pil,cv2 or tensor!!!" total_images = len(img_list) processed_images = 0 pbar = ProgressBar(total_images) images = map(load_image, img_list) try: prev_image = next(images) while True: next_image = next(images) yield prev_image processed_images += 1 pbar.update_absolute(processed_images, total_images) prev_image = next_image except StopIteration: pass if prev_image is not None: yield prev_image def load_images(img_list: list,): gen = images_generator(img_list) (width, height) = next(gen) images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3))))) if len(images) == 0: raise FileNotFoundError(f"No images could be loaded .") return images def tensor2pil(tensor): image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def cf_tensor2cv(tensor,width, height): d1, _, _, _ = tensor.size() if d1 > 1: tensor_list = list(torch.chunk(tensor, chunks=d1)) tensor = [tensor_list][0] cr_tensor=tensor_upscale(tensor,width, height) cv_img=tensor2cv(cr_tensor) return cv_img def center_resize_pad(img, new_width, new_height):#为简化，new已是正方形 h, w = img.shape[:2] if w == h: if w == new_width: return img else: return cv2.resize(img, (new_width, new_height)) else: #蒙版也有可能不是正方形 if h > w: # 竖直图左右填充 s = max(h, w) f = np.zeros((s, s, 3), np.uint8) ax, ay = (s - img.shape[1]) // 2, (s - img.shape[0]) // 2 f[ay:img.shape[0] + ay, ax:ax + img.shape[1]] = img else: f = center_crop(img, h, h) return cv2.resize(f, (new_width, new_height)) def img_padding(height,width,ref_image_pil): output_max = max(height, width) img = tensor2cv(ref_image_pil) h, w = img.shape[:2] if h==w: return cv2.resize(img, (output_max,output_max), interpolation=cv2.INTER_AREA) else: if h > w: #竖直图左右填充 s = max(h, w) f = np.zeros((s, s, 3), np.uint8) ax, ay = (s - w) // 2, (s - h) // 2 f[ay:h + ay, ax:ax + w] = img else: f=center_crop(img, h, h) return cv2.resize(f, (output_max,output_max), interpolation=cv2.INTER_AREA) def estimate_ratio(keypoint: list,box_xy,length=None): left_eye_y = [] left_shoulder_y = [] left_elbow_y = [] for i, (name, (x, y, conf)) in enumerate(keypoint[0].items()): if name == "left_eye": if conf > 0.3: left_eye_y.append(y) if name == "left_shoulder": if conf > 0.3: left_shoulder_y.append(y) if name == "left_elbow": if conf > 0.3: left_elbow_y.append(y) if left_eye_y and left_elbow_y: length=left_elbow_y[0] - left_eye_y[0] elif left_eye_y and left_shoulder_y and not left_elbow_y : length = left_shoulder_y[0] - left_eye_y[0] else: pass if left_eye_y: left_eye_y=[left_eye_y[0]+box_xy[0]] #眼部的实际高度要加上box的边界 return length,left_eye_y def align_img(input_length, first_length, height, width, input_frames_cv2_first, input_left_eye_y, first_left_eye_y): ratio = input_length / first_length # 82.0 [50.0] 76.0 [41.0] f f in in base_image=np.zeros((height, width,3), np.uint8) if input_length / height < first_length / height: # 输入图的人物占比要小，pose图需要缩小对齐,0.926 ratio input_frames_cv2_first=cv2.resize(input_frames_cv2_first, (int(height * ratio), int(height * ratio)), interpolation=cv2.INTER_AREA) #缩小 reduced_image,pad_size=center_paste(base_image, input_frames_cv2_first) #中心粘贴 move_ = -int(first_left_eye_y * ratio+pad_size[0] - input_left_eye_y) if first_left_eye_y * ratio +pad_size[0] >= \ input_left_eye_y else \ int(input_left_eye_y)- int(first_left_eye_y * ratio+pad_size[0]) #对齐眼睛 translation_matrix = np.float32([[1, 0, 0], [0, 1, move_]]) #y轴位移 shifted_image = cv2.warpAffine(reduced_image, translation_matrix, (width, height)) else: # pose图里人物的比例小于输入图，pose要放大 input_frames_cv2_first=cv2.resize(input_frames_cv2_first, (int(height / ratio), int(height / ratio)), interpolation=cv2.INTER_AREA) #放大 crpo_image=center_crop(input_frames_cv2_first, height, width) #中心裁切 h, w = input_frames_cv2_first.shape[:2] shift_y=(h-height)//2 move_ = -int(first_left_eye_y / ratio-shift_y - input_left_eye_y) if first_left_eye_y / ratio-shift_y >= \ input_left_eye_y else \ int(input_left_eye_y) - int(first_left_eye_y / ratio-shift_y) translation_matrix = np.float32([[1, 0, 0], [0, 1, move_]]) shifted_image = cv2.warpAffine(crpo_image, translation_matrix, (width, height)) return shifted_image def center_paste(img_b,img_f): b_h,b_w=img_b.shape[:2] f_h, f_w = img_f.shape[:2] x = (b_w - f_w) // 2 y = (b_h - f_h) // 2 # 确保坐标不会是负数 x = max(0, x) y = max(0, y) img_b[y:y + f_h, x:x + f_w] = img_f return img_b,(x,y) def affine_img(input_left_eye_y, first_left_eye_y, img): height,width=img.shape[:2] move_ = int(input_left_eye_y[0] - first_left_eye_y[0]) translation_matrix = np.float32([[1, 0, 0], [0, 1, move_]]) # y轴位移 shifted_image = cv2.warpAffine(img, translation_matrix, (width, height)) return shifted_image