From 0d46bb867f07d78a87c762ff791b3a8f38e200fd Mon Sep 17 00:00:00 2001 From: songkey-win Date: Tue, 10 Dec 2024 22:04:21 +0800 Subject: [PATCH] lc --- hellomeme/utils.py | 11 ++--------- meme.py | 31 +++++++++++++++++++++++-------- 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/hellomeme/utils.py b/hellomeme/utils.py index de3d046..ec06703 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -308,9 +308,9 @@ def crop_and_resize(frames, landmarks, save_size=512, crop=True): all_tl, all_br = np.min(landmarks, axis=1), np.max(landmarks, axis=1) mean_wh = np.mean(all_br - all_tl, axis=0) tl, br = np.min(all_tl, axis=0), np.max(all_br, axis=0) - new_size = min(max(mean_wh) * 2.5, min(H, W) - 1) + new_size = min(max(mean_wh) * 2.3, min(H, W) - 1) fcenter = (tl + br) * 0.5 - fcenter[1] -= new_size * 0.12 + fcenter[1] -= new_size * 0.115 ftl = fcenter - new_size * 0.5 fbr = fcenter + new_size * 0.5 @@ -333,15 +333,8 @@ def crop_and_resize(frames, landmarks, save_size=512, crop=True): opt_tl = (opt_center - opt_size * 0.5).astype(int) opt_br = (opt_center + opt_size * 0.5).astype(int) - # print('H, W', H, W) - # print('opt_center', opt_center) - # print('opt_size', opt_size) - # print('opt_tl', opt_tl) - # print('opt_br', opt_br) - frames = frames[:, opt_tl[1]:opt_br[1], opt_tl[0]:opt_br[0], :].copy() - # print('frames', frames.shape) landmarks = landmarks - opt_tl ratio = save_size / (opt_br - opt_tl) else: diff --git a/meme.py b/meme.py index 54b1325..eead5bb 100644 --- a/meme.py +++ b/meme.py @@ -13,6 +13,7 @@ from einops import rearrange import importlib.metadata import folder_paths +import logging cur_dir = osp.dirname(osp.abspath(__file__)) @@ -220,8 +221,10 @@ class CropPortrait: def crop_portrait(self, image, face_toolkits): image_np = cv2.cvtColor((image[0] * 255).cpu().numpy().astype(np.uint8), cv2.COLOR_BGR2RGB) - # image_np = cv2.resize(image_np, (512, 512)) - # print(image_np.shape) + + if min(image_np.shape[:2]) < 512: + raise Exception(f'Image size is too small -> min{image_np.shape[:2]} < 512') + face_toolkits['face_aligner'].reset_track() faces = face_toolkits['face_aligner'].forward(image_np) if len(faces) > 0: @@ -229,9 +232,11 @@ class CropPortrait: x['face_rect'][3] - x['face_rect'][1]))[-1] ref_landmark = face['pre_kpt_222'] - new_image = crop_and_resize(image_np[np.newaxis, :,:,:], ref_landmark[np.newaxis, :,:], 512, crop=True)[0] + new_image, new_landmark = crop_and_resize(image_np[np.newaxis, :,:,:], ref_landmark[np.newaxis, :,:], 512, crop=True) + # for x, y in new_landmark[0]: + # cv2.circle(new_image[0], (int(x), int(y)), 2, (0, 255, 0), -1) else: - new_image = image_np[np.newaxis, :,:,:] + raise Exception('No face detected') new_image = cv2.cvtColor(new_image[0], cv2.COLOR_RGB2BGR) return (torch.from_numpy(new_image[np.newaxis, :,:,:]).float() / 255., ) @@ -254,9 +259,11 @@ class GetFaceLandmarks: def get_face_landmarks(self, face_toolkits, images): frame_list = [cv2.cvtColor((frame * 255).cpu().numpy().astype(np.uint8), cv2.COLOR_BGR2RGB) for frame in images] frame_num = len(frame_list) - assert frame_num > 0 + if frame_num == 0: + raise Exception('No image detected') _, landmark_list = det_landmarks(face_toolkits['face_aligner'], frame_list) - assert len(frame_list) == frame_num + if len(frame_list) != frame_num: + raise Exception('Not all images have face detected!') return (torch.from_numpy(landmark_list).float(), ) @@ -372,12 +379,16 @@ class HMPipelineImage: device = get_torch_device(gpu_id) image_np = (ref_image[0] * 255).cpu().numpy().astype(np.uint8) + if min(image_np.shape[:2]) < 512: + raise Exception(f'Reference image size is too small -> min{image_np.shape[:2]} < 512') + image_np = cv2.resize(image_np, (512, 512)) image_pil = Image.fromarray(image_np) face_toolkits['face_aligner'].reset_track() faces = face_toolkits['face_aligner'].forward(image_np) - assert len(faces) > 0 + if len(faces) == 0: raise Exception('No face detected') + face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * ( x['face_rect'][3] - x['face_rect'][1]))[-1] ref_landmark = face['pre_kpt_222'] @@ -459,13 +470,17 @@ class HMPipelineVideo: device = get_torch_device(gpu_id) image_np = (ref_image[0] * 255).cpu().numpy().astype(np.uint8) + if min(image_np.shape[:2]) < 512: + raise Exception(f'Reference image size is too small -> min{image_np.shape[:2]} < 512') + image_np = cv2.resize(image_np, (512, 512)) image_pil = Image.fromarray(image_np) face_toolkits['face_aligner'].reset_track() faces = face_toolkits['face_aligner'].forward(image_np) face_toolkits['face_aligner'].reset_track() - assert len(faces) > 0 + if len(faces) == 0: raise Exception('No face detected') + face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * ( x['face_rect'][3] - x['face_rect'][1]))[-1] ref_landmark = face['pre_kpt_222']