feat: support 'step' in classify

williamfzc · Jul 19, 2019 · d828b40 · d828b40
1 parent 0ca7e9b
commit d828b40
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/example/classify.py b/example/classify.py
@@ -12,7 +12,12 @@
 # 如果你改动了，这里也要做相应修改
 cl.load('./cut_result')
 # 开始分析即可
-res = cl.classify('../test.mp4')
+res = cl.classify(
+    '../test.mp4',
+    # 步长，可以自行设置用于平衡效率与颗粒度
+    # 默认为1，即每帧都检测
+    step=1
+)
 
 # 分类出来的结果是一个 list，里面包含 ClassifierResult 对象
 # 你可以用它进行二次开发

diff --git a/stagesepx/classifier/ssim.py b/stagesepx/classifier/ssim.py
@@ -11,12 +11,15 @@ class SSIMClassifier(BaseClassifier):
     def classify(self,
                  video_path: str,
                  limit_range: typing.List[VideoCutRange] = None,
+                 step: int = None,
                  threshold: float = None) -> typing.List[ClassifierResult]:
         logger.debug(f'classify with {self.__class__.__name__}')
         assert self.data, 'should load data first'
 
         if not threshold:
             threshold = 0.85
+        if not step:
+            step = 1
 
         final_result: typing.List[ClassifierResult] = list()
         with toolbox.video_capture(video_path) as cap:
@@ -52,5 +55,6 @@ def classify(self,
                 logger.debug(f'frame {frame_id} ({frame_timestamp}) belongs to {result[0]}')
 
                 final_result.append(ClassifierResult(video_path, frame_id, frame_timestamp, result[0]))
+                toolbox.video_jump(cap, frame_id + step - 1)
                 ret, frame = cap.read()
         return final_result
diff --git a/stagesepx/classifier/svm.py b/stagesepx/classifier/svm.py
@@ -87,11 +87,15 @@ def predict_with_object(self, pic_object: np.ndarray) -> str:
 
     def classify(self,
                  video_path: str,
-                 limit_range: typing.List[VideoCutRange] = None) -> typing.List[ClassifierResult]:
+                 limit_range: typing.List[VideoCutRange] = None,
+                 step: int = None) -> typing.List[ClassifierResult]:
         logger.debug(f'classify with {self.__class__.__name__}')
         assert self.data, 'should load data first'
         assert self._model, 'should train before classify'
 
+        if not step:
+            step = 1
+
         final_result: typing.List[ClassifierResult] = list()
         with toolbox.video_capture(video_path) as cap:
             ret, frame = cap.read()
@@ -108,5 +112,6 @@ def classify(self,
                 result = self.predict_with_object(frame)
                 logger.debug(f'frame {frame_id} ({frame_timestamp}) belongs to {result}')
                 final_result.append(ClassifierResult(video_path, frame_id, frame_timestamp, result))
+                toolbox.video_jump(cap, frame_id + step - 1)
                 ret, frame = cap.read()
         return final_result