lc

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:

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']