Skip to content

Commit

Permalink
egocentric aligment
Browse files Browse the repository at this point in the history
  • Loading branch information
@sronilsson
sronilsson committed Oct 31, 2024
1 parent ff34807 commit 1226b32
Show file tree
Hide file tree
Showing 15 changed files with 509 additions and 239 deletions.
Binary file added docs/_static/img/plot_clf_cumcount.webp
View file
Binary file not shown.
2 changes: 1 addition & 1 deletion requirements-gpu.txt → requirements_gpu.txt
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
--extra-index-url https://pypi.nvidia.com
#--extra-index-url https://pypi.nvidia.com

cupy-cuda12x==13.3.0
shap==0.46.1.dev78
Expand Down
15 changes: 11 additions & 4 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,11 @@

import setuptools
import platform
from setuptools import setup, find_namespace_packages

REQUIREMENTS_PATH = 'requirements.txt'
ARM_REQUIREMENTS_PATH = 'requirements_arm.txt'
GPU_REQUIREMENTS_PATH = 'requirements_gpu.txt'

with open("docs/project_description.md", "r", encoding="utf-8") as fh:
long_description = fh.read()
Expand All @@ -21,18 +23,21 @@
with open(ARM_REQUIREMENTS_PATH, "r") as f:
arm_requirements = f.read().splitlines()

with open(GPU_REQUIREMENTS_PATH, "r") as f:
gpu_requirements = f.read().splitlines()

# Setup configuration
setuptools.setup(
name="Simba-UW-tf-dev",
version="2.2.4",
version="2.2.7",
author="Simon Nilsson, Jia Jie Choong, Sophia Hwang",
author_email="[email protected]",
description="Toolkit for computer classification and analysis of behaviors in experimental animals",
long_description=long_description,
long_description_content_type="text/markdown",
url="https://github.com/sgoldenlab/simba",
install_requires=requirements,
extras_require={'arm': [arm_requirements]},
extras_require={'arm': [arm_requirements], "gpu": [gpu_requirements]},
license='GNU General Public License v3 (GPLv3)',
packages=setuptools.find_packages(exclude=["*.tests",
"*.tests.*",
Expand All @@ -43,11 +48,13 @@
"sandbox",
"sandbox.*",
]),
# packages=find_namespace_packages(include=["simba*"],
# exclude=["*.tests", "*.tests.*", "tests.*", "tests", "__pycache__", "__pycache__.*", "pose_configurations_archive", "pose_configurations_archive.*", "sandbox", "sandbox.*"]),
include_package_data=True,
classifiers=(
classifiers=[
"Programming Language :: Python :: 3",
"License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
"Operating System :: OS Independent",
),
],
entry_points={'console_scripts':['simba=simba.SimBA:main'],}
)
5 changes: 3 additions & 2 deletions simba/SimBA.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import urllib.request
import webbrowser
from tkinter.filedialog import askdirectory
from tkinter.filedialog import askdirectory
from tkinter.messagebox import askyesno

import PIL.Image
Expand Down Expand Up @@ -416,8 +417,8 @@ def activate(box, *args):
button_runvalidmodel = SimbaButton(parent=label_model_validation, txt="RUN MODEL", txt_clr='blue', img='rocket', cmd=self.validate_model_first_step, thread=True)
button_generateplot = SimbaButton(parent=label_model_validation, txt="INTERACTIVE PROBABILITY PLOT", img='interactive_blue', txt_clr='blue', cmd=self.launch_interactive_plot, thread=False)

self.dis_threshold = Entry_Box(label_model_validation, "DISCRIMINATION THRESHOLD (0.0-1.0):", "30")
self.min_behaviorbout = Entry_Box(label_model_validation,"MINIMUM BOUT LENGTH (MS):","30",validation="numeric")
self.dis_threshold = Entry_Box(label_model_validation, "DISCRIMINATION THRESHOLD (0.0-1.0):", "35")
self.min_behaviorbout = Entry_Box(label_model_validation,"MINIMUM BOUT LENGTH (MS):","35",validation="numeric")
button_validate_model = SimbaButton(parent=label_model_validation, txt="CREATE VALIDATION VIDEO", txt_clr='blue', img='visualize_blue', cmd=ValidationVideoPopUp, cmd_kwargs={'config_path': lambda:config_path, 'simba_main_frm': lambda:self})

