simba-uw-tf-dev 4.6.1__py3-none-any.whl → 4.6.3__py3-none-any.whl

simba/plotting/spontaneous_alternation_plotter.py

@@ -14,7 +14,7 @@ from simba.data_processors.spontaneous_alternation_calculator import \
 from simba.mixins.config_reader import ConfigReader
 from simba.utils.checks import (check_file_exist_and_readable, check_float,
                                 check_int, check_str, check_valid_lst)
-from simba.utils.data import detect_bouts
+from simba.utils.data import detect_bouts, terminate_cpu_pool
 from simba.utils.enums import Formats, Paths, TextOptions
 from simba.utils.errors import AnimalNumberError, InvalidInputError
 from simba.utils.printing import stdout_success
@@ -296,8 +296,7 @@ class SpontaneousAlternationsPlotter(ConfigReader):
                 pool.imap(constants, frm_index, chunksize=self.multiprocess_chunksize)
             ):
                 print(f"Section {cnt} complete...")
-        pool.terminate()
-        pool.join()
+        terminate_cpu_pool(pool=pool, force=False)
         print(f"Joining {sa_computer.video_name} multiprocessed video...")
         concatenate_videos_in_folder(in_folder=self.temp_folder, save_path=save_path)
         self.timer.stop_timer()

simba/plotting/yolo_pose_track_visualizer.py

@@ -13,9 +13,10 @@ from simba.mixins.plotting_mixin import PlottingMixin
 from simba.utils.checks import (check_file_exist_and_readable, check_float,
                                 check_if_dir_exists, check_int,
                                 check_valid_boolean, check_valid_dataframe)
+from simba.utils.data import terminate_cpu_pool, get_cpu_pool
 from simba.utils.enums import Defaults, Formats
 from simba.utils.errors import InvalidFilepathError, NoFilesFoundError
-from simba.utils.lookups import get_random_color_palette, intermittent_palette
+from simba.utils.lookups import get_random_color_palette, intermittent_palette, get_current_time
 from simba.utils.printing import SimbaTimer, stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
                                     create_directory,
@@ -53,13 +54,12 @@ def _yolo_keypoint_track_visualizer(frm_ids: np.ndarray,
     video_save_path = os.path.join(save_dir, f'{batch_id}.mp4')
     video_writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], (video_meta_data["width"], video_meta_data["height"]))
     while current_frm <= end_frm:
-        print(f'Processing frame {current_frm}/{video_meta_data["frame_count"]} (batch: {batch_id})...')
+        print(f'[{get_current_time()}] Processing frame {current_frm}/{video_meta_data["frame_count"]} (batch: {batch_id}, video name: {video_meta_data["video_name"]})...')
         img = read_frm_of_video(video_path=video_path, frame_index=current_frm, raise_error=False)
         if img is not None:
             frm_data = data.loc[data[FRAME] == current_frm]
             frm_data = frm_data[frm_data[CONFIDENCE] > threshold]
             for cnt, (row, row_data) in enumerate(frm_data.iterrows()):
-
                 clrs = np.array(palettes[int(row_data[TRACK])]).astype(np.int32)
                 bbox_cords = row_data[BOX_CORD_FIELDS].values.astype(np.int32).reshape(-1, 2)
                 kp_coords = row_data.drop(EXPECTED_COLS).values.astype(np.int32).reshape(-1, 3)[:, :-1]
@@ -159,6 +159,8 @@ class YOLOPoseTrackVisualizer():
         self.threshold, self.circle_size, self.thickness, self.show_bbox, self.overwrite = threshold, circle_size, thickness, bbox, overwrite
 
     def run(self):
+        self.pool = get_cpu_pool(core_cnt=self.core_cnt, maxtasksperchild=Defaults.MAXIMUM_MAX_TASK_PER_CHILD.value, source=self.__class__.__name__)
+        self.timer = SimbaTimer(start=True)
         for video_cnt, (video_name, data_path) in enumerate(self.data_paths.items()):
             print(f'Visualizing YOLO pose tracks in video {video_name} ({video_cnt+1}/{len(self.data_paths.keys())}) ...')
             video_timer = SimbaTimer(start=True)
@@ -189,23 +191,25 @@ class YOLOPoseTrackVisualizer():
 
             frm_batches = np.array_split(np.array(list(range(0, df_frm_cnt))), self.core_cnt)
             frm_batches = [(i, j) for i, j in enumerate(frm_batches)]
