__author__ = "Simon Nilsson; sronilsson@gmail.com"
import os
from copy import deepcopy
from typing import Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
from PIL import Image
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_valid_rgb_tuple,
check_int, check_str, check_that_column_exist,
check_valid_boolean,
check_video_and_data_frm_count_align)
from simba.utils.data import create_color_palette, detect_bouts
from simba.utils.enums import (ConfigKey, Dtypes, Formats, Options, TagNames,
TextOptions)
from simba.utils.errors import (InvalidInputError, NoDataError,
NoSpecifiedOutputError)
from simba.utils.printing import log_event, stdout_success
from simba.utils.read_write import (find_all_videos_in_project, get_fn_ext,
get_video_meta_data, read_config_entry,
read_df)
from simba.utils.warnings import FrameRangeWarning
CIRCLE_SCALE = 'circle_scale'
FONT_SIZE = 'font_size'
SPACE_SCALE = 'spacing_scale'
TEXT_THICKNESS = 'text_thickness'
TEXT_SETTING_KEYS = ['circle_scale', 'font_size', 'spacing_scale', 'text_thickness']
CENTER_BP_TXT = ['centroid', 'center']
FOURCC = cv2.VideoWriter_fourcc(*Formats.MP4_CODEC.value)
[docs]class PlotSklearnResultsSingleCore(ConfigReader, TrainModelMixin, PlottingMixin):
"""
Plot classification results overlays on videos. Results are stored in the
`project_folder/frames/output/sklearn_results` directory of the SimBA project.
.. note::
For improved run-time, see :meth:`simba.plotting.plot_clf_results_mp.PlotSklearnResultsMultiProcess` for multiprocess class.
Scikit visualization documentation <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-10-sklearn-visualization__.
.. video:: _static/img/sklearn_visualization.webm
:width: 600
:autoplay:
:loop:
:muted:
:align: center
:param Union[str, os.PathLike] config_path: path to SimBA project config file in Configparser format
:param Optional[bool] video_setting: If True, SimBA will create compressed videos. Default True.
:param Optional[bool] frame_setting: If True, SimBA will create individual frames. Default True.
:param Optional[str] video_file_path: Path to video file to create classification visualizations for. If None, then all the videos in the csv/machine_results will be used. Default None.
:param Optional[Union[Dict[str, float], bool]] text_settings: Dictionary holding the circle size, font size, spacing size, and text thickness of the printed text. If None, then these are autocomputed.
:param Optional[bool] rotate: If True, the output video will be rotated 90 degrees from the input. Default False.
:param Optional[str] palette: The name of the palette used for the pose-estimation key-points. Default ``Set1``.
:param Optional[bool] print_timers: If True, the output video will have the cumulative time of the classified behaviours overlaid. Default True.
:param Optional[bool] show_bbox: If True, axis-aligned bounding boxes created from each anmals pose and displayed. Default True.
:example:
>>> text_settings = {'circle_scale': 5, 'font_size': 5, 'spacing_scale': 2, 'text_thickness': 10}
>>> test = PlotSklearnResultsSingleCore(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
>>> video_setting=True,
>>> frame_setting=False,
>>> video_file_path='Together_1.avi',
>>> print_timers=True,
>>> text_settings=text_settings,
>>> rotate=False)
>>> test.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_bbox: bool = False,
show_confidence: bool = False,
show_gantt: Optional[int] = None,
font_size: Optional[Union[int, float]] = None,
space_size: Optional[Union[int, float]] = None,
text_opacity: Optional[Union[int, float]] = None,
text_thickness: Optional[Union[int, float]] = None,
circle_size: Optional[Union[int, float]] = None,
pose_palette: Optional[str] = 'Set1',
print_timers: bool = True,
text_clr: Tuple[int, int,int] = (255, 255, 255),
text_bg_clr: Tuple[int, int,int] = (0, 0, 0)):
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, print_timers, animal_names, show_pose, show_bbox, 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)
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)
check_if_valid_rgb_tuple(data=text_bg_clr, source=f'{self.__class__.__name__} text_bg_clr')
check_if_valid_rgb_tuple(data=text_clr, source=f'{self.__class__.__name__} text_clr')
self.text_color, self.text_bg_color = text_clr, text_bg_clr
self.video_paths, self.print_timers, self.show_confidence = video_paths, print_timers, show_confidence
if self.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__)
if show_gantt is not None:
check_int(name=f"{self.__class__.__name__} show_gantt", value=show_gantt, max_value=2, min_value=1)
self.video_setting, self.frame_setting, self.rotate, self.text_opacity = video_setting, frame_setting, rotate, text_opacity
self.circle_size, self.font_size, self.animal_names, self.show_gantt, self.pose_palette = circle_size, font_size, animal_names, show_gantt, pose_palette
self.text_thickness, self.space_size, self.show_pose, self.show_bbox = text_thickness, space_size, show_pose, show_bbox
self.pose_threshold = read_config_entry(self.config, ConfigKey.THRESHOLD_SETTINGS.value, ConfigKey.SKLEARN_BP_PROB_THRESH.value, Dtypes.FLOAT.value, 0.00)
if not os.path.exists(self.sklearn_plot_dir):
os.makedirs(self.sklearn_plot_dir)
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)
self.conf_cols = [f'Probability_{x}' for x in self.clf_names]
self.clr_lst = create_color_palette(pallete_name=pose_palette, increments=len(self.body_parts_lst)+1)
if isinstance(self.video_paths, str): self.video_paths = [video_paths]
elif isinstance(self.video_paths, list): self.video_paths = video_paths
else:
raise InvalidInputError(msg=f'video_paths has to be a path of a list of paths. Got {type(video_paths)}', source=self.__class__.__name__)
for video_path in self.video_paths:
video_name = get_fn_ext(filepath=video_path)[1]
data_path = os.path.join(self.machine_results_dir, f'{video_name}.{self.file_type}')
if not os.path.isfile(data_path): raise NoDataError(msg=f'Cannot create classification videos for {video_name}. Expected classification data at location {data_path} but file does not exist', source=self.__class__.__name__)
def __get_print_settings(self):
self.video_text_thickness = TextOptions.TEXT_THICKNESS.value if self.text_thickness is None else int(max(self.text_thickness, 1))
longest_str = str(max(['TIMERS:', 'ENSEMBLE PREDICTION:'] + self.clf_names, key=len))
optimal_font_size, _, optimal_spacing_scale = self.get_optimal_font_scales(text=longest_str, accepted_px_width=int(self.video_meta_data["width"] / 3), accepted_px_height=int(self.video_meta_data["height"] / 10), text_thickness=self.video_text_thickness)
optimal_circle_size = self.get_optimal_circle_size(frame_size=(self.video_meta_data["width"], self.video_meta_data["height"]), circle_frame_ratio=100)
self.video_circle_size = optimal_circle_size if self.circle_size is None else int(self.circle_size)
self.video_font_size = optimal_font_size if self.font_size is None else self.font_size
self.video_space_size = optimal_spacing_scale 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):
for video_cnt, video_path in enumerate(self.video_paths):
_, self.video_name, _ = get_fn_ext(video_path)
self.data_path = os.path.join(self.machine_results_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)
height, width = deepcopy(self.video_meta_data["height"]), deepcopy(self.video_meta_data["width"])
self.save_path = os.path.join(self.sklearn_plot_dir, f"{self.video_name}.mp4")
if self.frame_setting:
self.video_frame_dir = os.path.join(self.sklearn_plot_dir, self.video_name)
if not os.path.exists(self.video_frame_dir): os.makedirs(self.video_frame_dir)
if self.rotate:
self.video_meta_data["height"], self.video_meta_data["width"] = (width, height)
self.clf_timers = {k: 0 for k in self.clf_names}
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=list(self.clf_timers.keys()), 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_rotation=90, 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
if self.show_gantt is None:
self.writer = cv2.VideoWriter(self.save_path, FOURCC, self.video_meta_data["fps"], (self.video_meta_data["width"], self.video_meta_data["height"]))
else:
self.writer = cv2.VideoWriter(self.save_path, FOURCC, self.video_meta_data["fps"], (int(self.video_meta_data["width"] + self.final_gantt_img.shape[1]), self.video_meta_data["height"]))
self.cap = cv2.VideoCapture(video_path)
frm_idx = 0
while self.cap.isOpened():
ret, self.frame = self.cap.read()
if ret:
clr_cnt = 0
for animal_cnt, (animal_name, animal_data) in enumerate(self.animal_bp_dict.items()):
if self.show_pose:
for bp_num in range(len(animal_data["X_bps"])):
x_bp, y_bp, p_bp = (animal_data["X_bps"][bp_num], animal_data["Y_bps"][bp_num], animal_data["P_bps"][bp_num])
bp_cords = self.data_df.loc[frm_idx, [x_bp, y_bp, p_bp]]
if bp_cords[p_bp] >= self.pose_threshold:
self.frame = cv2.circle(self.frame, (int(bp_cords[x_bp]), int(bp_cords[y_bp])), self.video_circle_size, self.clr_lst[clr_cnt], -1)
clr_cnt+=1
if self.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 = self.data_df.loc[frm_idx, [x_bp, y_bp, p_bp]]
cv2.putText(self.frame, animal_name, (int(bp_cords[x_bp]), int(bp_cords[y_bp])), self.font, self.video_font_size, self.clr_lst[0], self.video_text_thickness)
if self.show_bbox:
animal_headers = [val for pair in zip(animal_data["X_bps"], animal_data["Y_bps"]) for val in pair]
animal_cords = self.data_df.loc[frm_idx, animal_headers].values.reshape(-1, 2).astype(np.int32)
try:
bbox = GeometryMixin().keypoints_to_axis_aligned_bounding_box(keypoints=animal_cords.reshape(-1, len(animal_cords), 2).astype(np.int32))
cv2.polylines(self.frame, [bbox], True, self.clr_lst[animal_cnt], thickness=self.video_text_thickness, lineType=cv2.LINE_AA)
except:
pass
if self.rotate:
self.frame = np.array(Image.fromarray(self.frame).rotate(90, Image.BICUBIC, expand=True))
if self.show_gantt == 1:
self.frame = np.concatenate((self.frame, self.final_gantt_img), axis=1)
elif self.show_gantt == 2:
bout_rows = self.bouts_df.loc[self.bouts_df["End_frame"] <= frm_idx]
gantt_plot = PlottingMixin().make_gantt_plot(x_length=frm_idx + 1,
bouts_df=bout_rows,
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_rotation=90,
video_name=self.video_meta_data['video_name'],
save_path=None,
palette=self.gantt_clrs)
self.frame = np.concatenate((self.frame, gantt_plot), axis=1)
if self.print_timers:
self.frame = PlottingMixin().put_text(img=self.frame, text="TIMERS:", pos=(TextOptions.BORDER_BUFFER_Y.value, ((self.video_meta_data["height"] - self.video_meta_data["height"]) + self.video_space_size)), 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 = 2
for clf_name, clf_time in self.clf_timers.items():
frame_results = self.data_df.loc[frm_idx, clf_name]
self.clf_timers[clf_name] += (frame_results / self.video_meta_data['fps'])
if self.print_timers:
self.frame = PlottingMixin().put_text(img=self.frame, text=f"{clf_name} {round(self.clf_timers[clf_name], 2)}", 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
if self.show_confidence:
col_name = f'Probability_{clf_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)
self.add_spacer += 1
for clf_name, clf_time in self.clf_timers.items():
if self.data_df.loc[frm_idx, clf_name] == 1:
self.frame = PlottingMixin().put_text(img=self.frame, text=clf_name, 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_color=TextOptions.COLOR.value, text_bg_alpha=self.video_text_opacity)
self.add_spacer += 1
if self.video_setting:
self.writer.write(self.frame.astype(np.uint8))
if self.frame_setting:
frame_save_name = os.path.join(self.video_frame_dir, f"{frm_idx}.png")
cv2.imwrite(frame_save_name, self.frame)
frm_idx += 1
print(f'Frame: {frm_idx} / {self.video_meta_data["frame_count"]}. Video: {self.video_name} ({video_cnt + 1}/{len(self.video_paths)})')
else:
FrameRangeWarning(msg=f'Could not read frame {frm_idx} in video {video_path}. Stoping video creation.')
break
print(f"Video {self.video_name} saved at {self.save_path}...")
self.cap.release()
self.writer.release()
self.timer.stop_timer()
stdout_success(msg=f"{len(self.video_paths)} visualization(s) created in {self.sklearn_plot_dir} directory", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__)
# test = PlotSklearnResultsSingleCore(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=True,
# show_gantt=2)
# test.run()
# test = PlotSklearnResultsSingleCore(config_path=r"C:\troubleshooting\RAT_NOR\project_folder\project_config.ini",
# video_setting=True,
# frame_setting=False,
# video_paths=r"C:\troubleshooting\RAT_NOR\project_folder\videos\03152021_NOB_IOT_8.mp4",
# print_timers=True,
# rotate=True,
# animal_names=True)
# test.run()