import argparse
import os.path
from copy import deepcopy
import numpy as np
import pandas as pd
from simba.mixins.abstract_classes import AbstractFeatureExtraction
from simba.mixins.config_reader import ConfigReader
from simba.mixins.feature_extraction_mixin import FeatureExtractionMixin
from simba.mixins.geometry_mixin import GeometryMixin
from simba.plotting.geometry_plotter import GeometryPlotter
from simba.utils.errors import NoFilesFoundError
from simba.utils.printing import SimbaTimer, stdout_success
from simba.utils.read_write import find_core_cnt, get_fn_ext, read_df, write_df
WHITE = 'white'
BLACK = 'black'
CORE_CNT = -1
VISUALIZE = True
BUFFER = 20
RECTANGLES = True
[docs]class RatSocialFeaturizer(ConfigReader,
FeatureExtractionMixin,
GeometryMixin,
AbstractFeatureExtraction):
"""
Featurization class for analyzing rat social behaviors.
:example:
>>> extractor = RatSocialFeaturizer(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/piotr/project_folder/project_config.ini')
>>> extractor.run()
References
----------
.. [1] Popik, Piotr, Ewelina Cyrano, Joanna Golebiowska, Natalia Malikowska-Racia, Agnieszka Potasiewicz, and Agnieszka Nikiforuk. “Deep Learning Algorithms Reveal Increased Social Activity in Rats at the Onset of the Dark Phase of the Light/Dark Cycle.” Edited by Dragan Hrncic. PLOS ONE 19, no. 11 (November 8, 2024): e0307794. https://doi.org/10.1371/journal.pone.0307794.
.. [2] Popik, Piotr, Ewelina Cyrano, Diana Piotrowska, Malgorzata Holuj, Joanna Golebiowska, Natalia Malikowska-Racia, Agnieszka Potasiewicz, and Agnieszka Nikiforuk. “Effects of Ketamine on Rat Social Behavior as Analyzed by DeepLabCut and SimBA Deep Learning Algorithms.” Frontiers in Pharmacology 14 (January 10, 2024): 1329424. https://doi.org/10.3389/fphar.2023.1329424.
"""
def __init__(self,
config_path: str):
ConfigReader.__init__(self, config_path=config_path)
FeatureExtractionMixin.__init__(self)
GeometryMixin.__init__(self)
if len(self.outlier_corrected_paths) == 0:
raise NoFilesFoundError(msg=f'No files found in {self.outlier_corrected_dir}')
self.session_timer = SimbaTimer(start=True)
self.config_path = config_path
[docs] def run(self):
print(f'Used core count: {CORE_CNT}')
for file_cnt, file_path in enumerate(self.outlier_corrected_paths):
video_timer = SimbaTimer(start=True)
_, video_name, _ = get_fn_ext(filepath=file_path)
print(f'Processing {video_name}...')
_, pixels_per_mm, _ = self.read_video_info(video_name=video_name)
data_df = read_df(file_path=file_path, file_type=self.file_type)
data_df = data_df.apply(pd.to_numeric, errors='coerce').fillna(0).replace(0, np.nan)
data_df = data_df.interpolate(method='nearest').bfill().ffill()
nan_cols = data_df.columns[data_df.isnull().all(0)] # FIND COLUMN NAMES THAT ARE ALL NaNs
#data_df = data_df.drop(nan_cols, axis=1) # DROP COLUMNS THAT ARE ALL NaNs FROM THE DATA
results = deepcopy(data_df)
save_path = os.path.join(self.features_dir, f'{video_name}.{self.file_type}')
white_animal_bp_names, black_animal_bp_names = self.animal_bp_dict[WHITE], self.animal_bp_dict[BLACK]
white_animal_cols, black_animal_cols = [], []
for x, y in zip(white_animal_bp_names['X_bps'], white_animal_bp_names['Y_bps']): white_animal_cols.extend((x, y))
for x, y in zip(black_animal_bp_names['X_bps'], black_animal_bp_names['Y_bps']): black_animal_cols.extend((x, y))
black_animal_cols = [x for x in black_animal_cols if x not in nan_cols] # DROP COLUMN NAMES THAT ARE ALL NaNs FROM THE BLACK ANIMAL BODY-PART NAMES
white_animal_cols = [x for x in white_animal_cols if x not in nan_cols] # DROP COLUMN NAMES THAT ARE ALL NaNs FROM THE WHITE ANIMAL BODY-PART NAMES
white_animal_df, black_animal_df = data_df[white_animal_cols], data_df[black_animal_cols]
white_animal_df_arr = white_animal_df.values.reshape(len(white_animal_df), -1 , 2)
black_animal_df_arr = black_animal_df.values.reshape(len(black_animal_df), -1, 2)
white_animal_polygons = GeometryMixin().multiframe_bodyparts_to_polygon(data=white_animal_df_arr, pixels_per_mm=pixels_per_mm, parallel_offset=BUFFER, verbose=True, video_name=video_name, animal_name='white', core_cnt=CORE_CNT)
black_animal_polygons = GeometryMixin().multiframe_bodyparts_to_polygon(data=black_animal_df_arr, pixels_per_mm=pixels_per_mm, parallel_offset=BUFFER, verbose=True, video_name=video_name, animal_name='black', core_cnt=CORE_CNT)
if RECTANGLES:
white_animal_polygons = GeometryMixin().multiframe_minimum_rotated_rectangle(shapes=white_animal_polygons, video_name=video_name, animal_name='white', verbose=True, core_cnt=CORE_CNT)
black_animal_polygons = GeometryMixin().multiframe_minimum_rotated_rectangle(shapes=black_animal_polygons, video_name=video_name, animal_name='black', verbose=True, core_cnt=CORE_CNT)
results['polygon_pct_overlap'] = GeometryMixin().multiframe_compute_pct_shape_overlap(shape_1=white_animal_polygons, shape_2=black_animal_polygons, animal_names='black_white', video_name=video_name, verbose=True, core_cnt=CORE_CNT)
combined_list = [list(pair) for pair in list(zip(white_animal_polygons, black_animal_polygons))]
difference = GeometryMixin().multiframe_difference(shapes=combined_list, verbose=True, animal_names='white_black', video_name=video_name, core_cnt=CORE_CNT)
results['difference_area'] = GeometryMixin().multiframe_area(shapes=difference, pixels_per_mm=pixels_per_mm, verbose=True, video_name=video_name, core_cnt=CORE_CNT)
self.save(data=results, save_path=save_path)
if VISUALIZE:
geometry_plotter = GeometryPlotter(config_path=self.config_path, geometries=[white_animal_polygons, black_animal_polygons], video_name=video_name, core_cnt=CORE_CNT)
geometry_plotter.run()
video_timer.stop_timer()
stdout_success(msg=f'{video_name} complete!', elapsed_time=video_timer.elapsed_time_str)
self.session_timer.stop_timer()
stdout_success(msg=f'{len(self.outlier_corrected_paths)} data files saved in {self.features_dir}')
[docs] def save(self,
data: pd.DataFrame,
save_path: str):
write_df(df=data, file_type=self.file_type, save_path=save_path)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='SimBA Custom Feature Extractor')
parser.add_argument('--config_path', type=str, help='SimBA project config path')
args = parser.parse_args()
feature_extractor = RatSocialFeaturizer(config_path=args.config_path)
feature_extractor.run()
#
# feature_extractor = PiotrFeatureExtractor(config_path='/Users/simon/Desktop/envs/simba/troubleshooting/piotr/project_folder/project_config.ini')
# feature_extractor.run()