# !/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 pickle import gc import shutil 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 multiprocessing.pool import ThreadPool 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") MAX_SIZE = 768 def process_video_v2(infer_image_pil,ref_image_pil, uploaded_audio,face_locator_tensor, W_change, H_change, start_idx, seed, context_frames, context_overlap, cfg, steps, sample_rate, fps, pipe,mask_len,origin_h,origin_w, save_video, audio_file_prefix,L,whisper_chunks): # audio_clip = AudioFileClip(audio_clip) # # audio_clip = audio_clip.set_duration(L / fps) if seed is not None and seed > -1: generator = torch.manual_seed(seed) else: generator = torch.manual_seed(random.randint(100, 1000000)) video = pipe( infer_image_pil, #uploaded_audio, face_locator_tensor, width=W_change, height=H_change, video_length=L, num_inference_steps=steps, guidance_scale=cfg, generator=generator, audio_sample_rate=sample_rate, context_frames=context_frames, fps=fps, context_overlap=context_overlap, start_idx=start_idx, whisper_chunks=whisper_chunks, ).videos final_length = min(video.shape[2],mask_len, 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(ref_image_pil, uploaded_audio, width, height, length, seed, face_locator_tensor,context_frames, context_overlap, cfg, steps, sample_rate, fps, pipe,save_video,mask_len, audio_file_prefix,whisper_chunks): if seed is not None and seed > -1: generator = torch.manual_seed(seed) else: generator = torch.manual_seed(random.randint(100, 1000000)) video = pipe( ref_image_pil, #uploaded_audio, c_face_locator_tensor=face_locator_tensor, width=width, height=height, video_length=length, num_inference_steps=steps, guidance_scale=cfg, generator=generator, audio_sample_rate=sample_rate, context_frames=context_frames, fps=fps, context_overlap=context_overlap, whisper_chunks=whisper_chunks, ).videos final_length = min(video.shape[2], mask_len, 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): file_path = hf_hub_download( repo_id=repo_id, subfolder=subfolder, filename=pt_name, local_dir = file_dir, ) return file_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): file_path = hf_hub_download( repo_id=repo_id, filename=pt_name, local_dir=file_dir, ) return file_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 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 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): x1, y1, x2, y2=box_xy bbox_width, bbox_height = x2 - x1, y2 - y1 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: y_coord =int(y * bbox_height / 256) + y1 #获取坐标点在原图的绝对值 left_eye_y.append(y_coord) if name == "left_shoulder": if conf > 0.3: y_coord = int(y * bbox_height / 256) + y1 left_shoulder_y.append(y_coord) if name == "left_elbow": if conf > 0.3: y_coord = int(y * bbox_height / 256) + y1 left_elbow_y.append(y_coord) 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]] 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 #603 [201] 679 [220] input_frames_cv2_first = cv2.resize(input_frames_cv2_first, (int(height * ratio), int(height * ratio)), interpolation=cv2.INTER_AREA) # 缩小 base_image=np.zeros((height, width,3), np.uint8) if input_length / height < first_length / height: # 输入图的人物占比要小，pose图需要缩小0.88对齐,#603 [201] 679 [220] reduced_image,pad_size=center_paste(base_image, input_frames_cv2_first) #中心粘贴，pad为黑边尺寸 eye_y=int(first_left_eye_y*ratio+pad_size[1]) #图片缩放后的绝对值 else: # pose图里人物的比例小于输入图，pose要放大 reduced_image=center_crop(input_frames_cv2_first, height, width) #中心裁切 h, w = input_frames_cv2_first.shape[:2] shift_y=(h-height)//2 #图片放大后裁切掉的边界 eye_y= int(first_left_eye_y*ratio-shift_y) move_ = -int(eye_y - input_left_eye_y) if eye_y>= int(input_left_eye_y) else int(input_left_eye_y- eye_y) #对齐眼睛 translation_matrix = np.float32([[1, 0, 0], [0, 1, move_]]) shifted_image = cv2.warpAffine(reduced_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 def convert_fps(src_path, tgt_path, tgt_fps=24, tgt_sr=16000): clip = VideoFileClip(src_path) new_clip = clip.set_fps(tgt_fps) if tgt_fps is not None: audio = new_clip.audio audio = audio.set_fps(tgt_sr) new_clip = new_clip.set_audio(audio) if '.mov' in tgt_path: tgt_path = tgt_path.replace('.mov', '.mp4') new_clip.write_videofile(tgt_path, codec='libx264', audio_codec='aac') def get_video_pose(visualizer, frames, sample_stride: int = 1, max_frame=None): # read input video #vr = decord.VideoReader(video_path, ctx=decord.cpu(0)) #sample_stride *= max(1, int(vr.get_avg_fps() / 24)) #frames = vr.get_batch(list(range(0, len(vr), sample_stride))).asnumpy() #frames=vr.get_batch(list(range(0, len(vr), sample_stride))).asnumpy() # print(frames[0]) if max_frame is not None: frames = frames[:max_frame] height, width, _ = frames[0].shape print("start get video pose list") detected_poses = [visualizer(frm) for frm in frames] visualizer.release_memory() return detected_poses, height, width, frames def resize_and_pad(img, max_size): img_new = np.zeros((max_size, max_size, 3)).astype('uint8') imh, imw = img.shape[0], img.shape[1] half = max_size // 2 if imh > imw: imh_new = max_size imw_new = int(round(imw / imh * imh_new)) half_w = imw_new // 2 rb, re = 0, max_size cb = half - half_w ce = cb + imw_new else: imw_new = max_size imh_new = int(round(imh / imw * imw_new)) half_h = imh_new // 2 cb, ce = 0, max_size rb = half - half_h re = rb + imh_new img_resize = cv2.resize(img, (imw_new, imh_new)) img_new[rb:re, cb:ce, :] = img_resize return img_new def resize_and_pad_param(imh, imw, max_size): half = max_size // 2 if imh > imw: imh_new = max_size imw_new = int(round(imw / imh * imh_new)) half_w = imw_new // 2 rb, re = 0, max_size cb = half - half_w ce = cb + imw_new else: imw_new = max_size imh_new = int(round(imh / imw * imw_new)) imh_new = max_size half_h = imh_new // 2 cb, ce = 0, max_size rb = half - half_h re = rb + imh_new return imh_new, imw_new, rb, re, cb, ce def get_pose_params(detected_poses, max_size): print('get_pose_params...') # pose rescale height = 768 width=768 w_min_all, w_max_all, h_min_all, h_max_all = [], [], [], [] mid_all = [] for num, detected_pose in enumerate(detected_poses): detected_poses[num]['num'] = num candidate_body = detected_pose['bodies']['candidate'] score_body = detected_pose['bodies']['score'] candidate_face = detected_pose['faces'] score_face = detected_pose['faces_score'] candidate_hand = detected_pose['hands'] score_hand = detected_pose['hands_score'] # face if candidate_face.shape[0] > 1: index = 0 candidate_face = candidate_face[index] score_face = score_face[index] detected_poses[num]['faces'] = candidate_face.reshape(1, candidate_face.shape[0], candidate_face.shape[1]) detected_poses[num]['faces_score'] = score_face.reshape(1, score_face.shape[0]) else: candidate_face = candidate_face[0] score_face = score_face[0] # body if score_body.shape[0] > 1: tmp_score = [] for k in range(0, score_body.shape[0]): tmp_score.append(score_body[k].mean()) index = np.argmax(tmp_score) candidate_body = candidate_body[index * 18:(index + 1) * 18, :] score_body = score_body[index] score_hand = score_hand[(index * 2):(index * 2 + 2), :] candidate_hand = candidate_hand[(index * 2):(index * 2 + 2), :, :] else: score_body = score_body[0] all_pose = np.concatenate((candidate_body, candidate_face)) all_score = np.concatenate((score_body, score_face)) all_pose = all_pose[all_score > 0.8] body_pose = np.concatenate((candidate_body,)) mid_ = body_pose[1, 0] face_pose = candidate_face hand_pose = candidate_hand h_min, h_max = np.min(face_pose[:, 1]), np.max(body_pose[:7, 1]) h_ = h_max - h_min mid_w = mid_ w_min = mid_w - h_ // 2 w_max = mid_w + h_ // 2 w_min_all.append(w_min) w_max_all.append(w_max) h_min_all.append(h_min) h_max_all.append(h_max) mid_all.append(mid_w) w_min = np.min(w_min_all) w_max = np.max(w_max_all) h_min = np.min(h_min_all) h_max = np.max(h_max_all) mid = np.mean(mid_all) margin_ratio = 0.25 h_margin = (h_max - h_min) * margin_ratio h_min = max(h_min - h_margin * 0.8, 0) h_max = min(h_max + h_margin * 0.1, 1) h_new = h_max - h_min h_min_real = int(h_min * height) h_max_real = int(h_max * height) mid_real = int(mid * width) height_new = h_max_real - h_min_real + 1 width_new = height_new w_min_real = mid_real - width_new // 2 if w_min_real < 0: w_min_real = 0 width_new = mid_real * 2 w_max_real = w_min_real + width_new w_min = w_min_real / width w_max = w_max_real / width imh_new, imw_new, rb, re, cb, ce = resize_and_pad_param(height_new, width_new, max_size) res = {'draw_pose_params': [imh_new, imw_new, rb, re, cb, ce], 'pose_params': [w_min, w_max, h_min, h_max], 'video_params': [h_min_real, h_max_real, w_min_real, w_max_real], } return res def save_pose_params_item(input_items): detected_pose, pose_params, draw_pose_params, save_dir = input_items w_min, w_max, h_min, h_max = pose_params num = detected_pose['num'] candidate_body = detected_pose['bodies']['candidate'] candidate_face = detected_pose['faces'][0] candidate_hand = detected_pose['hands'] candidate_body[:, 0] = (candidate_body[:, 0] - w_min) / (w_max - w_min) candidate_body[:, 1] = (candidate_body[:, 1] - h_min) / (h_max - h_min) candidate_face[:, 0] = (candidate_face[:, 0] - w_min) / (w_max - w_min) candidate_face[:, 1] = (candidate_face[:, 1] - h_min) / (h_max - h_min) candidate_hand[:, :, 0] = (candidate_hand[:, :, 0] - w_min) / (w_max - w_min) candidate_hand[:, :, 1] = (candidate_hand[:, :, 1] - h_min) / (h_max - h_min) detected_pose['bodies']['candidate'] = candidate_body detected_pose['faces'] = candidate_face.reshape(1, candidate_face.shape[0], candidate_face.shape[1]) detected_pose['hands'] = candidate_hand detected_pose['draw_pose_params'] = draw_pose_params np.save(save_dir + '/' + str(num) + '.npy', detected_pose) def save_pose_params(detected_poses, pose_params, draw_pose_params, save_dir): #save_dir = ori_video_path.replace('video', 'pose/') if not os.path.exists(save_dir): os.makedirs(save_dir) input_list = [] for i, detected_pose in enumerate(detected_poses): input_list.append([detected_pose, pose_params, draw_pose_params, save_dir]) pool = ThreadPool(8) pool.map(save_pose_params_item, input_list) pool.close() pool.join() return save_dir def get_img_pose(visualizer, img_path: str, sample_stride: int = 1, max_frame=None): # read input img frame = cv2.imread(img_path) height, width, _ = frame.shape short_size = min(height, width) resize_ratio = max(MAX_SIZE / short_size, 1.0) frame = cv2.resize(frame, (int(resize_ratio * width), int(resize_ratio * height))) height, width, _ = frame.shape detected_poses = [visualizer(frame)] visualizer.release_memory() return detected_poses, height, width, frame def save_aligned_img(ori_frame, video_params, max_size): h_min_real, h_max_real, w_min_real, w_max_real = video_params img = ori_frame[h_min_real:h_max_real, w_min_real:w_max_real, :] img_aligened = resize_and_pad(img, max_size=max_size) print('aligned img shape:', img_aligened.shape) save_dir = './assets/refimg_aligned' os.makedirs(save_dir, exist_ok=True) save_path = os.path.join(save_dir, 'aligned.png') cv2.imwrite(save_path, img_aligened) return save_path