__author__ = "Simon Nilsson; sronilsson@gmail.com"
import os
import shutil
from typing import Any, Dict, List, Optional, Union
import cv2
import numpy as np
import pandas as pd
from numba import jit, prange
from simba.mixins.config_reader import ConfigReader
from simba.mixins.geometry_mixin import GeometryMixin
from simba.mixins.plotting_mixin import PlottingMixin
from simba.utils.checks import (
check_all_file_names_are_represented_in_video_log,
check_file_exist_and_readable, check_float, check_if_keys_exist_in_dict,
check_int, check_that_column_exist, check_valid_lst)
from simba.utils.enums import Formats, TagNames
from simba.utils.errors import NoSpecifiedOutputError
from simba.utils.printing import SimbaTimer, log_event, stdout_success
from simba.utils.read_write import get_fn_ext, read_df
STYLE_PALETTE = 'palette'
STYLE_SHADING = 'shading'
STYLE_BIN_SIZE = 'bin_size'
STYLE_MAX_SCALE = 'max_scale'
STYLE_ATTR = [STYLE_PALETTE, STYLE_SHADING, STYLE_BIN_SIZE, STYLE_MAX_SCALE]
[docs]class HeatmapperLocationSingleCore(ConfigReader, PlottingMixin):
"""
Create heatmaps representing the location where animals spend time.
For improved run-time, see :meth:`simba.heat_mapper_location_mp.HeatMapperLocationMultiprocess` for multiprocess class.
.. note::
`GitHub visualizations tutorial <https://github.com/sgoldenlab/simba/blob/master/docs/tutorial.md#step-11-visualizations>`__.
For improved run-time of videos, see :meth:`simba.heat_mapper_location_mp.HeatMapperLocationMultiprocess` for multiprocess class.
.. video:: _static/img/heatmap_location.webm
:width: 1000
:autoplay:
:loop:
:muted:
:align: center
:param str config_path: path to SimBA project config file in Configparser format
:param str bodypart: The name of the body-part used to infer the location of the animal.
:param int bin_size: The rectangular size of each heatmap location in millimeters. For example, `50` will divide the video frames into 5 centimeter rectangular spatial bins.
:param str palette: Heatmap pallette. Eg. 'jet', 'magma', 'inferno','plasma', 'viridis', 'gnuplot2'
:param int or 'auto' max_scale: The max value in the heatmap in seconds. E.g., with a value of `10`, if the classified behavior has occurred >= 10s within a rectangular bins, it will be filled with the same color.
:param bool final_img_setting: If True, create a single image representing the last frame of the input video
:param bool video_setting: If True, then create a video of heatmaps.
:param bool frame_setting: If True, then create individual heatmap frames.
:example:
>>> style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 100, 'max_scale': 'auto'}
>>> heatmapper = HeatmapperLocationSingleCore(config_path='/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini', style_attr = style_attr, final_img_setting=True, video_setting=True, frame_setting=False, bodypart='Nose_1', files_found=['/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/csv/machine_results/Together_1.csv'])
>>> heatmapper.run()
"""
def __init__(self,
config_path: Union[str, os.PathLike],
data_paths: List[Union[str, os.PathLike]],
bodypart: str,
style_attr: Dict[str, Any],
final_img_setting: Optional[bool] = True,
video_setting: Optional[bool] = False,
frame_setting: Optional[bool] = False):
log_event(logger_name=str(__class__.__name__), log_type=TagNames.CLASS_INIT.value, msg=self.create_log_msg_from_init_args(locals=locals()))
if (not frame_setting) and (not video_setting) and (not final_img_setting):
raise NoSpecifiedOutputError(msg="Please choose to select either heatmap videos, frames, and/or final image.", source=self.__class__.__name__)
check_file_exist_and_readable(file_path=config_path)
check_valid_lst(data=data_paths, valid_dtypes=(str,), min_len=1)
check_if_keys_exist_in_dict(data=style_attr, key=STYLE_ATTR, name=f'{self.__class__.__name__} style_attr')
ConfigReader.__init__(self, config_path=config_path)
PlottingMixin.__init__(self)
self.frame_setting, self.video_setting = frame_setting, video_setting
self.final_img_setting, self.bp = final_img_setting, bodypart
self.style_attr = style_attr
if not os.path.exists(self.heatmap_location_dir):
os.makedirs(self.heatmap_location_dir)
self.data_paths = data_paths
self.fourcc = cv2.VideoWriter_fourcc(*Formats.MP4_CODEC.value)
self.bp_lst = [f"{self.bp}_x", f"{self.bp}_y"]
print(f"Processing heatmaps for {len(self.data_paths)} video(s)...")
def run(self):
check_all_file_names_are_represented_in_video_log(video_info_df=self.video_info_df, data_paths=self.data_paths)
for file_cnt, file_path in enumerate(self.data_paths):
video_timer = SimbaTimer(start=True)
_, self.video_name, _ = get_fn_ext(file_path)
self.video_info, self.px_per_mm, self.fps = self.read_video_info(video_name=self.video_name)
self.width, self.height = int(self.video_info["Resolution_width"].values[0]), int(self.video_info["Resolution_height"].values[0])
if self.video_setting:
self.video_save_path = os.path.join(self.heatmap_location_dir, f"{self.video_name}.mp4")
self.writer = cv2.VideoWriter(self.video_save_path, self.fourcc, self.fps, (self.width, self.height),)
if self.frame_setting | self.final_img_setting:
self.save_video_folder = os.path.join(self.heatmap_location_dir, self.video_name)
if os.path.exists(self.save_video_folder):
shutil.rmtree(self.save_video_folder)
os.makedirs(self.save_video_folder)
self.data_df = read_df(file_path, self.file_type)
check_that_column_exist(df=self.data_df, column_name=self.bp_lst, file_name=file_path)
self.data_df = self.data_df[self.bp_lst]
squares, aspect_ratio = GeometryMixin().bucket_img_into_grid_square(bucket_grid_size_mm=self.style_attr[STYLE_BIN_SIZE], img_size=(self.width, self.height), px_per_mm=self.px_per_mm)
location_array = GeometryMixin().cumsum_coord_geometries(data=self.data_df.values, fps=self.fps, geometries=squares)
if self.style_attr[STYLE_MAX_SCALE] == "auto":
self.max_scale = np.round(np.max(np.max(location_array[-1], axis=0)), 3)
if self.max_scale == 0: self.max_scale = 1
else:
check_float(name=f'{self.__class__.__name__} style max scale', value=self.style_attr[STYLE_MAX_SCALE], min_value=10e-6)
self.max_scale = self.style_attr[STYLE_MAX_SCALE]
if self.final_img_setting:
final_img_save_name = os.path.join(self.heatmap_location_dir, f"{self.video_name}_final_frm.png")
self.make_location_heatmap_plot(frm_data=location_array[-1:, :, :][0],
max_scale=self.max_scale,
palette=self.style_attr[STYLE_PALETTE],
aspect_ratio=aspect_ratio,
file_name=final_img_save_name,
shading=self.style_attr[STYLE_SHADING],
img_size=(self.width, self.height))
print(f"Final heatmap image saved at {final_img_save_name}")
if (self.frame_setting) or (self.video_setting):
for frm_cnt, cumulative_frm in enumerate(range(location_array.shape[0])):
img = self.make_location_heatmap_plot(frm_data=location_array[cumulative_frm, :, :],
max_scale=self.max_scale,
palette=self.style_attr["palette"],
aspect_ratio=aspect_ratio,
file_name=None,
shading=self.style_attr["shading"],
img_size=(self.width, self.height))
if self.video_setting:
self.writer.write(img[:, :, :3])
if self.frame_setting:
frame_save_path = os.path.join(self.save_video_folder, f"{frm_cnt}.png")
cv2.imwrite(frame_save_path, img)
print(f"Heatmap frame: {frm_cnt + 1} / {len(self.data_df)}. Video: {self.video_name} ({file_cnt + 1}/{len(self.data_paths)})")
if self.video_setting:
self.writer.release()
video_timer.stop_timer()
print(f"Heatmap plot for video {self.video_name} saved (elapsed time: {video_timer.elapsed_time_str}s")
self.timer.stop_timer()
stdout_success(msg=f"Created heatmaps for {len(self.data_paths)} videos", elapsed_time=self.timer.elapsed_time_str, source=self.__class__.__name__,)
# test = HeatmapperLocationSingleCore(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/platea/project_folder/project_config.ini',
# style_attr = {'palette': 'jet', 'shading': 'flat', 'bin_size': 30, 'max_scale': 2},
# final_img_setting=True,
# video_setting=True,
# frame_setting=False,
# bodypart='CENTER',
# data_paths=['/Users/simon/Desktop/envs/simba/troubleshooting/platea/project_folder/csv/outlier_corrected_movement_location/Video_1.csv'])
# test.run()
# test = HeatmapperLocationSingleCore(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/open_field_below/project_folder/project_config.ini',
# style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 100, 'max_scale': 10},
# final_img_setting=True,
# video_setting=True,
# frame_setting=False,
# bodypart='Snout',
# data_paths=['/Users/simon/Desktop/envs/simba/troubleshooting/open_field_below/project_folder/csv/outlier_corrected_movement_location/raw_clip1.csv'])
# test.run()
#
# test = HeatmapperLocationSingleCore(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini',
# style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 100, 'max_scale': 'auto'},
# final_img_setting=True,
# video_setting=True,
# frame_setting=False,
# bodypart='Nose_1',
# files_found=['/Users/simon/Desktop/envs/simba/troubleshooting/two_black_animals_14bp/project_folder/csv/outlier_corrected_movement_location/Together_1.csv'])
# test.run()
# test = HeatmapperLocationSingleCore(config_path='/Users/simon/Desktop/envs/troubleshooting/locomotion/project_folder/project_config.ini',
# style_attr = {'palette': 'jet', 'shading': 'gouraud', 'bin_size': 50, 'max_scale': 'auto'},
# final_img_setting=False,
# video_setting=True,
# frame_setting=False,
# bodypart='Nose',
# files_found=['/Users/simon/Desktop/envs/troubleshooting/locomotion/project_folder/csv/outlier_corrected_movement_location/PD1406_2022-05-24_RVDG_GCaMP8s-Gi_Video_Day_22_Baseline.csv'])
# test.create_heatmaps()
# test = HeatmapperLocation(config_path='/Users/simon/Desktop/troubleshooting/train_model_project/project_folder/project_config.ini',
# bodypart='Nose_1',
# bin_size=50,
# palette='jet',
# max_scale='auto',
# final_img_setting=True,
# video_setting=False,
# frame_setting=False)
# test.create_heatmaps()