label_runmachinemodel = CreateLabelFrameWithIcon(parent=tab9,header="RUN MACHINE MODEL",icon_name=Keys.DOCUMENTATION.value,icon_link=Links.SCENARIO_2.value)
Expand Down
180 changes: 94 additions & 86 deletions simba/data_processors/egocentric_aligner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,7 @@
import pandas as pd

from simba.mixins.config_reader import ConfigReader
from simba.utils.checks import (
check_all_file_names_are_represented_in_video_log, check_int,
check_valid_dataframe)
from simba.utils.checks import (check_all_file_names_are_represented_in_video_log, check_int, check_valid_dataframe, check_if_dir_exists)
from simba.utils.enums import Formats
from simba.utils.printing import SimbaTimer
from simba.utils.read_write import (concatenate_videos_in_folder,
Expand All @@ -26,9 +24,11 @@
def _egocentric_aligner(frm_range: np.ndarray,
video_path: Union[str, os.PathLike],
temp_dir: Union[str, os.PathLike],
video_name: str,
centers: List[Tuple[int, int]],
rotation_vectors: np.ndarray,
target: Tuple[int, int]):
target: Tuple[int, int],
color: Tuple[int, int, int] = (255, 255, 255)):

video_meta = get_video_meta_data(video_path=video_path)
cap = cv2.VideoCapture(video_path)
Expand All @@ -41,13 +41,13 @@ def _egocentric_aligner(frm_range: np.ndarray,
img = read_frm_of_video(video_path=cap, frame_index=frm_id)
R, center = rotation_vectors[frm_id], centers[frm_id]
M_rotate = np.hstack([R, np.array([[-center[0] * R[0, 0] - center[1] * R[0, 1] + center[0]], [-center[0] * R[1, 0] - center[1] * R[1, 1] + center[1]]])])
rotated_frame = cv2.warpAffine(img, M_rotate, (video_meta['width'], video_meta['height']))
rotated_frame = cv2.warpAffine(img, M_rotate, (video_meta['width'], video_meta['height']), borderValue=color)
translation_x = target[0] - center[0]
translation_y = target[1] - center[1]
M_translate = np.float32([[1, 0, translation_x], [0, 1, translation_y]])
final_frame = cv2.warpAffine(rotated_frame, M_translate, (video_meta['width'], video_meta['height']))
final_frame = cv2.warpAffine(rotated_frame, M_translate, (video_meta['width'], video_meta['height']), borderValue=color)
writer.write(final_frame)
print(f'Creating frame {frm_id} (CPU core: {batch+1}).')
print(f'Creating frame {frm_id} ({video_name}, CPU core: {batch+1}).')

cap.release()
writer.release()
Expand Down Expand Up @@ -82,6 +82,7 @@ def __init__(self,
config_path: Union[str, os.PathLike],
save_dir: Union[str, os.PathLike],
data_dir: Optional[Union[str, os.PathLike]] = None,
videos_dir: Optional[Union[str, os.PathLike]] = None,
anchor_1: Optional[str] = 'tail_base',
anchor_2: Optional[str] = 'nose',
direction: int = 0,
Expand All @@ -106,91 +107,98 @@ def __init__(self,
self.cores = find_core_cnt()[0]
else:
self.cores = cores
if videos_dir is not None:
check_if_dir_exists(in_dir=videos_dir)
self.video_dir = videos_dir

def run(self):
for file_cnt, file_path in enumerate(self.data_paths):
_, self.video_name, _ = get_fn_ext(filepath=file_path)
save_path = os.path.join(self.save_dir, f'{self.video_name}.{self.file_type}')
df = read_df(file_path=file_path, file_type=self.file_type)
original_cols, self.file_path = list(df.columns), file_path
df.columns = [x.lower() for x in list(df.columns)]
check_valid_dataframe(df=df, source=self.__class__.__name__, valid_dtypes=Formats.NUMERIC_DTYPES.value, required_fields=self.anchor_1_cols + self.anchor_2_cols)
self.body_parts_lst = [x.lower() for x in self.body_parts_lst]
bp_cols = [x for x in df.columns if not x.endswith('_p')]
anchor_1_idx = self.body_parts_lst.index(self.anchor_1)
anchor_2_idx = self.body_parts_lst.index(self.anchor_2)
data_arr = df[bp_cols].values.reshape(len(df), len(self.body_parts_lst), 2).astype(np.int32)
results_arr = np.zeros_like(data_arr)
self.rotation_angles, self.rotation_vectors, self.centers, self.deltas = [], [], [], []
for frame_index in range(data_arr.shape[0]):
frame_points = data_arr[frame_index]
frame_anchor_1 = frame_points[anchor_1_idx]
self.centers.append(tuple(frame_anchor_1))
frame_anchor_2 = frame_points[anchor_2_idx]
delta_x, delta_y = frame_anchor_2[0] - frame_anchor_1[0], frame_anchor_2[1] - frame_anchor_1[1]
self.deltas.append((delta_x, delta_x))
current_angle = np.arctan2(delta_y, delta_x)
rotate_angle = self.target_angle - current_angle
self.rotation_angles.append(rotate_angle)
cos_theta, sin_theta = np.cos(rotate_angle), np.sin(rotate_angle)
R = np.array([[cos_theta, -sin_theta], [sin_theta, cos_theta]])
self.rotation_vectors.append(R)
keypoints_translated = frame_points - frame_anchor_1
keypoints_rotated = np.dot(keypoints_translated, R.T)
anchor_1_position_after_rotation = keypoints_rotated[anchor_1_idx]
translation_to_target = np.array(self.anchor_location) - anchor_1_position_after_rotation
keypoints_aligned = keypoints_rotated + translation_to_target
results_arr[frame_index] = keypoints_aligned

results_arr = results_arr.reshape(len(df), len(bp_cols))
self.out_df = pd.DataFrame(results_arr, columns=bp_cols)
df.update(self.out_df)
df.columns = original_cols
write_df(df=df, file_type=self.file_type, save_path=save_path)
if self.rotate_video:
self.run_video_rotation()
if not os.path.isfile(save_path):
df = read_df(file_path=file_path, file_type=self.file_type)
original_cols, self.file_path = list(df.columns), file_path
df.columns = [x.lower() for x in list(df.columns)]
check_valid_dataframe(df=df, source=self.__class__.__name__, valid_dtypes=Formats.NUMERIC_DTYPES.value, required_fields=self.anchor_1_cols + self.anchor_2_cols)
self.body_parts_lst = [x.lower() for x in self.body_parts_lst]
bp_cols = [x for x in df.columns if not x.endswith('_p')]
anchor_1_idx = self.body_parts_lst.index(self.anchor_1)
anchor_2_idx = self.body_parts_lst.index(self.anchor_2)
data_arr = df[bp_cols].values.reshape(len(df), len(self.body_parts_lst), 2).astype(np.int32)
results_arr = np.zeros_like(data_arr)
self.rotation_angles, self.rotation_vectors, self.centers, self.deltas = [], [], [], []
for frame_index in range(data_arr.shape[0]):
frame_points = data_arr[frame_index]
frame_anchor_1 = frame_points[anchor_1_idx]
self.centers.append(tuple(frame_anchor_1))
frame_anchor_2 = frame_points[anchor_2_idx]
delta_x, delta_y = frame_anchor_2[0] - frame_anchor_1[0], frame_anchor_2[1] - frame_anchor_1[1]
self.deltas.append((delta_x, delta_x))
current_angle = np.arctan2(delta_y, delta_x)
rotate_angle = self.target_angle - current_angle
self.rotation_angles.append(rotate_angle)
cos_theta, sin_theta = np.cos(rotate_angle), np.sin(rotate_angle)
R = np.array([[cos_theta, -sin_theta], [sin_theta, cos_theta]])
self.rotation_vectors.append(R)
keypoints_translated = frame_points - frame_anchor_1
keypoints_rotated = np.dot(keypoints_translated, R.T)
anchor_1_position_after_rotation = keypoints_rotated[anchor_1_idx]
translation_to_target = np.array(self.anchor_location) - anchor_1_position_after_rotation
keypoints_aligned = keypoints_rotated + translation_to_target
results_arr[frame_index] = keypoints_aligned

results_arr = results_arr.reshape(len(df), len(bp_cols))
self.out_df = pd.DataFrame(results_arr, columns=bp_cols)
df.update(self.out_df)
df.columns = original_cols
write_df(df=df, file_type=self.file_type, save_path=save_path)
if self.rotate_video:
self.run_video_rotation()

def run_video_rotation(self):
video_timer = SimbaTimer(start=True)
video_path = find_video_of_file(video_dir=self.video_dir, filename=self.video_name, raise_error=True)
video_meta = get_video_meta_data(video_path=video_path)
save_path = os.path.join(self.save_dir, f'{self.video_name}.mp4')
temp_dir = os.path.join(self.save_dir, 'temp')
if not os.path.isdir(temp_dir):
os.makedirs(temp_dir)
else:
remove_a_folder(folder_dir=temp_dir)
os.makedirs(temp_dir)
if video_meta['frame_count'] != len(self.out_df):
FrameRangeWarning(msg=f'The video {video_path} contains {video_meta["frame_count"]} frames while the file {self.file_path} contains {len(self.out_df)} frames', source=self.__class__.__name__)
frm_list = np.arange(0, video_meta['frame_count'])
frm_list = np.array_split(frm_list, self.cores)
frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)]
print(f"Creating rotated videos, multiprocessing (chunksize: {self.multiprocess_chunksize}, cores: {self.cores})...")
with multiprocessing.Pool(self.cores, maxtasksperchild=self.maxtasksperchild) as pool:
constants = functools.partial(_egocentric_aligner,
temp_dir=temp_dir,
video_path=video_path,
centers=self.centers,
rotation_vectors=self.rotation_vectors,
target=self.anchor_location)
for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=self.multiprocess_chunksize)):
print(f"Rotate batch {result}/{self.cores} complete...")
pool.terminate()
pool.join()

