import functools
import multiprocessing
import os
from typing import Optional, Tuple, Union
import cv2
import numpy as np
from simba.utils.checks import (check_file_exist_and_readable,
check_if_dir_exists, check_if_valid_rgb_tuple,
check_int, check_valid_array,
check_valid_boolean, check_valid_cpu_pool,
check_valid_tuple)
from simba.utils.data import (align_target_warpaffine_vectors,
center_rotation_warpaffine_vectors,
egocentrically_align_pose, get_cpu_pool,
terminate_cpu_pool)
from simba.utils.enums import Defaults, Formats
from simba.utils.printing import SimbaTimer, stdout_success
from simba.utils.read_write import (concatenate_videos_in_folder,
create_directory, find_core_cnt,
get_fn_ext, get_video_meta_data, read_df,
read_frm_of_video,
read_img_batch_from_video_gpu,
remove_a_folder)
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_meta["frame_count"]} ({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_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
writer.release()
return batch + 1
[docs]class EgocentricVideoRotator():
"""
Perform egocentric rotation of a video using CPU multiprocessing.
.. video:: _static/img/EgocentricalAligner_2.webm
:width: 800
:autoplay:
:loop:
:muted:
:align: center
.. 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,
pool: bool = None):
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
if pool is not None:
check_valid_cpu_pool(value=pool, source=self.__class__.__name__, max_cores=find_core_cnt()[0], min_cores=2, raise_error=True)
self.pool_termination_flag = True
else:
self.pool_termination_flag = False
self.pool = get_cpu_pool(core_cnt=self.core_cnt, source=self.__class__.__name__) if pool is None else pool
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(paths=temp_dir)
else:
remove_a_folder(folder_dir=temp_dir)
create_directory(paths=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})...")
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(self.pool.imap(constants, frm_list, chunksize=1)):
if self.verbose: print(f"Rotate batch {result}/{self.core_cnt} complete...")
if self.pool_termination_flag: terminate_cpu_pool(pool=self.pool, force=False)
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\desktop\rotate_ex\data\501_MA142_Gi_Saline_0513.csv"
# VIDEO_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513.mp4"
# SAVE_PATH = r"C:\Users\sroni\OneDrive\Desktop\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, core_cnt=16)
# rotater.run()