-            with multiprocessing.Pool(self.core_cnt, maxtasksperchild=Defaults.MAXIMUM_MAX_TASK_PER_CHILD.value) as pool:
-                constants = functools.partial(_yolo_keypoint_track_visualizer,
-                                              data=self.data_df,
-                                              threshold=self.threshold,
-                                              video_path=self.video_paths[video_name],
-                                              save_dir=video_temp_dir,
-                                              circle_size=video_circle_size,
-                                              thickness=video_thickness,
-                                              palettes=video_palettes,
-                                              show_bbox=self.show_bbox)
-                for cnt, result in enumerate(pool.imap(constants, frm_batches, chunksize=1)):
-                    print(f'Video batch {result+1}/{self.core_cnt} complete...')
-            pool.terminate()
-            pool.join()
+            constants = functools.partial(_yolo_keypoint_track_visualizer,
+                                          data=self.data_df,
+                                          threshold=self.threshold,
+                                          video_path=self.video_paths[video_name],
+                                          save_dir=video_temp_dir,
+                                          circle_size=video_circle_size,
+                                          thickness=video_thickness,
+                                          palettes=video_palettes,
+                                          show_bbox=self.show_bbox)
+            for cnt, result in enumerate(self.pool.imap(constants, frm_batches, chunksize=1)):
+                print(f'[{get_current_time()}] Video batch {result+1}/{self.core_cnt} complete...')
             video_timer.stop_timer()
             concatenate_videos_in_folder(in_folder=video_temp_dir, save_path=save_path, gpu=True)
             stdout_success(msg=f'YOLO track pose video saved at {save_path}', source=self.__class__.__name__, elapsed_time=video_timer.elapsed_time_str)
+        terminate_cpu_pool(pool=self.pool, force=False, source=self.__class__.__name__)
+        self.timer.stop_timer()
+        stdout_success(msg=f'YOLO track pose video data for {len(self.data_paths.keys())} videos saved in {self.save_dir}', source=self.__class__.__name__,  elapsed_time=self.timer.elapsed_time_str)
+
+
 #
 # if __name__ == "__main__" and not hasattr(sys, 'ps1'):
 #     parser = argparse.ArgumentParser(description="Visualize YOLO pose tracking CSV outputs on their source videos.")
@@ -247,13 +251,13 @@ class YOLOPoseTrackVisualizer():
 # #kp_vis.run()
 
 
-if __name__ == "__main__":
-    VIDEO_PATH = r"E:\netholabs_videos\primeintellect_100_videos"
-    DATA_PATH = r"E:\netholabs_videos\primeintellect_100_largest"
-    SAVE_DIR = r"E:\netholabs_videos\primeintellect_100_videos\out"
-    kp_vis = YOLOPoseTrackVisualizer(data_path=DATA_PATH,
-                                     video_path=VIDEO_PATH,
-                                     save_dir=SAVE_DIR,
-                                     core_cnt=8,
-                                     bbox=True)
-    kp_vis.run()
+# if __name__ == "__main__":
+#     VIDEO_PATH = r"E:\netholabs_videos\primeintellect_100_videos"
+#     DATA_PATH = r"E:\netholabs_videos\primeintellect_100_largest"
+#     SAVE_DIR = r"E:\netholabs_videos\primeintellect_100_videos\out"
+#     kp_vis = YOLOPoseTrackVisualizer(data_path=DATA_PATH,
+#                                      video_path=VIDEO_PATH,
+#                                      save_dir=SAVE_DIR,
+#                                      core_cnt=8,
+#                                      bbox=True)
+#     kp_vis.run()

simba/plotting/yolo_pose_visualizer.py

@@ -14,7 +14,7 @@ from simba.utils.checks import (check_file_exist_and_readable, check_float,
                                 check_if_dir_exists, check_int,
                                 check_valid_boolean, check_valid_dataframe,
                                 check_valid_lst, check_valid_tuple)
-from simba.utils.data import create_color_palette
+from simba.utils.data import create_color_palette, terminate_cpu_pool, get_cpu_pool
 from simba.utils.enums import Defaults, Options
 from simba.utils.errors import (CountError, DataHeaderError, FrameRangeError,
                                 InvalidInputError, NoDataError)
@@ -24,8 +24,7 @@ from simba.utils.read_write import (concatenate_videos_in_folder,
                                     find_files_of_filetypes_in_directory,
                                     get_fn_ext, get_video_meta_data,
                                     read_frm_of_video, recursive_file_search,
-                                    remove_a_folder)
-from simba.utils.warnings import InvalidValueWarning
+                                    remove_a_folder, get_current_time)
 
 FRAME = 'FRAME'
 CLASS_ID = 'CLASS_ID'
