import os
from copy import copy
from typing import Optional, Union
import numpy as np
import pandas as pd
from simba.data_processors.interpolate import Interpolate
from simba.data_processors.smoothing import Smoothing
from simba.mixins.config_reader import ConfigReader
from simba.utils.checks import (check_if_dir_exists,
check_if_keys_exist_in_dict, check_int,
check_str, check_valid_boolean,
check_valid_dataframe)
from simba.utils.enums import ConfigKey, Dtypes, Formats, Methods
from simba.utils.errors import InvalidInputError, NoFilesFoundError
from simba.utils.printing import SimbaTimer, stdout_success
from simba.utils.read_write import (find_files_of_filetypes_in_directory,
get_fn_ext, read_df, write_df)
REQUIRED_FIELDS = ['nose_x', 'nose_y', 'left_x', 'left_y', 'center_x', 'center_y', 'right_x', 'right_y', 'tail_x', 'tail_y']
BP_NAMES = ['nose', 'left', 'center', 'right', 'tail']
[docs]class SimBABlobImporter(ConfigReader):
"""
Import blob tracking data generated by SimBA's blob tracking tools into a SimBA project.
This class imports CSV files containing blob tracking data (with columns for nose, tail, center, left, and right body parts)
and converts them into SimBA's standard pose-estimation format. The imported data can optionally be interpolated and smoothed.
.. note::
The project must be configured as a SimBA blob project (pose setting must be set to 'simba_blob').
The input CSV files must contain the following required fields: nose_x, nose_y, left_x, left_y, center_x, center_y, right_x, right_y, tail_x, tail_y.
.. seealso::
To generate blob tracking data, see :func:`simba.video_processors.blob_tracking_executor.BlobTrackingExecutor`.
To visualize blob tracking results, see :func:`simba.plotting.blob_visualizer.BlobVisualizer`.
:param Union[str, os.PathLike] config_path: Path to the SimBA project configuration file.
:param Union[str, os.PathLike] data_path: Path to a directory containing blob tracking CSV files, or a single CSV file.
:param Optional[Union[str, os.PathLike]] save_dir: Directory where imported files will be saved. If None, saves to the project's outlier_corrected directory. Default: None.
:param Optional[dict] smoothing_settings: Dictionary with smoothing settings. Must contain 'method' ('savitzky-golay' or 'gaussian') and 'time_window' (positive integer). If None, no smoothing is applied. Default: None.
:param Optional[dict] interpolation_settings: Dictionary with interpolation settings. Must contain 'method' ('linear', 'quadratic', or 'nearest') and 'type' ('body-parts' or 'animals'). If None, no interpolation is applied. Default: None.
:param Optional[bool] verbose: If True, prints progress messages. Default: True.
:example:
>>> r = SimBABlobImporter(config_path=r"C:/troubleshooting/simba_blob_project/project_folder/project_config.ini", data_path=r'C:/troubleshooting/simba_blob_project/data')
>>> r.run()
>>> r = SimBABlobImporter(config_path=r"C:/troubleshooting/simba_blob_project/project_folder/project_config.ini",
... data_path=r'C:/troubleshooting/simba_blob_project/data',
... smoothing_settings={'method': 'savitzky-golay', 'time_window': 100},
... interpolation_settings={'method': 'nearest', 'type': 'body-parts'})
>>> r.run()
"""
def __init__(self,
config_path: Union[str, os.PathLike],
data_path: Union[str, os.PathLike],
save_dir: Optional[Union[str, os.PathLike]] = None,
smoothing_settings: Optional[dict] = None,
interpolation_settings: Optional[dict] = None,
verbose: Optional[bool] = True):
ConfigReader.__init__(self, config_path=config_path, read_video_info=False, create_logger=False)
pose_config_name = self.read_config_entry(config=self.config, section=ConfigKey.CREATE_ENSEMBLE_SETTINGS.value,
option=ConfigKey.POSE_SETTING.value, default_value=None,
data_type=Dtypes.STR.value).strip()
if pose_config_name != Methods.SIMBA_BLOB.value:
raise InvalidInputError(
msg=f'The project {config_path} is not a SimBA blob project. Cannot import SimBA blob data to a non SimBA blob project ({ConfigKey.POSE_SETTING.value}: {pose_config_name}, expected: {Methods.SIMBA_BLOB.value})',
source=self.__class__.__name__)
if os.path.isdir(data_path):
self.data_paths = find_files_of_filetypes_in_directory(directory=data_path, extensions=['.csv'],
raise_error=True)
elif os.path.isfile(data_path):
self.data_paths = [data_path]
else:
raise NoFilesFoundError(msg=f'{data_path} is not a valid file path or valid directory path',
source=self.__class__.__name__)
if interpolation_settings is not None:
check_if_keys_exist_in_dict(data=interpolation_settings, key=['method', 'type'], name=f'{self.__class__.__name__} interpolation_settings')
check_str(name=f'{self.__class__.__name__} interpolation_settings type', value=interpolation_settings['type'], options=('body-parts', 'animals'))
check_str(name=f'{self.__class__.__name__} interpolation_settings method', value=interpolation_settings['method'], options=('linear', 'quadratic', 'nearest'))
self.interpolation_type, self.interpolation_method = interpolation_settings['type'], interpolation_settings['method']
else:
self.interpolation_type, self.interpolation_method = None, None
if smoothing_settings is not None:
check_if_keys_exist_in_dict(data=smoothing_settings, key=['method', 'time_window'], name=f'{self.__class__.__name__} smoothing_settings')
check_str(name=f'{self.__class__.__name__} smoothing_settings method', value=smoothing_settings['method'], options=('savitzky-golay', 'gaussian'))
check_int(name=f'{self.__class__.__name__} smoothing_settings time_window', value=smoothing_settings['time_window'], min_value=1)
self.smoothing_time, self.smoothing_method = smoothing_settings['time_window'], smoothing_settings['method']
else:
self.smoothing_time, self.smoothing_method = None, None
check_valid_boolean(value=verbose, source=f'{self.__class__.__name__} verbose', raise_error=True)
if save_dir is not None:
check_if_dir_exists(in_dir=save_dir, source=f'{self.__class__.__name__} save_dir', raise_error=True)
else:
save_dir = copy(self.outlier_corrected_dir)
self.interpolation_settings, self.smoothing_settings, = (interpolation_settings, smoothing_settings)
self.save_dir, self.verbose = save_dir, verbose
def run(self):
for file_cnt, file_path in enumerate(self.data_paths):
file_timer = SimbaTimer(start=True)
df = read_df(file_path=file_path, file_type='csv')
df.columns = [x.strip().lower() for x in df.columns]
video_name = get_fn_ext(filepath=file_path)[1]
if self.verbose:
print(f'Importing SimBA blob data for video {video_name}...')
save_path = os.path.join(self.save_dir, f'{video_name}.csv')
check_valid_dataframe(df=df, source=f'{self.__class__.__name__} {file_path}',
valid_dtypes=Formats.NUMERIC_DTYPES.value, required_fields=REQUIRED_FIELDS)
df = df[REQUIRED_FIELDS].astype(np.int32)
df_out = pd.DataFrame()
for i in range(0, df.shape[1], 2):
df_out = pd.concat(
[df_out, df.iloc[:, i:i + 2], pd.DataFrame(1, index=df.index, columns=[f'{BP_NAMES[i // 2]}_p'])],
axis=1)
del df
write_df(df=df_out, file_type=self.file_type, save_path=save_path, multi_idx_header=False)
if self.interpolation_settings is not None:
interpolator = Interpolate(config_path=self.config_path, data_path=save_path, type=self.interpolation_type, method=self.interpolation_method, multi_index_df_headers=False, copy_originals=False)
interpolator.run()
if self.smoothing_settings is not None:
smoother = Smoothing(config_path=self.config_path, data_path=save_path, time_window=self.smoothing_time, method=self.smoothing_method, multi_index_df_headers=False, copy_originals=False)
smoother.run()
file_timer.stop_timer()
print(f'Imported data for video {video_name} (elapsed time: {file_timer.elapsed_time}s)')
self.timer.stop_timer()
stdout_success(
msg=f"{len(self.data_paths)} SimBA blob tracking files file(s) imported to the SimBA project {self.save_dir}",
source=self.__class__.__name__, elapsed_time=self.timer.elapsed_time)
#
# r = SimBABlobImporter(config_path=r"C:\troubleshooting\simba_blob_project\project_folder\project_config.ini",
# data_path=r'C:\troubleshooting\simba_blob_project\data',
# smoothing_settings={'method': 'savitzky-golay', 'time_window': 100},
# interpolation_settings='Body-parts: Nearest')
# r.run()