Add cv2 get(CAP_PROP_POS_FRAMES)

ixlab · Oct 28, 2024 · fafe087 · fafe087
1 parent f9b8786
commit fafe087
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/snake-pit/test_cv2.py b/snake-pit/test_cv2.py
@@ -14,6 +14,7 @@ def test_constants():
     assert ocv_cv2.CAP_PROP_FRAME_WIDTH == vf_cv2.CAP_PROP_FRAME_WIDTH
     assert ocv_cv2.CAP_PROP_FRAME_HEIGHT == vf_cv2.CAP_PROP_FRAME_HEIGHT
     assert ocv_cv2.CAP_PROP_FPS == vf_cv2.CAP_PROP_FPS
+    assert ocv_cv2.CAP_PROP_FRAME_COUNT == vf_cv2.CAP_PROP_FRAME_COUNT
 
     assert ocv_cv2.FONT_HERSHEY_SIMPLEX == vf_cv2.FONT_HERSHEY_SIMPLEX
     assert ocv_cv2.FONT_HERSHEY_PLAIN == vf_cv2.FONT_HERSHEY_PLAIN

diff --git a/vidformer-py/vidformer/cv2/vf_cv2.py b/vidformer-py/vidformer/cv2/vf_cv2.py
@@ -28,6 +28,7 @@
 CAP_PROP_FRAME_WIDTH = 3
 CAP_PROP_FRAME_HEIGHT = 4
 CAP_PROP_FPS = 5
+CAP_PROP_FRAME_COUNT = 7
 
 FONT_HERSHEY_SIMPLEX = 0
 FONT_HERSHEY_PLAIN = 1
@@ -128,6 +129,10 @@ def get(self, prop):
             return self._source.fmt()["width"]
         elif prop == CAP_PROP_FRAME_HEIGHT:
             return self._source.fmt()["height"]
+        elif prop == CAP_PROP_FRAME_COUNT:
+            return len(self._source.ts())
+        elif prop == CAP_PROP_POS_FRAMES:
+            return self._next_frame_idx
 
         raise Exception(f"Unknown property {prop}")
 

diff --git a/vidformer-py/vidformer/vf.py b/vidformer-py/vidformer/vf.py
@@ -540,7 +540,7 @@ def __init__(self, source):
 
     def __getitem__(self, idx):
         if type(idx) != int:
-            raise Exception("Source iloc index must be an integer")
+            raise Exception(f"Source iloc index must be an integer, got a {type(idx)}")
         return SourceExpr(self._source, idx, True)