feat(#78): extra args (optional) for contain_image

williamfzc · Dec 5, 2019 · ae04813 · ae04813
1 parent 16b70eb
commit ae04813
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git a/stagesepx/cutter/cut_range.py b/stagesepx/cutter/cut_range.py
@@ -90,7 +90,9 @@ def contain_image(
         operator = self.video.get_operator()
         frame = operator.get_frame_by_id(target_id)
 
-        return toolbox.match_template_with_object(image_object, frame.data)
+        return toolbox.match_template_with_object(
+            image_object, frame.data, *args, **kwargs
+        )
 
     def pick(
         self, frame_count: int = None, is_random: bool = None, *_, **__

diff --git a/stagesepx/toolbox.py b/stagesepx/toolbox.py
@@ -16,6 +16,7 @@
 from loguru import logger
 from findit import FindIt
 
+
 # DO NOT IMPORT ANYTHING FROM STAGESEPX HERE
 # MAKE TOOLBOX STATIC
 
@@ -275,20 +276,20 @@ def fps_convert(
 
 
 def match_template_with_object(
-    template: np.ndarray, target: np.ndarray
+    template: np.ndarray, target: np.ndarray, *args, **kwargs
 ) -> typing.Dict[str, typing.Any]:
     fi = FindIt(engine=["template"])
     # load template
     fi_template_name = "default"
-    fi.load_template(fi_template_name, pic_object=template)
+    fi.load_template(fi_template_name, pic_object=template, *args, **kwargs)
 
     result = fi.find("", target_pic_object=target)
     return result["data"][fi_template_name]["TemplateEngine"]
 
 
 def match_template_with_path(
-    template: str, target: np.ndarray
+    template: str, target: np.ndarray, *args, **kwargs
 ) -> typing.Dict[str, typing.Any]:
     assert os.path.isfile(template), f"image {template} not existed"
     template_object = turn_grey(imread(template))
-    return match_template_with_object(template_object, target)
+    return match_template_with_object(template_object, target, *args, **kwargs)
diff --git a/stagesepx/video.py b/stagesepx/video.py
@@ -30,15 +30,19 @@ def copy(self):
         return VideoFrame(self.frame_id, self.timestamp, self.data[:])
 
     def contain_image(
-        self, image_path: str = None, image_object: np.ndarray = None
+        self, image_path: str = None, image_object: np.ndarray = None, *args, **kwargs
     ) -> typing.Dict[str, typing.Any]:
         assert image_path or image_object, "should fill image_path or image_object"
 
         if image_path:
             logger.debug(f"found image path, use it first: {image_path}")
-            return toolbox.match_template_with_path(image_path, self.data)
+            return toolbox.match_template_with_path(
+                image_path, self.data, *args, **kwargs
+            )
         image_object = toolbox.turn_grey(image_object)
-        return toolbox.match_template_with_object(image_object, self.data)
+        return toolbox.match_template_with_object(
+            image_object, self.data, *args, **kwargs
+        )
 
 
 class _BaseFrameOperator(object):