From 362a3140fb1c76f59955be4b3e4bdbc42c8345bc Mon Sep 17 00:00:00 2001 From: songkey-win Date: Tue, 10 Dec 2024 21:05:08 +0800 Subject: [PATCH] lc --- meme.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/meme.py b/meme.py index 4fc3335..54b1325 100644 --- a/meme.py +++ b/meme.py @@ -224,12 +224,14 @@ class CropPortrait: # print(image_np.shape) face_toolkits['face_aligner'].reset_track() faces = face_toolkits['face_aligner'].forward(image_np) - assert len(faces) > 0 - 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'] + if len(faces) > 0: + 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'] - new_image = crop_and_resize(image_np[np.newaxis, :,:,:], ref_landmark[np.newaxis, :,:], 512, crop=True)[0] + new_image = crop_and_resize(image_np[np.newaxis, :,:,:], ref_landmark[np.newaxis, :,:], 512, crop=True)[0] + else: + new_image = image_np[np.newaxis, :,:,:] new_image = cv2.cvtColor(new_image[0], cv2.COLOR_RGB2BGR) return (torch.from_numpy(new_image[np.newaxis, :,:,:]).float() / 255., )