bug fix

jonahpearl · Feb 4, 2023 · e57b1ff · e57b1ff
1 parent d82ba08
commit e57b1ff
Showing 1 changed file with 17 additions and 6 deletions.
diff --git a/top_bottom_moseq/segmentation.py b/top_bottom_moseq/segmentation.py
@@ -47,22 +47,33 @@ def segment_session(prefix,
     if output_prefix is None:
         output_prefix = prefix
 
+    # Check if already done
+    matched_frames = load_matched_frames(prefix, camera_names)
+    if all(
+        [check_if_already_done(
+            out_movie,
+            len(matched_frames), 
+            overwrite=overwrite) 
+            for camera in camera_names 
+                for out_movie in [
+                    output_prefix+'.{}.mouse_mask.avi'.format(camera),
+                    output_prefix+'.{}.occl_mask.avi'.format(camera)
+                    ]
+                ]
+            ):
+        print('Movies already segmented, continuing...')
+        return
+
     # Load torch models and matched frames betw top/bottom cams
     mouse_model = load_segmentation_model(mouse_model_weights)
     occl_model = load_segmentation_model(occlusion_model_weights)
-    matched_frames = load_matched_frames(prefix, camera_names)
 
     for camera,frames in zip(camera_names,matched_frames.T):
 
         # Prepare names of files
         ir_reader_in = prefix+'.{}.ir.avi'.format(camera)
         mouse_mask_out = output_prefix+'.{}.mouse_mask.avi'.format(camera)
         occl_mask_out = output_prefix+'.{}.occl_mask.avi'.format(camera)
-
-        # Don't process if already done!
-        if all([check_if_already_done(out_movie, len(frames), overwrite=overwrite) for out_movie in [mouse_mask_out, occl_mask_out]]):
-            print('Movies already segmented, continuing...')
-            return
 
         # Segment the mouse in each frame
         with ExitStack() as stack: