simba-uw-tf-dev 4.6.2__py3-none-any.whl → 4.7.2__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,11 @@ 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 get_cpu_pool, terminate_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_current_time, get_random_color_palette,
+                                 intermittent_palette)
 from simba.utils.printing import SimbaTimer, stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
                                     create_directory,
@@ -53,13 +55,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 +160,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 +192,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 +252,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,8 @@ 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, get_cpu_pool,
+                              terminate_cpu_pool)
 from simba.utils.enums import Defaults, Options
 from simba.utils.errors import (CountError, DataHeaderError, FrameRangeError,
                                 InvalidInputError, NoDataError)
@@ -22,10 +23,9 @@ from simba.utils.printing import SimbaTimer, stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
                                     create_directory, find_core_cnt,
                                     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
+                                    get_current_time, get_fn_ext,
+                                    get_video_meta_data, read_frm_of_video,
+                                    recursive_file_search, remove_a_folder)
 
 FRAME = 'FRAME'
 CLASS_ID = 'CLASS_ID'
@@ -58,7 +58,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 +206,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 +249,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 +412,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/pose_importers/simba_blob_importer.py

@@ -44,10 +44,10 @@ class SimBABlobImporter(ConfigReader):
     :param Optional[bool] verbose: If True, prints progress messages. Default: True.
 
     :example:
-    >>> r = SimBABlobImporter(config_path=r"C:\troubleshooting\simba_blob_project\project_folder\project_config.ini", data_path=r'C:\troubleshooting\simba_blob_project\data')
+    >>> r = SimBABlobImporter(config_path=r"C:/troubleshooting/simba_blob_project/project_folder/project_config.ini", data_path=r'C:/troubleshooting/simba_blob_project/data')
     >>> r.run()
-    >>> r = SimBABlobImporter(config_path=r"C:\troubleshooting\simba_blob_project\project_folder\project_config.ini", 
-    ...                       data_path=r'C:\troubleshooting\simba_blob_project\data',
+    >>> r = SimBABlobImporter(config_path=r"C:/troubleshooting/simba_blob_project/project_folder/project_config.ini", 
+    ...                       data_path=r'C:/troubleshooting/simba_blob_project/data',
     ...                       smoothing_settings={'method': 'savitzky-golay', 'time_window': 100},
     ...                       interpolation_settings={'method': 'nearest', 'type': 'body-parts'})
     >>> r.run()

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
@@ -167,7 +168,7 @@ class ROIAggregateStatisticsAnalyzerMultiprocess(ConfigReader, FeatureExtraction
     :param save_path (str | os.PathLike, optional): Path to save summary statistics.
 
     :example:
-    >>> analyzer = ROIAggregateStatisticsAnalyzerMultiprocess(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini", body_parts=['Center'], first_entry_time=True, threshold=0.0, calculate_distances=True, transpose=False, detailed_bout_data=True)
+    >>> analyzer = ROIAggregateStatisticsAnalyzerMultiprocess(config_path=r"C:/troubleshooting/mitra/project_folder/project_config.ini", body_parts=['Center'], first_entry_time=True, threshold=0.0, calculate_distances=True, transpose=False, detailed_bout_data=True)
     >>> analyzer.run()
     >>> analyzer.save()
     """
@@ -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
@@ -149,7 +150,7 @@ class ROIClfCalculatorMultiprocess(ConfigReader):
        'GitHub tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/Scenario2.md#part-4--analyze-machine-results`__.
 
     :example:
-    >>> analyzer = ROIClfCalculatorMultiprocess(config_path=r"D:\troubleshooting\maplight_ri\project_folder\project_config.ini", bp_names=['resident_NOSE'], clf_names=['attack'], clf_time=True, started_bout_cnt=True, ended_bout_cnt=False, bout_table=True, transpose=True, core_cnt=20)
+    >>> analyzer = ROIClfCalculatorMultiprocess(config_path=r"D:/troubleshooting/maplight_ri/project_folder/project_config.ini", bp_names=['resident_NOSE'], clf_names=['attack'], clf_time=True, started_bout_cnt=True, ended_bout_cnt=False, bout_table=True, transpose=True, core_cnt=20)
     >>> analyzer.run()
     >>> analyzer.save()
     """
@@ -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()

simba/sandbox/analyze_runtimes.py

@@ -0,0 +1,30 @@
+"""Analyze runtime statistics for directionality_to_nonstatic_target"""
+import numpy as np
+from collections import defaultdict
+
+# Parse the runtime data
+data = {
+    10000: [0.4389, 0.0008, 0.0012],
+    100000: [0.0063, 0.0052, 0.0052],
+    1000000: [0.0768, 0.0306, 0.0239],
+    10000000: [0.2195, 0.2122, 0.2083],
+    50000000: [1.8936, 1.5664, 1.2548]
+}
+
+# Calculate statistics
+print("=" * 80)
+print(f"{'Observations':<15} {'Mean (s)':<12} {'Std (s)':<12} {'Min (s)':<12} {'Max (s)':<12} {'Throughput (M obs/s)':<20}")
+print("=" * 80)
+
+for obs_count in sorted(data.keys()):
+    times = np.array(data[obs_count])
+    mean_time = np.mean(times)
+    std_time = np.std(times)
+    min_time = np.min(times)
+    max_time = np.max(times)
+    throughput = obs_count / (mean_time * 1_000_000)  # Million observations per second
+    
+    print(f"{obs_count:<15,} {mean_time:<12.4f} {std_time:<12.4f} {min_time:<12.4f} {max_time:<12.4f} {throughput:<20.2f}")
+
+print("=" * 80)
+print("\nNote: First run typically includes JIT compilation overhead (especially for 10k observations)")