concatenate_videos_in_folder(in_folder=temp_dir, save_path=save_path, remove_splits=True, gpu=False)
video_timer.stop_timer()
print(f"Egocentric rotation video {save_path} complete (elapsed time: {video_timer.elapsed_time_str}s) ...")


# if __name__ == "__main__":
# aligner = EgocentricalAligner(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",
# rotate_video=True,
# anchor_1='tail_base',
# anchor_2='nose',
# data_dir=r'C:\troubleshooting\mitra\project_folder\csv\outlier_corrected_movement_location\test',
# save_dir=r"C:\troubleshooting\mitra\project_folder\csv\outlier_corrected_movement_location\test\bg_temp\rotated")
# aligner.run()
#
video_path = find_video_of_file(video_dir=self.video_dir, filename=self.video_name, raise_error=False)
if video_path is not None:
video_meta = get_video_meta_data(video_path=video_path)
save_path = os.path.join(self.save_dir, f'{self.video_name}.mp4')
temp_dir = os.path.join(self.save_dir, 'temp')
if not os.path.isdir(temp_dir):
os.makedirs(temp_dir)
else:
remove_a_folder(folder_dir=temp_dir)
os.makedirs(temp_dir)
if video_meta['frame_count'] != len(self.out_df):
FrameRangeWarning(msg=f'The video {video_path} contains {video_meta["frame_count"]} frames while the file {self.file_path} contains {len(self.out_df)} frames', source=self.__class__.__name__)
frm_list = np.arange(0, video_meta['frame_count'])
frm_list = np.array_split(frm_list, self.cores)
frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)]
print(f"Creating rotated videos, multiprocessing (chunksize: {self.multiprocess_chunksize}, cores: {self.cores})...")
with multiprocessing.Pool(self.cores, maxtasksperchild=self.maxtasksperchild) as pool:
constants = functools.partial(_egocentric_aligner,
temp_dir=temp_dir,
video_name=self.video_name,
video_path=video_path,
centers=self.centers,
rotation_vectors=self.rotation_vectors,
target=self.anchor_location)
for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=self.multiprocess_chunksize)):
print(f"Rotate batch {result}/{self.cores} complete...")
pool.terminate()
pool.join()

concatenate_videos_in_folder(in_folder=temp_dir, save_path=save_path, remove_splits=True, gpu=False)
video_timer.stop_timer()
print(f"Egocentric rotation video {save_path} complete (elapsed time: {video_timer.elapsed_time_str}s) ...")


if __name__ == "__main__":
aligner = EgocentricalAligner(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini",
rotate_video=True,
anchor_1='tail_base',
anchor_2='nose',
data_dir=r'C:\troubleshooting\mitra\project_folder\csv\outlier_corrected_movement_location',
videos_dir=r'C:\troubleshooting\mitra\project_folder\videos\bg_removed',
save_dir=r"C:\troubleshooting\mitra\project_folder\videos\bg_removed\rotated")
aligner.run()


Loading

0 comments on commit 1226b32

Please sign in to comment.