simba-uw-tf-dev 4.6.4__py3-none-any.whl → 4.6.6__py3-none-any.whl

simba/utils/read_write.py

@@ -2460,6 +2460,13 @@ def read_img_batch_from_video_gpu(video_path: Union[str, os.PathLike],
     """
     Reads a batch of frames from a video file using GPU acceleration.
 
+    .. csv-table::
+       :header: EXPECTED RUNTIMES
+       :file: ../../docs/tables/read_img_batch_from_video_gpu.csv
+       :widths: 10, 45, 45
+       :align: center
+       :header-rows: 1
+
     This function uses FFmpeg with CUDA acceleration to read frames from a specified range in a video file. It supports both RGB and greyscale video formats. Frames are returned as a dictionary where the keys are
     frame indices and the values are NumPy arrays representing the image data.
 
@@ -2468,7 +2475,7 @@ def read_img_batch_from_video_gpu(video_path: Union[str, os.PathLike],
        If you expect that the video you are reading in is black and white, set ``black_and_white`` to True to round any of these wonly value sto 0 and 255.
 
     .. seealso::
-       For CPU multicore acceleration, see :func:`simba.mixins.image_mixin.ImageMixin.read_img_batch_from_video`
+       For CPU multicore acceleration, see :func:`simba.mixins.image_mixin.ImageMixin.read_img_batch_from_video` or :func:`simba.utils.read_write.read_img_batch_from_video`.
 
     :param video_path: Path to the video file. Can be a string or an os.PathLike object.
     :param start_frm: The starting frame index to read. If None, starts from the beginning of the video.
@@ -2479,6 +2486,7 @@ def read_img_batch_from_video_gpu(video_path: Union[str, os.PathLike],
     :return: A dictionary where keys are frame indices (integers) and values are NumPy arrays containing the image data of each frame.
     """
 
+    timer = SimbaTimer(start=True)
     check_file_exist_and_readable(file_path=video_path)
     video_meta_data = get_video_meta_data(video_path=video_path, fps_as_int=False)
     if start_frm is not None:
@@ -2551,6 +2559,10 @@ def read_img_batch_from_video_gpu(video_path: Union[str, os.PathLike],
                 binary_frms[frm_id] = np.where(frames[frm_id] > 127, 255, 0).astype(np.uint8)
         frames = binary_frms
 
+    timer.stop_timer()
+    if verbose:
+        print(f'[{get_current_time()}] Read frames {start_frm}-{end_frm} (video: {video_name}, elapsed time: {timer.elapsed_time_str}s)')
+
     return frames
 
 
@@ -3156,7 +3168,7 @@ def _read_img_batch_from_video_helper(frm_idx: np.ndarray, video_path: Union[str
     cap.set(1, current_frm)
     while current_frm < end_frm:
         if verbose:
-            print(f'Reading frame {current_frm}/{video_meta_data["frame_count"]} ({video_meta_data["video_name"]})...')
+            print(f'[{get_current_time()}] Reading frame {current_frm} ({video_meta_data["video_name"]})...')
         img = cap.read()[1]
         if img is not None:
             if greyscale or black_and_white or clahe:
@@ -3188,6 +3200,14 @@ def read_img_batch_from_video(video_path: Union[str, os.PathLike],
     """
     Read a batch of frames from a video file. This method reads frames from a specified range of frames within a video file using multiprocessing.
 
+    .. csv-table::
+       :header: EXPECTED RUNTIMES
+       :file: ../../docs/tables/read_img_batch_from_video.csv
+       :widths: 10, 45, 45
+       :align: center
+       :header-rows: 1
+
+
     .. seealso::
        For GPU acceleration, see :func:`simba.utils.read_write.read_img_batch_from_video_gpu`
 
@@ -3209,6 +3229,8 @@ def read_img_batch_from_video(video_path: Union[str, os.PathLike],
     >>> read_img_batch_from_video(video_path='/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/videos/Together_1.avi', start_frm=0, end_frm=50)
     """
 
+
+    timer = SimbaTimer(start=True)
     if platform.system() == "Darwin":
         if not multiprocessing.get_start_method(allow_none=True):
             multiprocessing.set_start_method("fork", force=True)
@@ -3230,19 +3252,22 @@ def read_img_batch_from_video(video_path: Union[str, os.PathLike],
     if end_frm <= start_frm:
         FrameRangeError(msg=f"Start frame ({start_frm}) has to be before end frame ({end_frm})", source=read_img_batch_from_video.__name__)
     frm_lst = np.array_split(np.arange(start_frm, end_frm + 1), core_cnt)
+    pool = multiprocessing.Pool(core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value)
     results = {}
-    with multiprocessing.Pool(core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value) as pool:
-        constants = functools.partial(_read_img_batch_from_video_helper,
-                                      video_path=video_path,
-                                      greyscale=greyscale,
-                                      black_and_white=black_and_white,
-                                      clahe=clahe,
-                                      verbose=verbose)
-        for cnt, result in enumerate(pool.imap(constants, frm_lst, chunksize=1)):
-            results.update(result)
-    pool.join()
+    constants = functools.partial(_read_img_batch_from_video_helper,
+                                  video_path=video_path,
+                                  greyscale=greyscale,
+                                  black_and_white=black_and_white,
+                                  clahe=clahe,
+                                  verbose=verbose)
+    for cnt, result in enumerate(pool.imap(constants, frm_lst, chunksize=1)):
+        results.update(result)
     pool.close()
-    #terminate_cpu_pool(pool=pool, force=False)
+    pool.join()
+    pool.terminate()
+    timer.stop_timer()
+    if verbose:
+        print(f'[{get_current_time()}] Read frames {start_frm}-{end_frm} (video: {video_meta_data["video_name"]}, elapsed time: {timer.elapsed_time_str}s)')
     return results
 
 def read_yolo_bp_names_file(file_path: Union[str, os.PathLike]) -> Tuple[str]:

simba/video_processors/egocentric_video_rotator.py

@@ -9,10 +9,12 @@ import numpy as np
 from simba.utils.checks import (check_file_exist_and_readable,
                                 check_if_dir_exists, check_if_valid_rgb_tuple,
                                 check_int, check_valid_array,
-                                check_valid_boolean, check_valid_tuple)
+                                check_valid_boolean, check_valid_cpu_pool,
+                                check_valid_tuple)
 from simba.utils.data import (align_target_warpaffine_vectors,
                               center_rotation_warpaffine_vectors,
-                              egocentrically_align_pose, terminate_cpu_pool)
+                              egocentrically_align_pose, get_cpu_pool,
+                              terminate_cpu_pool)
 from simba.utils.enums import Defaults, Formats
 from simba.utils.printing import SimbaTimer, stdout_success
 from simba.utils.read_write import (concatenate_videos_in_folder,
@@ -114,7 +116,8 @@ class EgocentricVideoRotator():
                  fill_clr: Tuple[int, int, int] = (0, 0, 0),
                  core_cnt: int = -1,
                  save_path: Optional[Union[str, os.PathLike]] = None,
-                 gpu: Optional[bool] = True):
+                 gpu: Optional[bool] = True,
+                 pool: bool = None):
 
         check_file_exist_and_readable(file_path=video_path)
         self.video_meta_data = get_video_meta_data(video_path=video_path)
@@ -125,10 +128,14 @@ class EgocentricVideoRotator():
         check_valid_boolean(value=[verbose], source=f'{self.__class__.__name__} verbose')
         check_if_valid_rgb_tuple(data=fill_clr)
         check_int(name=f'{self.__class__.__name__} core_cnt', value=core_cnt, min_value=-1, unaccepted_vals=[0])
-        if core_cnt > find_core_cnt()[0] or core_cnt == -1:
-            self.core_cnt = find_core_cnt()[0]
+        if core_cnt > find_core_cnt()[0] or core_cnt == -1: self.core_cnt = find_core_cnt()[0]
+        else: self.core_cnt = core_cnt
+        if pool is not None:
+            check_valid_cpu_pool(value=pool, source=self.__class__.__name__, max_cores=find_core_cnt()[0], min_cores=2, raise_error=True)
+            self.pool_termination_flag = True
         else:
-            self.core_cnt = core_cnt
+            self.pool_termination_flag = False
+        self.pool = get_cpu_pool(core_cnt=self.core_cnt, source=self.__class__.__name__) if pool is None else pool
         video_dir, self.video_name, _ = get_fn_ext(filepath=video_path)
         if save_path is not None:
             self.save_dir = os.path.dirname(save_path)
@@ -151,37 +158,35 @@ class EgocentricVideoRotator():
         frm_list = np.arange(0, self.video_meta_data['frame_count'])
         frm_list = np.array_split(frm_list, self.core_cnt)
         frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)]
-        if self.verbose:
-            print(f"Creating rotated video {self.video_name}, multiprocessing (chunksize: {1}, cores: {self.core_cnt})...")
-        with multiprocessing.Pool(self.core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value) as pool:
-            constants = functools.partial(egocentric_video_aligner,
-                                          temp_dir=temp_dir,
-                                          video_name=self.video_name,
-                                          video_path=self.video_path,
-                                          centers=self.centers,
-                                          rotation_vectors=self.rotation_vectors,
-                                          target=self.anchor_loc,
-                                          verbose=self.verbose,
-                                          fill_clr=self.fill_clr,
-                                          gpu=self.gpu)
-            for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=1)):
-                if self.verbose:
-                    print(f"Rotate batch {result}/{self.core_cnt} complete...")
-        terminate_cpu_pool(pool=pool, force=False)
+        if self.verbose: print(f"Creating rotated video {self.video_name}, multiprocessing (chunksize: {1}, cores: {self.core_cnt})...")
+
+        constants = functools.partial(egocentric_video_aligner,
+                                      temp_dir=temp_dir,
+                                      video_name=self.video_name,
+                                      video_path=self.video_path,
+                                      centers=self.centers,
+                                      rotation_vectors=self.rotation_vectors,
+                                      target=self.anchor_loc,
+                                      verbose=self.verbose,
+                                      fill_clr=self.fill_clr,
+                                      gpu=self.gpu)
+        for cnt, result in enumerate(self.pool.imap(constants, frm_list, chunksize=1)):
+            if self.verbose: print(f"Rotate batch {result}/{self.core_cnt} complete...")
+        if self.pool_termination_flag: terminate_cpu_pool(pool=self.pool, force=False)
         concatenate_videos_in_folder(in_folder=temp_dir, save_path=self.save_path, remove_splits=True, gpu=self.gpu, verbose=self.verbose)
         video_timer.stop_timer()
         stdout_success(msg=f"Egocentric rotation video {self.save_path} complete", elapsed_time=video_timer.elapsed_time_str, source=self.__class__.__name__)
 
-if __name__ == "__main__":
-    DATA_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\data\501_MA142_Gi_Saline_0513.csv"
-    VIDEO_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513.mp4"
-    SAVE_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513_rotated.mp4"
-    ANCHOR_LOC = np.array([250, 250])
-
-    df = read_df(file_path=DATA_PATH, file_type='csv')
-    bp_cols = [x for x in df.columns if not x.endswith('_p')]
-    data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32)
-
-    _, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=5, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0)
-    rotater = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_vectors, anchor_location=(400, 100), save_path=SAVE_PATH, verbose=True, core_cnt=16)
-    rotater.run()
+# if __name__ == "__main__":
+#     DATA_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\data\501_MA142_Gi_Saline_0513.csv"
+#     VIDEO_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513.mp4"
+#     SAVE_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513_rotated.mp4"
+#     ANCHOR_LOC = np.array([250, 250])
+#
+#     df = read_df(file_path=DATA_PATH, file_type='csv')
+#     bp_cols = [x for x in df.columns if not x.endswith('_p')]
+#     data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32)
+#
+#     _, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=5, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0)
+#     rotater = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_vectors, anchor_location=(400, 100), save_path=SAVE_PATH, verbose=True, core_cnt=16)
+#     rotater.run()