feat(#78): support contain_image in VideoFrame

and SingleClassifierResult

stagesepx/toolbox.py

@@ -284,3 +284,11 @@ def match_template_with_object(
 
     result = fi.find("", target_pic_object=target)
     return result["data"][fi_template_name]["TemplateEngine"]
+
+
+def match_template_with_path(
+    template: str, target: np.ndarray
+) -> 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)

stagesepx/video.py

@@ -11,9 +11,9 @@
 
 class VideoFrame(object):
     def __init__(self, frame_id: int, timestamp: float, data: np.ndarray):
-        self.frame_id = frame_id
-        self.timestamp = timestamp
-        self.data = data
+        self.frame_id: int = frame_id
+        self.timestamp: float = timestamp
+        self.data: np.ndarray = data
 
     def __str__(self):
         return f"<VideoFrame id={self.frame_id} timestamp={self.timestamp}>"
@@ -29,6 +29,17 @@ def init(cls, cap: cv2.VideoCapture, frame: np.ndarray) -> "VideoFrame":
     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
+    ) -> 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)
+        image_object = toolbox.turn_grey(image_object)
+        return toolbox.match_template_with_object(image_object, self.data)
+
 
 class _BaseFrameOperator(object):
     def __init__(self, video: "VideoObject"):

test/test_toolbox.py

@@ -82,3 +82,6 @@ def test_match_template():
     image2 = toolbox.imread(IMAGE_PATH)
     ret = toolbox.match_template_with_object(image1, image2)
     assert ret["ok"]
+
+    ret = toolbox.match_template_with_path(IMAGE_PATH, image2)
+    assert ret["ok"]

test/test_video.py

@@ -4,6 +4,8 @@
 
 PROJECT_PATH = os.path.dirname(os.path.dirname(__file__))
 VIDEO_PATH = os.path.join(PROJECT_PATH, "demo.mp4")
+IMAGE_NAME = "demo.jpg"
+IMAGE_PATH = os.path.join(PROJECT_PATH, IMAGE_NAME)
 
 
 def test_read_from_file():
@@ -45,3 +47,10 @@ def test_custom_ffmpeg():
         VideoObject(VIDEO_PATH, fps=30)
     except FileNotFoundError:
         pass
+
+
+def test_contain_image():
+    v = VideoObject(VIDEO_PATH)
+    v.load_frames()
+    ret = v.data[0].contain_image(IMAGE_PATH)
+    assert ret["ok"]