fix cv2 params

2022-10-09 21:32:13 +08:00
parent 521a1e2858
commit f4fcece180
6 changed files with 37 additions and 1 deletions
--- a/lama_cleaner/model/opencv2.py
+++ b/lama_cleaner/model/opencv2.py
@@ -2,6 +2,11 @@ import cv2
 from lama_cleaner.model.base import InpaintModel
 from lama_cleaner.schema import Config

+flag_map = {
+    "INPAINT_NS": cv2.INPAINT_NS,
+    "INPAINT_TELEA": cv2.INPAINT_TELEA
+}
+
 class OpenCV2(InpaintModel):
    pad_mod = 1

@@ -15,5 +20,5 @@ class OpenCV2(InpaintModel):
        mask: [H, W, 1]
        return: BGR IMAGE
        """
-        cur_res = cv2.inpaint(image[:,:,::-1], mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
+        cur_res = cv2.inpaint(image[:,:,::-1], mask, inpaintRadius=config.cv2_radius, flags=flag_map[config.cv2_flag])
        return cur_res