@@ -58,7 +57,7 @@ def _yolo_keypoint_visualizer(frm_ids: np.ndarray,
     if TRACK in data.columns:
         data = data.drop([TRACK], axis=1)
     while current_frm <= end_frm:
-        print(f'Processing frame {current_frm}/{video_meta_data["frame_count"]} (batch: {batch_id}, video: {video_meta_data["video_name"]})...')
+        print(f'[{get_current_time()}] Processing frame {current_frm}/{video_meta_data["frame_count"]} (batch: {batch_id}, video: {video_meta_data["video_name"]})...')
         img = read_frm_of_video(video_path=video_path, frame_index=current_frm)
         frm_data = data.loc[data[FRAME] == current_frm]
         frm_data = frm_data[frm_data[CONFIDENCE] > threshold]
@@ -206,6 +205,7 @@ class YOLOPoseVisualizer():
         self.timer = SimbaTimer(start=True)
 
     def run(self):
+        self.pool = get_cpu_pool(core_cnt=self.core_cnt, maxtasksperchild=Defaults.MAXIMUM_MAX_TASK_PER_CHILD.value, verbose=True, source=self.__class__.__name__)
         for video_cnt, (video_name, data_path) in enumerate(self.data_paths.items()):
             video_timer = SimbaTimer(start=True)
             self.video_temp_dir = os.path.join(self.save_dir, video_name, "temp")
@@ -248,26 +248,24 @@ class YOLOPoseVisualizer():
                 thickness = deepcopy(self.thickness)
             frm_batches = np.array_split(np.array(list(range(0, self.df_frm_cnt))), self.core_cnt)
             frm_batches = [(i, j) for i, j in enumerate(frm_batches)]
-            if self.verbose: print(f'Visualizing video {self.video_meta_data["video_name"]} (frame count: {self.video_meta_data["frame_count"]})...')
-            with multiprocessing.Pool(self.core_cnt, maxtasksperchild=Defaults.MAXIMUM_MAX_TASK_PER_CHILD.value) as pool:
-                constants = functools.partial(_yolo_keypoint_visualizer,
-                                          data=self.data_df,
-                                          threshold=self.threshold,
-                                          video_path=self.video_paths[video_name],
-                                          save_dir=self.video_temp_dir,
-                                          circle_size=circle_size,
-                                          thickness=thickness,
-                                          palettes=self.clrs,
-                                          bbox=self.bbox,
-                                          skeleton=self.skeleton)
-                for cnt, result in enumerate(pool.imap(constants, frm_batches, chunksize=1)):
-                    print(f'Video batch {result+1}/{self.core_cnt} complete...')
-            pool.terminate()
-            pool.join()
+            if self.verbose: print(f'[{get_current_time()}] Visualizing video {self.video_meta_data["video_name"]} (frame count: {self.video_meta_data["frame_count"]})...')
+            constants = functools.partial(_yolo_keypoint_visualizer,
+                                      data=self.data_df,
+                                      threshold=self.threshold,
+                                      video_path=self.video_paths[video_name],
+                                      save_dir=self.video_temp_dir,
+                                      circle_size=circle_size,
+                                      thickness=thickness,
+                                      palettes=self.clrs,
+                                      bbox=self.bbox,
+                                      skeleton=self.skeleton)
+            for cnt, result in enumerate(self.pool.imap(constants, frm_batches, chunksize=1)):
+                print(f'[{get_current_time()}] Video batch {result+1}/{self.core_cnt} complete...')
             video_timer.stop_timer()
             concatenate_videos_in_folder(in_folder=self.video_temp_dir, save_path=self.save_path, gpu=True)
             stdout_success(msg=f'YOLO pose video saved at {self.save_path} (Video {video_cnt+1}/{len(list(self.data_paths.keys()))})', source=self.__class__.__name__, elapsed_time=video_timer.elapsed_time_str)
 
+        terminate_cpu_pool(pool=self.pool, force=False, source=self.__class__.__name__)
         self.timer.stop_timer()
         stdout_success(msg=f'{len(list(self.data_paths.keys()))} YOLO pose video saved in directory {self.save_dir}', source=self.__class__.__name__, elapsed_time=self.timer.elapsed_time_str)
 
@@ -413,18 +411,18 @@ class YOLOPoseVisualizer():
 #     kp_vis.run()
 
 
-# if __name__ == "__main__":
-#     video_path = r"E:\netholabs_videos\primeintellect_100_videos\cage_1_date_2025_08_28_hour_20_minute_21.avi"
-#     data_path = r"E:\netholabs_videos\primeintellect_100_largest\cage_1_date_2025_08_28_hour_20_minute_21.csv"
-#     save_dir = r'E:\netholabs_videos\test_order'
-#     kp_vis = YOLOPoseVisualizer(data_path=data_path,
-#                                 video_path=video_path,
-#                                 save_dir=save_dir,
-#                                 core_cnt=14,
-#                                 palettes=('tab20',),
-#                                 recursive=True,
-#                                 sample_n=None)
-#
-#
-#     kp_vis.run()
+if __name__ == "__main__":
+    video_path = r"E:\netholabs_videos\primeintellect_100_videos\cage_1_date_2025_08_28_hour_20_minute_21.avi"
+    data_path = r"E:\netholabs_videos\primeintellect_100_largest\cage_1_date_2025_08_28_hour_20_minute_21.csv"
+    save_dir = r'E:\netholabs_videos\test_order'
+    kp_vis = YOLOPoseVisualizer(data_path=data_path,
+                                video_path=video_path,
+                                save_dir=save_dir,
+                                core_cnt=14,
+                                palettes=('tab20',),
+                                recursive=True,
+                                sample_n=None)
+
+
+    kp_vis.run()
 

simba/plotting/yolo_seg_visualizer.py

@@ -12,7 +12,7 @@ from simba.utils.checks import (check_file_exist_and_readable, check_float,
                                 check_int, check_valid_boolean,
                                 check_valid_dataframe, check_valid_lst,
                                 check_valid_tuple)
-from simba.utils.data import create_color_palette
+from simba.utils.data import create_color_palette, terminate_cpu_pool
 from simba.utils.enums import Defaults, Options
 from simba.utils.errors import CountError, DataHeaderError, FrameRangeError
 from simba.utils.printing import SimbaTimer, stdout_success
@@ -140,8 +140,7 @@ class YOLOSegmentationVisualizer():
                                           shape_opacity=self.shape_opacity)
             for cnt, result in enumerate(pool.imap(constants, frm_batches, chunksize=1)):
                 print(f'Video batch {result+1}/{self.core_cnt} complete...')
