diff --git a/docs/_static/img/egocentric_nb_1.webm b/docs/_static/img/egocentric_nb_1.webm new file mode 100644 index 000000000..cbc7c28ac Binary files /dev/null and b/docs/_static/img/egocentric_nb_1.webm differ diff --git a/docs/_static/img/egocentric_nb_2.webm b/docs/_static/img/egocentric_nb_2.webm new file mode 100644 index 000000000..ba39e6ea5 Binary files /dev/null and b/docs/_static/img/egocentric_nb_2.webm differ diff --git a/docs/_static/img/egocentric_nb_3.webm.webm b/docs/_static/img/egocentric_nb_3.webm.webm new file mode 100644 index 000000000..b72d3d71c Binary files /dev/null and b/docs/_static/img/egocentric_nb_3.webm.webm differ diff --git a/docs/nb/egocentric_align.ipynb b/docs/nb/egocentric_align.ipynb new file mode 100644 index 000000000..6d5973b0f --- /dev/null +++ b/docs/nb/egocentric_align.ipynb @@ -0,0 +1,301 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4f019f1c", + "metadata": {}, + "source": [ + "# Egocentric data and video alignment" + ] + }, + { + "cell_type": "markdown", + "id": "cb011978", + "metadata": {}, + "source": [ + "In this notebook, we will egocentrically align pose estimation and associated video data.\n", + "\n", + "Egocentric alignment is a crucial preprocessing step with various applications.\n", + "\n", + "One primary use case is in unsupervised deep learning scenarios, where images of animals are used as inputs to train algorithms.\n", + "Standardizing the orientation of the subject (e.g., an animal) ensures that the model and subsequent analyses do not incorrectly interpret the same behavior in different orientations\n", + "(e.g., grooming while facing north versus south) as distinct behaviors. By aligning the animal to a consistent frame of reference, we eliminate orientation as a confounding variable.\n", + "\n", + "While egocentric alignment is an essential first step, it is often insufficient by itself for comprehensive analyses. Additional preprocessing steps are typically required, such as:\n", + "\n", + "* Background subtraction to isolate the animal from its surroundings (see other relevant methods and notebooks).\n", + "* Geometric segmentation to slice out and focus on the subject (again, see other relevant methods and notebooks).\n", + "\n", + "In this notebook, we will focus exclusively on performing egocentric alignment. Further preprocessing steps are outlined in related materials." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "4654f323", + "metadata": {}, + "outputs": [], + "source": [ + "from simba.data_processors.egocentric_aligner import EgocentricalAligner\n", + "from ipywidgets import Video\n", + "from IPython.display import HTML" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "4d4919e1", + "metadata": {}, + "outputs": [], + "source": [ + "# WE DEFINE HOW WE SHOULD EGOCENTRICALLY ALIGN THE DATA AND VIDEOs\n", + "\n", + "ANCHOR_POINT_1 = 'center' # Name of the body-part which is the \"primary\" anchor point around which the alignment centers. In rodents, this is often the center of the tail-base of the animal.\n", + "ANCHOR_POINT_2 = 'nose' # The name of the secondary anchor point defining the alignment direction. This is often the anterior body-part, in rodents it can be the nose or nape.\n", + "DIRECTION = 0 # The egocentric alignment angle, in degrees. For example, `0` and the animals `ANCHOR_POINT_2` is directly to the east (right) of `ANCHOR_POINT_1`. `180` and the animals `ANCHOR_POINT_2` is directly to the west (left) of `ANCHOR_POINT_1`.\n", + "ANCHOR_LOCATION = (250, 250) # The pixel location in the video where `ANCHOR_POINT_1` should be placed. For example, if the videos are 500x500, 250x250 will place the anchor right in the middle.\n", + "GPU = True # If we have an NVIDEA GPU availalable, we can use it to speed up processing. Otherwise set this to `False`.\n", + "FILL_COLOR = (0, 0, 0) # We are rotating videos, while at the same time retaining the original video size. Therefore, there will be some \"new\" areas exposed in the video (see below for more info). This is what color to color these new areas.\n", + "VERBOSE = False # If True, prints progress (like which frame and video is being processed etc). However, this information will flood this notebook is I am going to turn it off." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "8627686d", + "metadata": {}, + "outputs": [], + "source": [ + "# WE DEFINE THE PATHS TO THE DIRECTORIES HOLDING THE DATA AND VIDEOS, AND DIRECTORY WHERE WE SHOULD STORE THE RESULTS.\n", + "\n", + "DATA_DIRECTORY = r'C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\data' #DIRECTORY WHICH IS HOLDING POSE-ESTIMATION DATA\n", + "VIDEOS_DIRECTORY = r'C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\videos' #DIRECTORY WHICH IS VIDEOS, ONE FOR EACH FILE IN THE DATA_DIRECTORY. NOTE: IF YOU SET THIS TO None, THEN THE ALIGNMENT WILL BE PERFORMED ON THE DATA ONLY.\n", + "SAVE_DIRECTORY = r\"C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\" #DIRECTORY WHERE WE SHOULD SAVE THE ROTATED POSE-ESTIMATION AND ROTATED VIDEOS." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a3680603", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "SIMBA COMPLETE: Video concatenated (elapsed time: 16.214s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\501_MA142_Gi_Saline_0513.mp4 complete (elapsed time: 102.2651s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 15.0453s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\501_MA142_Gi_Saline_0515.mp4 complete (elapsed time: 101.7265s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 14.7956s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\501_MA142_Gi_Saline_0517.mp4 complete (elapsed time: 105.1418s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 16.5156s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\502_MA141_Gi_Saline_0513.mp4 complete (elapsed time: 107.4481s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 14.7832s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\502_MA141_Gi_Saline_0515.mp4 complete (elapsed time: 103.6864s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 14.9907s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\502_MA141_Gi_Saline_0517.mp4 complete (elapsed time: 106.1234s) \tcomplete\n", + "SIMBA COMPLETE: Egocentrically aligned data for 6 files saved in C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated (elapsed time: 627.5204s) \tcomplete\n" + ] + } + ], + "source": [ + "# NOW WE ARE GOOD TO GO, USING THE INFORMATION ABOVE, WE DEFINE AN INSTANCE OF AN SimBA EGOCENTRIALIGNER AND RUN IT\n", + "# ON THE 6 VIDEOS AND VIDEO DATA INSIDE THE DATA AND VIDEO DIRECTORIES, RESPECTIVELY.\n", + "aligner = EgocentricalAligner(anchor_1=ANCHOR_POINT_1,\n", + " anchor_2=ANCHOR_POINT_2,\n", + " data_dir=DATA_DIRECTORY,\n", + " videos_dir=VIDEOS_DIRECTORY,\n", + " save_dir=SAVE_DIRECTORY,\n", + " direction=DIRECTION,\n", + " gpu=GPU,\n", + " anchor_location=ANCHOR_LOCATION,\n", + " fill_clr=FILL_COLOR,\n", + " verbose=VERBOSE)\n", + "aligner.run()" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "54298ed3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + " \n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "#EXAMPLE VIDEO EXPECTED RESULTS SNIPPET\n", + "video_url = 'https://raw.githubusercontent.com/sgoldenlab/simba/master/docs/_static/img/egocentric_nb_1.webm'\n", + "HTML(f''' \n", + "''')" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "e67378a7", + "metadata": {}, + "outputs": [], + "source": [ + "#NOW, LET'S CHANGE A FEW SETTINGS, TO GET A FEELING FOR HOW IT BEHAVES …\n", + "ANCHOR_LOCATION = (500, 100)\n", + "FILL_COLOR = (255, 0, 0)\n", + "DIRECTION = 180" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "8ead691a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "SIMBA COMPLETE: Video concatenated (elapsed time: 16.2393s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\501_MA142_Gi_Saline_0513.mp4 complete (elapsed time: 108.1662s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 15.0821s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\501_MA142_Gi_Saline_0515.mp4 complete (elapsed time: 101.4825s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 14.7555s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\501_MA142_Gi_Saline_0517.mp4 complete (elapsed time: 102.9438s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 16.4037s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\502_MA141_Gi_Saline_0513.mp4 complete (elapsed time: 109.8081s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 14.8025s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\502_MA141_Gi_Saline_0515.mp4 complete (elapsed time: 103.8362s) \tcomplete\n", + "SIMBA COMPLETE: Video concatenated (elapsed time: 15.0489s) \tcomplete\n", + "SIMBA COMPLETE: Egocentric rotation video C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated\\502_MA141_Gi_Saline_0517.mp4 complete (elapsed time: 107.5951s) \tcomplete\n", + "SIMBA COMPLETE: Egocentrically aligned data for 6 files saved in C:\\Users\\sroni\\OneDrive\\Desktop\\rotate_ex\\rotated (elapsed time: 634.8647s) \tcomplete\n" + ] + } + ], + "source": [ + "# ... WE CREATE A NEW INSTANCE BASED ON THE UPDATED INFORMATION ABOVE, AND RUN IT.\n", + "aligner = EgocentricalAligner(anchor_1=ANCHOR_POINT_1,\n", + " anchor_2=ANCHOR_POINT_2,\n", + " data_dir=DATA_DIRECTORY,\n", + " videos_dir=VIDEOS_DIRECTORY,\n", + " save_dir=SAVE_DIRECTORY,\n", + " direction=DIRECTION,\n", + " gpu=GPU,\n", + " anchor_location=ANCHOR_LOCATION,\n", + " fill_clr=FILL_COLOR,\n", + " verbose=VERBOSE)\n", + "aligner.run()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "dd458e84", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + " \n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "#EXAMPLE VIDEO EXPECTED RESULTS SNIPPET\n", + "video_url = 'https://raw.githubusercontent.com/sgoldenlab/simba/master/docs/_static/img/egocentric_nb_2.webm'\n", + "HTML(f''' \n", + "''')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0f659939", + "metadata": {}, + "outputs": [], + "source": [ + "# FINALLY, LET'S USE A DIFFERENT EXPERIMENT, WITH DIFFERENT BODY-PART NAMES TO SEE HOW IT BEHAVES\n", + "ANCHOR_POINT_1 = 'center' \n", + "ANCHOR_POINT_2 = 'nose' \n", + "DIRECTION = 0 \n", + "ANCHOR_LOCATION = (600, 300)\n", + "FILL_COLOR = (128, 120, 128) \n", + "\n", + "DATA_DIRECTORY = r'C:\\troubleshooting\\open_field_below\\project_folder\\csv\\outlier_corrected_movement_location'\n", + "VIDEOS_DIRECTORY = r'C:\\troubleshooting\\open_field_below\\project_folder\\videos'\n", + "SAVE_DIRECTORY = r\"C:\\troubleshooting\\open_field_below\\project_folder\\videos\\rotated\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c654f6a", + "metadata": {}, + "outputs": [], + "source": [ + "# ... AGAIN WE CREATE A NEW INSTANCE BASED ON THE UPDATED INFORMATION ABOVE, AND RUN IT.\n", + "aligner = EgocentricalAligner(anchor_1=ANCHOR_POINT_1,\n", + " anchor_2=ANCHOR_POINT_2,\n", + " data_dir=DATA_DIRECTORY,\n", + " videos_dir=VIDEOS_DIRECTORY,\n", + " save_dir=SAVE_DIRECTORY,\n", + " direction=DIRECTION,\n", + " gpu=GPU,\n", + " anchor_location=ANCHOR_LOCATION,\n", + " fill_clr=FILL_COLOR,\n", + " verbose=VERBOSE)\n", + "aligner.run()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a70d3f8", + "metadata": {}, + "outputs": [], + "source": [ + "#EXAMPLE VIDEO EXPECTED RESULTS SNIPPET\n", + "video_url = 'https://raw.githubusercontent.com/sgoldenlab/simba/master/docs/_static/img/egocentric_nb_2.webm'\n", + "HTML(f''' \n", + "''')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "simba", + "language": "python", + "name": "simba" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/nb/geometry_example_1.ipynb b/docs/nb/geometry_example_1.ipynb index a649f6cb8..6515ee4f6 100644 --- a/docs/nb/geometry_example_1.ipynb +++ b/docs/nb/geometry_example_1.ipynb @@ -9,7 +9,6 @@ ] }, { - "attachments": {}, "cell_type": "markdown", "id": "85aedd1a", "metadata": {}, @@ -335,7 +334,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3", "language": "python", "name": "python3" }, @@ -349,7 +348,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.15" + "version": "3.6.13" } }, "nbformat": 4, diff --git a/docs/nb/geometry_example_6.ipynb b/docs/nb/geometry_example_6.ipynb index 400cc35b8..21be1b74a 100644 --- a/docs/nb/geometry_example_6.ipynb +++ b/docs/nb/geometry_example_6.ipynb @@ -503,7 +503,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3", "language": "python", "name": "python3" }, @@ -517,7 +517,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.15" + "version": "3.6.13" } }, "nbformat": 4, diff --git a/docs/nb/geometry_example_7.ipynb b/docs/nb/geometry_example_7.ipynb index 230a28a68..be51c9010 100644 --- a/docs/nb/geometry_example_7.ipynb +++ b/docs/nb/geometry_example_7.ipynb @@ -690,7 +690,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3", "language": "python", "name": "python3" }, @@ -704,7 +704,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.15" + "version": "3.6.13" } }, "nbformat": 4, diff --git a/docs/simba.data_processors.rst b/docs/simba.data_processors.rst index 867f04a61..8bd548684 100644 --- a/docs/simba.data_processors.rst +++ b/docs/simba.data_processors.rst @@ -180,3 +180,12 @@ Heuristic freezing detector :members: :undoc-members: + +Data GPU methods +---------------------------------------------- + +.. automodule:: simba.data_processors.cuda.data + :members: + :undoc-members: + :show-inheritance: + diff --git a/docs/simba.video_processing.rst b/docs/simba.video_processing.rst index 3f18563de..d807866cd 100644 --- a/docs/simba.video_processing.rst +++ b/docs/simba.video_processing.rst @@ -51,9 +51,16 @@ Interactive brightness / contrast :show-inheritance: -Batch process executor +Batch video process executor ----------------------------------------------------------------------- .. automodule:: simba.video_processors.batch_process_create_ffmpeg_commands :members: :show-inheritance: + +Egocentrically rotate videos +----------------------------------------------------------------------- + +.. automodule:: simba.video_processors.egocentric_video_rotator + :members: + :show-inheritance: diff --git a/docs/tables/egocentrically_align_pose_cuda.csv b/docs/tables/egocentrically_align_pose_cuda.csv new file mode 100644 index 000000000..9146f28c3 --- /dev/null +++ b/docs/tables/egocentrically_align_pose_cuda.csv @@ -0,0 +1,12 @@ +FRAMES (MILLIONS),CUDA TIME (S),CUDA TIME (STEV) +0.25,0.1882001,0.15372434 +0.5,0.1221498,0.00574819 +1,0.24,0.0307 +2,0.38,0.1092 +4,0.505,0.01590969 +8,1.037,0.0346 +16,3.42,1.53194867 +7 BODY-PARTS PER FRAME,, +3 ITERATIONS,, +batch size 16M,, +DIDN'T TEST HIGHER N AS I DON'T HAVE THE NON-GPU RAM,, diff --git a/setup.py b/setup.py index d2a4a5ec9..1bb88002b 100644 --- a/setup.py +++ b/setup.py @@ -28,8 +28,8 @@ # Setup configuration setuptools.setup( - name="Simba-UW-tf-dev", - version="2.3.9", + name="simba-uw-tf-dev", + version="2.4.3", author="Simon Nilsson, Jia Jie Choong, Sophia Hwang", author_email="sronilsson@gmail.com", description="Toolkit for computer classification and analysis of behaviors in experimental animals", diff --git a/simba/data_processors/cuda/data.py b/simba/data_processors/cuda/data.py new file mode 100644 index 000000000..89b3fb0a5 --- /dev/null +++ b/simba/data_processors/cuda/data.py @@ -0,0 +1,175 @@ +import math +import time + +from simba.utils.checks import check_int, check_valid_array +import numpy as np +from typing import Tuple +from simba.utils.errors import InvalidInputError +from simba.utils.enums import Formats +from simba.utils.read_write import read_df +from numba import cuda +from simba.data_processors.cuda.utils import _cuda_matrix_multiplication, _cuda_2d_transpose, _cuda_subtract_2d, _cuda_add_2d + +THREADS_PER_BLOCK = 1024 + + +@cuda.jit() +def _egocentric_align_kernel(data, centers, rotation_vectors, results, target_angle, anchor_idx, transposed_rotation_vectors, matrix_multiplier_arr, anchor_loc): + frm_idx = cuda.grid(1) + if frm_idx >= data.shape[0]: + return + else: + frm_points = data[frm_idx] + frm_anchor_1, frm_anchor_2 = frm_points[anchor_idx[0]], frm_points[anchor_idx[1]] + centers[frm_idx][0], centers[frm_idx][1] = frm_anchor_1[0], frm_anchor_1[1] + delta_x, delta_y = frm_anchor_2[0] - frm_anchor_1[0], frm_anchor_2[1] - frm_anchor_1[1] + frm_angle = math.atan2(delta_y, delta_x) + frm_rotation_angle = target_angle[0] - frm_angle + frm_cos_theta, frm_sin_theta = math.cos(frm_rotation_angle), math.sin(frm_rotation_angle) + rotation_vectors[frm_idx][0][0], rotation_vectors[frm_idx][0][1] = frm_cos_theta, -frm_sin_theta + rotation_vectors[frm_idx][1][0], rotation_vectors[frm_idx][1][1] = frm_sin_theta, frm_cos_theta + keypoints_rotated = _cuda_subtract_2d(frm_points, frm_anchor_1) + r_transposed = _cuda_2d_transpose(rotation_vectors[frm_idx], transposed_rotation_vectors[frm_idx]) + keypoints_rotated = _cuda_matrix_multiplication(keypoints_rotated, r_transposed, matrix_multiplier_arr[frm_idx]) + anchor_1_position_after_rotation = keypoints_rotated[anchor_idx[0]] + anchor_1_position_after_rotation[0] = anchor_loc[0] - anchor_1_position_after_rotation[0] + anchor_1_position_after_rotation[1] = anchor_loc[1] - anchor_1_position_after_rotation[1] + + frm_results = _cuda_add_2d(keypoints_rotated, anchor_1_position_after_rotation) + for i in range(frm_results.shape[0]): + for j in range(frm_results.shape[1]): + results[frm_idx][i][j] = frm_results[i][j] + + +def egocentrically_align_pose_cuda(data: np.ndarray, + anchor_1_idx: int, + anchor_2_idx: int, + anchor_location: np.ndarray, + direction: int, + batch_size: int = int(10e+5)) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + + """ + Aligns a set of 2D points egocentrically based on two anchor points and a target direction using GPU acceleration. + + Rotates and translates a 3D array of 2D points (e.g., time-series of frame-wise data) such that + one anchor point is aligned to a specified location, and the direction between the two anchors is aligned + to a target angle. + + .. video:: _static/img/EgocentricalAligner.webm + :width: 600 + :autoplay: + :loop: + + .. seealso:: + For numpy function, see :func:`simba.utils.data.egocentrically_align_pose`. + For numba alternative, see :func:`simba.utils.data.egocentrically_align_pose_numba`. + To align both pose and video, see :func:`simba.data_processors.egocentric_aligner.EgocentricalAligner`. + To egocentrically rotate video, see :func:`simba.video_processors.egocentric_video_rotator.EgocentricVideoRotator` + + .. csv-table:: + :header: EXPECTED RUNTIMES + :file: ../../../docs/tables/egocentrically_align_pose_cuda.csv + :widths: 10, 45, 45 + :align: center + :class: simba-table + :header-rows: 1 + + :param np.ndarray data: A 3D array of shape `(num_frames, num_points, 2)` containing 2D points for each frame. Each frame is represented as a 2D array of shape `(num_points, 2)`, where each row corresponds to a point's (x, y) coordinates. + :param int anchor_1_idx: The index of the first anchor point in `data` used as the center of alignment. This body-part will be placed in the center of the image. + :param int anchor_2_idx: The index of the second anchor point in `data` used to calculate the direction vector. This bosy-part will be located `direction` degrees from the anchor_1 body-part. + :param int direction: The target direction in degrees to which the vector between the two anchors will be aligned. + :param np.ndarray anchor_location: A 1D array of shape `(2,)` specifying the target (x, y) location for `anchor_1_idx` after alignment. + :param int batch_size: Size of data that is processed on each iteration on GPU. default 1m. Increase if GPU allows. + + :return: A tuple containing the rotated data, and variables required for also rotating the video using the same rules: + - `aligned_data`: A 3D array of shape `(num_frames, num_points, 2)` with the aligned 2D points. + - `centers`: A 2D array of shape `(num_frames, 2)` containing the original locations of `anchor_1_idx` in each frame before alignment. + - `rotation_vectors`: A 3D array of shape `(num_frames, 2, 2)` containing the rotation matrices applied to each frame. + :rtype: Tuple[np.ndarray, np.ndarray, np.ndarray] + + :example: + >>> DATA_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/data/501_MA142_Gi_Saline_0513.csv" + >>> VIDEO_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/videos/501_MA142_Gi_Saline_0513.mp4" + >>> SAVE_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/videos/501_MA142_Gi_Saline_0513_rotated.mp4" + >>> ANCHOR_LOC = np.array([300, 300]) + >>> + >>> df = read_df(file_path=DATA_PATH, file_type='csv') + >>> bp_cols = [x for x in df.columns if not x.endswith('_p')] + >>> data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int64) + >>> data, centers, rotation_matrices = egocentrically_align_pose_cuda(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=180,batch_size=36000000) + """ + + check_valid_array(data=data, source=egocentrically_align_pose_cuda.__name__, accepted_ndims=(3,), accepted_dtypes=Formats.NUMERIC_DTYPES.value) + check_int(name=f'{egocentrically_align_pose_cuda.__name__} anchor_1_idx', min_value=0, max_value=data.shape[1], value=anchor_1_idx) + check_int(name=f'{egocentrically_align_pose_cuda.__name__} anchor_2_idx', min_value=0, max_value=data.shape[1], value=anchor_2_idx) + if anchor_1_idx == anchor_2_idx: raise InvalidInputError(msg=f'Anchor 1 index ({anchor_1_idx}) cannot be the same as Anchor 2 index ({anchor_2_idx})', source=egocentrically_align_pose_cuda.__name__) + check_int(name=f'{egocentrically_align_pose_cuda.__name__} direction', value=direction, min_value=0, max_value=360) + check_valid_array(data=anchor_location, source=egocentrically_align_pose_cuda.__name__, accepted_ndims=(1,), accepted_axis_0_shape=[2,], accepted_dtypes=Formats.NUMERIC_DTYPES.value) + check_int(name=f'{egocentrically_align_pose_cuda.__name__} batch_size', value=batch_size, min_value=1) + results = np.full_like(a=data, fill_value=-1, dtype=np.int64) + results_centers = np.full((data.shape[0], 2), fill_value=-1, dtype=np.int64) + results_rotation_vectors = np.full((data.shape[0], 2, 2), fill_value=-1, dtype=np.float64) + transposed_results_rotation_vectors = np.full((data.shape[0], 2, 2), fill_value=np.nan, dtype=np.float64) + matrix_multiplier_arr = np.full((data.shape[0], data.shape[1], 2), fill_value=-1, dtype=np.int64) + target_angle = np.deg2rad(direction) + target_angle_dev = cuda.to_device(np.array([target_angle])) + anchor_idx_dev = cuda.to_device(np.array([anchor_1_idx, anchor_2_idx])) + anchor_loc_dev = cuda.to_device(anchor_location) + + for l in range(0, data.shape[0], batch_size): + r = l + batch_size + sample_data = np.ascontiguousarray(data[l:r]).astype(np.float64) + sample_centers = np.ascontiguousarray(results_centers[l:r]).astype(np.int64) + sample_rotation_vectors = np.ascontiguousarray(results_rotation_vectors[l:r].astype(np.float64)) + sample_transposed_rotation_vectors = np.ascontiguousarray(transposed_results_rotation_vectors[l:r]) + sample_matrix_multiplier_arr = np.ascontiguousarray(matrix_multiplier_arr[l:r]) + sample_results = np.ascontiguousarray(results[l:r].astype(np.float64)) + sample_data_dev = cuda.to_device(sample_data) + sample_centers_dev = cuda.to_device(sample_centers) + sample_matrix_multiplier_arr_dev = cuda.to_device(sample_matrix_multiplier_arr) + sample_transposed_rotation_vectors_dev = cuda.to_device(sample_transposed_rotation_vectors) + sample_rotation_vectors_dev = cuda.to_device(sample_rotation_vectors) + sample_results_dev = cuda.to_device(sample_results) + bpg = (sample_data.shape[0] + (THREADS_PER_BLOCK - 1)) // THREADS_PER_BLOCK + _egocentric_align_kernel[bpg, THREADS_PER_BLOCK](sample_data_dev, + sample_centers_dev, + sample_rotation_vectors_dev, + sample_results_dev, + target_angle_dev, + anchor_idx_dev, + sample_transposed_rotation_vectors_dev, + sample_matrix_multiplier_arr_dev, + anchor_loc_dev) + results[l:r] = sample_results_dev.copy_to_host() + results_centers[l:r] = sample_centers_dev.copy_to_host() + results_rotation_vectors[l:r] = sample_rotation_vectors_dev.copy_to_host() + + return results, results_centers, results_rotation_vectors + + +# DATA_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/data/501_MA142_Gi_Saline_0513.csv" +# VIDEO_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/videos/501_MA142_Gi_Saline_0513.mp4" +# SAVE_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/videos/501_MA142_Gi_Saline_0513_rotated.mp4" +# ANCHOR_LOC = np.array([300, 300]) +# # +# # df = read_df(file_path=DATA_PATH, file_type='csv') +# # bp_cols = [x for x in df.columns if not x.endswith('_p')] +# # data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int64) +# # data, centers, rotation_matrices = egocentrically_align_pose_cuda(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=180,batch_size=36000000) +# # +# for i in [250000, 500000, 1000000, 2000000, 4000000, 8000000, 16000000]: +# data = np.random.randint(0, 500, (i, 6, 2)) +# times = [] +# for j in range(3): +# start_t = time.perf_counter() +# data, centers, rotation_matrices = egocentrically_align_pose_cuda(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=180, batch_size=36000000) +# times.append(time.perf_counter() - start_t) +# print(i, '\t' * 4, np.mean(times), '\t' * 4, np.std(times)) + + +# from simba.video_processors.egocentric_video_rotator import EgocentricVideoRotator +# +# runner = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_matrices, anchor_location=(300, 300)) +# runner.run() + +#_, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0) \ No newline at end of file diff --git a/simba/data_processors/cuda/utils.py b/simba/data_processors/cuda/utils.py index cbb73a438..c95c5adaa 100644 --- a/simba/data_processors/cuda/utils.py +++ b/simba/data_processors/cuda/utils.py @@ -5,6 +5,7 @@ from numba import cuda, float64, guvectorize + @cuda.jit(device=True) def _cuda_sum(x: np.ndarray): s = 0 @@ -82,6 +83,40 @@ def _cuda_digital_pixel_to_grey(r: int, g: int, b: int): def _euclid_dist(x, y): return math.sqrt(((y[0] - x[0]) ** 2) + ((y[1] - x[1]) ** 2)) +@cuda.jit(device=True) +def _cuda_matrix_multiplication(mA, mB, out): + """ Matrix multiplication""" + for i in range(mA.shape[0]): + for j in range(mB.shape[1]): + for k in range(mA.shape[1]): + out[i][j] += mA[i][k] * mB[k][j] + return out + +@cuda.jit(device=True) +def _cuda_2d_transpose(x, y): + """ Transpose a 2d array """ + for i in range(x.shape[0]): + for j in range(x.shape[1]): + y[j][i] = x[i][j] + return y + +@cuda.jit(device=True) +def _cuda_subtract_2d(x: np.ndarray, vals: np.ndarray) -> np.ndarray: + """ Subtract 1d array values for every row in a 2d array""" + for i in range(x.shape[0]): + for j in range(x.shape[1]): + x[i][j] = x[i][j] - vals[j] + return x + + +@cuda.jit(device=True) +def _cuda_add_2d(x: np.ndarray, vals: np.ndarray) -> np.ndarray: + """ Add 1d array values for every row in a 2d array""" + for i in range(x.shape[0]): + for j in range(x.shape[1]): + x[i][j] = x[i][j] + vals[j] + return x + def _cuda_available() -> Tuple[bool, Dict[int, Any]]: """ Check if GPU available. If True, returns the GPUs, the model, physical slots and compute capabilitie(s). diff --git a/simba/data_processors/egocentric_aligner.py b/simba/data_processors/egocentric_aligner.py index 08379cd93..a26576c17 100644 --- a/simba/data_processors/egocentric_aligner.py +++ b/simba/data_processors/egocentric_aligner.py @@ -1,62 +1,15 @@ -import functools -import multiprocessing import os -from typing import List, Optional, Tuple, Union +from typing import Optional, Tuple, Union -import cv2 import numpy as np import pandas as pd -from simba.utils.checks import ( - check_all_file_names_are_represented_in_video_log, check_if_dir_exists, - check_if_valid_rgb_tuple, check_instance, check_int, check_str, - check_valid_boolean, check_valid_dataframe, check_valid_tuple) -from simba.utils.data import egocentrically_align_pose +from simba.utils.checks import (check_if_dir_exists, check_if_valid_rgb_tuple, check_int, check_str, check_valid_dataframe, check_valid_tuple, check_valid_boolean) +from simba.utils.data import egocentrically_align_pose_numba from simba.utils.enums import Formats, Options -from simba.utils.printing import SimbaTimer -from simba.utils.read_write import (bgr_to_rgb_tuple, - concatenate_videos_in_folder, - find_core_cnt, - find_files_of_filetypes_in_directory, - find_video_of_file, get_fn_ext, - get_video_meta_data, read_df, - read_frm_of_video, read_video_info_csv, - remove_a_folder, write_df) -from simba.utils.warnings import FrameRangeWarning - - -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], - fill_clr: Tuple[int, int, int] = (255, 255, 255), - verbose: bool = False): - - video_meta = get_video_meta_data(video_path=video_path) - cap = cv2.VideoCapture(video_path) - batch, frm_range = frm_range[0], frm_range[1] - save_path = os.path.join(temp_dir, f'{batch}.mp4') - fourcc = cv2.VideoWriter_fourcc(*f'{Formats.MP4_CODEC.value}') - writer = cv2.VideoWriter(save_path, fourcc, video_meta['fps'], (video_meta['width'], video_meta['height'])) - - for frm_cnt, frm_id in enumerate(frm_range): - 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']), borderValue=fill_clr) - 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']), borderValue=fill_clr) - writer.write(final_frame) - if verbose: - print(f'Creating frame {frm_id} ({video_name}, CPU core: {batch+1}).') - - writer.release() - return batch+1 +from simba.utils.printing import SimbaTimer, stdout_success +from simba.utils.read_write import (bgr_to_rgb_tuple, find_core_cnt, find_files_of_filetypes_in_directory, find_video_of_file, get_fn_ext, read_df, write_df) +from simba.video_processors.egocentric_video_rotator import EgocentricVideoRotator class EgocentricalAligner(): @@ -73,16 +26,21 @@ class EgocentricalAligner(): :autoplay: :loop: + .. seealso:: + To produce rotation vectors, uses :func:`~simba.utils.data.egocentrically_align_pose_numba` or :func:`~simba.utils.data.egocentrically_align_pose`. + To rotate video only, see :func:`~simba.video_processors.egocentric_video_rotator.EgocentricVideoRotator` + + :param Union[str, os.PathLike] config_path: Path to the configuration file. :param Union[str, os.PathLike] save_dir: Directory where the processed output will be saved. :param Optional[Union[str, os.PathLike]] data_dir: Directory containing CSV files with movement data. :param Optional[str] anchor_1: Primary anchor point (e.g., 'tail_base') around which the alignment centers. :param Optional[str] anchor_2: Secondary anchor point (e.g., 'nose') defining the alignment direction. - :param int direction: Target angle, in degrees, for alignment; e.g., `0` aligns along the x-axis. + :param int direction: Target angle, in degrees, for alignment; e.g., `0` aligns east :param Optional[Tuple[int, int]] anchor_location: Pixel location in the output where `anchor_1` should appear; default is `(250, 250)`. :param Tuple[int, int, int] fill_clr: If rotating the videos, the color of the additional pixels. :param Optional[bool] rotate_video: Whether to rotate the video to align with the specified direction. - :param Optional[int] cores: Number of CPU cores to use for video rotation; `-1` uses all available cores. + :param Optional[int] core_cnt: Number of CPU cores to use for video rotation; `-1` uses all available cores. :example: >>> aligner = EgocentricalAligner(rotate_video=True, anchor_1='tail_base', anchor_2='nose', data_dir=r"/data_dir", videos_dir=r'/videos_dir', save_dir=r"/save_dir", video_info=r"C:\troubleshooting\mitra\project_folder\logs\video_info.csv", direction=0, anchor_location=(250, 250), fill_clr=(0, 0, 0)) @@ -95,13 +53,12 @@ def __init__(self, anchor_1: str = 'tail_base', anchor_2: str = 'nose', direction: int = 0, - anchor_location: Tuple[int, int] = (250, 250), + anchor_location: Union[Tuple[int, int], str] = (250, 250), core_cnt: int = -1, - rotate_video: bool = False, fill_clr: Tuple[int, int, int] = (250, 250, 255), verbose: bool = True, - videos_dir: Optional[Union[str, os.PathLike]] = None, - video_info: Optional[Union[str, os.PathLike, pd.DataFrame]] = None): + gpu: bool = False, + videos_dir: Optional[Union[str, os.PathLike]] = None): self.data_paths = find_files_of_filetypes_in_directory(directory=data_dir, extensions=['.csv']) check_if_dir_exists(in_dir=save_dir, source=f'{self.__class__.__name__} save_dir') @@ -111,26 +68,25 @@ def __init__(self, if core_cnt == -1: self.core_cnt = find_core_cnt()[0] check_int(name=f'{self.__class__.__name__} direction', value=direction, min_value=0, max_value=360) check_valid_tuple(x=anchor_location, source=f'{self.__class__.__name__} anchor_location', accepted_lengths=(2,), valid_dtypes=(int,)) - for i in anchor_location: check_int(name=f'{self.__class__.__name__} anchor_location', value=i, min_value=1) - check_valid_boolean(value=[rotate_video, verbose], source=f'{self.__class__.__name__} rotate_video') - if rotate_video: + check_valid_boolean(value=[gpu], source=f'{self.__class__.__name__} gpu') + for i in anchor_location: + check_int(name=f'{self.__class__.__name__} anchor_location', value=i, min_value=1) + if videos_dir is not None: check_if_valid_rgb_tuple(data=fill_clr) fill_clr = bgr_to_rgb_tuple(value=fill_clr) check_if_dir_exists(in_dir=videos_dir, source=f'{self.__class__.__name__} videos_dir') - check_instance(source=f'{self.__class__.__name__} video_info', accepted_types=(str, pd.DataFrame), instance=video_info) - if isinstance(video_info, str): video_info = read_video_info_csv(file_path=video_info) - else: check_valid_dataframe(df=video_info, source=f'{self.__class__.__name__} video_info', required_fields=Formats.EXPECTED_VIDEO_INFO_COLS.value) self.video_paths = find_files_of_filetypes_in_directory(directory=videos_dir, extensions=Options.ALL_VIDEO_FORMAT_OPTIONS.value) for file_path in self.data_paths: find_video_of_file(video_dir=videos_dir, filename=get_fn_ext(file_path)[1], raise_error=True) - check_all_file_names_are_represented_in_video_log(video_info_df=video_info, data_paths=self.data_paths) self.anchor_1_cols = [f'{anchor_1}_x'.lower(), f'{anchor_1}_y'.lower()] self.anchor_2_cols = [f'{anchor_2}_x'.lower(), f'{anchor_2}_y'.lower()] - self.anchor_1, self.anchor_2, self.videos_dir = anchor_1, anchor_2, videos_dir - self.rotate_video, self.save_dir, self.verbose = rotate_video, save_dir, verbose - self.anchor_location, self.direction, self.fill_clr = np.array(anchor_location), direction, fill_clr + self.anchor_1, self.anchor_2, self.videos_dir = anchor_1.lower(), anchor_2.lower(), videos_dir + self.save_dir, self.verbose, self.gpu = save_dir, verbose, gpu + self.anchor_location, self.direction, self.fill_clr = anchor_location, direction, fill_clr + self.rotate_video = [True if videos_dir is not None else False] def run(self): + timer = SimbaTimer(start=True) for file_cnt, file_path in enumerate(self.data_paths): video_timer = SimbaTimer(start=True) _, self.video_name, _ = get_fn_ext(filepath=file_path) @@ -147,66 +103,55 @@ def run(self): _= [body_parts_lst.append(x[:-2]) for x in bp_cols if x[:-2] not in body_parts_lst] anchor_1_idx, anchor_2_idx = body_parts_lst.index(self.anchor_1), body_parts_lst.index(self.anchor_2) data_arr = df[bp_cols].values.reshape(len(df), len(body_parts_lst), 2).astype(np.int32) - results_arr, self.centers, self.rotation_vectors = egocentrically_align_pose(data=data_arr, anchor_1_idx=anchor_1_idx, anchor_2_idx=anchor_2_idx, direction=self.direction, anchor_location=self.anchor_location) + results_arr, self.centers, self.rotation_vectors = egocentrically_align_pose_numba(data=data_arr, anchor_1_idx=anchor_1_idx, anchor_2_idx=anchor_2_idx, direction=self.direction, anchor_location=np.array(self.anchor_location).astype(np.int64)) 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=Formats.CSV.value, save_path=save_path) video_timer.stop_timer() - print(f'{self.video_name} complete, saved at {save_path} (elapsed time: {video_timer.elapsed_time_str}s)') + if self.verbose: + print(f'{self.video_name} complete, saved at {save_path} (elapsed time: {video_timer.elapsed_time_str}s)') if self.rotate_video: - #self.out_df = self.out_df.head(500) - self.run_video_rotation() - - - - def run_video_rotation(self): - video_timer = SimbaTimer(start=True) - video_path = find_video_of_file(video_dir=self.videos_dir, filename=self.video_name, raise_error=False) - 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.arange(0, 500) - frm_list = np.array_split(frm_list, self.core_cnt) - frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)] - print(f"Creating rotated video {self.video_name}, multiprocessing (chunksize: {1}, cores: {self.core_cnt})...") - with multiprocessing.Pool(self.core_cnt, maxtasksperchild=100) 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, - verbose=self.verbose, - fill_clr=self.fill_clr) - for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=1)): - print(f"Rotate batch {result}/{self.core_cnt} 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 self.verbose: + print(f'Rotating video {self.video_name}...') + video_path = find_video_of_file(video_dir=self.videos_dir, filename=self.video_name, raise_error=False) + save_path = os.path.join(self.save_dir, f'{self.video_name}.mp4') + video_rotator = EgocentricVideoRotator(video_path=video_path, + centers=self.centers, + rotation_vectors=self.rotation_vectors, + anchor_location=self.anchor_location, + verbose=self.verbose, + gpu=self.gpu, + fill_clr=self.fill_clr, + core_cnt=self.core_cnt, + save_path=save_path) + video_rotator.run() + if self.verbose: + print(f'Rotated data for video {self.video_name} ({file_cnt+1}/{len(self.data_paths)}) saved in {self.save_dir}.') + timer.stop_timer() + stdout_success(msg=f'Egocentrically aligned data for {len(self.data_paths)} files saved in {self.save_dir}', elapsed_time=timer.elapsed_time_str) + # if __name__ == "__main__": -# aligner = EgocentricalAligner(rotate_video=True, -# anchor_1='tail_base', +# aligner = EgocentricalAligner(anchor_1='butt/proximal tail', +# anchor_2='snout', +# data_dir=r'C:\troubleshooting\open_field_below\project_folder\csv\outlier_corrected_movement_location', +# videos_dir=r'C:\troubleshooting\open_field_below\project_folder\videos', +# save_dir=r"C:\troubleshooting\open_field_below\project_folder\videos\rotated", +# direction=0, +# gpu=True, +# anchor_location=(600, 300), +# fill_clr=(128,128,128)) +# aligner.run() + +# aligner = EgocentricalAligner(anchor_1='tail_base', # anchor_2='nose', # data_dir=r'C:\Users\sroni\OneDrive\Desktop\rotate_ex\data', # videos_dir=r'C:\Users\sroni\OneDrive\Desktop\rotate_ex\videos', -# save_dir=r"C:\troubleshooting\mitra\project_folder\videos\additional\examples\rotated", -# video_info=r"C:\troubleshooting\mitra\project_folder\logs\video_info.csv", +# save_dir=r"C:\Users\sroni\OneDrive\Desktop\rotate_ex\rotated", # direction=0, +# gpu=True, # anchor_location=(250, 250), # fill_clr=(0, 0, 0)) # aligner.run() diff --git a/simba/labelling/labelling_interface.py b/simba/labelling/labelling_interface.py index 2ef5d9674..d59f562fa 100644 --- a/simba/labelling/labelling_interface.py +++ b/simba/labelling/labelling_interface.py @@ -158,7 +158,7 @@ def __init__(self, self.checkboxes[target] = {} self.checkboxes[target]["name"] = target if self.current_frm_n.get() not in list(self.data_df_targets.index): - raise FrameRangeError(msg=f'Frame data {self.current_frm_n.get()} could not be found in video {self.video_name}. Modify the [Last saved frames] section in the project config if needed.', source=self.__class__.__name__) + raise FrameRangeError(msg=f'Frame pose-estimation data for frame {self.current_frm_n.get()} could not be found for video {self.video_name}. This suggests that the video ({self.video_path}) has a different frame number count than the rows in the data files (e.g., {self.targets_inserted_file_path}, {self.features_extracted_file_path}, {self.machine_results_file_path}, where they exist). Alternatively, modify the [Last saved frames] section in the {self.config_path} if needed.', source=self.__class__.__name__) self.checkboxes[target]["var"] = IntVar(value=self.data_df_targets[target].iloc[self.current_frm_n.get()]) self.checkboxes[target]["cb"] = Checkbutton(self.check_frame, text=target, diff --git a/simba/mixins/train_model_mixin.py b/simba/mixins/train_model_mixin.py index 1ed491de7..c1905b9ac 100644 --- a/simba/mixins/train_model_mixin.py +++ b/simba/mixins/train_model_mixin.py @@ -585,10 +585,11 @@ def create_clf_report(self, save_dir: Union[str, os.PathLike], digits: Optional[int] = 4, clf_name: Optional[str] = None, - img_size: Optional[tuple] = (13.7, 8.27), + img_size: Optional[tuple] = (2500, 4500), #width by height cmap: Optional[str] = "coolwarm", threshold: Optional[int] = 0.5, - save_file_no: Optional[int] = None) -> None: + save_file_no: Optional[int] = None, + dpi: Optional[int] = 300) -> None: """ Create classifier truth table report. @@ -631,18 +632,14 @@ def create_clf_report(self, y_pred = np.where(y_pred > threshold, 1, 0) plt.figure() - clf_report = classification_report(y_true=y_df.values, y_pred=y_pred, target_names=class_names, digits=digits, - output_dict=True, zero_division=0) + clf_report = classification_report(y_true=y_df.values, y_pred=y_pred, target_names=class_names, digits=digits, output_dict=True, zero_division=0) clf_report = pd.DataFrame.from_dict({key: clf_report[key] for key in class_names}) - img = sns.heatmap(pd.DataFrame(clf_report).T, annot=True, cmap=cmap, vmin=0.0, vmax=1.0, linewidth=2.0, - linecolor='black', fmt='g', annot_kws={"size": 20}) + plt.figure(figsize=(round((img_size[1] / dpi), 2), round((img_size[0] / dpi), 2)), dpi=dpi) + img = sns.heatmap(pd.DataFrame(clf_report).T, annot=True, cmap=cmap, vmin=0.0, vmax=1.0, linewidth=2.0, linecolor='black', fmt='g', annot_kws={"size": 20, "weight": "bold", "color": "white", "family": "sans-serif"}) img.set_xticklabels(img.get_xticklabels(), size=16) img.set_yticklabels(img.get_yticklabels(), size=16) - - img.figure.set_size_inches(img_size) - plt.savefig(save_path, dpi=300) + plt.savefig(save_path, dpi=dpi) plt.close("all") - timer.stop_timer() print(f'Classification report saved at {save_path} (elapsed time: {timer.elapsed_time_str}s)') diff --git a/simba/model/train_rf.py b/simba/model/train_rf.py index d208a9ecd..2b88f0bca 100644 --- a/simba/model/train_rf.py +++ b/simba/model/train_rf.py @@ -361,7 +361,7 @@ def save(self) -> None: stdout_success(msg=f"Evaluation files are in models/generated_models/model_evaluations folders", source=self.__class__.__name__) - +# # test = TrainRandomForestClassifier(config_path=r"C:\troubleshooting\mitra\project_folder\project_config.ini") # test.run() # test.save() diff --git a/simba/utils/data.py b/simba/utils/data.py index 05fc03b00..c8ef3c7cc 100644 --- a/simba/utils/data.py +++ b/simba/utils/data.py @@ -1436,6 +1436,10 @@ def egocentrically_align_pose(data: np.ndarray, one anchor point is aligned to a specified location, and the direction between the two anchors is aligned to a target angle. + .. seealso:: + For numba acceleration, see :func:`simba.utils.data.egocentrically_align_pose_numba`. + To align both pose and video, see :func:`simba.data_processors.egocentric_aligner.EgocentricalAligner`. + To egocentrically rotate video, see :func:`simba.video_processors.egocentric_video_rotator.EgocentricVideoRotator` .. video:: _static/img/EgocentricalAligner_2.webm :width: 800 @@ -1492,7 +1496,7 @@ def egocentrically_align_pose(data: np.ndarray, return results, centers, rotation_vectors -@njit("(int32[:, :, :], int64, int64, int64, int32[:])") +@njit("(int32[:, :, :], int64, int64, int64, int64[:])") def egocentrically_align_pose_numba(data: np.ndarray, anchor_1_idx: int, anchor_2_idx: int, @@ -1513,7 +1517,7 @@ def egocentrically_align_pose_numba(data: np.ndarray, .. csv-table:: :header: EXPECTED RUNTIMES - :file: ../../../docs/tables/egocentrically_align_pose_numba.csv + :file: ../../docs/tables/egocentrically_align_pose_numba.csv :widths: 12, 22, 22, 22, 22 :align: center :class: simba-table @@ -1521,7 +1525,8 @@ def egocentrically_align_pose_numba(data: np.ndarray, .. seealso:: For numpy function, see :func:`simba.utils.data.egocentrically_align_pose`. - To align both pose and video, see :func:`simba.data_processors.egocentric_aligner.EgocentricalAligner` + To align both pose and video, see :func:`simba.data_processors.egocentric_aligner.EgocentricalAligner`. + To egocentrically rotate video, see :func:`simba.video_processors.egocentric_video_rotator.EgocentricVideoRotator` :param np.ndarray data: A 3D array of shape `(num_frames, num_points, 2)` containing 2D points for each frame. Each frame is represented as a 2D array of shape `(num_points, 2)`, where each row corresponds to a point's (x, y) coordinates. :param int anchor_1_idx: The index of the first anchor point in `data` used as the center of alignment. This body-part will be placed in the center of the image. @@ -1565,6 +1570,116 @@ def egocentrically_align_pose_numba(data: np.ndarray, return results, centers, rotation_vectors +@jit(nopython=True) +def center_rotation_warpaffine_vectors(rotation_vectors: np.ndarray, centers: np.ndarray): + """ + Create WarpAffine vectors for rotating a video around the center. These are used for egocentric alignment of video. + + .. note:: + `rotation_vectors` and `centers` are returned by :func:`simba.utils.data.egocentrically_align_pose`, or :func:`simba.utils.data.egocentrically_align_pose_numba` + `results are used within e.g., :func:`simba.video_processors.egocentric_video_rotator.EgocentricVideoRotator` + + """ + results = np.full((rotation_vectors.shape[0], 2, 3), fill_value=np.nan, dtype=np.float64) + for idx in range(rotation_vectors.shape[0]): + R, center = rotation_vectors[idx], centers[idx] + top = np.hstack((R[0, :], np.array([-center[0] * R[0, 0] - center[1] * R[0, 1] + center[0]]))) + bottom = np.hstack((R[1, :], np.array([-center[0] * R[1, 0] - center[1] * R[1, 1] + center[1]]))) + results[idx] = np.vstack((top, bottom)) + return results + + +@jit(nopython=True) +def align_target_warpaffine_vectors(centers: np.ndarray, target: np.ndarray): + """ + Create WarpAffine for placing original center at new target position. These are used for egocentric alignment of video. + + .. note:: + `centers` are returned by :func:`simba.utils.data.egocentrically_align_pose`, or :func:`simba.utils.data.egocentrically_align_pose_numba` + `target` in the location in the image where the anchor body-part should be placed. + `results are used within e.g., :func:`simba.video_processors.egocentric_video_rotator.EgocentricVideoRotator` + + """ + results = np.full((centers.shape[0], 2, 3), fill_value=np.nan, dtype=np.float64) + for idx in range(centers.shape[0]): + translation_x = target[0] - centers[idx][0] + translation_y = target[1] - centers[idx][1] + results[idx] = np.array([[1, 0, translation_x], + [0, 1, translation_y]]) + return results + + +@jit(nopython=True, cache=True) +def _bilinear_interpolate(image: np.ndarray, x: int, y: int): + """ + Helper called by :func:`simba.sandbox.warp_numba.egocentric_frm_rotator` Perform bilinear interpolation on an image at fractional coordinates (x, y). Assumes coordinates (x, y) are within the bounds of the image. + """ + x0, y0 = int(np.floor(x)), int(np.floor(y)) + dx, dy = x - x0, y - y0 + if x0 < 0 or x0 + 1 >= image.shape[1] or y0 < 0 or y0 + 1 >= image.shape[0]: + return 0 + I00, I01 = image[y0, x0], image[y0, x0+1] + I10, I11 = image[y0+1, x0], image[y0+1, x0+1] + return (I00 * (1 - dx) * (1 - dy) + I01 * dx * (1 - dy) + I10 * (1 - dx) * dy + I11 * dx * dy) + +@jit(nopython=True) +def egocentric_frm_rotator(frames: np.ndarray, + rotation_matrices: np.ndarray, + interpolate: Optional[bool] = True) -> np.ndarray: + """ + Rotates a sequence of frames using the provided rotation matrices in an egocentric manner using acceleration through numba JIT.. + + Applies a geometric transformation to each frame in the input sequence based on + its corresponding rotation matrix. The transformation includes rotation and translation, + followed by bilinear interpolation to map pixel values from the source frame to the output frame. + + .. note:: + To create rotation matrices, see :func:`simba.utils.data.center_rotation_warpaffine_vectors` and :func:`simba.utils.data.align_target_warpaffine_vectors` + + :param np.ndarray frames: A 4D array of shape (N, H, W, C) + :param np.ndarray rotation_matrices: A 3D array of shape (N, 3, 3), where each 3x3 matrix represents an affine transformation for a corresponding frame. The matrix should include rotation and translation components. + :return: A 4D array of shape (N, H, W, C), representing the warped frames after applying the transformations. The shape matches the input `frames`. + :rtype: np.ndarray + + :example: + >>> DATA_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/data/501_MA142_Gi_Saline_0513.csv" + >>> VIDEO_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/videos/501_MA142_Gi_Saline_0513.mp4" + >>> SAVE_PATH = r"/mnt/c/Users/sroni/OneDrive/Desktop/rotate_ex/videos/501_MA142_Gi_Saline_0513_rotated.mp4" + >>> ANCHOR_LOC = np.array([300, 300]) + >>> + >>> df = read_df(file_path=DATA_PATH, file_type='csv') + >>> bp_cols = [x for x in df.columns if not x.endswith('_p')] + >>> data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int64) + >>> data, centers, rotation_matrices = egocentrically_align_pose(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=180) + >>> imgs = read_img_batch_from_video_gpu(video_path=VIDEO_PATH, start_frm=0, end_frm=100) + >>> imgs = np.stack(list(imgs.values()), axis=0) + >>> + >>> rot_matrices_center = center_rotation_warpaffine_vectors(rotation_vectors=rotation_matrices, centers=centers) + >>> rot_matrices_align = align_target_warpaffine_vectors(centers=centers, target=ANCHOR_LOC) + >>> + >>> imgs_centered = egocentric_frm_rotator(frames=imgs, rotation_matrices=rot_matrices_center) + >>> imgs_out = egocentric_frm_rotator(frames=imgs_centered, rotation_matrices=rot_matrices_align) + """ + + N, H, W, C = frames.shape + warped_frames = np.zeros_like(frames) + for i in prange(N): + frame = frames[i] + rotation_matrix = rotation_matrices[i] + affine_matrix = rotation_matrix[:2, :2] + translation = np.ascontiguousarray(rotation_matrix[:2, 2]) + inverse_affine_matrix = np.ascontiguousarray(np.linalg.inv(affine_matrix)) + inverse_translation = -np.dot(inverse_affine_matrix, translation) + for r in range(H): + for c in range(W): + src_x = inverse_affine_matrix[0, 0] * c + inverse_affine_matrix[0, 1] * r + inverse_translation[0] + src_y = inverse_affine_matrix[1, 0] * c + inverse_affine_matrix[1, 1] * r + inverse_translation[1] + for ch in range(C): + val = _bilinear_interpolate(frame[:, :, ch], src_x, src_y) + warped_frames[i, r, c, ch] = val + return warped_frames + + # run_user_defined_feature_extraction_class(config_path='/Users/simon/Desktop/envs/troubleshooting/circular_features_zebrafish/project_folder/project_config.ini', file_path='/Users/simon/Desktop/fish_feature_extractor_2023_version_5.py') diff --git a/simba/utils/read_write.py b/simba/utils/read_write.py index ffed4839e..4d8bac7bc 100644 --- a/simba/utils/read_write.py +++ b/simba/utils/read_write.py @@ -49,7 +49,7 @@ check_int, check_nvidea_gpu_available, check_str, check_valid_array, check_valid_boolean, check_valid_dataframe, - check_valid_lst) + check_valid_lst, is_video_color) from simba.utils.enums import ConfigKey, Dtypes, Formats, Keys, Options from simba.utils.errors import (DataHeaderError, DuplicationError, FFMPEGCodecGPUError, FileExistError, @@ -464,10 +464,13 @@ def get_video_meta_data(video_path: Union[str, os.PathLike, cv2.VideoCapture], f return video_data -def remove_a_folder(folder_dir: Union[str, os.PathLike]) -> None: +def remove_a_folder(folder_dir: Union[str, os.PathLike], ignore_errors: Optional[bool] = True) -> None: """Helper to remove a directory""" check_if_dir_exists(in_dir=folder_dir, source=remove_a_folder.__name__) - shutil.rmtree(folder_dir, ignore_errors=True) + try: + shutil.rmtree(folder_dir, ignore_errors=ignore_errors) + except Exception as e: + raise PermissionError(msg=f'Could not delete directory: {folder_dir}. is the directory or its content beeing used by anothe process?', source=remove_a_folder.__name__) def concatenate_videos_in_folder(in_folder: Union[str, os.PathLike], @@ -477,7 +480,8 @@ def concatenate_videos_in_folder(in_folder: Union[str, os.PathLike], substring: Optional[str] = None, remove_splits: Optional[bool] = True, gpu: Optional[bool] = False, - fps: Optional[Union[int, str]] = None) -> None: + fps: Optional[Union[int, str]] = None, + verbose: bool = True) -> None: """ Concatenate (temporally) all video files in a folder into a single video. @@ -540,9 +544,11 @@ def concatenate_videos_in_folder(in_folder: Union[str, os.PathLike], if check_nvidea_gpu_available() and gpu: if fps is None: - returned = os.system(f'ffmpeg -hwaccel auto -c:v h264_cuvid -f concat -safe 0 -i "{temp_txt_path}" -c:v h264_nvenc -c:a copy -hide_banner -loglevel info "{save_path}" -y') + returned = os.system(f"ffmpeg -f concat -safe 0 -i \"{temp_txt_path}\" -c:v h264_nvenc -pix_fmt yuv420p -c:a copy -hide_banner -loglevel info \"{save_path}\" -y") + # returned = os.system(f'ffmpeg -hwaccel auto -c:v h264_cuvid -f concat -safe 0 -i "{temp_txt_path}" -c:v h264_nvenc -c:a copy -hide_banner -loglevel info "{save_path}" -y') else: - returned = os.system(f'ffmpeg -hwaccel auto -c:v h264_cuvid -f concat -safe 0 -i "{temp_txt_path}" -r {out_fps} -c:v h264_nvenc -c:a copy -hide_banner -loglevel info "{save_path}" -y') + returned = os.system(f"ffmpeg -f concat -safe 0 -i \"{temp_txt_path}\" -r {out_fps} -c:v h264_nvenc -pix_fmt yuv420p -c:a copy -hide_banner -loglevel info \"{save_path}\" -y") + #returned = os.system(f'ffmpeg -hwaccel auto -c:v h264_cuvid -f concat -safe 0 -i "{temp_txt_path}" -r {out_fps} -c:v h264_nvenc -c:a copy -hide_banner -loglevel info "{save_path}" -y') #returned = os.system(f'ffmpeg -hwaccel cuda -hwaccel_output_format cuda -c:v h264_cuvid -f concat -safe 0 -i "{temp_txt_path}" -vf scale_cuda=1280:720,format=nv12 -r {out_fps} -c:v h264_nvenc -c:a copy -hide_banner -loglevel info "{save_path}" -y') else: if fps is None: @@ -557,11 +563,8 @@ def concatenate_videos_in_folder(in_folder: Union[str, os.PathLike], remove_a_folder(folder_dir=Path(in_folder)) break timer.stop_timer() - stdout_success( - msg="Video concatenated", - elapsed_time=timer.elapsed_time_str, - source=concatenate_videos_in_folder.__name__, - ) + if verbose: + stdout_success(msg="Video concatenated", elapsed_time=timer.elapsed_time_str, source=concatenate_videos_in_folder.__name__) def get_bp_headers(body_parts_lst: List[str]) -> list: @@ -696,7 +699,8 @@ def read_frm_of_video(video_path: Union[str, os.PathLike, cv2.VideoCapture], opacity: Optional[float] = None, size: Optional[Tuple[int, int]] = None, greyscale: Optional[bool] = False, - clahe: Optional[bool] = False) -> np.ndarray: + clahe: Optional[bool] = False, + use_ffmpeg: Optional[bool] = False) -> np.ndarray: """ Reads single image from video file. @@ -721,51 +725,55 @@ def read_frm_of_video(video_path: Union[str, os.PathLike, cv2.VideoCapture], check_file_exist_and_readable(file_path=video_path) video_meta_data = get_video_meta_data(video_path=video_path) else: - video_meta_data = {"frame_count": int(video_path.get(cv2.CAP_PROP_FRAME_COUNT))} + video_meta_data = {"frame_count": int(video_path.get(cv2.CAP_PROP_FRAME_COUNT)), + "fps": video_path.get(cv2.CAP_PROP_FPS), + 'width': int(video_path.get(cv2.CAP_PROP_FRAME_WIDTH)), + 'height': int(video_path.get(cv2.CAP_PROP_FRAME_HEIGHT))} + check_int(name='frame_index', value=frame_index, min_value=-1) - if frame_index == -1: frame_index = video_meta_data["frame_count"] - 1 + if frame_index == -1: + frame_index = video_meta_data["frame_count"] - 1 if (frame_index > video_meta_data["frame_count"]) or (frame_index < 0): raise FrameRangeError(msg=f'Frame {frame_index} is out of range: The video {video_path} contains {video_meta_data["frame_count"]} frames.', source=read_frm_of_video.__name__) - if type(video_path) == str: - capture = cv2.VideoCapture(video_path) + if not use_ffmpeg: + if type(video_path) == str: + capture = cv2.VideoCapture(video_path) + else: + capture = video_path + capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index) + ret, img = capture.read() + if not ret: + raise FrameRangeError(msg=f"Frame {frame_index} for video {video_path} could not be read.") else: - capture = video_path - capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index) - ret, img = capture.read() - if ret: - if opacity: - opacity = float(opacity / 100) - check_float( - name="Opacity", - value=opacity, - min_value=0.00, - max_value=1.00, - raise_error=True, - ) - opacity = 1 - opacity - h, w, clr = img.shape[:3] - opacity_image = np.ones((h, w, clr), dtype=np.uint8) * int(255 * opacity) - img = cv2.addWeighted( - img.astype(np.uint8), - 1 - opacity, - opacity_image.astype(np.uint8), - opacity, - 0, - ) - if size: - img = cv2.resize(img, size, interpolation=cv2.INTER_LINEAR) - if greyscale: - if len(img.shape) > 2: - img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) - if clahe: - if len(img.shape) > 2: - img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) - img = cv2.createCLAHE(clipLimit=2, tileGridSize=(16, 16)).apply(img) + if not isinstance(video_path, str): + raise NoDataError(msg='When using FFMpeg, pass video path', source=read_frm_of_video.__name__) + is_color = is_video_color(video=video_path) + timestamp = frame_index / video_meta_data['fps'] + if is_color: + cmd = f"ffmpeg -hwaccel cuda -ss {timestamp:.10f} -i {video_path} -vframes 1 -f rawvideo -pix_fmt bgr24 -v error -" + result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + img = np.frombuffer(result.stdout, np.uint8).reshape((video_meta_data["height"], video_meta_data["width"], 3)) + else: + cmd = f"ffmpeg -hwaccel cuda -ss {timestamp:.10f} -i {video_path} -vframes 1 -f rawvideo -pix_fmt gray -v error -" + result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + img = np.frombuffer(result.stdout, np.uint8).reshape((video_meta_data["height"], video_meta_data["width"])) + if opacity: + opacity = float(opacity / 100) + check_float(name="Opacity", value=opacity, min_value=0.00, max_value=1.00, raise_error=True) + opacity = 1 - opacity + h, w, clr = img.shape[:3] + opacity_image = np.ones((h, w, clr), dtype=np.uint8) * int(255 * opacity) + img = cv2.addWeighted( img.astype(np.uint8), 1 - opacity, opacity_image.astype(np.uint8), opacity, 0) + if size: + img = cv2.resize(img, size, interpolation=cv2.INTER_LINEAR) + if greyscale: + if len(img.shape) > 2: + img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + if clahe: + if len(img.shape) > 2: + img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + img = cv2.createCLAHE(clipLimit=2, tileGridSize=(16, 16)).apply(img) - else: - NoDataFoundWarning( - msg=f"Frame {frame_index} for video {video_path} could not be read." - ) return img @@ -1540,7 +1548,10 @@ def read_roi_data(roi_path: Union[str, os.PathLike]) -> Tuple[pd.DataFrame, pd.D def create_directory(path: Union[str, os.PathLike]): if not os.path.exists(path): - os.makedirs(path) + try: + os.makedirs(path) + except FileExistsError: + raise PermissionError(msg=f'SimBA is not allowed to create the directory {path}.', source=create_directory.__name__) else: pass @@ -2019,11 +2030,14 @@ def read_data_paths(path: Union[str, os.PathLike, None], def img_stack_to_greyscale(imgs: np.ndarray): """ Jitted conversion of a 4D stack of color images (RGB format) to grayscale. + .. image:: _static/img/img_stack_to_greyscale.png :width: 600 :align: center + :parameter np.ndarray imgs: A 4D array representing color images. It should have the shape (num_images, height, width, 3) where the last dimension represents the color channels (R, G, B). :returns np.ndarray: A 3D array containing the grayscale versions of the input images. The shape of the output array is (num_images, height, width). + :example: >>> imgs = ImageMixin().read_img_batch_from_video( video_path='/Users/simon/Desktop/envs/troubleshooting/two_black_animals_14bp/videos/Together_1.avi', start_frm=0, end_frm=100) >>> imgs = np.stack(list(imgs.values())) @@ -2083,7 +2097,6 @@ def read_img_batch_from_video_gpu(video_path: Union[str, os.PathLike], '-pix_fmt', f'{color_format}', '-'] - ffmpeg_process = subprocess.Popen(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if out_format == 'dict': frames = {} diff --git a/simba/video_processors/egocentric_video_rotator.py b/simba/video_processors/egocentric_video_rotator.py new file mode 100644 index 000000000..fd1f23aaf --- /dev/null +++ b/simba/video_processors/egocentric_video_rotator.py @@ -0,0 +1,179 @@ +import os +from typing import Union, Tuple, Optional +import numpy as np +import functools +import multiprocessing +import cv2 + +from simba.utils.checks import check_if_valid_rgb_tuple, check_valid_boolean, check_int, check_file_exist_and_readable, check_if_dir_exists, check_valid_array, check_valid_tuple +from simba.utils.enums import Formats, Defaults +from simba.utils.read_write import get_video_meta_data, get_fn_ext, find_core_cnt, remove_a_folder, read_frm_of_video, concatenate_videos_in_folder, read_df, read_img_batch_from_video_gpu, create_directory +from simba.utils.printing import SimbaTimer, stdout_success +from simba.utils.data import egocentrically_align_pose, center_rotation_warpaffine_vectors, align_target_warpaffine_vectors + + +def egocentric_video_aligner(frm_range: np.ndarray, + video_path: Union[str, os.PathLike], + temp_dir: Union[str, os.PathLike], + video_name: str, + centers: np.ndarray, + rotation_vectors: np.ndarray, + target: Tuple[int, int], + fill_clr: Tuple[int, int, int] = (255, 255, 255), + verbose: bool = False, + gpu: bool = True): + + video_meta = get_video_meta_data(video_path=video_path) + + batch, frm_range = frm_range[0], frm_range[1] + save_path = os.path.join(temp_dir, f'{batch}.mp4') + fourcc = cv2.VideoWriter_fourcc(*f'{Formats.MP4_CODEC.value}') + writer = cv2.VideoWriter(save_path, fourcc, video_meta['fps'], (video_meta['width'], video_meta['height'])) + batch_rotation_vectors = rotation_vectors[frm_range[0]: frm_range[-1]+1] + batch_centers = centers[frm_range[0]: frm_range[-1]+1] + m_rotates = center_rotation_warpaffine_vectors(rotation_vectors=batch_rotation_vectors, centers=batch_centers) + m_translations = align_target_warpaffine_vectors(centers=batch_centers, target=np.array(target)) + + if gpu: + img_counter = 0 + frm_batches = np.array_split(frm_range, (len(frm_range) + 30 - 1) // 30) + for frm_batch_cnt, frm_ids in enumerate(frm_batches): + frms = read_img_batch_from_video_gpu(video_path=video_path, start_frm=frm_ids[0], end_frm=frm_ids[-1]) + frms = np.stack(list(frms.values()), axis=0) + for img_cnt, img in enumerate(frms): + frame_id = img_counter * (batch+1) + rotated_frame = cv2.warpAffine(img, m_rotates[img_counter], (video_meta['width'], video_meta['height']), borderValue=fill_clr) + final_frame = cv2.warpAffine(rotated_frame, m_translations[img_counter],(video_meta['width'], video_meta['height']), borderValue=fill_clr) + writer.write(final_frame) + if verbose: + print(f'Creating frame {frame_id} ({video_name}, CPU core: {batch + 1}).') + img_counter+=1 + else: + cap = cv2.VideoCapture(video_path) + for frm_idx, frm_id in enumerate(frm_range): + img = read_frm_of_video(video_path=cap, frame_index=frm_id) + rotated_frame = cv2.warpAffine(img, m_rotates[frm_idx], (video_meta['width'], video_meta['height']), borderValue=fill_clr) + final_frame = cv2.warpAffine(rotated_frame, m_translations[frm_idx], (video_meta['width'], video_meta['height']), borderValue=fill_clr) + writer.write(final_frame) + if verbose: + print(f'Creating frame {frm_id} ({video_name}, CPU core: {batch + 1}).') + writer.release() + return batch + 1 + +class EgocentricVideoRotator(): + """ + Perform egocentric rotation of a video using CPU multiprocessing. + + .. video:: _static/img/EgocentricalAligner_2.webm + :width: 800 + :autoplay: + :loop: + + .. seealso:: + To perform joint egocentric alignment of both pose and video, or pose only, use :func:`~simba.data_processors.egocentric_aligner.EgocentricalAligner`. + To produce rotation vectors, use :func:`~simba.utils.data.egocentrically_align_pose_numba` or :func:`~simba.utils.data.egocentrically_align_pose`. + + :param Union[str, os.PathLike] video_path: Path to a video file. + :param np.ndarray centers: A 2D array of shape `(num_frames, 2)` containing the original locations of `anchor_1_idx` in each frame before alignment. Returned by :func:`~simba.utils.data.egocentrically_align_pose_numba` or :func:`~simba.utils.data.egocentrically_align_pose`. + :param np.ndarray rotation_vectors: A 3D array of shape `(num_frames, 2, 2)` containing the rotation matrices applied to each frame. Returned by :func:`~simba.utils.data.egocentrically_align_pose_numba` or :func:`~simba.utils.data.egocentrically_align_pose`. + :param bool verbose: If True, prints progress. Deafult True. + :param Tuple[int, int, int] fill_clr: The color of the additional pixels. Deafult black. (0, 0, 0). + :param int core_cnt: Number of CPU cores to use for video rotation; `-1` uses all available cores. + :param Optional[Union[str, os.PathLike]] save_path: The location where to store the rotated video. If None, saves the video as the same dir as the input video with the `_rotated` suffix. + + :example: + >>> DATA_PATH = "C:\501_MA142_Gi_Saline_0513.csv" + >>> VIDEO_PATH = "C:\501_MA142_Gi_Saline_0513.mp4" + >>> SAVE_PATH = "C:\501_MA142_Gi_Saline_0513_rotated.mp4" + >>> ANCHOR_LOC = np.array([250, 250]) + + >>> df = read_df(file_path=DATA_PATH, file_type='csv') + >>> bp_cols = [x for x in df.columns if not x.endswith('_p')] + >>> data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32) + >>> _, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0) + >>> rotater = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_vectors, anchor_location=ANCHOR_LOC, save_path=SAVE_PATH) + >>> rotater.run() + """ + + def __init__(self, + video_path: Union[str, os.PathLike], + centers: np.ndarray, + rotation_vectors: np.ndarray, + anchor_location: Tuple[int, int], + verbose: bool = True, + fill_clr: Tuple[int, int, int] = (0, 0, 0), + core_cnt: int = -1, + save_path: Optional[Union[str, os.PathLike]] = None, + gpu: Optional[bool] = True): + + check_file_exist_and_readable(file_path=video_path) + self.video_meta_data = get_video_meta_data(video_path=video_path) + check_valid_array(data=centers, source=f'{self.__class__.__name__} centers', accepted_ndims=(2,), accepted_axis_1_shape=[2, ], accepted_axis_0_shape=[self.video_meta_data['frame_count']], accepted_dtypes=Formats.NUMERIC_DTYPES.value) + check_valid_array(data=rotation_vectors, source=f'{self.__class__.__name__} rotation_vectors', accepted_ndims=(3,), accepted_axis_0_shape=[self.video_meta_data['frame_count']], accepted_dtypes=Formats.NUMERIC_DTYPES.value) + check_valid_tuple(x=anchor_location, source=f'{self.__class__.__name__} anchor_location', accepted_lengths=(2,), valid_dtypes=(int,)) + for i in anchor_location: check_int(name=f'{self.__class__.__name__} anchor_location', value=i, min_value=1) + check_valid_boolean(value=[verbose], source=f'{self.__class__.__name__} verbose') + check_if_valid_rgb_tuple(data=fill_clr) + check_int(name=f'{self.__class__.__name__} core_cnt', value=core_cnt, min_value=-1, unaccepted_vals=[0]) + if core_cnt > find_core_cnt()[0] or core_cnt == -1: + self.core_cnt = find_core_cnt()[0] + else: + self.core_cnt = core_cnt + video_dir, self.video_name, _ = get_fn_ext(filepath=video_path) + if save_path is not None: + self.save_dir = os.path.dirname(save_path) + check_if_dir_exists(in_dir=self.save_dir, source=f'{self.__class__.__name__} save_path') + else: + self.save_dir = video_dir + save_path = os.path.join(video_dir, f'{self.video_name}_rotated.mp4') + self.video_path, self.save_path = video_path, save_path + self.centers, self.rotation_vectors, self.gpu = centers, rotation_vectors, gpu + self.verbose, self.fill_clr, self.anchor_loc = verbose, fill_clr, anchor_location + + def run(self): + video_timer = SimbaTimer(start=True) + temp_dir = os.path.join(self.save_dir, 'temp') + if not os.path.isdir(temp_dir): + create_directory(path=temp_dir) + else: + remove_a_folder(folder_dir=temp_dir) + create_directory(path=temp_dir) + frm_list = np.arange(0, self.video_meta_data['frame_count']) + frm_list = np.array_split(frm_list, self.core_cnt) + frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)] + if self.verbose: + print(f"Creating rotated video {self.video_name}, multiprocessing (chunksize: {1}, cores: {self.core_cnt})...") + with multiprocessing.Pool(self.core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value) as pool: + constants = functools.partial(egocentric_video_aligner, + temp_dir=temp_dir, + video_name=self.video_name, + video_path=self.video_path, + centers=self.centers, + rotation_vectors=self.rotation_vectors, + target=self.anchor_loc, + verbose=self.verbose, + fill_clr=self.fill_clr, + gpu=self.gpu) + for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=1)): + if self.verbose: + print(f"Rotate batch {result}/{self.core_cnt} complete...") + pool.terminate() + pool.join() + + concatenate_videos_in_folder(in_folder=temp_dir, save_path=self.save_path, remove_splits=True, gpu=self.gpu, verbose=self.verbose) + video_timer.stop_timer() + stdout_success(msg=f"Egocentric rotation video {self.save_path} complete", elapsed_time=video_timer.elapsed_time_str, source=self.__class__.__name__) + +if __name__ == "__main__": + DATA_PATH = r"C:\Users\sroni\OneDrive\Desktop\rotate_ex\data\501_MA142_Gi_Saline_0513.csv" + VIDEO_PATH = r"C:\Users\sroni\OneDrive\Desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513.mp4" + SAVE_PATH = r"C:\Users\sroni\OneDrive\Desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513_rotated.mp4" + ANCHOR_LOC = np.array([250, 250]) + + df = read_df(file_path=DATA_PATH, file_type='csv') + bp_cols = [x for x in df.columns if not x.endswith('_p')] + data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32) + + _, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=5, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0) + rotater = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_vectors, anchor_location=(400, 100), save_path=SAVE_PATH, verbose=True) + rotater.run()