simba-uw-tf-dev 4.7.2__py3-none-any.whl → 4.7.5__py3-none-any.whl

simba/mixins/image_mixin.py

@@ -18,6 +18,7 @@ from collections import ChainMap
 import cv2
 import pandas as pd
 from numba import float64, int64, jit, njit, prange, uint8
+from PIL import Image, ImageDraw, ImageFont
 from shapely.geometry import Polygon
 from skimage.metrics import structural_similarity
 
@@ -30,11 +31,12 @@ from simba.utils.checks import (check_file_exist_and_readable, check_float,
 from simba.utils.data import terminate_cpu_pool
 from simba.utils.enums import Defaults, Formats, GeometryEnum, Options
 from simba.utils.errors import ArrayError, FrameRangeError, InvalidInputError
+from simba.utils.lookups import get_fonts
 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_frm_of_video, seconds_to_timestamp)
 
 
 class ImageMixin(object):
@@ -2052,18 +2054,141 @@ class ImageMixin(object):
 
         return denoised_img
 
+    @staticmethod
+    def get_timelapse_img(video_path: Union[str, os.PathLike],
+                          frame_cnt: int = 25,
+                          size: Optional[int] = None,
+                          crop_ratio: int = 50) -> np.ndarray:
 
+        """
+        Creates timelapse image from video.
 
+        .. image:: _static/img/get_timelapse_img.png
+           :width: 600
+           :align: center
 
+        :param Union[str, os.PathLike] video_path: Path to the video to cerate the timelapse image from.
+        :param int frame_cnt: Number of frames to grab from the video. There will be an even interval between each frame.
+        :param Optional[int] size: The total width in pixels of the final timelapse image. If None, uses the video width (adjusted for crop_ratio).
+        :param int crop_ratio: The percent of each original video (from the left) to show.
+        :return np.ndarray: The timelapse image as a numpy array
 
+        :example:
+        >>> img = ImageMixin.get_timelapse_img(video_path=r"E:\troubleshooting\mitra_emergence\project_folder\clip_test\Box1_180mISOcontrol_Females_clipped_progress_bar.mp4", size=100)
+        """
 
+        video_meta = get_video_meta_data(video_path=video_path, raise_error=True)
+        frm_ids = [int(i * video_meta['frame_count'] / frame_cnt) for i in range(frame_cnt)]
+        cap = cv2.VideoCapture(video_path)
+        frms = [read_frm_of_video(video_path=cap, frame_index=x, use_ffmpeg=False) for x in frm_ids]
+
+        effective_video_width = int(video_meta['width'] * (crop_ratio / 100))
+        if size is None:
+            size = effective_video_width
+        per_frame_width_after_crop = size / frame_cnt
+        per_frame_width_before_crop = per_frame_width_after_crop / (crop_ratio / 100)
+        scale_factor = per_frame_width_before_crop / frms[0].shape[1]
+        scaled_frms = [cv2.resize(x, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR) for x in frms]
+        if crop_ratio is not None:
+            scaled_frms = [ImageMixin.segment_img_vertical(img=x, pct=crop_ratio, left=True) for x in scaled_frms]
+        
+        return cv2.hconcat(scaled_frms)
 
+    @staticmethod
+    def create_time_ruler(width: int,
+                          video_path: Union[str, os.PathLike],
+                          height: int = 60,
+                          num_divisions: int = 6,
+                          font: str = 'Arial',
+                          bg_color: Tuple[int, int, int] = (255, 255, 255),
+                          line_color: Tuple[int, int, int] = (128, 128, 128),
+                          text_color: Tuple[int, int, int] = (0, 0, 0),
+                          padding: int = 60,
+                          show_time: bool = True) -> np.ndarray:
+        """
+        Create a horizontal ruler/scale bar with tick marks and labels.
+
+        .. image:: _static/img/create_time_ruler.png
+           :width: 600
+           :align: center
 
+        :param int width: Width of the ruler in pixels (should match timelapse image width if one is used)
+        :param Union[str, os.PathLike] video_path: Path to video file to get metadata from
+        :param int height: Height of the ruler in pixels. Default 60.
+        :param int num_divisions: Number of major divisions on the ruler. Default 6.
+        :param str font: Font name to use for labels. Default 'Algerian'.
+        :param Tuple[int, int, int] bg_color: Background color (R, G, B). Default white.
+        :param Tuple[int, int, int] line_color: Color for tick marks and lines (R, G, B). Default grey.
+        :param Tuple[int, int, int] text_color: Color for text labels (R, G, B). Default black.
+        :param bool show_time: If True, show time labels, else show frame numbers. Default True.
+        :return: Ruler image as numpy array (BGR format for OpenCV compatibility)
+        :rtype: np.ndarray
 
+        :example:
+        >>> ruler = ImageMixin.create_time_ruler(width=1920, video_path='path/to/video.mp4', height=60, num_divisions=6)
+        """
 
-
-
-
+        check_file_exist_and_readable(file_path=video_path)
+        check_int(name='width', value=width, min_value=1, raise_error=True)
+        check_int(name='height', value=height, min_value=1, raise_error=True)
+        check_int(name='num_divisions', value=num_divisions, min_value=1, raise_error=True)
+        check_int(name='padding', value=padding, min_value=0, raise_error=True)
+        check_str(name='font', value=font, allow_blank=False, raise_error=True)
+        check_if_valid_rgb_tuple(data=bg_color, raise_error=True, source=ImageMixin.create_time_ruler.__name__)
+        check_if_valid_rgb_tuple(data=line_color, raise_error=True, source=ImageMixin.create_time_ruler.__name__)
+        check_if_valid_rgb_tuple(data=text_color, raise_error=True, source=ImageMixin.create_time_ruler.__name__)
+
+        video_meta = get_video_meta_data(video_path=video_path, raise_error=True)
+        total_width = width + (2 * padding)
+
+        img = Image.new('RGB', (total_width, height), color=bg_color)
+        draw = ImageDraw.Draw(img)
+        font_dict = get_fonts()
+        try:
+            font_path = font_dict[font]
+            pil_font = ImageFont.truetype(font_path, size=12)
+        except (KeyError, OSError):
+            pil_font = ImageFont.load_default()
+        major_tick_height, half_tick_height = height * 0.6, height * 0.4
+        quarter_tick_height, eighth_tick_height = height * 0.25, height * 0.15
+
+        for i in range(num_divisions + 1):
+            x = padding + int(i * width / num_divisions)
+            draw.line([(x, 0), (x, major_tick_height)], fill=line_color, width=2)
+            if show_time and video_meta['video_length_s'] is not None:
+                seconds_at_division = i * video_meta['video_length_s'] / num_divisions
+                label = seconds_to_timestamp(seconds=seconds_at_division)
+            elif video_meta['frame_count'] is not None:
+                label = str(int(i * video_meta['frame_count'] / num_divisions))
+            else:
+                label = str(i)
+            bbox = draw.textbbox((0, 0), label, font=pil_font)
+            text_width = bbox[2] - bbox[0]
+            if i == 0:
+                draw.text((x, major_tick_height + 5), label, fill=text_color, font=pil_font)
+            elif i == num_divisions:
+                draw.text((x - text_width, major_tick_height + 5), label, fill=text_color, font=pil_font)
+            else:
+                draw.text((x - text_width // 2, major_tick_height + 5), label, fill=text_color, font=pil_font)
+            if i < num_divisions:
+                x_half = padding + int((i + 0.5) * width / num_divisions)
+                draw.line([(x_half, 0), (x_half, half_tick_height)], fill=line_color, width=1)
+                for q in [0.25, 0.75]:
+                    x_quarter = padding + int((i + q) * width / num_divisions)
+                    draw.line([(x_quarter, 0), (x_quarter, quarter_tick_height)], fill=line_color, width=1)
+                for e in [0.125, 0.375, 0.625, 0.875]:
+                    x_eighth = padding + int((i + e) * width / num_divisions)
+                    draw.line([(x_eighth, 0), (x_eighth, eighth_tick_height)], fill=line_color, width=1)
+
+        draw.line([(0, height - 1), (total_width, height - 1)], fill=line_color, width=1)
+        img_bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+        return img_bgr
+
+
+# img = ImageMixin.create_time_ruler(width=1920, video_path=r"E:\troubleshooting\mitra_emergence\project_folder\clip_test\Box1_180mISOcontrol_Females_clipped_progress_bar.mp4", height=60, num_divisions=6)
+#
+# cv2.imshow('sadasdas', img)
+# cv2.waitKey(40000)
 
 #x = ImageMixin.get_blob_locations(video_path=r"C:\troubleshooting\RAT_NOR\project_folder\videos\2022-06-20_NOB_DOT_4_downsampled_bg_subtracted.mp4", gpu=True)
 # imgs = ImageMixin().read_all_img_in_dir(dir='/Users/simon/Desktop/envs/simba/troubleshooting/RAT_NOR/project_folder/videos/examples')

simba/model/yolo_fit.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from contextlib import redirect_stderr, redirect_stdout
 
 os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
 import argparse
@@ -21,7 +22,8 @@ from simba.utils.checks import (check_file_exist_and_readable,
                                 check_valid_boolean, check_valid_device)
 from simba.utils.enums import Options
 from simba.utils.errors import SimBAGPUError, SimBAPAckageVersionError
-from simba.utils.read_write import find_core_cnt
+from simba.utils.printing import stdout_information
+from simba.utils.read_write import find_core_cnt, get_current_time
 from simba.utils.yolo import load_yolo_model
 
 
@@ -108,20 +110,25 @@ class FitYolo():
 
 
     def run(self):
-        model = load_yolo_model(weights_path=self.weights_path,
-                                verbose=self.verbose,
-                                format=self.format,
-                                device=self.device)
-
-        model.train(data=self.model_yaml,
-                    epochs=self.epochs,
-                    project=self.save_path,
-                    batch=self.batch,
-                    plots=self.plots,
-                    imgsz=self.imgsz,
-                    workers=self.workers,
-                    device=self.device,
-                    patience=self.patience)
+        # Temporarily redirect stdout/stderr to terminal to ensure ultralytics output goes to terminal
+        # sys.__stdout__ and sys.__stderr__ are the original terminal streams
+        stdout_information(msg=f'[{get_current_time()}] Please follow the YOLO pose model training in the terminal from where SimBA was launched ...', source=self.__class__.__name__)
+        stdout_information(msg=f'[{get_current_time()}] Results will be stored in the {self.save_path} directory ..', source=self.__class__.__name__)
+        with redirect_stdout(sys.__stdout__), redirect_stderr(sys.__stderr__):
+            model = load_yolo_model(weights_path=self.weights_path,
+                                    verbose=self.verbose,
+                                    format=self.format,
+                                    device=self.device)
+
+            model.train(data=self.model_yaml,
+                        epochs=self.epochs,
+                        project=self.save_path,
+                        batch=self.batch,
+                        plots=self.plots,
+                        imgsz=self.imgsz,
+                        workers=self.workers,
+                        device=self.device,
+                        patience=self.patience)
 
 
 if __name__ == "__main__" and not hasattr(sys, 'ps1'):

simba/model/yolo_pose_inference.py

@@ -34,7 +34,7 @@ from simba.utils.errors import (CountError, InvalidFilepathError,
                                 InvalidFileTypeError, SimBAGPUError,
                                 SimBAPAckageVersionError)
 from simba.utils.lookups import get_current_time
-from simba.utils.printing import SimbaTimer, stdout_success
+from simba.utils.printing import SimbaTimer, stdout_information, stdout_success
 from simba.utils.read_write import (find_files_of_filetypes_in_directory,
                                     get_video_meta_data, recursive_file_search)
 from simba.utils.warnings import FileExistWarning, NoDataFoundWarning
@@ -182,7 +182,12 @@ class YOLOPoseInference():
         results = {}
         class_dict = self.model.names
         timer = SimbaTimer(start=True)
-        print(f'Starting tracking inference for {len(self.video_path)} video(s) ({get_current_time()})... ')
+        if self.save_dir is not None:
+            msg = f'[{get_current_time()}] Starting tracking inference for {len(self.video_path)} video(s). Results will be saved in {self.save_dir} ... '
+        else:
+            msg = f'[{get_current_time()}] Starting tracking inference for {len(self.video_path)} video(s) ... '
+        stdout_information(msg=msg, source=self.__class__.__name__)
+        stdout_information(msg='Follow progress in OS terminal window ...', source=self.__class__.__name__)
         for video_cnt, path in enumerate(self.video_path):
             video_timer = SimbaTimer(start=True)
             _, video_name, _ = get_fn_ext(filepath=path)

simba/outlier_tools/skip_outlier_correction.py

@@ -5,7 +5,7 @@ from typing import Union
 
 from simba.mixins.config_reader import ConfigReader
 from simba.utils.checks import check_if_filepath_list_is_empty
-from simba.utils.printing import SimbaTimer, stdout_success
+from simba.utils.printing import SimbaTimer, stdout_information, stdout_success
 from simba.utils.read_write import get_fn_ext, read_df, write_df
 
 
@@ -43,7 +43,7 @@ class OutlierCorrectionSkipper(ConfigReader):
             save_path = os.path.join(self.outlier_corrected_dir, f"{video_name}.{self.file_type}")
             write_df(df=data_df, file_type=self.file_type, save_path=save_path)
             video_timer.stop_timer()
-            print(f"Skipped outlier correction for video {video_name} (elapsed time {video_timer.elapsed_time_str}s)...")
+            stdout_information(msg=f"Skipped outlier correction for video {video_name} (Video {file_cnt+1}/{len(self.input_csv_paths)})", elapsed_time=video_timer.elapsed_time_str)
         self.timer.stop_timer()
         stdout_success(msg=f"Skipped outlier correction for {len(self.input_csv_paths)} files", elapsed_time=self.timer.elapsed_time_str)
 

simba/plotting/heat_mapper_clf_mp.py

@@ -17,7 +17,7 @@ from simba.utils.checks import (
     check_all_file_names_are_represented_in_video_log,
     check_filepaths_in_iterable_exist, check_int, check_str,
     check_valid_boolean, check_valid_dataframe, check_valid_dict)
-from simba.utils.data import terminate_cpu_pool
+from simba.utils.data import get_cpu_pool, terminate_cpu_pool
 from simba.utils.enums import Formats
 from simba.utils.errors import InvalidInputError, NoSpecifiedOutputError
 from simba.utils.printing import SimbaTimer, stdout_success
@@ -149,6 +149,7 @@ class HeatMapperClfMultiprocess(ConfigReader, PlottingMixin):
     def run(self):
         print(f"Processing {len(self.data_paths)} video(s)...")
         check_all_file_names_are_represented_in_video_log(video_info_df=self.video_info_df, data_paths=self.data_paths)
+        pool = get_cpu_pool(core_cnt=self.core_cnt, source=self.__class__.__name__)
         for file_cnt, file_path in enumerate(self.data_paths):
             video_timer = SimbaTimer(start=True)
             _, self.video_name, _ = get_fn_ext(file_path)
@@ -173,7 +174,8 @@ class HeatMapperClfMultiprocess(ConfigReader, PlottingMixin):
             if len(np.unique(clf_data)) == 1:
                 raise InvalidInputError(msg=f'Cannot plot heatmap for behavior {self.clf_name} in video {self.video_name}. The behavior is classified as {np.unique(clf_data)} in every single frame.')
             grid, aspect_ratio = GeometryMixin.bucket_img_into_grid_square(img_size=(self.width, self.height), bucket_grid_size_mm=self.bin_size, px_per_mm=self.px_per_mm, add_correction=False, verbose=False)
-            clf_data = GeometryMixin().cumsum_bool_geometries(data=bp_data, geometries=grid, bool_data=clf_data, fps=self.fps, verbose=False)
+
+            clf_data = GeometryMixin().cumsum_bool_geometries(data=bp_data, geometries=grid, bool_data=clf_data, fps=self.fps, verbose=False, core_cnt=self.core_cnt, pool=pool)
             if self.max_scale == "auto":
                 self.max_scale = max(1, self.__calculate_max_scale(clf_array=clf_data))
             if self.final_img_setting:
@@ -197,32 +199,29 @@ class HeatMapperClfMultiprocess(ConfigReader, PlottingMixin):
                     frm_per_core_w_batch.append((batch_cnt, frm_range, frame_arrays[batch_cnt]))
                 del frame_arrays
                 print(f"Creating heatmaps, multiprocessing (chunksize: {self.multiprocess_chunksize}, cores: {self.core_cnt})...")
-                with multiprocessing.Pool(self.core_cnt, maxtasksperchild=self.maxtasksperchild) as pool:
-                    constants = functools.partial(_heatmap_multiprocessor,
-                                                  video_setting=self.video_setting,
-                                                  frame_setting=self.frame_setting,
-                                                  style_attr=self.style_attr,
-                                                  fps=self.fps,
-                                                  video_temp_dir=self.temp_folder,
-                                                  frame_dir=self.frames_save_dir,
-                                                  max_scale=self.max_scale,
-                                                  aspect_ratio=aspect_ratio,
-                                                  clf_name=self.clf_name,
-                                                  size=(self.width, self.height),
-                                                  video_name=self.video_name,
-                                                  make_clf_heatmap_plot=self.make_clf_heatmap_plot)
-
-                    for cnt, batch in enumerate(pool.imap(constants, frm_per_core_w_batch, chunksize=self.multiprocess_chunksize)):
-                        print(f'Batch core {batch+1}/{self.core_cnt} complete (Video {self.video_name})... ')
-                    terminate_cpu_pool(pool=pool, force=False)
-
+                constants = functools.partial(_heatmap_multiprocessor,
+                                              video_setting=self.video_setting,
+                                              frame_setting=self.frame_setting,
+                                              style_attr=self.style_attr,
+                                              fps=self.fps,
+                                              video_temp_dir=self.temp_folder,
+                                              frame_dir=self.frames_save_dir,
+                                              max_scale=self.max_scale,
+                                              aspect_ratio=aspect_ratio,
+                                              clf_name=self.clf_name,
+                                              size=(self.width, self.height),
+                                              video_name=self.video_name,
+                                              make_clf_heatmap_plot=self.make_clf_heatmap_plot)
+
+
+                for cnt, batch in enumerate(pool.imap(constants, frm_per_core_w_batch, chunksize=self.multiprocess_chunksize)):
+                    print(f'Batch core {batch+1}/{self.core_cnt} complete (Video {self.video_name})... ')
                 if self.video_setting:
                     print(f"Joining {self.video_name} multiprocessed video...")
                     concatenate_videos_in_folder(in_folder=self.temp_folder, save_path=self.save_video_path)
-
                 video_timer.stop_timer()
                 print(f"Heatmap video {self.video_name} complete, (elapsed time: {video_timer.elapsed_time_str}s) ...")
-
+        terminate_cpu_pool(pool=pool, force=False, source=self.__class__.__name__)
         self.timer.stop_timer()
         stdout_success(msg=f"Heatmap visualizations for {len(self.data_paths)} video(s) created in {self.heatmap_clf_location_dir} directory", elapsed_time=self.timer.elapsed_time_str)
 
@@ -261,3 +260,26 @@ class HeatMapperClfMultiprocess(ConfigReader, PlottingMixin):
 #                                core_cnt=5,
 #                      files_found=['/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/csv/machine_results/Together_1.csv'])
 # test.create_heatmaps()
+# if __name__ == "__main__":
+#     x = HeatMapperClfMultiprocess(config_path=r"E:\troubleshooting\mitra_emergence\project_folder\project_config.ini",
+#                               bodypart='nose',
+#                               clf_name='GROOMING',
+#                               style_attr={'palette': 'jet', 'shading': 'gouraud', 'bin_size': 25, 'max_scale': 'auto'},
+#                               final_img_setting=True,
+#                               video_setting=False,
+#                               frame_setting=False,
+#                               core_cnt=12,
+#                               data_paths=[r"E:\troubleshooting\mitra_emergence\project_folder\csv\machine_results\Box1_180mISOcontrol_Females.csv"])
+#
+#     x.run()
+
+    # def __init__(self,
+    #              config_path: Union[str, os.PathLike],
+    #              bodypart: str,
+    #              clf_name: str,
+    #              data_paths: List[str],
+    #              style_attr: dict,
+    #              final_img_setting: bool = True,
+    #              video_setting: bool = False,
+    #              frame_setting: bool = False,
+    #              core_cnt: int = -1):

simba/plotting/plot_clf_results.py

@@ -236,7 +236,8 @@ class PlotSklearnResultsSingleCore(ConfigReader, TrainModelMixin, PlottingMixin)
                             self.add_spacer += 1
                         if self.show_confidence:
                             col_name = f'Probability_{clf_name}'
-                            conf_txt = f'{clf_name} CONFIDENCE {self.data_df.loc[frm_idx, col_name]}'
+                            conf = round(self.data_df.loc[frm_idx, col_name], 4)
+                            conf_txt = f'{clf_name} CONFIDENCE {conf:.4f}'
                             self.frame = PlottingMixin().put_text(img=self.frame, text=conf_txt, pos=(TextOptions.BORDER_BUFFER_Y.value, ((self.video_meta_data["height"] - self.video_meta_data["height"]) + self.video_space_size * self.add_spacer)), font_size=self.video_font_size, font_thickness=self.video_text_thickness, font=self.font, text_bg_alpha=self.video_text_opacity, text_color_bg=self.text_bg_color, text_color=self.text_color)
                             self.add_spacer += 1
                     self.frame = PlottingMixin().put_text(img=self.frame, text="ENSEMBLE PREDICTION:", pos=(TextOptions.BORDER_BUFFER_Y.value, ((self.video_meta_data["height"] - self.video_meta_data["height"]) + self.video_space_size * self.add_spacer)), font_size=self.video_font_size, font_thickness=self.video_text_thickness, font=self.font, text_bg_alpha=self.video_text_opacity, text_color_bg=self.text_bg_color, text_color=self.text_color)