From 2e40fe38205cd5c41571b7fdade6a4780a2c7209 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Tue, 9 Jul 2024 21:23:10 +0300 Subject: [PATCH] Update live_portrait_pipeline.py --- liveportrait/live_portrait_pipeline.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/liveportrait/live_portrait_pipeline.py b/liveportrait/live_portrait_pipeline.py index 07d6ed6..246c1fb 100644 --- a/liveportrait/live_portrait_pipeline.py +++ b/liveportrait/live_portrait_pipeline.py @@ -79,9 +79,10 @@ class LivePortraitPipeline(object): source_np, i, total_frames, mismatch_method ) driving_frame = driving_images_np[i] - - source_lmk = crop_info[i]["lmk_crop"] - img_crop_256x256 = crop_info[i]["img_crop_256x256"] + + safe_index = min(i, len(crop_info) - 1) + source_lmk = crop_info[safe_index]["lmk_crop"] + img_crop_256x256 = crop_info[safe_index]["img_crop_256x256"] if inference_cfg.flag_do_crop: I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256) @@ -252,7 +253,7 @@ class LivePortraitPipeline(object): # Transform and blend I_p_i_to_ori = _transform_img( I_p_i, - crop_info[i]["M_c2o"], + crop_info[safe_index]["M_c2o"], dsize=(source_frame_rgb.shape[1], source_frame_rgb.shape[0]), ) @@ -261,7 +262,7 @@ class LivePortraitPipeline(object): inference_cfg.mask_crop = cv2.imread(os.path.join(script_directory, "utils", "resources", "mask_template.png"), cv2.IMREAD_COLOR) mask_ori = _transform_img( inference_cfg.mask_crop, - crop_info[i]["M_c2o"], + crop_info[safe_index]["M_c2o"], dsize=(source_frame_rgb.shape[1], source_frame_rgb.shape[0]), ) mask_ori = mask_ori.astype(np.float32) / 255.0