-        pool.terminate()
-        pool.join()
+        terminate_cpu_pool(pool=pool, force=False)
         video_timer.stop_timer()
         concatenate_videos_in_folder(in_folder=self.video_temp_dir, save_path=self.save_path, gpu=True)
         stdout_success(msg=f'YOLO pose video saved at {self.save_path}', source=self.__class__.__name__, elapsed_time=video_timer.elapsed_time_str)

simba/roi_tools/roi_aggregate_stats_mp.py

@@ -17,7 +17,8 @@ from simba.utils.checks import (
     check_all_file_names_are_represented_in_video_log,
     check_file_exist_and_readable, check_float, check_if_dir_exists, check_int,
     check_that_column_exist, check_valid_boolean, check_valid_lst)
-from simba.utils.data import detect_bouts, slice_roi_dict_for_video
+from simba.utils.data import (detect_bouts, slice_roi_dict_for_video,
+                              terminate_cpu_pool)
 from simba.utils.enums import ROI_SETTINGS, Formats, Keys
 from simba.utils.errors import CountError, ROICoordinatesNotFoundError
 from simba.utils.printing import SimbaTimer, stdout_success
@@ -297,8 +298,8 @@ class ROIAggregateStatisticsAnalyzerMultiprocess(ConfigReader, FeatureExtraction
                 self.results.append(result); self.detailed_dfs.append(detailed_dfs)
                 print(f"Data batch core {batch_id} / {self.core_cnt} complete...")
         self.results = pd.concat(self.results, axis=0).reset_index(drop=True)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
+
 
     def save(self):
         self.__clean_results()

simba/roi_tools/roi_clf_calculator_mp.py

@@ -16,7 +16,8 @@ from simba.utils.checks import (
     check_all_file_names_are_represented_in_video_log,
     check_file_exist_and_readable, check_if_dir_exists, check_int,
     check_valid_boolean, check_valid_dataframe, check_valid_lst)
-from simba.utils.data import detect_bouts, slice_roi_dict_for_video
+from simba.utils.data import (detect_bouts, slice_roi_dict_for_video,
+                              terminate_cpu_pool)
 from simba.utils.enums import ROI_SETTINGS, Keys
 from simba.utils.errors import InvalidInputError, NoROIDataError
 from simba.utils.lookups import get_current_time
@@ -237,8 +238,7 @@ class ROIClfCalculatorMultiprocess(ConfigReader):
                 self.bouts_results.append(batch_bout_results)
                 print(f"Data batch core {batch_id + 1} / {self.core_cnt} complete...")
         self.bouts_results = pd.concat(self.bouts_results, axis=0).reset_index(drop=True) if len(self.bouts_results) > 0 else None
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
 
     def save(self):
         self.timer.stop_timer()