improves: video improves (#2)

cereja-project · Jul 14, 2022 · 8cf9fc9 · 8cf9fc9
1 parent a579c4a
commit 8cf9fc9
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/calango/__init__.py b/calango/__init__.py
@@ -26,5 +26,5 @@
 from .devices import Mouse
 from .media import Image, VideoWriter, Video
 
-VERSION = "2.1.8.final.0"
+VERSION = "2.1.9.final.0"
 __version__ = get_version_pep440_compliant(VERSION)
diff --git a/calango/media.py b/calango/media.py
@@ -981,11 +981,30 @@ def get_frames(self, n_frames=None):
                 self.stop()
                 break
 
+    def _cut(self, start, end=None, step=1):
+        assert not self._cap.is_webcam, 'Not available for webcam'
+        assert not self._th_show_running, "The video is showing, so you can't get frames"
+        if not self.is_opened:
+            self._build()
+        filter_map = set(range(start, end or self.total_frames, step))
+        max_frame = max(filter_map)
+        for frame in self.get_frames():
+            if self.current_number_frame > max_frame:
+                self.stop()
+            if self.current_number_frame in filter_map:
+                yield frame
+
+    def cut(self, start, end, step=1):
+        return Video(list(self._cut(start, end, step=step)), fps=self.fps)
+
     def _save(self, file_path, n_frames=None, fourcc=None):
         VideoWriter.write_frames(file_path, self.get_frames(n_frames=n_frames), fps=self._cap.fps, fourcc=fourcc)
 
-    def save(self, file_path, n_frames=None, fourcc=None):
-        threading.Thread(target=self._save, args=(file_path, n_frames, fourcc)).start()
+    def save(self, file_path, n_frames=None, fourcc=None, use_thread=True):
+        if use_thread:
+            threading.Thread(target=self._save, args=(file_path, n_frames, fourcc)).start()
+            return
+        self._save(file_path, n_frames, fourcc)
 
     def show(self):
         if ON_COLAB_JUPYTER:
@@ -1017,7 +1036,7 @@ def save_frames(self, p: str, start=1, end=None, step=1, img_format='jpg', limit
         size_number = len(str(_frame_count))
         max_frame = max(filter_map)
         for frame in self.get_frames():
-            if not (self.current_number_frame < max_frame):
+            if self.current_number_frame > max_frame:
                 break
             prefix = cj.get_zero_mask(self.current_number_frame, size_number)
             if self.current_number_frame in filter_map: