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

simba/SimBA.py

@@ -966,8 +966,8 @@ class App(object):
         video_process_menu.add_cascade(label="Drop body-parts from tracking data", compound="left", image=self.menu_icons["trash"]["img"], command=DropTrackingDataPopUp, font=Formats.FONT_REGULAR.value)
         extract_frames_menu = Menu(video_process_menu, font=Formats.FONT_REGULAR.value)
         extract_frames_menu.add_command(label="Extract defined frames", command=ExtractSpecificFramesPopUp, font=Formats.FONT_REGULAR.value, image=self.menu_icons["frames"]["img"], compound="left")
-        extract_frames_menu.add_command(label="Extract frames from single video", command=SingleVideo2FramesPopUp, font=Formats.FONT_REGULAR.value, image=self.menu_icons["video"]["img"], compound="left")
-        extract_frames_menu.add_command(label="Extract frames from multiple video", command=MultipleVideos2FramesPopUp, font=Formats.FONT_REGULAR.value, image=self.menu_icons["stack"]["img"], compound="left")
+        extract_frames_menu.add_command(label="Extract frames from single video", command=SingleVideo2FramesPopUp, font=Formats.FONT_REGULAR.value, image=self.menu_icons["frames_2"]["img"], compound="left")
+        extract_frames_menu.add_command(label="Extract frames from multiple videos", command=MultipleVideos2FramesPopUp, font=Formats.FONT_REGULAR.value, image=self.menu_icons["stack"]["img"], compound="left")
         extract_frames_menu.add_command(label="Extract frames from seq files", command=ExtractSEQFramesPopUp, font=Formats.FONT_REGULAR.value, image=self.menu_icons["fire"]["img"], compound="left")
         video_process_menu.add_cascade(label="Extract frames...", compound="left", image=self.menu_icons["frames"]["img"], menu=extract_frames_menu, font=Formats.FONT_REGULAR.value)
 

simba/data_processors/agg_clf_counter_mp.py

@@ -20,7 +20,7 @@ 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
+from simba.utils.data import detect_bouts, terminate_cpu_pool
 from simba.utils.enums import TagNames
 from simba.utils.errors import NoChoosenMeasurementError
 from simba.utils.printing import SimbaTimer, log_event, stdout_success
@@ -210,8 +210,7 @@ class AggregateClfCalculatorMultiprocess(ConfigReader):
                 self.bouts_df_lst.append(batch_bouts_df_lst)
                 print(f"Data batch core {batch_id+1} / {self.core_cnt} complete...")
         self.bouts_df_lst = [df for sub in self.bouts_df_lst for df in sub]
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
 
     def save(self) -> None:
         """
@@ -242,56 +241,56 @@ class AggregateClfCalculatorMultiprocess(ConfigReader):
         self.timer.stop_timer()
         stdout_success(msg=f"Data aggregate log saved at {self.save_path}", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__)
 
-if __name__ == "__main__" and not hasattr(sys, 'ps1'):
-    parser = argparse.ArgumentParser(description='Compute aggregate descriptive statistics from classification data.')
-    parser.add_argument('--config_path', type=str, required=True, help='Path to SimBA project config file')
-    parser.add_argument('--classifiers', type=str, nargs='+', required=True, help='List of classifier names to analyze')
-    parser.add_argument('--data_dir', type=str, default=None, help='Directory containing machine results CSV files (default: project machine_results directory)')
-    parser.add_argument('--detailed_bout_data', action='store_true', help='Save detailed bout data for each bout')
-    parser.add_argument('--transpose', action='store_true', help='Create output with one video per row')
-    parser.add_argument('--no_first_occurrence', action='store_true', help='Disable first occurrence calculation')
-    parser.add_argument('--no_event_count', action='store_true', help='Disable event count calculation')
-    parser.add_argument('--no_total_event_duration', action='store_true', help='Disable total event duration calculation')
-    parser.add_argument('--no_mean_event_duration', action='store_true', help='Disable mean event duration calculation')
-    parser.add_argument('--no_median_event_duration', action='store_true', help='Disable median event duration calculation')
-    parser.add_argument('--no_mean_interval_duration', action='store_true', help='Disable mean interval duration calculation')
-    parser.add_argument('--no_median_interval_duration', action='store_true', help='Disable median interval duration calculation')
-    parser.add_argument('--frame_count', action='store_true', help='Include frame count in output')
-    parser.add_argument('--video_length', action='store_true', help='Include video length in output')
-
-    args = parser.parse_args()
-
-    clf_calculator = AggregateClfCalculatorMultiprocess(
-        config_path=args.config_path,
-        classifiers=args.classifiers,
-        data_dir=args.data_dir,
-        detailed_bout_data=args.detailed_bout_data,
-        transpose=args.transpose,
-        first_occurrence=not args.no_first_occurrence,
-        event_count=not args.no_event_count,
-        total_event_duration=not args.no_total_event_duration,
-        mean_event_duration=not args.no_mean_event_duration,
-        median_event_duration=not args.no_median_event_duration,
-        mean_interval_duration=not args.no_mean_interval_duration,
-        median_interval_duration=not args.no_median_interval_duration,
-        frame_count=args.frame_count,
-        video_length=args.video_length
-    )
-    clf_calculator.run()
-    clf_calculator.save()
-
-# if __name__ == "__main__":
-#     test = AggregateClfCalculatorMultiprocess(config_path=r"D:\troubleshooting\maplight_ri\project_folder\project_config.ini",
-#                                   classifiers=['attack'],
-#                                   transpose=True,
-#                                   mean_event_duration = True,
-#                                   median_event_duration = True,
-#                                   mean_interval_duration = True,
-#                                   median_interval_duration = True,
-#                                   detailed_bout_data=True,
-#                                   core_cnt=12)
-#     test.run()
-#     test.save()
+# if __name__ == "__main__" and not hasattr(sys, 'ps1'):
+#     parser = argparse.ArgumentParser(description='Compute aggregate descriptive statistics from classification data.')
+#     parser.add_argument('--config_path', type=str, required=True, help='Path to SimBA project config file')
+#     parser.add_argument('--classifiers', type=str, nargs='+', required=True, help='List of classifier names to analyze')
+#     parser.add_argument('--data_dir', type=str, default=None, help='Directory containing machine results CSV files (default: project machine_results directory)')
+#     parser.add_argument('--detailed_bout_data', action='store_true', help='Save detailed bout data for each bout')
+#     parser.add_argument('--transpose', action='store_true', help='Create output with one video per row')
+#     parser.add_argument('--no_first_occurrence', action='store_true', help='Disable first occurrence calculation')
+#     parser.add_argument('--no_event_count', action='store_true', help='Disable event count calculation')
+#     parser.add_argument('--no_total_event_duration', action='store_true', help='Disable total event duration calculation')
+#     parser.add_argument('--no_mean_event_duration', action='store_true', help='Disable mean event duration calculation')
+#     parser.add_argument('--no_median_event_duration', action='store_true', help='Disable median event duration calculation')
+#     parser.add_argument('--no_mean_interval_duration', action='store_true', help='Disable mean interval duration calculation')
+#     parser.add_argument('--no_median_interval_duration', action='store_true', help='Disable median interval duration calculation')
+#     parser.add_argument('--frame_count', action='store_true', help='Include frame count in output')
+#     parser.add_argument('--video_length', action='store_true', help='Include video length in output')
+#
+#     args = parser.parse_args()
+#
+#     clf_calculator = AggregateClfCalculatorMultiprocess(
+#         config_path=args.config_path,
+#         classifiers=args.classifiers,
+#         data_dir=args.data_dir,
+#         detailed_bout_data=args.detailed_bout_data,
+#         transpose=args.transpose,
+#         first_occurrence=not args.no_first_occurrence,
+#         event_count=not args.no_event_count,
+#         total_event_duration=not args.no_total_event_duration,
+#         mean_event_duration=not args.no_mean_event_duration,
+#         median_event_duration=not args.no_median_event_duration,
+#         mean_interval_duration=not args.no_mean_interval_duration,
+#         median_interval_duration=not args.no_median_interval_duration,
+#         frame_count=args.frame_count,
+#         video_length=args.video_length
+#     )
+#     clf_calculator.run()
+#     clf_calculator.save()
+
+if __name__ == "__main__":
+    test = AggregateClfCalculatorMultiprocess(config_path=r"D:\troubleshooting\maplight_ri\project_folder\project_config.ini",
+                                  classifiers=['attack'],
+                                  transpose=True,
+                                  mean_event_duration = True,
+                                  median_event_duration = True,
+                                  mean_interval_duration = True,
+                                  median_interval_duration = True,
+                                  detailed_bout_data=True,
+                                  core_cnt=12)
+    test.run()
+    test.save()
 
 
 # test = AggregateClfCalculator(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",

simba/data_processors/cuda/image.py

@@ -213,7 +213,6 @@ def _average_3d_stack_cuda(image_stack: np.ndarray) -> np.ndarray:
     return results
 
 
-
 def create_average_frm_cuda(video_path: Union[str, os.PathLike],
                             start_frm: Optional[int] = None,
                             end_frm: Optional[int] = None,
@@ -1512,6 +1511,9 @@ def pose_plotter(data: Union[str, os.PathLike, np.ndarray],
         stdout_success(msg=f'Pose-estimation video saved at {save_path}.', elapsed_time=total_timer.elapsed_time_str)
 
 
+
+#x = create_average_frm_cuda(video_path=r"D:\troubleshooting\mitra\project_folder\videos\average_cpu_test\20min.mp4", verbose=True, batch_size=500, async_frame_read=False)
+
 # VIDEO_PATH = "/mnt/d/troubleshooting/maplight_ri/project_folder/blob/videos/Trial_1_C24_D1_1.mp4"
 # #
 #

simba/data_processors/cue_light_analyzer.py

@@ -1,13 +1,9 @@
 __author__ = "Simon Nilsson; sronilsson@gmail.com"
 
 import functools
-import glob
-import itertools
 import multiprocessing
 import os
-import platform
-import time
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Union
 
 import cv2
 import numpy as np
@@ -17,9 +13,9 @@ from simba.mixins.config_reader import ConfigReader
 from simba.mixins.statistics_mixin import Statistics
 from simba.utils.checks import (
     check_all_file_names_are_represented_in_video_log, check_if_dir_exists,
-    check_if_valid_img, check_int, check_nvidea_gpu_available,
-    check_valid_boolean, check_valid_lst)
-from simba.utils.data import detect_bouts, slice_roi_dict_from_attribute
+    check_if_valid_img, check_int, check_valid_boolean, check_valid_lst)
+from simba.utils.data import (detect_bouts, slice_roi_dict_from_attribute,
+                              terminate_cpu_pool)
 from simba.utils.enums import Defaults, Keys
 from simba.utils.errors import NoROIDataError
 from simba.utils.printing import SimbaTimer, stdout_success
@@ -220,7 +216,7 @@ class CueLightAnalyzer(ConfigReader):
                         else: self.intensities[key] = subdict
                         if self.verbose:
                             print(f'Batch {int(np.ceil(cnt + 1 / self.core_cnt))} complete...')
-            pool.terminate(); pool.join()
+            terminate_cpu_pool(pool=pool, force=False)
             kmeans = self._get_kmeans(intensities=self.intensities)
             self.data_df = self._append_light_data(data_df=self.data_df, kmeans_data=kmeans)
             self.data_df = self._remove_outlier_events(data_df=self.data_df)

simba/mixins/geometry_mixin.py

@@ -1339,8 +1339,7 @@ class GeometryMixin(object):
             )
             for cnt, result in enumerate(pool.imap(constants, data, chunksize=1)):
                 results.append(result)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         if data_ndim == 2:
             return [i for s in results for i in s]
         else:
@@ -1370,8 +1369,7 @@ class GeometryMixin(object):
                                           cap_style=cap_style)
             for cnt, mp_return in enumerate(pool.imap(constants, geomety_lst, chunksize=1)):
                 results.append(mp_return)
-            pool.join()
-            pool.terminate()
+            terminate_cpu_pool(pool=pool, force=False)
         return [l for ll in results for l in ll]
 
     def multiframe_bodyparts_to_circle(self,
@@ -1524,8 +1522,7 @@ class GeometryMixin(object):
             )
             for cnt, result in enumerate(pool.imap(constants, data, chunksize=1)):
                 results.append(result)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_compute_pct_shape_overlap(self,
@@ -1798,8 +1795,7 @@ class GeometryMixin(object):
         timer.stop_timer()
         if verbose:
             stdout_success(msg="Rotated rectangles complete.", elapsed_time=timer.elapsed_time_str)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     @staticmethod
@@ -2003,8 +1999,7 @@ class GeometryMixin(object):
             )
             for cnt, result in enumerate(pool.imap(constants, shapes, chunksize=1)):
                 results.append(result)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_union(self, shapes: Iterable[Union[LineString, MultiLineString, Polygon]], core_cnt: int = -1) -> \
@@ -2043,8 +2038,7 @@ class GeometryMixin(object):
         with multiprocessing.Pool(core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value) as pool:
             for cnt, result in enumerate(pool.imap(GeometryMixin().union, shapes, chunksize=1)):
                 results.append(result)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_symmetric_difference(self, shapes: Iterable[Union[LineString, MultiLineString, Polygon]],
@@ -2084,8 +2078,7 @@ class GeometryMixin(object):
                     pool.imap(GeometryMixin().symmetric_difference, shapes, chunksize=1)
             ):
                 results.append(result)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_delaunay_triangulate_keypoints(self, data: np.ndarray, core_cnt: int = -1) -> List[List[Polygon]]:
@@ -2132,8 +2125,7 @@ class GeometryMixin(object):
             ):
                 results.append(result)
 
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_difference(
@@ -2221,8 +2213,7 @@ class GeometryMixin(object):
             msg="Multi-frame difference compute complete",
             elapsed_time=timer.elapsed_time_str,
         )
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_area(self,
@@ -2276,8 +2267,7 @@ class GeometryMixin(object):
 
         timer.stop_timer()
         stdout_success(msg="Multi-frame area compute complete", elapsed_time=timer.elapsed_time_str)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     def multiframe_bodyparts_to_multistring_skeleton(
@@ -2619,8 +2609,7 @@ class GeometryMixin(object):
                     pool.imap(GeometryMixin.is_shape_covered, shapes, chunksize=1)
             ):
                 results.append(mp_return)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     @staticmethod
@@ -3321,8 +3310,7 @@ class GeometryMixin(object):
             for cnt, result in enumerate(pool.imap(constants, data, chunksize=1)):
                 if result[1] != -1:
                     img_arr[result[0], result[2] - 1, result[1] - 1] = 1
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         timer.stop_timer()
         stdout_success(
             msg="Cumulative coordinates in geometries complete",
@@ -3415,8 +3403,7 @@ class GeometryMixin(object):
             for cnt, result in enumerate(pool.imap(constants, data, chunksize=1)):
                 if result[1] != -1:
                     img_arr[result[0], result[2] - 1, result[1] - 1] = 1
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         if fps is None:
             return np.cumsum(img_arr, axis=0)
         else:
@@ -3559,8 +3546,7 @@ class GeometryMixin(object):
             constants = functools.partial(GeometryMixin._compute_framewise_geometry_idx, grid=grid, verbose=verbose)
             for cnt, result in enumerate(pool.imap(constants, data, chunksize=1)):
                 results.append(result)
-        pool.join();
-        pool.terminate();
+        terminate_cpu_pool(pool=pool, force=False)
         del data
 
         results = np.vstack(results)[:, 1:].astype(np.int32)

simba/mixins/image_mixin.py

@@ -18,7 +18,7 @@ from collections import ChainMap
 import cv2
 import pandas as pd
 from numba import float64, int64, jit, njit, prange, uint8
-from shapely.geometry import MultiPolygon, Polygon
+from shapely.geometry import Polygon
 from skimage.metrics import structural_similarity
 
 from simba.utils.checks import (check_file_exist_and_readable, check_float,
@@ -27,16 +27,14 @@ from simba.utils.checks import (check_file_exist_and_readable, check_float,
                                 check_int, check_str, check_valid_array,
                                 check_valid_boolean, check_valid_lst,
                                 check_valid_tuple, is_img_bw, is_img_greyscale)
+from simba.utils.data import terminate_cpu_pool
 from simba.utils.enums import Defaults, Formats, GeometryEnum, Options
-from simba.utils.errors import (ArrayError, FFMPEGCodecGPUError,
-                                FrameRangeError, InvalidInputError,
-                                NotDirectoryError)
+from simba.utils.errors import ArrayError, FrameRangeError, InvalidInputError
 from simba.utils.printing import SimbaTimer, stdout_success
 from simba.utils.read_write import (find_core_cnt,
                                     find_files_of_filetypes_in_directory,
                                     get_fn_ext, get_video_meta_data,
-                                    read_frm_of_video,
-                                    read_img_batch_from_video_gpu, write_df)
+                                    read_frm_of_video)
 
 
 class ImageMixin(object):
@@ -546,8 +544,8 @@ class ImageMixin(object):
                 pool.imap(constants, split_frm_idx, chunksize=1)
             ):
                 results.append(result)
-        pool.terminate()
-        pool.join()
+
+        terminate_cpu_pool(pool=pool, force=False)
         results = dict(ChainMap(*results))
 
         max_value, max_frm = -np.inf, None
@@ -876,8 +874,7 @@ class ImageMixin(object):
                 pool.imap(ImageMixin()._image_reader_helper, file_paths, chunksize=1)
             ):
                 imgs.update(result)
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return imgs
 
     @staticmethod
@@ -1027,8 +1024,7 @@ class ImageMixin(object):
             for cnt, result in enumerate(pool.imap(constants, frm_lst, chunksize=1)):
                 results.update(result)
 
-        pool.join()
-        pool.terminate()
+        terminate_cpu_pool(pool=pool, force=False)
         return results
 
     @staticmethod
@@ -1509,8 +1505,8 @@ class ImageMixin(object):
                 for cnt, result in enumerate(pool.imap(constants, shapes, chunksize=1)):
                     results.append(result)
                 results = dict(ChainMap(*results))
-        pool.join()
-        pool.terminate()
+
+        terminate_cpu_pool(pool=pool, force=False)
         results = dict(sorted(results.items(), key=lambda item: int(item[0])))
         timer.stop_timer()
         stdout_success(msg="Geometry image slicing complete.", elapsed_time=timer.elapsed_time_str, source=self.__class__.__name__)

simba/mixins/train_model_mixin.py

@@ -67,7 +67,7 @@ from simba.utils.checks import (check_all_dfs_in_list_has_same_cols,
                                 check_valid_boolean, check_valid_dataframe,
                                 check_valid_lst, is_lxc_container)
 from simba.utils.data import (detect_bouts, detect_bouts_multiclass,
-                              get_library_version)
+                              get_library_version, terminate_cpu_pool)
 from simba.utils.enums import (OS, ConfigKey, Defaults, Dtypes, Formats, Links,
                                Methods, MLParamKeys, Options)
 from simba.utils.errors import (ClassifierInferenceError, CorruptedFileError,
@@ -1859,7 +1859,7 @@ class TrainModelMixin(object):
                     shap_raw.append(shap_data[result[1]][1].drop(clf_name, axis=1))
                     if verbose: print(f"Completed SHAP care batch (Batch {result[1] + 1}/{len(shap_data)}).")
 
-            pool.terminate(); pool.join()
+            terminate_cpu_pool(pool=pool, force=False)
             shap_df = pd.DataFrame(data=np.row_stack(shap_results), columns=list(x_names) + ["Expected_value", "Sum", "Prediction_probability", clf_name])
             raw_df = pd.DataFrame(data=np.row_stack(shap_raw), columns=list(x_names))
             out_shap_path, out_raw_path, img_save_path, df_save_paths, summary_dfs, img = None, None, None, None, None, None

simba/plotting/ROI_feature_visualizer_mp.py

@@ -24,10 +24,9 @@ from simba.utils.checks import (check_file_exist_and_readable,
                                 check_if_valid_rgb_tuple, check_int, check_str,
                                 check_valid_boolean, check_valid_lst,
                                 check_video_and_data_frm_count_align)
-from simba.utils.data import slice_roi_dict_for_video
+from simba.utils.data import slice_roi_dict_for_video, terminate_cpu_pool
 from simba.utils.enums import Formats, TextOptions
-from simba.utils.errors import (BodypartColumnNotFoundError, NoFilesFoundError,
-                                ROICoordinatesNotFoundError)
+from simba.utils.errors import BodypartColumnNotFoundError, NoFilesFoundError
 from simba.utils.printing import stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
                                     find_core_cnt, get_fn_ext,
@@ -315,8 +314,7 @@ class ROIfeatureVisualizerMultiprocess(ConfigReader):
             print(f"Joining {self.video_name} multi-processed video...")
             concatenate_videos_in_folder(in_folder=self.save_temp_dir, save_path=self.save_path, video_format="mp4", remove_splits=True, gpu=self.gpu)
             self.timer.stop_timer()
-            pool.terminate()
-            pool.join()
+            terminate_cpu_pool(pool=pool, force=False)
             stdout_success(msg=f"Video {self.video_name} complete. Video saved in directory {self.roi_features_save_dir}.", elapsed_time=self.timer.elapsed_time_str)
 
 

simba/plotting/clf_validator_mp.py

@@ -14,9 +14,9 @@ from simba.mixins.plotting_mixin import PlottingMixin
 from simba.utils.checks import (check_float, check_if_valid_rgb_tuple,
                                 check_int, check_str, check_that_column_exist,
                                 check_valid_lst)
-from simba.utils.data import detect_bouts
-from simba.utils.enums import Formats, TagNames, TextOptions
-from simba.utils.errors import NoFilesFoundError, NoSpecifiedOutputError
+from simba.utils.data import detect_bouts, terminate_cpu_pool
+from simba.utils.enums import Formats, TextOptions
+from simba.utils.errors import NoSpecifiedOutputError
 from simba.utils.printing import SimbaTimer, log_event, stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
                                     find_core_cnt, get_fn_ext,
@@ -218,8 +218,7 @@ class ClassifierValidationClipsMultiprocess(ConfigReader):
                 for cnt, result in enumerate(
                     pool.imap(constants, clip_data, chunksize=self.multiprocess_chunksize)):
                     print(f"Bout {cnt+1} complete...")
-                pool.terminate()
-                pool.join()
+                terminate_cpu_pool(pool=pool, force=False)
 
             if self.concat_video:
                 print(f"Joining {file_name} multiprocessed video...")

simba/plotting/cue_light_visualizer.py

@@ -1,10 +1,9 @@
 __author__ = "Simon Nilsson; sronilsson@gmail.com"
 
 import functools
-import itertools
 import multiprocessing
 import os
-from typing import List, Optional, Union
+from typing import List, Union
 
 import cv2
 import numpy as np
@@ -16,14 +15,15 @@ from simba.utils.checks import (check_file_exist_and_readable, check_int,
                                 check_valid_boolean, check_valid_dataframe,
                                 check_valid_lst)
 from simba.utils.data import (create_color_palettes, detect_bouts,
-                              slice_roi_dict_from_attribute)
-from simba.utils.enums import Defaults, Formats, Keys, TextOptions
+                              slice_roi_dict_from_attribute,
+                              terminate_cpu_pool)
+from simba.utils.enums import Defaults, Formats, TextOptions
 from simba.utils.errors import NoROIDataError, NoSpecifiedOutputError
 from simba.utils.printing import stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
                                     create_directory, find_core_cnt,
                                     get_fn_ext, get_video_meta_data, read_df,
-                                    read_frm_of_video, remove_a_folder)
+                                    read_frm_of_video)
 
 
 def _plot_cue_light_data(frm_idxs: list,
@@ -197,8 +197,7 @@ class CueLightVisualizer(ConfigReader):
             for cnt, result in enumerate(pool.imap(constants, self.frame_chunks, chunksize=self.multiprocess_chunksize)):
                 if self.verbose:
                     print(f'Batch {int(result+1/self.core_cnt)} complete...')
-            pool.terminate()
-            pool.join()
+            terminate_cpu_pool(pool=pool, force=False)
         self.timer.stop_timer()
         if self.video_setting:
             print(f"Joining {self.video_name} multiprocessed video...")

simba/plotting/directing_animals_visualizer_mp.py

@@ -19,7 +19,7 @@ from simba.utils.checks import (check_file_exist_and_readable,
                                 check_if_valid_rgb_tuple, check_int,
                                 check_valid_lst,
                                 check_video_and_data_frm_count_align)
-from simba.utils.data import create_color_palettes
+from simba.utils.data import create_color_palettes, terminate_cpu_pool
 from simba.utils.enums import OS, Formats, Keys, TextOptions
 from simba.utils.errors import (AnimalNumberError, InvalidInputError,
                                 NoFilesFoundError)
@@ -226,8 +226,7 @@ class DirectingOtherAnimalsVisualizerMultiprocess(ConfigReader, PlottingMixin):
             print(f"Joining {self.video_name} multi-processed video...")
             concatenate_videos_in_folder(in_folder=self.save_temp_path, save_path=self.save_path, video_format="mp4", remove_splits=True)
             self.timer.stop_timer()
-            pool.terminate()
-            pool.join()
+            terminate_cpu_pool(pool=pool, force=False)
             stdout_success(msg=f"Video {self.video_name} complete. Video saved in {self.directing_animals_video_output_path} directory", elapsed_time=self.timer.elapsed_time_str)