simba-uw-tf-dev 4.5.8__py3-none-any.whl → 4.7.1__py3-none-any.whl

simba/plotting/frame_mergerer_ffmpeg.py

@@ -1,196 +1,137 @@
-__author__ = "Simon Nilsson; sronilsson@gmail.com"
-
-import os
-import shutil
-from datetime import datetime
-from typing import List, Optional, Union
-
-try:
-    from typing import Literal
-except ImportError:
-    from typing_extensions import Literal
-
-from simba.mixins.config_reader import ConfigReader
-from simba.utils.checks import (check_ffmpeg_available,
-                                check_file_exist_and_readable, check_int,
-                                check_nvidea_gpu_available, check_str,
-                                check_valid_lst)
-from simba.utils.enums import Paths, TagNames
-from simba.utils.errors import FFMPEGCodecGPUError
-from simba.utils.printing import SimbaTimer, log_event, stdout_success
-from simba.utils.read_write import copy_files_to_directory, get_fn_ext
-from simba.video_processors.video_processing import (
-    horizontal_video_concatenator, mixed_mosaic_concatenator,
-    mosaic_concatenator, vertical_video_concatenator)
-
-HORIZONTAL = "horizontal"
-VERTICAL = "vertical"
-MOSAIC = "mosaic"
-MIXED_MOSAIC = "mixed_mosaic"
-ACCEPTED_TYPES = [HORIZONTAL, VERTICAL, MOSAIC, MIXED_MOSAIC]
-
-
-class FrameMergererFFmpeg(ConfigReader):
-    """
-    Merge separate visualizations of classifications, descriptive statistics etc., into  single video mosaic.
-
-    .. note::
-       `GitHub tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-12-merge-frames>`_.
-
-        .. image:: _static/img/mosaic_videos.gif
-          :width: 600
-          :align: center
-
-    :parameter str config_path: Optional path to SimBA project config file in Configparser format.
-    :parameter Literal["horizontal", "vertical", "mosaic", "mixed_mosaic"] concat_type: Type of concatenation. OPTIONS: 'horizontal', 'vertical', 'mosaic', 'mixed_mosaic'.
-    :parameter List[Union[str, os.PathLike]] video_paths: List with videos to concatenate.
-    :parameter Optional[int] video_height: Optional height of the canatenated videos. Required if concat concat_type is not mixed_mosaic.
-    :parameter int video_width: Optional wisth of the canatenated videos. Required if concat concat_type is not mixed_mosaic.
-    :parameter Optional[bool] gpu: If True, use NVIDEA FFMpeg GPU codecs. Default False.
-
-    :example:
-    >>> video_paths = ['/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4', '/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4']
-    >>> merger = FrameMergererFFmpeg(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini', video_paths=videos, video_height=600, video_width=600, concat_type='mosaic')
-    >>> merger.run()
-    """
-
-    def __init__(
-        self,
-        concat_type: Literal["horizontal", "vertical", "mosaic", "mixed_mosaic"],
-        video_paths: List[Union[str, os.PathLike]],
-        video_height: Optional[int] = None,
-        video_width: Optional[int] = None,
-        config_path: Optional[str] = None,
-        gpu: Optional[bool] = False,
-    ):
-
-        if gpu and not check_nvidea_gpu_available():
-            raise FFMPEGCodecGPUError(
-                msg="NVIDEA GPU not available (as evaluated by nvidea-smi returning None",
-                source=self.__class__.__name__,
-            )
-        check_ffmpeg_available()
-        check_str(
-            name=f"{FrameMergererFFmpeg.__name__} concat_type",
-            value=concat_type,
-            options=ACCEPTED_TYPES,
-        )
-        check_valid_lst(
-            data=video_paths,
-            source=f"{self.__class__.__name__} video_paths",
-            valid_dtypes=(str,),
-            min_len=2,
-        )
-        for i in video_paths:
-            check_file_exist_and_readable(file_path=i)
-        if concat_type != MIXED_MOSAIC:
-            check_int(
-                name=f"{FrameMergererFFmpeg.__name__} video_height",
-                value=video_height,
-                min_value=0,
-            )
-            check_int(
-                name=f"{FrameMergererFFmpeg.__name__} video_width",
-                value=video_height,
-                min_value=0,
-            )
-        if config_path is not None:
-            ConfigReader.__init__(self, config_path=config_path)
-            log_event(
-                logger_name=str(__class__.__name__),
-                log_type=TagNames.CLASS_INIT.value,
-                msg=self.create_log_msg_from_init_args(locals=locals()),
-            )
-            self.output_dir = os.path.join(
-                self.project_path, Paths.CONCAT_VIDEOS_DIR.value
-            )
-            self.output_path = os.path.join(
-                self.project_path,
-                Paths.CONCAT_VIDEOS_DIR.value,
-                f"merged_video_{self.datetime}.mp4",
-            )
-        else:
-            self.timer = SimbaTimer(start=True)
-            self.datetime = datetime.now().strftime("%Y%m%d%H%M%S")
-            self.output_dir, _, _ = get_fn_ext(filepath=video_paths[0])
-            self.output_path = os.path.join(
-                self.output_dir, f"merged_video_{self.datetime}.mp4"
-            )
-
-        self.video_height, self.video_width, self.gpu = video_height, video_width, gpu
-        self.video_paths, self.concat_type = video_paths, concat_type
-        if not os.path.exists(self.output_dir):
-            os.makedirs(self.output_dir)
-
-    def run(self):
-        if self.concat_type == HORIZONTAL:
-            _ = horizontal_video_concatenator(
-                video_paths=self.video_paths,
-                save_path=self.output_path,
-                height_px=self.video_height,
-                gpu=self.gpu,
-                verbose=True,
-            )
-        elif self.concat_type == VERTICAL:
-            _ = vertical_video_concatenator(
-                video_paths=self.video_paths,
-                save_path=self.output_path,
-                width_px=self.video_width,
-                gpu=self.gpu,
-                verbose=True,
-            )
-        elif self.concat_type == MOSAIC:
-            _ = mosaic_concatenator(
-                video_paths=self.video_paths,
-                save_path=self.output_path,
-                width_px=self.video_width,
-                height_px=self.video_height,
-                gpu=self.gpu,
-                verbose=True,
-            )
-        else:
-            _ = mixed_mosaic_concatenator(
-                video_paths=self.video_paths,
-                save_path=self.output_path,
-                gpu=self.gpu,
-                verbose=True,
-            )
-        self.timer.stop_timer()
-        stdout_success(
-            msg=f"Merged video saved at {self.output_path}",
-            source=self.__class__.__name__,
-            elapsed_time=self.timer.elapsed_time_str,
-        )
-
-
-# videos = ['/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4', '/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4']
-#
-# merger = FrameMergererFFmpeg(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
-#                     video_paths=videos,
-#                     video_height=600,
-#                     video_width=600,
-#                     concat_type='vertical') #horizontal, vertical, mosaic, mixed_mosaic
-# merger.run()
-
-
-#
-# FrameMergererFFmpeg(config_path=None,
-#                     frame_types={'Video 1': '/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.avi',
-#                                  'Video 2': '/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.avi'},
-#                     video_height=640,
-#                     video_width=480,
-#                     concat_type='vertical') #horizontal, vertical, mosaic, mixed_mosaic
-#
-#
-
-
-# FrameMergererFFmpeg(config_path=None,
-#                     frame_types={'Video 1': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4',
-#                                  'Video 2': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4',
-#                                  'Video 3': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4',
-#                                  'Video 4': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4'},
-#                     video_height=640,
-#                     video_width=480,
-#                     concat_type='mixed_mosaic',
-#                     gpu=False) #horizontal, vertical, mosaic, mixed_mosaic
-#
+__author__ = "Simon Nilsson; sronilsson@gmail.com"
+
+import os
+import shutil
+from datetime import datetime
+from typing import List, Optional, Union
+
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+from simba.mixins.config_reader import ConfigReader
+from simba.utils.checks import (check_ffmpeg_available,
+                                check_file_exist_and_readable, check_int,
+                                check_nvidea_gpu_available, check_str,
+                                check_valid_boolean, check_valid_lst)
+from simba.utils.enums import Paths, TagNames
+from simba.utils.errors import FFMPEGCodecGPUError
+from simba.utils.printing import SimbaTimer, log_event, stdout_success
+from simba.utils.read_write import copy_files_to_directory, get_fn_ext
+from simba.video_processors.video_processing import (
+    horizontal_video_concatenator, mixed_mosaic_concatenator,
+    mosaic_concatenator, vertical_video_concatenator)
+
+HORIZONTAL = "horizontal"
+VERTICAL = "vertical"
+MOSAIC = "mosaic"
+MIXED_MOSAIC = "mixed_mosaic"
+ACCEPTED_TYPES = [HORIZONTAL, VERTICAL, MOSAIC, MIXED_MOSAIC]
+
+
+class FrameMergererFFmpeg(ConfigReader):
+    """
+    Merge separate visualizations of classifications, descriptive statistics etc., into  single video mosaic.
+
+    .. note::
+       `GitHub tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-12-merge-frames>`_.
+
+        .. image:: _static/img/mosaic_videos.gif
+          :width: 600
+          :align: center
+
+    :parameter str config_path: Optional path to SimBA project config file in Configparser format.
+    :parameter Literal["horizontal", "vertical", "mosaic", "mixed_mosaic"] concat_type: Type of concatenation. OPTIONS: 'horizontal', 'vertical', 'mosaic', 'mixed_mosaic'.
+    :parameter List[Union[str, os.PathLike]] video_paths: List with videos to concatenate.
+    :param int quality: Video quality (CRF value). Lower values = higher quality. Range 0-52. Default 23.
+    :parameter Optional[int] video_height: Optional height of the canatenated videos. Required if concat concat_type is not mixed_mosaic.
+    :parameter int video_width: Optional wisth of the canatenated videos. Required if concat concat_type is not mixed_mosaic.
+    :parameter Optional[bool] gpu: If True, use NVIDEA FFMpeg GPU codecs. Default False.
+
+    :example:
+    >>> video_paths = ['/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4', '/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4']
+    >>> merger = FrameMergererFFmpeg(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini', video_paths=videos, video_height=600, video_width=600, concat_type='mosaic')
+    >>> merger.run()
+    """
+
+    def __init__( self,
+                 concat_type: Literal["horizontal", "vertical", "mosaic", "mixed_mosaic"],
+                 video_paths: List[Union[str, os.PathLike]],
+                 video_height: Optional[int] = None,
+                 video_width: Optional[int] = None,
+                 config_path: Optional[str] = None,
+                 quality: int = 23,
+                 gpu: bool = False):
+
+        if gpu and not check_nvidea_gpu_available():
+            raise FFMPEGCodecGPUError(msg="NVIDEA GPU not available (as evaluated by nvidea-smi returning None", source=self.__class__.__name__)
+        check_ffmpeg_available()
+        check_str(name=f"{FrameMergererFFmpeg.__name__} concat_type", value=concat_type, options=ACCEPTED_TYPES)
+        check_valid_lst(data=video_paths, source=f"{self.__class__.__name__} video_paths", valid_dtypes=(str,), min_len=2)
+        check_valid_boolean(value=gpu, source=f'{self.__class__.__name__} gpu', raise_error=True)
+        self.quality = 23 if not check_int(name='quality', value=quality, min_value=0, max_value=52, raise_error=False)[0] else int(quality)
+        for i in video_paths:
+            check_file_exist_and_readable(file_path=i)
+        if concat_type != MIXED_MOSAIC:
+            check_int(name=f"{FrameMergererFFmpeg.__name__} video_height", value=video_height, min_value=0)
+            check_int(name=f"{FrameMergererFFmpeg.__name__} video_width", value=video_height, min_value=0)
+        if config_path is not None:
+            ConfigReader.__init__(self, config_path=config_path)
+            log_event(logger_name=str(__class__.__name__), log_type=TagNames.CLASS_INIT.value, msg=self.create_log_msg_from_init_args(locals=locals()))
+            self.output_dir = os.path.join(self.project_path, Paths.CONCAT_VIDEOS_DIR.value)
+            self.output_path = os.path.join(self.project_path, Paths.CONCAT_VIDEOS_DIR.value, f"merged_video_{self.datetime}.mp4")
+        else:
+            self.timer = SimbaTimer(start=True)
+            self.datetime = datetime.now().strftime("%Y%m%d%H%M%S")
+            self.output_dir, _, _ = get_fn_ext(filepath=video_paths[0])
+            self.output_path = os.path.join(self.output_dir, f"merged_video_{self.datetime}.mp4")
+
+        self.video_height, self.video_width, self.gpu = video_height, video_width, gpu
+        self.video_paths, self.concat_type = video_paths, concat_type
+        if not os.path.exists(self.output_dir): os.makedirs(self.output_dir)
+
+    def run(self):
+        if self.concat_type == HORIZONTAL:
+            _ = horizontal_video_concatenator(video_paths=self.video_paths, save_path=self.output_path, height_px=self.video_height, gpu=self.gpu, quality=self.quality, verbose=True)
+        elif self.concat_type == VERTICAL:
+            _ = vertical_video_concatenator(video_paths=self.video_paths, save_path=self.output_path, width_px=self.video_width, gpu=self.gpu, quality=self.quality, verbose=True)
+        elif self.concat_type == MOSAIC:
+            _ = mosaic_concatenator(video_paths=self.video_paths, save_path=self.output_path, width_px=self.video_width, height_px=self.video_height, gpu=self.gpu, quality=self.quality, verbose=True)
+        else:
+            _ = mixed_mosaic_concatenator(video_paths=self.video_paths, save_path=self.output_path, gpu=self.gpu, verbose=True)
+        self.timer.stop_timer()
+        stdout_success(msg=f"Merged video saved at {self.output_path}", source=self.__class__.__name__, elapsed_time=self.timer.elapsed_time_str)
+
+
+# videos = ['/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4', '/Users/simon/Desktop/envs/simba/troubleshooting/mouse_open_field/project_folder/videos/SI_DAY3_308_CD1_PRESENT_downsampled.mp4']
+#
+# merger = FrameMergererFFmpeg(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
+#                     video_paths=videos,
+#                     video_height=600,
+#                     video_width=600,
+#                     concat_type='vertical') #horizontal, vertical, mosaic, mixed_mosaic
+# merger.run()
+
+
+#
+# FrameMergererFFmpeg(config_path=None,
+#                     frame_types={'Video 1': '/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.avi',
+#                                  'Video 2': '/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.avi'},
+#                     video_height=640,
+#                     video_width=480,
+#                     concat_type='vertical') #horizontal, vertical, mosaic, mixed_mosaic
+#
+#
+
+
+# FrameMergererFFmpeg(config_path=None,
+#                     frame_types={'Video 1': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4',
+#                                  'Video 2': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4',
+#                                  'Video 3': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4',
+#                                  'Video 4': r'C:\Users\Nape_Computer_2\Desktop\test_videos\Box1_PM2_day_5_20211104T171021.mp4'},
+#                     video_height=640,
+#                     video_width=480,
+#                     concat_type='mixed_mosaic',
+#                     gpu=False) #horizontal, vertical, mosaic, mixed_mosaic
+#

simba/plotting/gantt_creator.py

@@ -2,6 +2,7 @@ __author__ = "Simon Nilsson; sronilsson@gmail.com"
 
 import os
 import shutil
+from copy import deepcopy
 from typing import List, Optional, Union
 
 import cv2
@@ -16,7 +17,7 @@ from simba.utils.checks import (
 from simba.utils.data import create_color_palette, detect_bouts
 from simba.utils.enums import Formats, Options
 from simba.utils.errors import NoSpecifiedOutputError
-from simba.utils.lookups import get_named_colors
+from simba.utils.lookups import get_fonts, get_named_colors
 from simba.utils.printing import stdout_success
 from simba.utils.read_write import get_fn_ext, read_df
 
@@ -60,16 +61,18 @@ class GanttCreatorSingleProcess(ConfigReader, PlottingMixin):
 
     def __init__(self,
                  config_path: Union[str, os.PathLike],
-                 data_paths: List[Union[str, os.PathLike]],
+                 data_paths: Optional[Union[Union[str, os.PathLike], List[Union[str, os.PathLike]]]] = None,
                  width: int = 640,
                  height: int = 480,
                  font_size: int = 8,
                  font_rotation: int = 45,
+                 font: Optional[str] = None,
                  palette: str = 'Set1',
-                 frame_setting: Optional[bool] = False,
-                 video_setting: Optional[bool] = False,
-                 last_frm_setting: Optional[bool] = True,
-                 hhmmss: Optional[bool] = True):
+                 frame_setting: bool = False,
+                 video_setting: bool = False,
+                 last_frm_setting: bool = True,
+                 hhmmss: bool = True,
+                 clf_names: Optional[List[str]] = None):
 
         if ((frame_setting != True) and (video_setting != True) and (last_frm_setting != True)):
             raise NoSpecifiedOutputError(msg="SIMBA ERROR: Please select gantt videos, frames, and/or last frame.")
@@ -78,7 +81,13 @@ class GanttCreatorSingleProcess(ConfigReader, PlottingMixin):
         check_int(value=height, min_value=1, name=f'{self.__class__.__name__} height')
         check_int(value=font_size, min_value=1, name=f'{self.__class__.__name__} font_size')
         check_int(value=font_rotation, min_value=0, max_value=180, name=f'{self.__class__.__name__} font_rotation')
-        check_valid_lst(data=data_paths, source=f'{self.__class__.__name__} data_paths', valid_dtypes=(str,), min_len=1)
+        if isinstance(data_paths, list):
+            check_valid_lst(data=data_paths, source=f'{self.__class__.__name__} data_paths', valid_dtypes=(str,), min_len=1)
+        elif isinstance(data_paths, str):
+            check_file_exist_and_readable(file_path=data_paths)
+            data_paths = [data_paths]
+        else:
+            data_paths = deepcopy(self.machine_results_paths)
         check_valid_boolean(value=hhmmss, source=f'{self.__class__.__name__} hhmmss', raise_error=False)
         palettes = Options.PALETTE_OPTIONS_CATEGORICAL.value + Options.PALETTE_OPTIONS.value
         check_str(name=f'{self.__class__.__name__} palette', value=palette, options=palettes)
@@ -90,7 +99,12 @@ class GanttCreatorSingleProcess(ConfigReader, PlottingMixin):
         if not os.path.exists(self.gantt_plot_dir): os.makedirs(self.gantt_plot_dir)
         self.frame_setting, self.video_setting, self.last_frm_setting = frame_setting, video_setting, last_frm_setting
         self.width, self.height, self.font_size, self.font_rotation = width, height, font_size, font_rotation
-        self.data_paths, self.hhmmss = data_paths, hhmmss
+        if font is not None:
+            check_str(name=f'{self.__class__.__name__} font', value=font, options=list(get_fonts().keys()), raise_error=True)
+        if clf_names is not None:
+            check_valid_lst(data=clf_names, source=f'{self.__class__.__name__} clf_names', valid_dtypes=(str,), valid_values=self.clf_names, min_len=1, raise_error=True)
+            self.clf_names = clf_names
+        self.data_paths, self.hhmmss, self.font = data_paths, hhmmss, font
         self.colours = get_named_colors()
         self.colour_tuple_x = list(np.arange(3.5, 203.5, 5))
 
@@ -121,6 +135,7 @@ class GanttCreatorSingleProcess(ConfigReader, PlottingMixin):
                                      font_size=self.font_size,
                                      font_rotation=self.font_rotation,
                                      video_name=self.video_name,
+                                     font=self.font,
                                      save_path=os.path.join(self.gantt_plot_dir, f"{self.video_name }_final_image.png"),
                                      palette=self.clr_lst,
                                      hhmmss=self.hhmmss)
@@ -135,6 +150,7 @@ class GanttCreatorSingleProcess(ConfigReader, PlottingMixin):
                                                 width=self.width,
                                                 height=self.height,
                                                 font_size=self.font_size,
+                                                font=self.font,
                                                 font_rotation=self.font_rotation,
                                                 video_name=self.video_name,
                                                 palette=self.clr_lst,
@@ -156,13 +172,16 @@ class GanttCreatorSingleProcess(ConfigReader, PlottingMixin):
 # test = GanttCreatorSingleProcess(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",
 #                                 frame_setting=False,
 #                                 video_setting=False,
-#                                 data_paths=[r"C:\troubleshooting\mitra\project_folder\csv\machine_results\592_MA147_Gq_CNO_0515.csv"],
+#                                 data_paths=[r"C:\troubleshooting\mitra\project_folder\csv\machine_results\501_MA142_Gi_CNO_0516.csv"],
 #                                 last_frm_setting=True,
 #                                 width=640,
 #                                 height= 480,
 #                                 font_size=10,
+#                                  font=None,
 #                                 font_rotation=45,
-#                                 palette='Set1')
+#                                 hhmmss=False,
+#                                 palette='Set1',
+#                                 clf_names=['straub_tail'])
 # test.run()