chore: refine mask_crop loading

app.py

@@ -16,8 +16,8 @@ import gdown
 import os
 import spaces
 
-folder_url = f"https://drive.google.com/drive/folders/1UtKgzKjFAOmZkhNK-OYT0caJ_w2XAnib"
-gdown.download_folder(url=folder_url, output="pretrained_weights", quiet=False)
+# folder_url = f"https://drive.google.com/drive/folders/1UtKgzKjFAOmZkhNK-OYT0caJ_w2XAnib"
+# gdown.download_folder(url=folder_url, output="pretrained_weights", quiet=False)
 
 def partial_fields(target_class, kwargs):
     return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})
@@ -175,4 +175,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[eye_retargeting_slider, lip_retargeting_slider, retargeting_input_image]
     )
 
-demo.launch()
+demo.launch()

src/config/inference_config.py

@@ -5,6 +5,8 @@ config dataclass used for inference
 """
 
 import os.path as osp
+import cv2
+from numpy import ndarray
 from dataclasses import dataclass
 from typing import Literal, Tuple
 from .base_config import PrintableConfig, make_abs_path
@@ -38,7 +40,7 @@ class InferenceConfig(PrintableConfig):
 
     flag_write_result: bool = True  # whether to write output video
     flag_pasteback: bool = True  # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
-    mask_crop = None
+    mask_crop: ndarray = cv2.imread(make_abs_path('../utils/resources/mask_template.png'), cv2.IMREAD_COLOR)
     flag_write_gif: bool = False
     size_gif: int = 256
     ref_max_shape: int = 1280

src/utils/crop.py

@@ -409,4 +409,4 @@ def paste_back(image_to_processed, crop_M_c2o, rgb_ori, mask_ori):
     dsize = (rgb_ori.shape[1], rgb_ori.shape[0])
     result = _transform_img(image_to_processed, crop_M_c2o, dsize=dsize)
     result = np.clip(mask_ori * result + (1 - mask_ori) * rgb_ori, 0, 255).astype(np.uint8)
-    return result
+    return result