__author__ = "Simon Nilsson; sronilsson@gmail.com"
import functools
import multiprocessing
import os
import platform
from copy import deepcopy
from typing import Dict, List, Optional, Tuple, Union
try:
from typing import Literal
except:
from typing_extensions import Literal
import cv2
import numpy as np
import pandas as pd
from simba.mixins.config_reader import ConfigReader
from simba.mixins.geometry_mixin import GeometryMixin
from simba.mixins.plotting_mixin import PlottingMixin
from simba.mixins.train_model_mixin import TrainModelMixin
from simba.utils.checks import (check_float, check_if_dir_exists,
check_if_valid_rgb_tuple, check_int,
check_nvidea_gpu_available, check_str,
check_that_column_exist,
check_that_hhmmss_start_is_before_end,
check_valid_boolean, check_valid_dict,
check_valid_lst,
check_video_and_data_frm_count_align)
from simba.utils.data import (check_if_string_value_is_valid_video_timestamp,
create_color_palette, detect_bouts,
find_frame_numbers_from_time_stamp, get_cpu_pool,
terminate_cpu_pool)
from simba.utils.enums import ConfigKey, Dtypes, Options, TagNames, TextOptions
from simba.utils.errors import (InvalidInputError, NoDataError,
NoSpecifiedOutputError)
from simba.utils.lookups import get_simba_font_name_and_path
from simba.utils.printing import (SimbaTimer, log_event, stdout_information,
stdout_success)
from simba.utils.read_write import (concatenate_videos_in_folder,
create_directory,
find_all_videos_in_project, find_core_cnt,
get_fn_ext, get_video_meta_data,
read_config_entry, read_df,
seconds_to_timestamp)
from simba.utils.warnings import (FileExistWarning, FrameRangeWarning,
NoDataFoundWarning)
START_TIME, END_TIME = 'start_time', 'end_time'
SECONDS, HHMMSSSSSS = ['seconds', 'hh:mm:ss.ssss']
def _multiprocess_sklearn_video(data: pd.DataFrame,
bp_dict: dict,
video_save_dir: str,
frame_save_dir: str,
clf_cumsum: dict,
rotate: bool,
video_path: str,
video_setting: bool,
frame_setting: bool,
pose_threshold: float,
clf_confidence: Union[dict, None],
show_pose: bool,
show_animal_names: bool,
bbox: Optional[Literal['axis-aligned', 'animal-aligned']],
circle_size: int,
font_size: int,
font_path: Optional[str],
font_size_px: Optional[int],
space_size: int,
text_thickness: int,
timer: str,
text_opacity: float,
text_bg_clr: Tuple[int, int, int],
text_color: Tuple[int, int, int],
pose_clr_lst: List[Tuple[int, int, int]],
show_gantt: Optional[int],
bouts_df: Optional[pd.DataFrame],
final_gantt: Optional[np.ndarray],
gantt_clrs: List[Tuple[float, float, float]],
clf_names: List[str],
verbose:bool):
fourcc, font = cv2.VideoWriter_fourcc(*"mp4v"), cv2.FONT_HERSHEY_DUPLEX
font_size = font_size_px if font_path is not None else font_size
gantt_font = os.path.splitext(os.path.basename(font_path))[0] if font_path is not None else None
video_meta_data = get_video_meta_data(video_path=video_path)
if rotate:
video_meta_data["height"], video_meta_data["width"] = (video_meta_data['width'], video_meta_data['height'])
cap = cv2.VideoCapture(video_path)
batch, data = data
start_frm, current_frm, end_frm = (data["index"].iloc[0], data["index"].iloc[0], data["index"].iloc[-1])
if video_setting:
video_save_path = os.path.join(video_save_dir, f"{batch}.mp4")
if show_gantt is None:
video_writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], (video_meta_data["width"], video_meta_data["height"]))
else:
video_writer = cv2.VideoWriter(video_save_path, fourcc, video_meta_data["fps"], (int(video_meta_data["width"] + final_gantt.shape[1]), video_meta_data["height"]))
cap.set(1, start_frm)
while current_frm < end_frm:
ret, img = cap.read()
if ret:
clr_cnt = 0
for animal_cnt, (animal_name, animal_data) in enumerate(bp_dict.items()):
if show_pose:
for bp_no in range(len(animal_data["X_bps"])):
x_bp, y_bp, p_bp = (animal_data["X_bps"][bp_no], animal_data["Y_bps"][bp_no], animal_data["P_bps"][bp_no])
bp_cords = data.loc[current_frm, [x_bp, y_bp, p_bp]]
if bp_cords[p_bp] >= pose_threshold:
img = cv2.circle(img, (int(bp_cords[x_bp]), int(bp_cords[y_bp])), circle_size, pose_clr_lst[clr_cnt], -1)
clr_cnt += 1
if show_animal_names:
x_bp, y_bp, p_bp = (animal_data["X_bps"][0], animal_data["Y_bps"][0], animal_data["P_bps"][0])
bp_cords = data.loc[current_frm, [x_bp, y_bp, p_bp]]
img = PlottingMixin().put_text(img=img, text=animal_name, pos=(int(bp_cords[x_bp]), int(bp_cords[y_bp])), font_size=font_size, font=font, font_thickness=text_thickness, font_path=font_path, text_color=tuple(int(v) for v in pose_clr_lst[0]), text_bg_alpha=0.0)
if bbox is not None:
animal_headers = [val for pair in zip(animal_data["X_bps"], animal_data["Y_bps"]) for val in pair]
animal_cords = data.loc[current_frm, animal_headers].values.reshape(-1, 2).astype(np.int32)
try:
if bbox == Options.AXIS_ALIGNED.value:
animal_bbox = GeometryMixin().keypoints_to_axis_aligned_bounding_box(keypoints=animal_cords.reshape(-1, len(animal_cords), 2).astype(np.int32))
else:
animal_bbox = GeometryMixin().minimum_rotated_rectangle(shape=animal_cords, buffer=None)
animal_bbox = np.round(np.array(animal_bbox.exterior.coords)).astype(np.int32)
img = cv2.polylines(img, [animal_bbox], True, pose_clr_lst[animal_cnt], thickness=circle_size, lineType=cv2.LINE_AA)
except Exception as e:
pass
if rotate:
img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
if show_gantt == 1:
img = np.concatenate((img, final_gantt), axis=1)
elif show_gantt == 2:
bout_rows = bouts_df.loc[bouts_df["End_frame"] <= current_frm]
gantt_plot = PlottingMixin().make_gantt_plot(x_length=current_frm + 1,
bouts_df=bout_rows,
clf_names=clf_names,
fps=video_meta_data['fps'],
width=video_meta_data['width'],
height=video_meta_data['height'],
font_size=12,
font=gantt_font,
font_rotation=45,
x_tick_lbl_rotation=45,
hhmmss=True,
video_name=video_meta_data['video_name'],
save_path=None,
palette=gantt_clrs)
img = np.concatenate((img, gantt_plot), axis=1)
if timer is not None:
img = PlottingMixin().put_text(img=img, text="TIMERS:", pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size)), font_size=font_size, font_path=font_path, font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
add_spacer = 2
for clf_name, clf_time_df in clf_cumsum.items():
frame_results = clf_time_df.loc[current_frm]
clf_time = round(frame_results / video_meta_data['fps'], 2)
if timer is not None:
if timer == HHMMSSSSSS:
clf_time = seconds_to_timestamp(seconds=clf_time, hh_mm_ss_sss=True)
img = PlottingMixin().put_text(img=img, text=f"{clf_name} {clf_time}",pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer)), font_size=font_size, font_path=font_path, font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
add_spacer += 1
if clf_confidence is not None:
conf = round(clf_confidence[clf_name][current_frm], 4)
frm_clf_conf_txt = f'{clf_name} CONFIDENCE: {conf:.4f}'
img = PlottingMixin().put_text(img=img, text=frm_clf_conf_txt,pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer)), font_size=font_size, font_path=font_path, font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
add_spacer += 1
img = PlottingMixin().put_text(img=img, text="ENSEMBLE PREDICTION:", pos=(TextOptions.BORDER_BUFFER_Y.value, ((video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer)), font_size=font_size, font_path=font_path, font_thickness=text_thickness, font=font, text_bg_alpha=text_opacity, text_color_bg=text_bg_clr, text_color=text_color)
add_spacer += 1
for clf_name in clf_cumsum.keys():
if data.loc[current_frm, clf_name] == 1:
img = PlottingMixin().put_text(img=img, text=clf_name, pos=(TextOptions.BORDER_BUFFER_Y.value, (video_meta_data["height"] - video_meta_data["height"]) + space_size * add_spacer), font_size=font_size, font_path=font_path, font_thickness=text_thickness, font=font, text_color=TextOptions.FLAMINGO.value, text_bg_alpha=text_opacity)
add_spacer += 1
if video_setting:
video_writer.write(img.astype(np.uint8))
if frame_setting:
frame_save_name = os.path.join(frame_save_dir, f"{current_frm}.png")
cv2.imwrite(frame_save_name, img)
current_frm += 1
time = seconds_to_timestamp(seconds= current_frm/ video_meta_data['fps'])
if verbose: stdout_information(msg=f"Multi-processing video frame {current_frm} (time-stamp: {time}, core batch: {batch}, video name: {video_meta_data['video_name']})...")
else:
FrameRangeWarning(msg=f'Could not read frame {current_frm} in video {video_path}. Stopping video creation.')
break
cap.release()
if video_setting:
video_writer.release()
return batch
[docs]class PlotSklearnResultsMultiProcess(ConfigReader, TrainModelMixin, PlottingMixin):
"""
Plot classification results on videos using multiprocessing. Results are stored in the
`project_folder/frames/output/sklearn_results` directory of the SimBA project.
This class creates annotated videos/frames showing classifier predictions overlaid on pose-estimation data,
with optional Gantt charts, timers, and bounding boxes. Processing is parallelized across multiple CPU cores
for faster rendering of large video datasets.
.. seealso::
`Tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-10-sklearn-visualization>`__.
For single-core processing, see :meth:`simba.plotting.plot_clf_results.PlotSklearnResultsSingleCore`.
.. video:: _static/img/sklearn_visualization.webm
:width: 600
:autoplay:
:loop:
:muted:
:align: center
.. video:: _static/img/T1.webm
:width: 1000
:autoplay:
:loop:
:muted:
:align: center
.. youtube:: Frq6mMcaHBc
:width: 640
:height: 480
:align: center
:param Union[str, os.PathLike] config_path: Path to SimBA project config file in Configparser format.
:param bool video_setting: If True, creates compressed MP4 videos. Default True.
:param bool frame_setting: If True, saves individual annotated frames as PNG images. Default False.
:param Optional[Union[List[Union[str, os.PathLike]], Union[str, os.PathLike]]] video_paths: Path(s) to video file(s) to process. Accepts a single video file path, a list of video file paths, or a path to a directory containing videos (in which case all videos in that directory are processed). If None, processes all videos found in the project's video directory. Default None.
:param bool rotate: If True, rotates output videos 90 degrees clockwise. Default False.
:param bool animal_names: If True, displays animal names on the video frames. Default False.
:param bool show_pose: If True, overlays pose-estimation keypoints on the video. Default True.
:param bool show_confidence: If True, displays per-frame classifier confidence values (probabilities) for each behavior. Default False.
:param Optional[Union[int, float]] font_size: Font size for text overlays. If None, auto-computed based on video resolution. Default None.
:param Optional[Union[int, float]] space_size: Vertical spacing between text lines. If None, auto-computed. Default None.
:param Optional[Union[int, float]] text_thickness: Thickness of text characters. If None, uses default. Default None.
:param Optional[Union[int, float]] text_opacity: Opacity of text background (0.0-1.0). If None, defaults to 0.8. Default None.
:param Optional[Union[int, float]] circle_size: Radius of pose keypoint circles. If None, auto-computed based on video resolution. Default None.
:param Optional[str] pose_palette: Name of color palette for pose keypoints. Must be from :class:`simba.utils.enums.Options.PALETTE_OPTIONS_CATEGORICAL` or :class:`simba.utils.enums.Options.PALETTE_OPTIONS`. Default 'Set1'.
:param Optional[Literal['seconds', 'hh:mm:ss.ssss']] print_timer: Timer display mode for cumulative classifier durations. ``'seconds'`` shows numeric seconds, ``'hh:mm:ss.ssss'`` shows clock-style timestamps with sub-second precision, and ``None`` disables timer display. Default ``'seconds'``.
:param bool overwrite: If True, existing output files in the target output location may be overwritten. Default True.
:param Optional[Literal['axis-aligned', 'animal-aligned']] bbox: If not None, draw bounding boxes around each animal. ``'axis-aligned'`` = rectangle aligned with video axes; ``'animal-aligned'`` = minimum rotated rectangle aligned with the animal's orientation. Default None (no bounding boxes).
:param Optional[Dict[str, str]] time_slice: Optional time interval to render in ``{'start_time': 'HH:MM:SS', 'end_time': 'HH:MM:SS'}`` format. If None, processes full video duration. Default None.
:param Optional[int] show_gantt: If 1, appends static Gantt chart to video. If 2, appends dynamic Gantt chart that updates per frame. If None, no Gantt chart. Default None.
:param Tuple[int, int, int] text_clr: RGB color tuple for text foreground. Default (255, 255, 255) (white).
:param Tuple[int, int, int] text_bg_clr: RGB color tuple for text background. Default (0, 0, 0) (black).
:param bool gpu: If True, uses GPU acceleration for video concatenation (requires CUDA-capable GPU). Default False.
:param bool verbose: If True, prints progress and status messages during processing. Default True.
:param int core_cnt: Number of CPU cores to use for parallel processing. Pass -1 to use all available cores. Default -1.
:param Optional[Union[str, os.PathLike]] data_dir: Directory containing the classification data files (one per video, named ``<video_name>.<file_type>``) to overlay on the videos. If None, defaults to the project's ``machine_results_dir``. Default None.
:param Optional[Union[str, os.PathLike]] save_dir: Directory in which to write annotated videos/frames. Created if it does not exist. If None, defaults to the project's ``sklearn_plot_dir``. Default None.
:param Optional[Union[List[str], Tuple[str, ...]]] clf_names: Optional subset of classifier names to visualize. Each entry must match a classifier defined in the project config. If None, all project classifiers are plotted (current behavior). Default None.
:example:
>>> clf_plotter = PlotSklearnResultsMultiProcess(
... config_path='/Users/simon/Desktop/envs/simba/troubleshooting/beepboop174/project_folder/project_config.ini',
... video_setting=True,
... frame_setting=False,
... video_paths='Trial_10.mp4',
... rotate=False,
... show_pose=True,
... bbox='axis-aligned',
... print_timers=True,
... show_gantt=1,
... core_cnt=5
... )
>>> clf_plotter.run()
"""
def __init__(self,
config_path: Union[str, os.PathLike],
video_setting: bool = True,
frame_setting: bool = False,
video_paths: Optional[Union[List[Union[str, os.PathLike]], Union[str, os.PathLike]]] = None,
rotate: bool = False,
animal_names: bool = False,
show_pose: bool = True,
show_confidence: bool = False,
font_size: Optional[Union[int, float]] = None,
font: Optional[str] = None,
space_size: Optional[Union[int, float]] = None,
text_thickness: Optional[Union[int, float]] = None,
text_opacity: Optional[Union[int, float]] = None,
circle_size: Optional[Union[int, float]] = None,
pose_palette: Optional[str] = 'Set1',
print_timer: Optional[Literal['seconds', 'hh:mm:ss.ssss']] = 'seconds',
overwrite: bool = True,
bbox: Optional[Literal['axis-aligned', 'animal-aligned']] = None,
time_slice: Optional[Dict[str, str]] = None,
show_gantt: Optional[int] = None,
text_clr: Tuple[int, int, int] = (255, 255, 255),
text_bg_clr: Tuple[int, int, int] = (0, 0, 0),
gpu: bool = False,
verbose: bool = True,
core_cnt: int = -1,
data_dir: Optional[Union[str, os.PathLike]] = None,
save_dir: Optional[Union[str, os.PathLike]] = None,
clf_names: Optional[Union[List[str], Tuple[str, ...]]] = None):
ConfigReader.__init__(self, config_path=config_path)
TrainModelMixin.__init__(self)
PlottingMixin.__init__(self)
log_event(logger_name=str(__class__.__name__), log_type=TagNames.CLASS_INIT.value, msg=self.create_log_msg_from_init_args(locals=locals()))
for i in [video_setting, frame_setting, rotate, animal_names, show_pose, gpu, show_confidence]:
check_valid_boolean(value=i, source=self.__class__.__name__, raise_error=True)
if (not video_setting) and (not frame_setting):
raise NoSpecifiedOutputError(msg="Please choose to create a video and/or frames. SimBA found that you ticked neither video and/or frames", source=self.__class__.__name__)
if font_size is not None: check_float(name=f'{self.__class__.__name__} font_size', value=font_size, min_value=0.1)
self.font_path, self.font = None, None
if font is not None:
self.font, self.font_path = get_simba_font_name_and_path(font=font)
if space_size is not None: check_float(name=f'{self.__class__.__name__} space_size', value=space_size, min_value=0.1)
if text_thickness is not None: check_float(name=f'{self.__class__.__name__} text_thickness', value=text_thickness, min_value=0.1)
if circle_size is not None: check_float(name=f'{self.__class__.__name__} text_thickness', value=circle_size, min_value=0.1)
if circle_size is not None: check_float(name=f'{self.__class__.__name__} text_thickness', value=circle_size, min_value=0.1)
if text_opacity is not None: check_float(name=f'{self.__class__.__name__} text_opacity', value=text_opacity, min_value=0.1)
pose_palettes = Options.PALETTE_OPTIONS_CATEGORICAL.value + Options.PALETTE_OPTIONS.value
check_str(name=f'{self.__class__.__name__} pose_palette', value=pose_palette, options=pose_palettes)
if print_timer is not None: check_str(name=f'{self.__class__.__name__} timer', value=print_timer, options=(SECONDS, HHMMSSSSSS,))
self.clr_lst = create_color_palette(pallete_name=pose_palette, increments=len(self.body_parts_lst)+1)
check_if_valid_rgb_tuple(data=text_clr, source=f'{self.__class__.__name__} text_clr')
check_if_valid_rgb_tuple(data=text_bg_clr, source=f'{self.__class__.__name__} text_bg_clr')
check_valid_boolean(value=verbose, source=f'{self.__class__.__name__} verbose', raise_error=True)
check_valid_boolean(value=overwrite, source=f'{self.__class__.__name__} overwrite', raise_error=True)
if show_gantt is not None: check_int(name=f"{self.__class__.__name__} show_gantt", value=show_gantt, max_value=2, min_value=1)
if time_slice is not None:
check_valid_dict(x=time_slice, valid_key_dtypes=(str,), valid_values_dtypes=(str,), valid_keys=(START_TIME, END_TIME), required_keys=(START_TIME, END_TIME),)
check_if_string_value_is_valid_video_timestamp(value=time_slice[START_TIME], name='START TIME', raise_error=True)
check_if_string_value_is_valid_video_timestamp(value=time_slice[END_TIME], name='END TIME', raise_error=True)
check_that_hhmmss_start_is_before_end(start_time=time_slice[START_TIME], end_time=time_slice[END_TIME], name=f'TIME SLICE', raise_error=True)
if bbox is not None:
check_str(name=f'{self.__class__.__name__} bbox', value=bbox, options=Options.BBOX_OPTIONS.value, allow_blank=False, raise_error=True)
self.video_setting, self.frame_setting, self.rotate = video_setting, frame_setting, rotate
self.circle_size, self.font_size, self.animal_names, self.text_opacity = circle_size, font_size, animal_names, text_opacity
self.text_thickness, self.space_size, self.show_pose, self.pose_palette, self.verbose = text_thickness, space_size, show_pose, pose_palette, verbose
self.text_color, self.text_bg_color, self.bbox, self.show_gantt, self.show_confidence = text_clr, text_bg_clr, bbox, show_gantt, show_confidence
self.gpu = True if check_nvidea_gpu_available() and gpu else False
self.pose_threshold = read_config_entry(self.config, ConfigKey.THRESHOLD_SETTINGS.value, ConfigKey.SKLEARN_BP_PROB_THRESH.value, Dtypes.FLOAT.value, 0.00)
self.time_slice, self.overwrite, self.print_timer = time_slice, overwrite, print_timer
if save_dir is not None:
check_if_dir_exists(in_dir=save_dir, source=f'{self.__class__.__name__} save_dir', create_if_not_exist=True, raise_error=True)
self.save_dir = save_dir
else:
self.save_dir = self.sklearn_plot_dir
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
if isinstance(video_paths, str):
if os.path.isdir(video_paths):
self.video_paths = find_all_videos_in_project(videos_dir=video_paths, raise_error=False)
if len(self.video_paths) == 0:
raise NoDataError(msg=f'Cannot create classification videos. No videos exist in {video_paths} directory', source=self.__class__.__name__)
else:
self.video_paths = [video_paths]
elif isinstance(video_paths, list): self.video_paths = video_paths
elif video_paths is None:
self.video_paths = find_all_videos_in_project(videos_dir=self.video_dir)
if len(self.video_paths) == 0:
raise NoDataError(msg=f'Cannot create classification videos. No videos exist in {self.video_dir} directory', source=self.__class__.__name__)
else:
raise InvalidInputError(msg=f'video_paths has to be a path, directory, or a list of paths. Got {type(video_paths)}', source=self.__class__.__name__)
if data_dir is not None:
check_if_dir_exists(in_dir=data_dir, source=f'{self.__class__.__name__} data_dir', raise_error=True)
self.data_dir = data_dir
else:
self.data_dir = self.machine_results_dir
# Keep only videos that have corresponding classification data in data_dir; warn about (and skip) the rest, and error only if none have data.
videos_with_data, videos_missing_data = [], []
for video_path in self.video_paths:
video_name = get_fn_ext(filepath=video_path)[1]
data_path = os.path.join(self.data_dir, f'{video_name}.{self.file_type}')
if os.path.isfile(data_path):
videos_with_data.append(video_path)
else:
videos_missing_data.append(video_name)
if len(videos_with_data) == 0:
raise NoDataError(msg=f'Cannot create classification videos: none of the {len(self.video_paths)} video(s) have corresponding classification data in {self.data_dir} (expected files named <video_name>.{self.file_type}).', source=self.__class__.__name__)
if len(videos_missing_data) > 0:
NoDataFoundWarning(msg=f'Skipping {len(videos_missing_data)} video(s) with no classification data in {self.data_dir}: {videos_missing_data}', source=self.__class__.__name__)
self.video_paths = videos_with_data
check_int(name=f'{self.__class__.__name__} core_cnt', value=core_cnt, min_value=-1, unaccepted_vals=[0])
self.core_cnt = find_core_cnt()[0] if int(core_cnt) == -1 or int(core_cnt) > find_core_cnt()[0] else int(core_cnt)
if clf_names is not None:
if isinstance(clf_names, tuple): clf_names = list(clf_names)
check_valid_lst(data=clf_names, source=f'{self.__class__.__name__} clf_names', valid_dtypes=(str,), valid_values=self.clf_names, min_len=1)
self.clf_names = clf_names
self.conf_cols = [f'Probability_{x}' for x in self.clf_names]
if platform.system() == "Darwin":
multiprocessing.set_start_method("spawn", force=True)
def __get_print_settings(self):
optimal_circle_size = self.get_optimal_circle_size(frame_size=(self.video_meta_data["width"], self.video_meta_data["height"]), circle_frame_ratio=100)
longest_str = str(max(['TIMERS:', 'ENSEMBLE PREDICTION:'] + self.clf_names, key=len))
auto_thickness = max(TextOptions.TEXT_THICKNESS.value, int(round(min(self.video_meta_data["width"], self.video_meta_data["height"]) / 500)))
self.video_text_thickness = auto_thickness if self.text_thickness is None else int(max(self.text_thickness, 1))
optimal_font_size, _, optimal_spacing_scale = self.get_optimal_font_scales(text=longest_str, accepted_px_width=int(self.video_meta_data["width"] / 2.5), accepted_px_height=int(self.video_meta_data["height"] / 5), text_thickness=self.video_text_thickness)
self.video_circle_size = optimal_circle_size if self.circle_size is None else int(max(1, self.circle_size))
self.video_font_size = optimal_font_size if self.font_size is None else self.font_size
if self.font_path is not None:
self.video_font_size_px, _, _ = self.get_optimal_font_size_ttf(text=longest_str, font_path=self.font_path, accepted_px_width=int(self.video_meta_data["width"] / 2.5), accepted_px_height=int(self.video_meta_data["height"] / 5))
default_space_size = self.get_optimal_font_spacing_ttf(font_path=self.font_path, size_px=self.video_font_size_px, text=['TIMERS:', 'ENSEMBLE PREDICTION:'] + list(self.clf_names) + [f'{c} CONFIDENCE: 0.0000' for c in self.clf_names], gap=0)
else:
self.video_font_size_px, default_space_size = None, optimal_spacing_scale
self.video_space_size = default_space_size if self.space_size is None else int(max(self.space_size, 1))
self.video_text_opacity = 0.8 if self.text_opacity is None else float(self.text_opacity)
def run(self):
if self.verbose: stdout_information(msg=f'Creating {len(self.video_paths)} classification visualization(s) using {self.core_cnt} cores...')
self.pool = get_cpu_pool(core_cnt=self.core_cnt, source=self.__class__.__name__, )
for video_cnt, video_path in enumerate(self.video_paths):
video_timer = SimbaTimer(start=True)
_, self.video_name, _ = get_fn_ext(video_path)
if self.verbose: stdout_information(msg=f"Creating classification visualization for video {self.video_name}...")
self.data_path = os.path.join(self.data_dir, f'{self.video_name}.{self.file_type}')
self.data_df = read_df(self.data_path, self.file_type).reset_index(drop=True).fillna(0)
if self.show_pose: check_that_column_exist(df=self.data_df, column_name=self.bp_col_names, file_name=self.data_path)
if self.show_confidence: check_that_column_exist(df=self.data_df, column_name=self.conf_cols, file_name=self.data_path)
self.video_meta_data = get_video_meta_data(video_path=video_path)
if self.time_slice is not None:
check_if_string_value_is_valid_video_timestamp(value=self.time_slice[START_TIME], name=f'START TIME {START_TIME}')
check_if_string_value_is_valid_video_timestamp(value=self.time_slice[END_TIME], name=f'END TIME {END_TIME}')
frm_ids = find_frame_numbers_from_time_stamp(start_time=self.time_slice[START_TIME], end_time=self.time_slice[END_TIME], fps=int(self.video_meta_data['fps']))
self.data_df = self.data_df.loc[frm_ids]
height, width = deepcopy(self.video_meta_data["height"]), deepcopy(self.video_meta_data["width"])
self.save_path = os.path.join(self.save_dir, f"{self.video_name}.mp4")
self.video_frame_dir, self.video_temp_dir = None, None
if self.video_setting:
if os.path.isfile(self.save_path) and not self.overwrite:
FileExistWarning(msg=f'Skipping video {self.video_name}: The output video exist ({self.save_path}) and overwrite is FALSE', source=self.__class__.__name__)
continue
self.video_save_path = os.path.join(self.save_dir, f"{self.video_name}.mp4")
self.video_temp_dir = os.path.join(self.save_dir, self.video_name, "temp")
create_directory(paths=self.video_temp_dir, overwrite=True)
if self.frame_setting:
self.video_frame_dir = os.path.join(self.save_dir, self.video_name)
create_directory(paths=self.video_temp_dir, overwrite=True)
if self.rotate:
self.video_meta_data["height"], self.video_meta_data["width"] = (width, height)
if self.time_slice is None: check_video_and_data_frm_count_align(video=video_path, data=self.data_df, name=self.video_name, raise_error=False)
check_that_column_exist(df=self.data_df, column_name=self.clf_names, file_name=self.data_path)
self.__get_print_settings()
if self.show_gantt is not None:
self.gantt_clrs = create_color_palette(pallete_name=self.pose_palette, increments=len(self.clf_names) + 1, as_int=True, as_rgb_ratio=True)
self.bouts_df = detect_bouts(data_df=self.data_df, target_lst=list(self.clf_names), fps=int(self.video_meta_data["fps"]))
self.final_gantt_img = PlottingMixin().make_gantt_plot(x_length=len(self.data_df) + 1,
bouts_df=self.bouts_df,
clf_names=self.clf_names,
fps=self.video_meta_data["fps"],
width=self.video_meta_data["width"],
height=self.video_meta_data["height"],
font_size=12,
font=self.font,
font_rotation=45,
x_tick_lbl_rotation=45,
hhmmss=True,
video_name=self.video_meta_data["video_name"],
save_path=None,
palette=self.gantt_clrs)
self.final_gantt_img = self.resize_gantt(self.final_gantt_img, self.video_meta_data["height"])
else:
self.bouts_df, self.final_gantt_img, self.gantt_clrs = None, None, None
self.clf_cumsums, self.clf_p = {}, {} if self.show_confidence else None
for clf_name in self.clf_names:
self.clf_cumsums[clf_name] = self.data_df[clf_name].cumsum()
if self.show_confidence: self.clf_p[clf_name] = np.round(self.data_df[f'Probability_{clf_name}'].values.reshape(-1), 4)
self.data_df["index"] = self.data_df.index
data = np.array_split(self.data_df, self.core_cnt)
data = [(cnt, x) for (cnt, x) in enumerate(data)]
constants = functools.partial(_multiprocess_sklearn_video,
bp_dict=self.animal_bp_dict,
video_save_dir=self.video_temp_dir,
frame_save_dir=self.video_frame_dir,
clf_cumsum=self.clf_cumsums,
rotate=self.rotate,
video_path=video_path,
clf_confidence=self.clf_p,
timer=self.print_timer,
video_setting=self.video_setting,
frame_setting=self.frame_setting,
pose_threshold=self.pose_threshold,
show_pose=self.show_pose,
show_animal_names=self.animal_names,
circle_size=self.video_circle_size,
font_size=self.video_font_size,
font_path=self.font_path,
font_size_px=self.video_font_size_px,
space_size=self.video_space_size,
text_thickness=self.video_text_thickness,
text_opacity=self.video_text_opacity,
text_bg_clr=self.text_bg_color,
text_color=self.text_color,
pose_clr_lst=self.clr_lst,
bbox=self.bbox,
show_gantt=self.show_gantt,
bouts_df=self.bouts_df,
final_gantt=self.final_gantt_img,
gantt_clrs=self.gantt_clrs,
clf_names=self.clf_names,
verbose=self.verbose)
for cnt, result in enumerate(self.pool.imap(constants, data, chunksize=self.multiprocess_chunksize)):
if self.verbose: stdout_information(f"Image batch {result} complete, Video {(video_cnt + 1)}/{len(self.video_paths)}...")
if self.video_setting:
if self.verbose: stdout_information(msg=f"Joining {self.video_name} multiprocessed video...")
concatenate_videos_in_folder(in_folder=self.video_temp_dir, save_path=self.video_save_path, gpu=self.gpu, verbose=self.verbose)
video_timer.stop_timer()
stdout_information(msg=f"Video {self.video_name} complete (elapsed time: {video_timer.elapsed_time_str}s)...")
terminate_cpu_pool(pool=self.pool, force=False, source=self.__class__.__name__)
self.timer.stop_timer()
if self.video_setting:
stdout_success(msg=f"{len(self.video_paths)} video(s) saved in {self.save_dir} directory", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__)
if self.frame_setting:
stdout_success(f"Frames for {len(self.video_paths)} videos saved in sub-folders within {self.save_dir} directory", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__)
if __name__ == "__main__":
clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"H:\projects\jason_zhang\jason_project\project_folder\project_config.ini",
video_paths=r"H:\projects\jason_zhang\jason_project\project_folder\videos",
data_dir=r"H:\projects\jason_zhang\jason_project\project_folder\csv\REARING\200_0.30_smoothing_500ms\csvs",
save_dir=r'H:\projects\jason_zhang\jason_project\project_folder\csv\REARING\200_0.30_smoothing_500ms\videos',
clf_names=('REARING',),
video_setting=True,
frame_setting=False,
rotate=False,
show_confidence=True,
core_cnt=8,
show_pose=True,
animal_names=False,
font='poppins regular',
print_timer=HHMMSSSSSS,
overwrite=False,
time_slice=None, # {START_TIME: '00:00:00', END_TIME: '00:00:10'}, #{START_TIME: '00:00:00', END_TIME: '00:01:00'}, #,{START_TIME: '00:00:00', END_TIME: '00:01:00'}
bbox=None, #'animal-aligned'
text_opacity=0.6,
show_gantt=2)
clf_plotter.run()
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"H:\projects\brainwide_trap\brainwide_trap\project_folder\project_config.ini",
# video_paths=r"H:\projects\brainwide_trap\brainwide_trap\project_folder\videos",
# data_dir=r"H:\projects\brainwide_trap\brainwide_trap\project_folder\csv\machine_results",
# save_dir=r'H:\projects\brainwide_trap\brainwide_trap\project_folder\frames\output\freezing_immbility',
# clf_names=('FREEZING', 'IMMOBILITY'),
# video_setting=True,
# frame_setting=False,
# rotate=False,
# show_confidence=False,
# core_cnt=8,
# show_pose=True,
# animal_names=False,
# print_timer=HHMMSSSSSS,
# overwrite=True,
# time_slice=None, # {START_TIME: '00:00:00', END_TIME: '00:00:10'}, #{START_TIME: '00:00:00', END_TIME: '00:01:00'}, #,{START_TIME: '00:00:00', END_TIME: '00:01:00'}
# bbox=None, #'animal-aligned'
# text_opacity=0.6,
# show_gantt=None)
# clf_plotter.run()
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"E:\troubleshooting\mitra_emergence_hour\project_folder\project_config.ini",
# video_paths=r"E:\troubleshooting\mitra_emergence_hour\project_folder\videos\Box3_180mISOcontrol_Females.mp4", #,#,
# video_setting=True,
# frame_setting=False,
# rotate=False,
# show_confidence=False,
# core_cnt=12,
# show_pose=False,
# animal_names=False,
# print_timer=HHMMSSSSSS,
# overwrite=False,
# time_slice=None, #,{START_TIME: '00:00:00', END_TIME: '00:01:00'}
# bbox=None,
# text_opacity=0.6,
# show_gantt=None)
# clf_plotter.run()
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"E:\troubleshooting\mitra_emergence_hour\project_folder\project_config.ini",
# video_paths=r"E:\troubleshooting\mitra_emergence_hour\project_folder\videos\Box3_180mISOcontrol_Females.mp4", #,#,
# video_setting=True,
# frame_setting=False,
# rotate=False,
# show_confidence=False,
# core_cnt=12,
# show_pose=False,
# animal_names=False,
# print_timer=HHMMSSSSSS,
# overwrite=False,
# time_slice=None, #,{START_TIME: '00:00:00', END_TIME: '00:01:00'}
# bbox=None,
# text_opacity=0.6,
# show_gantt=None)
# clf_plotter.run()
#
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"F:\troubleshooting\sam\sam\project_folder\project_config.ini",
# video_paths=None,#r"F:\troubleshooting\sam\sam\project_folder\videos\20251212_150730_380_cam1_CROPPED.mp4",
# video_setting=True,
# frame_setting=False,
# rotate=False,
# show_confidence=False,
# core_cnt=8,
# animal_names=False,
# print_timer=HHMMSSSSSS,
# overwrite=False,
# time_slice=None,#{START_TIME: '00:00:00', END_TIME: '00:01:00'},
# bbox=None,
# text_opacity=0.6,
# show_gantt=2)
# clf_plotter.run()
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"F:\troubleshooting\sophiaa\project_folder\project_config.ini",
# video_setting=True,
# frame_setting=False,
# #video_paths=r"F:\troubleshooting\sophiaa\project_folder\videos\Choice222.mp4",
# print_timers=True,
# rotate=False,
# animal_names=True,
# core_cnt=6,
# circle_size=4,
# time_slice={START_TIME: '00:00:00', END_TIME: '00:01:00'},
# bbox=None,#'animal-aligned',
# show_gantt=None)
# clf_plotter.run()
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini",
# video_setting=True,
# frame_setting=False,
# video_paths=r"/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.mp4",
# print_timers=True,
# rotate=False,
# animal_names=False,
# show_bbox=True,
# show_gantt=None)
# clf_plotter.run()
#text_settings = {'circle_scale': 5, 'font_size': 0.528, 'spacing_scale': 28, 'text_thickness': 2}
# clf_plotter = PlotSklearnResultsMultiProcess(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/beepboop174/project_folder/project_config.ini',
# video_setting=True,
# frame_setting=False,
# rotate=False,
# video_file_path='592_MA147_Gq_CNO_0515.mp4',
# cores=-1,
# text_settings=False)
# clf_plotter.run()
#
# clf_plotter = PlotSklearnResultsMultiProcess(config_path='/Users/simon/Desktop/envs/troubleshooting/DLC_2_Black_animals/project_folder/project_config.ini', video_setting=True, frame_setting=False, rotate=False, video_file_path='Together_1.avi', cores=5)
# clf_plotter.run()
# if __name__ == "__main__":
# clf_plotter = PlotSklearnResultsMultiProcess(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",
# video_setting = True,
# frame_setting = False,
# rotate = False,
# core_cnt = 6,
# show_confidence=True,
# video_paths=r"C:\troubleshooting\mitra\project_folder\videos\501_MA142_Gi_CNO_0521.mp4")
# clf_plotter.run()