From 7cae83134bdfde60ec26a6476b034bfe7c822e41 Mon Sep 17 00:00:00 2001 From: Simon Nilsson Date: Fri, 29 Sep 2023 14:11:45 -0400 Subject: [PATCH] Add files via upload --- simba/roi_tools/ROI_size_standardizer.py | 118 +++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 simba/roi_tools/ROI_size_standardizer.py diff --git a/simba/roi_tools/ROI_size_standardizer.py b/simba/roi_tools/ROI_size_standardizer.py new file mode 100644 index 000000000..8b776b564 --- /dev/null +++ b/simba/roi_tools/ROI_size_standardizer.py @@ -0,0 +1,118 @@ +from simba.mixins.config_reader import ConfigReader +from simba.mixins.feature_extraction_mixin import FeatureExtractionMixin +from simba.utils.errors import NoROIDataError, ParametersFileError +from shapely.geometry import Polygon +from scipy.spatial import ConvexHull +import pandas as pd +import numpy as np +import math +from shapely.affinity import scale + +from simba.utils.printing import stdout_success + +class ROISizeStandardizer(ConfigReader, + FeatureExtractionMixin): + + + def __init__(self, + config_path: str, + reference_video: str): + + ConfigReader.__init__(self, config_path=config_path) + FeatureExtractionMixin.__init__(self) + self.reference_video_name = reference_video + self.read_roi_data() + self.checks() + + def checks(self): + if self.reference_video_name not in self.video_names_w_rois: + raise NoROIDataError(msg=f'The reference video {self.reference_video_name} does not have any defined ROIs.', source=self.__class__.__name__) + for video_name in self.video_names_w_rois: + if video_name not in list(self.video_info_df['Video']): + raise ParametersFileError(msg=f'Found defined ROIs for video {video_name}, but this video does not have pixels per millimeter defined in the project_folder/logs/video_info.csv file.', source=self.__class__.__name__) + + def compute_correction_factors(self): + self.correction_factors = {} + for video_name in self.video_names_w_rois: + px_per_mm = self.read_video_info(video_name=video_name)[1] + self.correction_factors[video_name] = round((((px_per_mm / self.reference_px_per_mm) - 1) * 100), 4) + + def find_scale_factor(self, video_name: str, correction_factor: float): + if self.correction_factors[video_name] < 0: + return math.sqrt(1 - abs(correction_factor) / 100) + else: + return math.sqrt(1 + correction_factor / 100) + + def run(self): + _, self.reference_px_per_mm, _ = self.read_video_info(video_name=self.reference_video_name) + self.compute_correction_factors() + self.video_names_to_change = [x for x in self.video_names_w_rois if x != self.reference_video_name] + + updated_rectangles, updated_polygons, updated_circles = pd.DataFrame(columns=self.rectangles_df.columns), pd.DataFrame(columns=self.polygon_df.columns), pd.DataFrame(columns=self.circles_df.columns) + for video_name in self.video_names_to_change: + rectangles = self.rectangles_df[self.rectangles_df['Video'] == video_name] + circles = self.circles_df[self.circles_df['Video'] == video_name] + polygons = self.polygon_df[self.polygon_df['Video'] == video_name] + + #UPDATE RECTANGLES + for idx, rectangle in rectangles.iterrows(): + if self.correction_factors[video_name] != 0: + vertices = np.array([rectangle['Tags']['Top left tag'], rectangle['Tags']['Top right tag'],rectangle['Tags']['Bottom left tag'], rectangle['Tags']['Bottom right tag']]) + scale_factor = self.find_scale_factor(video_name=video_name, correction_factor=self.correction_factors[video_name]) + new_vertices = np.unique(np.array(scale(Polygon(vertices), xfact=scale_factor, yfact=scale_factor).exterior.coords).astype(np.int16), axis=0) + + top_left = np.min(new_vertices, axis=0) + bottom_right = np.max(new_vertices, axis=0) + + rectangle['topLeftX'], rectangle['topLeftY'] = top_left[0], top_left[1] + rectangle['Bottom_right_X'], rectangle['Bottom_right_Y'] = bottom_right[0], bottom_right[1] + rectangle['width'], rectangle['height'] = int(rectangle['Bottom_right_X'] - rectangle['topLeftX']), int(rectangle['Bottom_right_Y'] - rectangle['topLeftY']) + rectangle['Tags']['Top left tag'], rectangle['Tags']['Top right tag'] = (int(rectangle['topLeftX']), int(rectangle['topLeftY'])), (int(rectangle['Bottom_right_X']), int(rectangle['topLeftY'])) + rectangle['Tags']['Bottom left tag'], rectangle['Tags']['Bottom right tag'] = (rectangle['topLeftX'], rectangle['Bottom_right_Y']), (rectangle['Bottom_right_X'], rectangle['Bottom_right_Y']) + rectangle['Tags']['Left tag'] = (int(rectangle['topLeftX']), int(rectangle['topLeftY'] + rectangle['height'] / 2)) + rectangle['Tags']['Right tag'] = (int(rectangle['topLeftX'] + rectangle['width']), int(rectangle['topLeftY'] + rectangle['height'] / 2)) + rectangle['Tags']['Top tag'] = (int(rectangle['topLeftX'] + rectangle['width']/2), int(rectangle['topLeftY'])) + rectangle['Tags']['Bottom tag'] = (int(rectangle['topLeftX'] + rectangle['width'] / 2), int(rectangle['topLeftY'] + rectangle['height'])) + updated_rectangles.loc[idx] = rectangle + + # UPDATE POLYGONS + for idx, polygon in polygons.iterrows(): + if self.correction_factors[video_name] != 0: + scale_factor = self.find_scale_factor(video_name=video_name, correction_factor=self.correction_factors[video_name]) + new_vertices = np.array(scale(Polygon(polygon['vertices']), xfact=scale_factor, yfact=scale_factor).exterior.coords) + new_vertices = ConvexHull(new_vertices).points.astype(np.int32) + polygon['vertices'] = new_vertices + for cnt, i in enumerate(polygon['vertices']): + polygon['Tags'][f'Tag_{cnt}'] = i + updated_polygons.loc[idx] = polygon + + # UPDATE CIRCLES + for idx, circle in circles.iterrows(): + if self.correction_factors[video_name] != 0: + circle['radius'] = int(circle['radius'] * (self.correction_factors[video_name] / 100)) + circle['Tags']['Border tag'] = (int(circle['centerX'] - circle['radius']), circle['centerY']) + updated_circles.loc[idx] = circle + + self.rectangles_df.update(updated_rectangles) + self.polygon_df.update(updated_polygons) + self.circles_df.update(updated_circles) + + def save(self): + store = pd.HDFStore(self.roi_coordinates_path, mode='w') + store['rectangles'] = self.rectangles_df + store['circleDf'] = self.circles_df + store['polygons'] = self.polygon_df + store.close() + stdout_success(msg=f'ROI size definitions standardized according to pixels per millimeter in video {self.reference_video_name}', source=self.__class__.__name__) + +# test = ROISizeStandardizer(config_path='/Users/simon/Desktop/envs/troubleshooting/DLC_2_Black_animals/project_folder/project_config.ini', reference_video='Together_1') +# test.run() +# test.save() + + + + +# +# test = ROISizeStandardizer(config_path='/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/project_folder/project_config.ini', reference_video='Together_1') +# test.run() +# test.save() \ No newline at end of file