__author__ = "Simon Nilsson; sronilsson@gmail.com"
import functools
import multiprocessing
import os
import platform
from typing import Any, Dict, Optional, Tuple, Union
import cv2
import numpy as np
import pandas as pd
from simba.data_processors.directing_other_animals_calculator import \
DirectingOtherAnimalsAnalyzer
from simba.mixins.config_reader import ConfigReader
from simba.mixins.plotting_mixin import PlottingMixin
from simba.utils.checks import (check_file_exist_and_readable, check_float,
check_if_keys_exist_in_dict,
check_if_string_value_is_valid_video_timestamp,
check_if_valid_rgb_tuple, check_int,
check_that_hhmmss_start_is_before_end,
check_valid_lst,
check_video_and_data_frm_count_align)
from simba.utils.data import (create_color_palettes,
find_frame_numbers_from_time_stamp, get_cpu_pool,
terminate_cpu_pool)
from simba.utils.enums import OS, Formats, Keys, TextOptions
from simba.utils.errors import (AnimalNumberError, InvalidInputError,
NoFilesFoundError)
from simba.utils.printing import stdout_information, stdout_success
from simba.utils.read_write import (
check_if_hhmmss_timestamp_is_valid_part_of_video,
concatenate_videos_in_folder, find_core_cnt, get_fn_ext,
get_video_meta_data, read_df, seconds_to_timestamp)
from simba.utils.warnings import NoDataFoundWarning
DIRECTION_THICKNESS = "direction_thickness"
DIRECTIONALITY_COLOR = "directionality_color"
CIRCLE_SIZE = "circle_size"
HIGHLIGHT_ENDPOINTS = "highlight_endpoints"
SHOW_POSE = "show_pose"
ANIMAL_NAMES = "animal_names"
START_TIME = 'start_time'
END_TIME = 'end_time'
STYLE_ATTR = [DIRECTION_THICKNESS, DIRECTIONALITY_COLOR, CIRCLE_SIZE, HIGHLIGHT_ENDPOINTS, SHOW_POSE, ANIMAL_NAMES]
FOURCC = cv2.VideoWriter_fourcc(*Formats.AVI_CODEC.value)
X_BPS, Y_BPS = Keys.X_BPS.value, Keys.Y_BPS.value
def _directing_animals_mp(frm_range: Tuple[int, np.ndarray],
directionality_data: pd.DataFrame,
pose_data: pd.DataFrame,
style_attr: dict,
animal_bp_dict: dict,
save_temp_dir: str,
line_opacity: float,
video_path: str,
video_meta_data: dict,
colors: list):
batch = frm_range[0]
start_frm, current_frm, end_frm = frm_range[1][0], frm_range[1][0], frm_range[1][-1]
save_path = os.path.join(save_temp_dir, f"{batch}.avi")
writer = cv2.VideoWriter(save_path, FOURCC, video_meta_data["fps"], (video_meta_data["width"], video_meta_data["height"]))
cap = cv2.VideoCapture(video_path)
cap.set(1, start_frm)
while current_frm <= end_frm:
ret, img = cap.read()
if ret:
frm_data = pose_data.iloc[current_frm]
if style_attr[SHOW_POSE]:
for animal_cnt, (animal_name, animal_bps) in enumerate(animal_bp_dict.items()):
for bp_cnt, bp in enumerate(zip(animal_bps[X_BPS], animal_bps[Y_BPS])):
x_bp, y_bp = frm_data[bp[0]], frm_data[bp[1]]
cv2.circle(img, (int(x_bp), int(y_bp)), style_attr[CIRCLE_SIZE], animal_bp_dict[animal_name]["colors"][bp_cnt], -1)
if style_attr[ANIMAL_NAMES]:
for animal_name, bp_data in animal_bp_dict.items():
headers = [bp_data[X_BPS][-1], bp_data[Y_BPS][-1]]
bp_cords = pose_data.loc[current_frm, headers].values.astype(np.int64)
cv2.putText(img, animal_name, (bp_cords[0], bp_cords[1]), TextOptions.FONT.value, 2, animal_bp_dict[animal_name]["colors"][0], 1)
if current_frm in list(directionality_data["Frame_#"].unique()):
img_data = directionality_data[directionality_data["Frame_#"] == current_frm]
for animal_name in img_data["Animal_1"].unique():
animal_img_data = img_data[img_data["Animal_1"] == animal_name].reset_index(drop=True)
img = PlottingMixin.draw_lines_on_img(img=img, start_positions=animal_img_data[["Eye_x", "Eye_y"]].values.astype(np.int64), end_positions=animal_img_data[["Animal_2_bodypart_x", "Animal_2_bodypart_y"]].values.astype(np.int64), color=tuple(colors[animal_name]), highlight_endpoint=style_attr[HIGHLIGHT_ENDPOINTS], thickness=style_attr[DIRECTION_THICKNESS], circle_size=style_attr[CIRCLE_SIZE], opacity=line_opacity)
current_frm += 1
writer.write(img.astype(np.uint8))
stdout_information(msg=f"Created directing frame: {current_frm}, frame timestamp: {seconds_to_timestamp(current_frm / video_meta_data['fps'])} (core batch: {batch}) ...")
else:
break
writer.release()
return batch
[docs]class DirectingOtherAnimalsVisualizerMultiprocess(ConfigReader, PlottingMixin):
"""
Create videos visualizing when animals direct their gaze toward body parts of other animals using multiprocessing.
Draws directional lines from eye positions (calculated from nose and ear coordinates) to target body parts.
Uses parallel processing across CPU cores for faster video creation.
.. important::
Requires pose-estimation data for left ear, right ear, and nose of each animal. Project must contain at least 2 animals.
.. note::
`Tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/Scenario2.md#visualizing-data-tables>`__.
.. image:: _static/img/directing_other_animals.png
:alt: Directing other animals
:width: 450
:align: center
.. youtube:: tsOJCOYZRAA
:width: 640
:height: 480
:align: center
.. video:: _static/img/DirectingOtherAnimalsVisualizerMultiprocess.mp4
:width: 900
:loop:
:muted:
:align: center
.. seealso::
For single core function, see :func:`simba.plotting.directing_animals_visualizer.DirectingOtherAnimalsVisualizer`.
:param Union[str, os.PathLike] config_path: Path to SimBA project config file.
:param Union[str, os.PathLike] video_path: Path to video file. Corresponding pose data must exist in outlier_corrected_movement_location directory.
:param Dict[str, Any] style_attr: Style attributes with keys: 'show_pose', 'animal_names', 'circle_size', 'directionality_color', 'direction_thickness', 'highlight_endpoints'. See example.
:param Optional[int] core_cnt: Number of CPU cores. -1 = all available. Default -1.
:param Optional[Dict[str, str]] time_slice: If set, restrict to time period. Dict with keys 'start_time' and 'end_time' (HH:MM:SS). Default None.
:param Optional[str] left_ear_name: Left ear body-part name. If None, auto-detected. Must provide all three body parts or none.
:param Optional[str] right_ear_name: Right ear body-part name. If None, auto-detected.
:param Optional[str] nose_name: Nose body-part name. If None, auto-detected.
:param float line_opacity: Opacity of direction lines (0.0–1.0). Default 1.0.
:raises AnimalNumberError: If project contains fewer than 2 animals.
:raises NoFilesFoundError: If pose-estimation data file not found.
:raises InvalidInputError: If body part names partially provided.
:example:
>>> style_attr = {'show_pose': True, 'animal_names': False, 'circle_size': 3, 'directionality_color': [(255, 0, 0), (0, 0, 255)], 'direction_thickness': None, 'highlight_endpoints': True}
>>> visualizer = DirectingOtherAnimalsVisualizerMultiprocess(config_path='project_config.ini', video_path='video.avi', style_attr=style_attr, core_cnt=-1)
>>> visualizer.run()
"""
def __init__(self,
config_path: Union[str, os.PathLike],
video_path: Union[str, os.PathLike],
style_attr: Dict[str, Any],
core_cnt: int = -1,
time_slice: Optional[Dict[str, str]] = None,
left_ear_name: Optional[str] = None,
line_opacity: float = 1.0,
right_ear_name: Optional[str] = None,
nose_name: Optional[str] = None):
if platform.system() == OS.MAC.value:
multiprocessing.set_start_method("spawn", force=True)
check_file_exist_and_readable(file_path=video_path)
check_file_exist_and_readable(file_path=config_path)
check_if_keys_exist_in_dict(data=style_attr, key=STYLE_ATTR, name=f"{self.__class__.__name__} style_attr")
check_int(name=f"{self.__class__.__name__} core_cnt", value=core_cnt, min_value=-1, max_value=find_core_cnt()[0], unaccepted_vals=[0])
check_float(name=f"{self.__class__.__name__} line_opacity", value=line_opacity, min_value=0.0, max_value=1.0, raise_error=True)
if core_cnt == -1: core_cnt = find_core_cnt()[0]
if time_slice is not None:
check_if_keys_exist_in_dict(data=time_slice, key=[START_TIME, END_TIME], name=f'{self.__class__.__name__} slicing')
check_if_string_value_is_valid_video_timestamp(value=time_slice[START_TIME], name="Video slicing START TIME")
check_if_string_value_is_valid_video_timestamp(value=time_slice[END_TIME], name="Video slicing END TIME")
check_that_hhmmss_start_is_before_end(start_time=time_slice[START_TIME], end_time=time_slice[END_TIME], name="SLICE TIME STAMPS")
ConfigReader.__init__(self, config_path=config_path)
PlottingMixin.__init__(self)
if self.animal_cnt < 2:
raise AnimalNumberError("Cannot analyze directionality between animals in a project with less than two animals.", source=self.__class__.__name__)
passed_bps = [left_ear_name, right_ear_name, nose_name]
if sum(p is None for p in passed_bps) not in (0, len(passed_bps)):
raise InvalidInputError(msg="left_ear_name, right_ear_name, and nose_name must either all be None or all be provided as strings", source=self.__class__.__name__)
self.animal_names = [k for k in self.animal_bp_dict.keys()]
_, self.video_name, _ = get_fn_ext(video_path)
self.data_path = os.path.join(self.outlier_corrected_dir, f"{self.video_name}.{self.file_type}")
if not os.path.isfile(self.data_path):
raise NoFilesFoundError(msg=f"SIMBA ERROR: Could not find the file at path {self.data_path}. Make sure the data file exist to create directionality visualizations", source=self.__class__.__name__)
self.direction_analyzer = DirectingOtherAnimalsAnalyzer(config_path=config_path, bool_tables=False, summary_tables=False, aggregate_statistics_tables=False, data_paths=[self.data_path], left_ear_name=left_ear_name, right_ear_name=right_ear_name, nose_name=nose_name, verbose=False)
self.direction_analyzer.run()
self.direction_analyzer.transpose_results()
self.style_attr, self.direction_colors = style_attr, {}
if isinstance(self.style_attr[DIRECTIONALITY_COLOR], list):
check_valid_lst(data=self.style_attr[DIRECTIONALITY_COLOR], source=f"{self.__class__.__name__} colors", valid_dtypes=(tuple,), min_len=self.animal_cnt)
for i in range(len(self.animal_names)):
check_if_valid_rgb_tuple(data=self.style_attr[DIRECTIONALITY_COLOR][i])
self.direction_colors[self.animal_names[i]] = self.style_attr[DIRECTIONALITY_COLOR][i]
if isinstance(self.style_attr[DIRECTIONALITY_COLOR], tuple):
check_if_valid_rgb_tuple(self.style_attr[DIRECTIONALITY_COLOR])
for i in range(len(self.animal_names)):
self.direction_colors[self.animal_names[i]] = self.style_attr[DIRECTIONALITY_COLOR]
else:
self.random_colors = create_color_palettes(1, int(self.animal_cnt**2))[0]
self.random_colors = [[int(item) for item in sublist] for sublist in self.random_colors]
for cnt in range(len(self.animal_names)):
self.direction_colors[self.animal_names[cnt]] = self.random_colors[cnt]
self.data_dict = self.direction_analyzer.directionality_df_dict
if not os.path.exists(self.directing_animals_video_output_path):
os.makedirs(self.directing_animals_video_output_path)
self.data_df = read_df(self.data_path, file_type=self.file_type)
self.video_save_path = os.path.join(self.directing_animals_video_output_path, f"{self.video_name}.mp4")
self.cap = cv2.VideoCapture(video_path)
self.video_meta_data = get_video_meta_data(video_path)
if style_attr[CIRCLE_SIZE] is None:
style_attr[CIRCLE_SIZE] = PlottingMixin().get_optimal_circle_size(frame_size=(self.video_meta_data['width'], self.video_meta_data['height']), circle_frame_ratio=100)
if style_attr[DIRECTION_THICKNESS] is None:
style_attr[DIRECTION_THICKNESS] = PlottingMixin().get_optimal_circle_size(frame_size=(self.video_meta_data['width'], self.video_meta_data['height']), circle_frame_ratio=80)
check_video_and_data_frm_count_align(video=video_path, data=self.data_path, name=video_path, raise_error=False)
if not os.path.exists(self.directing_animals_video_output_path):
os.makedirs(self.directing_animals_video_output_path)
self.save_path = os.path.join(self.directing_animals_video_output_path, f"{self.video_name}.mp4")
self.save_temp_path = os.path.join(self.directing_animals_video_output_path, "temp")
if os.path.exists(self.save_temp_path):
self.remove_a_folder(folder_dir=self.save_temp_path)
os.makedirs(self.save_temp_path)
self.core_cnt, self.video_path, self.time_slice, self.video_path, self.line_opacity = core_cnt, video_path, time_slice, video_path, line_opacity
stdout_information(msg=f"Processing video {self.video_name}...")
def run(self):
video_data = self.data_dict[self.video_name]
if len(video_data) < 1:
NoDataFoundWarning(msg=f"SimBA skipping video {self.video_name}: No animals are directing each other in the video.")
else:
pool = get_cpu_pool(core_cnt=self.core_cnt, maxtasksperchild=self.maxtasksperchild, source=self.__class__.__name__)
if self.time_slice is None:
frm_data = np.array_split(list(range(0, self.video_meta_data["frame_count"] + 1)), self.core_cnt)
else:
check_if_hhmmss_timestamp_is_valid_part_of_video(timestamp=self.time_slice[START_TIME], video_path=self.video_path)
check_if_hhmmss_timestamp_is_valid_part_of_video(timestamp=self.time_slice[END_TIME], video_path=self.video_path)
frm_numbers = find_frame_numbers_from_time_stamp(start_time=self.time_slice[START_TIME], end_time=self.time_slice[END_TIME], fps=self.video_meta_data["fps"])
frm_data = np.array_split(list(range(min(frm_numbers), max(frm_numbers) + 1)), self.core_cnt)
frm_data = [(i, j) for i, j in enumerate(frm_data)]
stdout_information(msg=f"Creating directing images, multiprocessing (chunksize: {self.multiprocess_chunksize}, cores: {self.core_cnt})...")
constants = functools.partial(_directing_animals_mp,
directionality_data=video_data,
pose_data=self.data_df,
video_meta_data=self.video_meta_data,
style_attr=self.style_attr,
save_temp_dir=self.save_temp_path,
video_path=self.video_path,
line_opacity=self.line_opacity,
animal_bp_dict=self.animal_bp_dict,
colors=self.direction_colors)
for cnt, result in enumerate(pool.imap(constants, frm_data, chunksize=self.multiprocess_chunksize)):
stdout_information(msg=f"Core batch {result+1} complete...")
stdout_information(msg=f"Joining {self.video_name} multi-processed video...")
concatenate_videos_in_folder(in_folder=self.save_temp_path, save_path=self.save_path, video_format="avi", fps=self.video_meta_data["fps"], remove_splits=True)
self.timer.stop_timer()
terminate_cpu_pool(pool=pool, force=False, source=self.__class__.__name__)
stdout_success(msg=f"Video {self.video_name} complete. Video saved in {self.directing_animals_video_output_path} directory", elapsed_time=self.timer.elapsed_time_str)
# if __name__ == "__main__":
# style_attr = {SHOW_POSE: True,
# ANIMAL_NAMES: False,
# CIRCLE_SIZE: 3,
# DIRECTIONALITY_COLOR: [(255, 0, 0), (0, 0, 255)],
# DIRECTION_THICKNESS: 10,
# HIGHLIGHT_ENDPOINTS: True}
# test = DirectingOtherAnimalsVisualizerMultiprocess(config_path=r"D:\troubleshooting\two_animals_sleap\project_folder\project_config.ini",
# video_path=r"D:\troubleshooting\two_animals_sleap\project_folder\videos\Test2025-05-14_11-15-49.mp4",
# style_attr=style_attr,
# core_cnt=5,
# time_slice={START_TIME: '00:01:00', END_TIME: '00:02:00'},
# left_ear_name='earL1',
# right_ear_name='earR1',
# nose_name='nose1')
# test.run()
# style_attr = {'Show_pose': True,
# 'Pose_circle_size': 3,
# "Direction_color": 'Random',
# 'Direction_thickness': 4,
# 'Highlight_endpoints': True,
# 'Polyfill': True}
# test = DirectingOtherAnimalsVisualizerMultiprocess(config_path='/Users/simon/Desktop/envs/troubleshooting/sleap_5_animals/project_folder/project_config.ini',
# data_path='/Users/simon/Desktop/envs/troubleshooting/sleap_5_animals/project_folder/csv/outlier_corrected_movement_location/Testing_Video_3.csv',
# style_attr=style_attr,
# core_cnt=5)
#
# test.run()
# style_attr = {'Show_pose': True, 'Pose_circle_size': 3, "Direction_color": 'Random', 'Direction_thickness': 4, 'Highlight_endpoints': True, 'Polyfill': True}
# test = DirectingOtherAnimalsVisualizerMultiprocess(config_path='/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
# data_path='/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/csv/outlier_corrected_movement_location/Together_1.csv',
# style_attr=style_attr,
# core_cnt=5)
#
# test.run()
# if __name__ == "__main__":
# style_attr = {SHOW_POSE: True,
# ANIMAL_NAMES: False,
# CIRCLE_SIZE: 3,
# DIRECTIONALITY_COLOR: [(255, 0, 0), (0, 0, 255)],
# DIRECTION_THICKNESS: 10,
# HIGHLIGHT_ENDPOINTS: True}
#
# test = DirectingOtherAnimalsVisualizerMultiprocess(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
# video_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/videos/Together_1.avi',
# style_attr=style_attr,
# core_cnt=-1)
#
# test.run()