Merge pull request #91 from nvbln/fix_progressbar

Fix progress indicator in offline video processing
portugueslab · Mar 5, 2024 · 6acca5d · 6acca5d
2 parents df21fb6 + d6fc1a0
commit 6acca5d
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/stytra/experiments/tracking_experiments.py b/stytra/experiments/tracking_experiments.py
@@ -161,7 +161,8 @@ def wrap_up(self, *args, **kwargs) -> None:
         for q in [self.camera.frame_queue]:
             q.clear()
 
-        self.camera.join()
+        if self.camera.is_alive():
+            self.camera.join()
 
     def _setup_frame_dispatcher(self, recording_event: Event = None) -> DispatchProcess:
         """

diff --git a/stytra/hardware/video/__init__.py b/stytra/hardware/video/__init__.py
@@ -331,8 +331,11 @@ def run(self):
             container.streams.video[0].thread_count = 1
 
             prt = None
-            while self.loop:
+            while not self.kill_event.is_set():
                 for framedata in container.decode(video=0):
+                    if self.kill_event.is_set():
+                        break
+
                     messages = []
                     if self.paused:
                         frame = self.old_frame

diff --git a/stytra/offline/track_video.py b/stytra/offline/track_video.py
@@ -65,15 +65,16 @@ def track(self):
         fileformat = self.cmb_fmt.currentText()
 
         self.exp.camera.kill_event.set()
-        reader = imageio.get_reader(str(self.input_path))
+        reader = imageio.get_reader(str(self.input_path), "ffmpeg")
         data = []
         self.exp.window_main.stream_plot.toggle_freeze()
 
         output_name = str(self.output_path) + "." + fileformat
         self.diag_track.show()
-        l = reader.get_length()
-        if not (0 < l < 100000):
-            l = 1
+        if hasattr(reader, "count_frames"):
+            l = reader.count_frames()
+        else:
+            l = reader.get_length()
         self.diag_track.prog_track.setMaximum(l)
         self.diag_track.lbl_status.setText("Tracking to " + output_name)