simba-uw-tf-dev 4.6.4__py3-none-any.whl → 4.6.6__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- simba/data_processors/cuda/geometry.py +45 -27
- simba/data_processors/cuda/image.py +1620 -1600
- simba/data_processors/cuda/statistics.py +17 -9
- simba/data_processors/egocentric_aligner.py +24 -6
- simba/data_processors/kleinberg_calculator.py +6 -2
- simba/feature_extractors/feature_subsets.py +12 -5
- simba/feature_extractors/straub_tail_analyzer.py +0 -2
- simba/mixins/statistics_mixin.py +9 -2
- simba/sandbox/analyze_runtimes.py +30 -0
- simba/sandbox/cuda/egocentric_rotator.py +374 -374
- simba/sandbox/proboscis_to_tip.py +28 -0
- simba/sandbox/test_directionality.py +47 -0
- simba/sandbox/test_nonstatic_directionality.py +27 -0
- simba/sandbox/test_pycharm_cuda.py +51 -0
- simba/sandbox/test_simba_install.py +41 -0
- simba/sandbox/test_static_directionality.py +26 -0
- simba/sandbox/test_static_directionality_2d.py +26 -0
- simba/sandbox/verify_env.py +42 -0
- simba/ui/pop_ups/fsttc_pop_up.py +27 -25
- simba/ui/pop_ups/kleinberg_pop_up.py +3 -2
- simba/utils/data.py +0 -1
- simba/utils/errors.py +441 -440
- simba/utils/lookups.py +1203 -1203
- simba/utils/read_write.py +38 -13
- simba/video_processors/egocentric_video_rotator.py +41 -36
- simba/video_processors/video_processing.py +5247 -5233
- simba/video_processors/videos_to_frames.py +41 -31
- {simba_uw_tf_dev-4.6.4.dist-info → simba_uw_tf_dev-4.6.6.dist-info}/METADATA +2 -2
- {simba_uw_tf_dev-4.6.4.dist-info → simba_uw_tf_dev-4.6.6.dist-info}/RECORD +33 -24
- {simba_uw_tf_dev-4.6.4.dist-info → simba_uw_tf_dev-4.6.6.dist-info}/LICENSE +0 -0
- {simba_uw_tf_dev-4.6.4.dist-info → simba_uw_tf_dev-4.6.6.dist-info}/WHEEL +0 -0
- {simba_uw_tf_dev-4.6.4.dist-info → simba_uw_tf_dev-4.6.6.dist-info}/entry_points.txt +0 -0
- {simba_uw_tf_dev-4.6.4.dist-info → simba_uw_tf_dev-4.6.6.dist-info}/top_level.txt +0 -0
|
@@ -1,374 +1,374 @@
|
|
|
1
|
-
import functools
|
|
2
|
-
import multiprocessing
|
|
3
|
-
import os
|
|
4
|
-
from typing import Optional, Tuple, Union
|
|
5
|
-
|
|
6
|
-
import cv2
|
|
7
|
-
import numpy as np
|
|
8
|
-
|
|
9
|
-
try:
|
|
10
|
-
import cupy as cp
|
|
11
|
-
from cupyx.scipy.ndimage import affine_transform
|
|
12
|
-
CUPY_AVAILABLE = True
|
|
13
|
-
except ImportError:
|
|
14
|
-
import numpy as cp
|
|
15
|
-
from scipy.ndimage import affine_transform
|
|
16
|
-
CUPY_AVAILABLE = False
|
|
17
|
-
|
|
18
|
-
from simba.utils.checks import (check_file_exist_and_readable,
|
|
19
|
-
check_if_dir_exists, check_if_valid_rgb_tuple,
|
|
20
|
-
check_int, check_valid_array,
|
|
21
|
-
check_valid_boolean, check_valid_tuple)
|
|
22
|
-
from simba.utils.data import (align_target_warpaffine_vectors,
|
|
23
|
-
center_rotation_warpaffine_vectors,
|
|
24
|
-
egocentrically_align_pose)
|
|
25
|
-
from simba.utils.enums import Defaults, Formats
|
|
26
|
-
from simba.utils.printing import SimbaTimer, stdout_success
|
|
27
|
-
from simba.utils.read_write import (concatenate_videos_in_folder,
|
|
28
|
-
create_directory, find_core_cnt,
|
|
29
|
-
get_fn_ext, get_video_meta_data, read_df,
|
|
30
|
-
read_frm_of_video,
|
|
31
|
-
read_img_batch_from_video_gpu,
|
|
32
|
-
remove_a_folder,
|
|
33
|
-
_read_img_batch_from_video_helper)
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
def egocentric_video_aligner(frm_range: np.ndarray,
|
|
37
|
-
video_path: Union[str, os.PathLike],
|
|
38
|
-
temp_dir: Union[str, os.PathLike],
|
|
39
|
-
video_name: str,
|
|
40
|
-
centers: np.ndarray,
|
|
41
|
-
rotation_vectors: np.ndarray,
|
|
42
|
-
target: Tuple[int, int],
|
|
43
|
-
fill_clr: Tuple[int, int, int] = (255, 255, 255),
|
|
44
|
-
verbose: bool = False,
|
|
45
|
-
gpu: bool = True):
|
|
46
|
-
|
|
47
|
-
video_meta = get_video_meta_data(video_path=video_path)
|
|
48
|
-
|
|
49
|
-
batch, frm_range = frm_range[0], frm_range[1]
|
|
50
|
-
save_path = os.path.join(temp_dir, f'{batch}.mp4')
|
|
51
|
-
fourcc = cv2.VideoWriter_fourcc(*f'{Formats.MP4_CODEC.value}')
|
|
52
|
-
writer = cv2.VideoWriter(save_path, fourcc, video_meta['fps'], (video_meta['width'], video_meta['height']))
|
|
53
|
-
batch_rotation_vectors = rotation_vectors[frm_range[0]: frm_range[-1]+1]
|
|
54
|
-
batch_centers = centers[frm_range[0]: frm_range[-1]+1]
|
|
55
|
-
m_rotates = center_rotation_warpaffine_vectors(rotation_vectors=batch_rotation_vectors, centers=batch_centers)
|
|
56
|
-
m_translations = align_target_warpaffine_vectors(centers=batch_centers, target=np.array(target))
|
|
57
|
-
|
|
58
|
-
if gpu:
|
|
59
|
-
# Combine rotation and translation matrices into single transform
|
|
60
|
-
# This reduces two sequential operations to one
|
|
61
|
-
batch_size = len(frm_range)
|
|
62
|
-
m_combined = np.zeros((batch_size, 2, 3), dtype=np.float32)
|
|
63
|
-
|
|
64
|
-
for i in range(batch_size):
|
|
65
|
-
# Convert rotation matrix (2x3) to 3x3 homogeneous
|
|
66
|
-
m_rot_3x3 = np.eye(3, dtype=np.float32)
|
|
67
|
-
m_rot_3x3[:2, :] = m_rotates[i].astype(np.float32)
|
|
68
|
-
|
|
69
|
-
# Convert translation matrix (2x3) to 3x3 homogeneous
|
|
70
|
-
m_trans_3x3 = np.eye(3, dtype=np.float32)
|
|
71
|
-
m_trans_3x3[:2, :] = m_translations[i].astype(np.float32)
|
|
72
|
-
|
|
73
|
-
# Combine: translation after rotation (matches sequential cv2.warpAffine order)
|
|
74
|
-
m_combined_3x3 = m_trans_3x3 @ m_rot_3x3
|
|
75
|
-
|
|
76
|
-
# Convert back to 2x3 for warpAffine compatibility
|
|
77
|
-
m_combined[i] = m_combined_3x3[:2, :]
|
|
78
|
-
|
|
79
|
-
# Process frames in batches using GPU reading
|
|
80
|
-
# Use same batch size as original (30) for optimal I/O overlap
|
|
81
|
-
# Main optimization: combined matrix (one warpAffine instead of two)
|
|
82
|
-
img_counter = 0
|
|
83
|
-
frm_batches = np.array_split(frm_range, (len(frm_range) + 30 - 1) // 30)
|
|
84
|
-
for frm_batch_cnt, frm_ids in enumerate(frm_batches):
|
|
85
|
-
frms = read_img_batch_from_video_gpu(video_path=video_path, start_frm=frm_ids[0], end_frm=frm_ids[-1], verbose=False)
|
|
86
|
-
frms = np.stack(list(frms.values()), axis=0)
|
|
87
|
-
for img_cnt, img in enumerate(frms):
|
|
88
|
-
# Use combined matrix for single warpAffine (faster than two separate calls)
|
|
89
|
-
m = m_combined[img_counter].astype(np.float32)
|
|
90
|
-
final_frame = cv2.warpAffine(img, m, (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
91
|
-
writer.write(final_frame)
|
|
92
|
-
if verbose:
|
|
93
|
-
frame_id = frm_ids[img_cnt]
|
|
94
|
-
print(f'Creating frame {frame_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
95
|
-
img_counter += 1
|
|
96
|
-
|
|
97
|
-
# Legacy CuPy code (commented out - CPU is faster for this use case)
|
|
98
|
-
if False and CUPY_AVAILABLE:
|
|
99
|
-
# Pre-compute all inverse matrices upfront (much faster than per-frame)
|
|
100
|
-
# For CuPy affine_transform, we need inverse matrices
|
|
101
|
-
m_inv_matrices = []
|
|
102
|
-
m_offsets = []
|
|
103
|
-
for i in range(batch_size):
|
|
104
|
-
m = m_combined[i]
|
|
105
|
-
matrix_2x2 = m[:2, :2].astype(np.float32)
|
|
106
|
-
offset = m[:2, 2].astype(np.float32)
|
|
107
|
-
m_inv_matrices.append(cp.asarray(matrix_2x2))
|
|
108
|
-
m_offsets.append(cp.asarray(offset))
|
|
109
|
-
# Batch invert all matrices at once
|
|
110
|
-
m_inv_matrices_gpu = cp.stack(m_inv_matrices)
|
|
111
|
-
m_inv_matrices_gpu = cp.linalg.inv(m_inv_matrices_gpu)
|
|
112
|
-
m_offsets_gpu = cp.stack(m_offsets)
|
|
113
|
-
|
|
114
|
-
# Create async reader for GPU
|
|
115
|
-
async_reader = AsyncVideoFrameReader(
|
|
116
|
-
video_path=video_path,
|
|
117
|
-
batch_size=batch_size_gpu,
|
|
118
|
-
max_que_size=3,
|
|
119
|
-
start_idx=frm_range[0],
|
|
120
|
-
end_idx=frm_range[-1] + 1,
|
|
121
|
-
gpu=True, # Use GPU reading
|
|
122
|
-
verbose=False
|
|
123
|
-
)
|
|
124
|
-
async_reader.start()
|
|
125
|
-
|
|
126
|
-
# Process batches as they become available from async reader
|
|
127
|
-
# Batch process and transfer to minimize GPU->CPU overhead
|
|
128
|
-
processed_frames_batch = []
|
|
129
|
-
frame_ids_batch = []
|
|
130
|
-
|
|
131
|
-
while True:
|
|
132
|
-
batch_result = get_async_frame_batch(batch_reader=async_reader, timeout=10)
|
|
133
|
-
if batch_result is None:
|
|
134
|
-
# Write any remaining frames
|
|
135
|
-
if processed_frames_batch:
|
|
136
|
-
for frame in processed_frames_batch:
|
|
137
|
-
writer.write(frame)
|
|
138
|
-
break
|
|
139
|
-
|
|
140
|
-
start_idx, end_idx, frms = batch_result
|
|
141
|
-
batch_len = end_idx - start_idx + 1
|
|
142
|
-
frms_gpu = cp.asarray(frms)
|
|
143
|
-
|
|
144
|
-
# Process all frames in batch on GPU first
|
|
145
|
-
batch_transformed = []
|
|
146
|
-
batch_frame_indices = []
|
|
147
|
-
|
|
148
|
-
for i in range(batch_len):
|
|
149
|
-
# Map frame index from video to frm_range index
|
|
150
|
-
frame_id = start_idx + i
|
|
151
|
-
frame_idx_in_range = np.where(frm_range == frame_id)[0]
|
|
152
|
-
if len(frame_idx_in_range) == 0:
|
|
153
|
-
continue
|
|
154
|
-
frame_idx_in_range = frame_idx_in_range[0]
|
|
155
|
-
batch_frame_indices.append((i, frame_idx_in_range))
|
|
156
|
-
|
|
157
|
-
# Process all frames in this batch on GPU
|
|
158
|
-
for i, frame_idx_in_range in batch_frame_indices:
|
|
159
|
-
img_gpu = frms_gpu[i]
|
|
160
|
-
matrix_inv = m_inv_matrices_gpu[frame_idx_in_range]
|
|
161
|
-
offset = m_offsets_gpu[frame_idx_in_range]
|
|
162
|
-
|
|
163
|
-
if len(img_gpu.shape) == 3: # Multi-channel
|
|
164
|
-
transformed_channels = []
|
|
165
|
-
for c in range(img_gpu.shape[2]):
|
|
166
|
-
transformed_ch = affine_transform(
|
|
167
|
-
img_gpu[:, :, c],
|
|
168
|
-
matrix=matrix_inv,
|
|
169
|
-
offset=offset,
|
|
170
|
-
output_shape=(video_meta['height'], video_meta['width']),
|
|
171
|
-
order=1,
|
|
172
|
-
mode='constant',
|
|
173
|
-
cval=fill_clr[c] if c < len(fill_clr) else fill_clr[0],
|
|
174
|
-
prefilter=False
|
|
175
|
-
)
|
|
176
|
-
transformed_channels.append(transformed_ch)
|
|
177
|
-
transformed = cp.stack(transformed_channels, axis=2)
|
|
178
|
-
else: # Single channel
|
|
179
|
-
transformed = affine_transform(
|
|
180
|
-
img_gpu,
|
|
181
|
-
matrix=matrix_inv,
|
|
182
|
-
offset=offset,
|
|
183
|
-
output_shape=(video_meta['height'], video_meta['width']),
|
|
184
|
-
order=1,
|
|
185
|
-
mode='constant',
|
|
186
|
-
cval=fill_clr[0] if len(fill_clr) > 0 else 0,
|
|
187
|
-
prefilter=False
|
|
188
|
-
)
|
|
189
|
-
batch_transformed.append(transformed)
|
|
190
|
-
|
|
191
|
-
# Batch transfer all frames from GPU to CPU at once
|
|
192
|
-
if batch_transformed:
|
|
193
|
-
# Stack all transformed frames and transfer in one go
|
|
194
|
-
batch_transformed_stack = cp.stack(batch_transformed)
|
|
195
|
-
batch_cpu = cp.asnumpy(batch_transformed_stack).astype(np.uint8)
|
|
196
|
-
|
|
197
|
-
# Write all frames from this batch
|
|
198
|
-
for frame_idx, (i, frame_idx_in_range) in enumerate(batch_frame_indices):
|
|
199
|
-
final_frame = batch_cpu[frame_idx]
|
|
200
|
-
writer.write(final_frame)
|
|
201
|
-
|
|
202
|
-
if verbose:
|
|
203
|
-
frame_id = start_idx + i
|
|
204
|
-
print(f'Creating frame {frame_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
205
|
-
|
|
206
|
-
async_reader.kill()
|
|
207
|
-
|
|
208
|
-
else:
|
|
209
|
-
# Fallback to CPU with combined matrix and batch reading
|
|
210
|
-
# Process frames in batches
|
|
211
|
-
# Use helper function directly to avoid nested multiprocessing (we're already in a worker process)
|
|
212
|
-
# Larger batch size reduces overhead
|
|
213
|
-
batch_size_gpu = 500
|
|
214
|
-
frm_batches = np.array_split(frm_range, (len(frm_range) + batch_size_gpu - 1) // batch_size_gpu)
|
|
215
|
-
|
|
216
|
-
# Create a mapping from frame_id to index in frm_range for fast lookup
|
|
217
|
-
frm_id_to_idx = {frame_id: idx for idx, frame_id in enumerate(frm_range)}
|
|
218
|
-
|
|
219
|
-
for frm_batch_cnt, frm_ids in enumerate(frm_batches):
|
|
220
|
-
# Read batch of frames directly using helper (no multiprocessing)
|
|
221
|
-
frm_idx_array = np.array(frm_ids)
|
|
222
|
-
frms_dict = _read_img_batch_from_video_helper(
|
|
223
|
-
frm_idx=frm_idx_array,
|
|
224
|
-
video_path=video_path,
|
|
225
|
-
greyscale=False,
|
|
226
|
-
verbose=False,
|
|
227
|
-
black_and_white=False,
|
|
228
|
-
clahe=False
|
|
229
|
-
)
|
|
230
|
-
frms = np.stack([frms_dict[f] for f in frm_ids], axis=0)
|
|
231
|
-
|
|
232
|
-
# Process all frames in batch using optimized CPU cv2.warpAffine with combined matrices
|
|
233
|
-
for i, frame_id in enumerate(frm_ids):
|
|
234
|
-
# Fast dictionary lookup instead of np.where
|
|
235
|
-
frame_idx_in_range = frm_id_to_idx.get(frame_id)
|
|
236
|
-
if frame_idx_in_range is None:
|
|
237
|
-
continue
|
|
238
|
-
|
|
239
|
-
img = frms[i]
|
|
240
|
-
m = m_combined[frame_idx_in_range].astype(np.float32)
|
|
241
|
-
final_frame = cv2.warpAffine(img, m, (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
242
|
-
writer.write(final_frame)
|
|
243
|
-
|
|
244
|
-
if verbose:
|
|
245
|
-
print(f'Creating frame {frame_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
246
|
-
else:
|
|
247
|
-
cap = cv2.VideoCapture(video_path)
|
|
248
|
-
for frm_idx, frm_id in enumerate(frm_range):
|
|
249
|
-
img = read_frm_of_video(video_path=cap, frame_index=frm_id)
|
|
250
|
-
rotated_frame = cv2.warpAffine(img, m_rotates[frm_idx], (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
251
|
-
final_frame = cv2.warpAffine(rotated_frame, m_translations[frm_idx], (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
252
|
-
writer.write(final_frame)
|
|
253
|
-
if verbose:
|
|
254
|
-
print(f'Creating frame {frm_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
255
|
-
writer.release()
|
|
256
|
-
return batch + 1
|
|
257
|
-
|
|
258
|
-
class EgocentricVideoRotator():
|
|
259
|
-
"""
|
|
260
|
-
Perform egocentric rotation of a video using CPU multiprocessing.
|
|
261
|
-
|
|
262
|
-
.. video:: _static/img/EgocentricalAligner_2.webm
|
|
263
|
-
:width: 800
|
|
264
|
-
:autoplay:
|
|
265
|
-
:loop:
|
|
266
|
-
|
|
267
|
-
.. seealso::
|
|
268
|
-
To perform joint egocentric alignment of both pose and video, or pose only, use :func:`~simba.data_processors.egocentric_aligner.EgocentricalAligner`.
|
|
269
|
-
To produce rotation vectors, use :func:`~simba.utils.data.egocentrically_align_pose_numba` or :func:`~simba.utils.data.egocentrically_align_pose`.
|
|
270
|
-
|
|
271
|
-
:param Union[str, os.PathLike] video_path: Path to a video file.
|
|
272
|
-
: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`.
|
|
273
|
-
: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`.
|
|
274
|
-
:param bool verbose: If True, prints progress. Deafult True.
|
|
275
|
-
:param Tuple[int, int, int] fill_clr: The color of the additional pixels. Deafult black. (0, 0, 0).
|
|
276
|
-
:param int core_cnt: Number of CPU cores to use for video rotation; `-1` uses all available cores.
|
|
277
|
-
: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.
|
|
278
|
-
|
|
279
|
-
:example:
|
|
280
|
-
>>> DATA_PATH = "C:\501_MA142_Gi_Saline_0513.csv"
|
|
281
|
-
>>> VIDEO_PATH = "C:\501_MA142_Gi_Saline_0513.mp4"
|
|
282
|
-
>>> SAVE_PATH = "C:\501_MA142_Gi_Saline_0513_rotated.mp4"
|
|
283
|
-
>>> ANCHOR_LOC = np.array([250, 250])
|
|
284
|
-
|
|
285
|
-
>>> df = read_df(file_path=DATA_PATH, file_type='csv')
|
|
286
|
-
>>> bp_cols = [x for x in df.columns if not x.endswith('_p')]
|
|
287
|
-
>>> data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32)
|
|
288
|
-
>>> _, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0)
|
|
289
|
-
>>> rotater = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_vectors, anchor_location=ANCHOR_LOC, save_path=SAVE_PATH)
|
|
290
|
-
>>> rotater.run()
|
|
291
|
-
"""
|
|
292
|
-
|
|
293
|
-
def __init__(self,
|
|
294
|
-
video_path: Union[str, os.PathLike],
|
|
295
|
-
centers: np.ndarray,
|
|
296
|
-
rotation_vectors: np.ndarray,
|
|
297
|
-
anchor_location: Tuple[int, int],
|
|
298
|
-
verbose: bool = True,
|
|
299
|
-
fill_clr: Tuple[int, int, int] = (0, 0, 0),
|
|
300
|
-
core_cnt: int = -1,
|
|
301
|
-
save_path: Optional[Union[str, os.PathLike]] = None,
|
|
302
|
-
gpu: Optional[bool] = True):
|
|
303
|
-
|
|
304
|
-
check_file_exist_and_readable(file_path=video_path)
|
|
305
|
-
self.video_meta_data = get_video_meta_data(video_path=video_path)
|
|
306
|
-
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)
|
|
307
|
-
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)
|
|
308
|
-
check_valid_tuple(x=anchor_location, source=f'{self.__class__.__name__} anchor_location', accepted_lengths=(2,), valid_dtypes=(int,))
|
|
309
|
-
for i in anchor_location: check_int(name=f'{self.__class__.__name__} anchor_location', value=i, min_value=1)
|
|
310
|
-
check_valid_boolean(value=[verbose], source=f'{self.__class__.__name__} verbose')
|
|
311
|
-
check_if_valid_rgb_tuple(data=fill_clr)
|
|
312
|
-
check_int(name=f'{self.__class__.__name__} core_cnt', value=core_cnt, min_value=-1, unaccepted_vals=[0])
|
|
313
|
-
if core_cnt > find_core_cnt()[0] or core_cnt == -1:
|
|
314
|
-
self.core_cnt = find_core_cnt()[0]
|
|
315
|
-
else:
|
|
316
|
-
self.core_cnt = core_cnt
|
|
317
|
-
video_dir, self.video_name, _ = get_fn_ext(filepath=video_path)
|
|
318
|
-
if save_path is not None:
|
|
319
|
-
self.save_dir = os.path.dirname(save_path)
|
|
320
|
-
check_if_dir_exists(in_dir=self.save_dir, source=f'{self.__class__.__name__} save_path')
|
|
321
|
-
else:
|
|
322
|
-
self.save_dir = video_dir
|
|
323
|
-
save_path = os.path.join(video_dir, f'{self.video_name}_rotated.mp4')
|
|
324
|
-
self.video_path, self.save_path = video_path, save_path
|
|
325
|
-
self.centers, self.rotation_vectors, self.gpu = centers, rotation_vectors, gpu
|
|
326
|
-
self.verbose, self.fill_clr, self.anchor_loc = verbose, fill_clr, anchor_location
|
|
327
|
-
|
|
328
|
-
def run(self):
|
|
329
|
-
video_timer = SimbaTimer(start=True)
|
|
330
|
-
temp_dir = os.path.join(self.save_dir, 'temp')
|
|
331
|
-
if not os.path.isdir(temp_dir):
|
|
332
|
-
create_directory(paths=temp_dir)
|
|
333
|
-
else:
|
|
334
|
-
remove_a_folder(folder_dir=temp_dir)
|
|
335
|
-
create_directory(paths=temp_dir)
|
|
336
|
-
frm_list = np.arange(0, self.video_meta_data['frame_count'])
|
|
337
|
-
frm_list = np.array_split(frm_list, self.core_cnt)
|
|
338
|
-
frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)]
|
|
339
|
-
if self.verbose:
|
|
340
|
-
print(f"Creating rotated video {self.video_name}, multiprocessing (chunksize: {1}, cores: {self.core_cnt})...")
|
|
341
|
-
with multiprocessing.Pool(self.core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value) as pool:
|
|
342
|
-
constants = functools.partial(egocentric_video_aligner,
|
|
343
|
-
temp_dir=temp_dir,
|
|
344
|
-
video_name=self.video_name,
|
|
345
|
-
video_path=self.video_path,
|
|
346
|
-
centers=self.centers,
|
|
347
|
-
rotation_vectors=self.rotation_vectors,
|
|
348
|
-
target=self.anchor_loc,
|
|
349
|
-
verbose=self.verbose,
|
|
350
|
-
fill_clr=self.fill_clr,
|
|
351
|
-
gpu=self.gpu)
|
|
352
|
-
for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=1)):
|
|
353
|
-
if self.verbose:
|
|
354
|
-
print(f"Rotate batch {result}/{self.core_cnt} complete...")
|
|
355
|
-
pool.terminate()
|
|
356
|
-
pool.join()
|
|
357
|
-
|
|
358
|
-
concatenate_videos_in_folder(in_folder=temp_dir, save_path=self.save_path, remove_splits=True, gpu=self.gpu, verbose=self.verbose)
|
|
359
|
-
video_timer.stop_timer()
|
|
360
|
-
stdout_success(msg=f"Egocentric rotation video {self.save_path} complete", elapsed_time=video_timer.elapsed_time_str, source=self.__class__.__name__)
|
|
361
|
-
|
|
362
|
-
if __name__ == "__main__":
|
|
363
|
-
DATA_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\data\501_MA142_Gi_Saline_0513.csv"
|
|
364
|
-
VIDEO_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513.mp4"
|
|
365
|
-
SAVE_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513_rotated.mp4"
|
|
366
|
-
ANCHOR_LOC = np.array([250, 250])
|
|
367
|
-
|
|
368
|
-
df = read_df(file_path=DATA_PATH, file_type='csv')
|
|
369
|
-
bp_cols = [x for x in df.columns if not x.endswith('_p')]
|
|
370
|
-
data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32)
|
|
371
|
-
|
|
372
|
-
_, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=5, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0)
|
|
373
|
-
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, gpu=True)
|
|
374
|
-
rotater.run()
|
|
1
|
+
import functools
|
|
2
|
+
import multiprocessing
|
|
3
|
+
import os
|
|
4
|
+
from typing import Optional, Tuple, Union
|
|
5
|
+
|
|
6
|
+
import cv2
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
try:
|
|
10
|
+
import cupy as cp
|
|
11
|
+
from cupyx.scipy.ndimage import affine_transform
|
|
12
|
+
CUPY_AVAILABLE = True
|
|
13
|
+
except ImportError:
|
|
14
|
+
import numpy as cp
|
|
15
|
+
from scipy.ndimage import affine_transform
|
|
16
|
+
CUPY_AVAILABLE = False
|
|
17
|
+
|
|
18
|
+
from simba.utils.checks import (check_file_exist_and_readable,
|
|
19
|
+
check_if_dir_exists, check_if_valid_rgb_tuple,
|
|
20
|
+
check_int, check_valid_array,
|
|
21
|
+
check_valid_boolean, check_valid_tuple)
|
|
22
|
+
from simba.utils.data import (align_target_warpaffine_vectors,
|
|
23
|
+
center_rotation_warpaffine_vectors,
|
|
24
|
+
egocentrically_align_pose)
|
|
25
|
+
from simba.utils.enums import Defaults, Formats
|
|
26
|
+
from simba.utils.printing import SimbaTimer, stdout_success
|
|
27
|
+
from simba.utils.read_write import (concatenate_videos_in_folder,
|
|
28
|
+
create_directory, find_core_cnt,
|
|
29
|
+
get_fn_ext, get_video_meta_data, read_df,
|
|
30
|
+
read_frm_of_video,
|
|
31
|
+
read_img_batch_from_video_gpu,
|
|
32
|
+
remove_a_folder,
|
|
33
|
+
_read_img_batch_from_video_helper)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def egocentric_video_aligner(frm_range: np.ndarray,
|
|
37
|
+
video_path: Union[str, os.PathLike],
|
|
38
|
+
temp_dir: Union[str, os.PathLike],
|
|
39
|
+
video_name: str,
|
|
40
|
+
centers: np.ndarray,
|
|
41
|
+
rotation_vectors: np.ndarray,
|
|
42
|
+
target: Tuple[int, int],
|
|
43
|
+
fill_clr: Tuple[int, int, int] = (255, 255, 255),
|
|
44
|
+
verbose: bool = False,
|
|
45
|
+
gpu: bool = True):
|
|
46
|
+
|
|
47
|
+
video_meta = get_video_meta_data(video_path=video_path)
|
|
48
|
+
|
|
49
|
+
batch, frm_range = frm_range[0], frm_range[1]
|
|
50
|
+
save_path = os.path.join(temp_dir, f'{batch}.mp4')
|
|
51
|
+
fourcc = cv2.VideoWriter_fourcc(*f'{Formats.MP4_CODEC.value}')
|
|
52
|
+
writer = cv2.VideoWriter(save_path, fourcc, video_meta['fps'], (video_meta['width'], video_meta['height']))
|
|
53
|
+
batch_rotation_vectors = rotation_vectors[frm_range[0]: frm_range[-1]+1]
|
|
54
|
+
batch_centers = centers[frm_range[0]: frm_range[-1]+1]
|
|
55
|
+
m_rotates = center_rotation_warpaffine_vectors(rotation_vectors=batch_rotation_vectors, centers=batch_centers)
|
|
56
|
+
m_translations = align_target_warpaffine_vectors(centers=batch_centers, target=np.array(target))
|
|
57
|
+
|
|
58
|
+
if gpu:
|
|
59
|
+
# Combine rotation and translation matrices into single transform
|
|
60
|
+
# This reduces two sequential operations to one
|
|
61
|
+
batch_size = len(frm_range)
|
|
62
|
+
m_combined = np.zeros((batch_size, 2, 3), dtype=np.float32)
|
|
63
|
+
|
|
64
|
+
for i in range(batch_size):
|
|
65
|
+
# Convert rotation matrix (2x3) to 3x3 homogeneous
|
|
66
|
+
m_rot_3x3 = np.eye(3, dtype=np.float32)
|
|
67
|
+
m_rot_3x3[:2, :] = m_rotates[i].astype(np.float32)
|
|
68
|
+
|
|
69
|
+
# Convert translation matrix (2x3) to 3x3 homogeneous
|
|
70
|
+
m_trans_3x3 = np.eye(3, dtype=np.float32)
|
|
71
|
+
m_trans_3x3[:2, :] = m_translations[i].astype(np.float32)
|
|
72
|
+
|
|
73
|
+
# Combine: translation after rotation (matches sequential cv2.warpAffine order)
|
|
74
|
+
m_combined_3x3 = m_trans_3x3 @ m_rot_3x3
|
|
75
|
+
|
|
76
|
+
# Convert back to 2x3 for warpAffine compatibility
|
|
77
|
+
m_combined[i] = m_combined_3x3[:2, :]
|
|
78
|
+
|
|
79
|
+
# Process frames in batches using GPU reading
|
|
80
|
+
# Use same batch size as original (30) for optimal I/O overlap
|
|
81
|
+
# Main optimization: combined matrix (one warpAffine instead of two)
|
|
82
|
+
img_counter = 0
|
|
83
|
+
frm_batches = np.array_split(frm_range, (len(frm_range) + 30 - 1) // 30)
|
|
84
|
+
for frm_batch_cnt, frm_ids in enumerate(frm_batches):
|
|
85
|
+
frms = read_img_batch_from_video_gpu(video_path=video_path, start_frm=frm_ids[0], end_frm=frm_ids[-1], verbose=False)
|
|
86
|
+
frms = np.stack(list(frms.values()), axis=0)
|
|
87
|
+
for img_cnt, img in enumerate(frms):
|
|
88
|
+
# Use combined matrix for single warpAffine (faster than two separate calls)
|
|
89
|
+
m = m_combined[img_counter].astype(np.float32)
|
|
90
|
+
final_frame = cv2.warpAffine(img, m, (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
91
|
+
writer.write(final_frame)
|
|
92
|
+
if verbose:
|
|
93
|
+
frame_id = frm_ids[img_cnt]
|
|
94
|
+
print(f'Creating frame {frame_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
95
|
+
img_counter += 1
|
|
96
|
+
|
|
97
|
+
# Legacy CuPy code (commented out - CPU is faster for this use case)
|
|
98
|
+
if False and CUPY_AVAILABLE:
|
|
99
|
+
# Pre-compute all inverse matrices upfront (much faster than per-frame)
|
|
100
|
+
# For CuPy affine_transform, we need inverse matrices
|
|
101
|
+
m_inv_matrices = []
|
|
102
|
+
m_offsets = []
|
|
103
|
+
for i in range(batch_size):
|
|
104
|
+
m = m_combined[i]
|
|
105
|
+
matrix_2x2 = m[:2, :2].astype(np.float32)
|
|
106
|
+
offset = m[:2, 2].astype(np.float32)
|
|
107
|
+
m_inv_matrices.append(cp.asarray(matrix_2x2))
|
|
108
|
+
m_offsets.append(cp.asarray(offset))
|
|
109
|
+
# Batch invert all matrices at once
|
|
110
|
+
m_inv_matrices_gpu = cp.stack(m_inv_matrices)
|
|
111
|
+
m_inv_matrices_gpu = cp.linalg.inv(m_inv_matrices_gpu)
|
|
112
|
+
m_offsets_gpu = cp.stack(m_offsets)
|
|
113
|
+
|
|
114
|
+
# Create async reader for GPU
|
|
115
|
+
async_reader = AsyncVideoFrameReader(
|
|
116
|
+
video_path=video_path,
|
|
117
|
+
batch_size=batch_size_gpu,
|
|
118
|
+
max_que_size=3,
|
|
119
|
+
start_idx=frm_range[0],
|
|
120
|
+
end_idx=frm_range[-1] + 1,
|
|
121
|
+
gpu=True, # Use GPU reading
|
|
122
|
+
verbose=False
|
|
123
|
+
)
|
|
124
|
+
async_reader.start()
|
|
125
|
+
|
|
126
|
+
# Process batches as they become available from async reader
|
|
127
|
+
# Batch process and transfer to minimize GPU->CPU overhead
|
|
128
|
+
processed_frames_batch = []
|
|
129
|
+
frame_ids_batch = []
|
|
130
|
+
|
|
131
|
+
while True:
|
|
132
|
+
batch_result = get_async_frame_batch(batch_reader=async_reader, timeout=10)
|
|
133
|
+
if batch_result is None:
|
|
134
|
+
# Write any remaining frames
|
|
135
|
+
if processed_frames_batch:
|
|
136
|
+
for frame in processed_frames_batch:
|
|
137
|
+
writer.write(frame)
|
|
138
|
+
break
|
|
139
|
+
|
|
140
|
+
start_idx, end_idx, frms = batch_result
|
|
141
|
+
batch_len = end_idx - start_idx + 1
|
|
142
|
+
frms_gpu = cp.asarray(frms)
|
|
143
|
+
|
|
144
|
+
# Process all frames in batch on GPU first
|
|
145
|
+
batch_transformed = []
|
|
146
|
+
batch_frame_indices = []
|
|
147
|
+
|
|
148
|
+
for i in range(batch_len):
|
|
149
|
+
# Map frame index from video to frm_range index
|
|
150
|
+
frame_id = start_idx + i
|
|
151
|
+
frame_idx_in_range = np.where(frm_range == frame_id)[0]
|
|
152
|
+
if len(frame_idx_in_range) == 0:
|
|
153
|
+
continue
|
|
154
|
+
frame_idx_in_range = frame_idx_in_range[0]
|
|
155
|
+
batch_frame_indices.append((i, frame_idx_in_range))
|
|
156
|
+
|
|
157
|
+
# Process all frames in this batch on GPU
|
|
158
|
+
for i, frame_idx_in_range in batch_frame_indices:
|
|
159
|
+
img_gpu = frms_gpu[i]
|
|
160
|
+
matrix_inv = m_inv_matrices_gpu[frame_idx_in_range]
|
|
161
|
+
offset = m_offsets_gpu[frame_idx_in_range]
|
|
162
|
+
|
|
163
|
+
if len(img_gpu.shape) == 3: # Multi-channel
|
|
164
|
+
transformed_channels = []
|
|
165
|
+
for c in range(img_gpu.shape[2]):
|
|
166
|
+
transformed_ch = affine_transform(
|
|
167
|
+
img_gpu[:, :, c],
|
|
168
|
+
matrix=matrix_inv,
|
|
169
|
+
offset=offset,
|
|
170
|
+
output_shape=(video_meta['height'], video_meta['width']),
|
|
171
|
+
order=1,
|
|
172
|
+
mode='constant',
|
|
173
|
+
cval=fill_clr[c] if c < len(fill_clr) else fill_clr[0],
|
|
174
|
+
prefilter=False
|
|
175
|
+
)
|
|
176
|
+
transformed_channels.append(transformed_ch)
|
|
177
|
+
transformed = cp.stack(transformed_channels, axis=2)
|
|
178
|
+
else: # Single channel
|
|
179
|
+
transformed = affine_transform(
|
|
180
|
+
img_gpu,
|
|
181
|
+
matrix=matrix_inv,
|
|
182
|
+
offset=offset,
|
|
183
|
+
output_shape=(video_meta['height'], video_meta['width']),
|
|
184
|
+
order=1,
|
|
185
|
+
mode='constant',
|
|
186
|
+
cval=fill_clr[0] if len(fill_clr) > 0 else 0,
|
|
187
|
+
prefilter=False
|
|
188
|
+
)
|
|
189
|
+
batch_transformed.append(transformed)
|
|
190
|
+
|
|
191
|
+
# Batch transfer all frames from GPU to CPU at once
|
|
192
|
+
if batch_transformed:
|
|
193
|
+
# Stack all transformed frames and transfer in one go
|
|
194
|
+
batch_transformed_stack = cp.stack(batch_transformed)
|
|
195
|
+
batch_cpu = cp.asnumpy(batch_transformed_stack).astype(np.uint8)
|
|
196
|
+
|
|
197
|
+
# Write all frames from this batch
|
|
198
|
+
for frame_idx, (i, frame_idx_in_range) in enumerate(batch_frame_indices):
|
|
199
|
+
final_frame = batch_cpu[frame_idx]
|
|
200
|
+
writer.write(final_frame)
|
|
201
|
+
|
|
202
|
+
if verbose:
|
|
203
|
+
frame_id = start_idx + i
|
|
204
|
+
print(f'Creating frame {frame_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
205
|
+
|
|
206
|
+
async_reader.kill()
|
|
207
|
+
|
|
208
|
+
else:
|
|
209
|
+
# Fallback to CPU with combined matrix and batch reading
|
|
210
|
+
# Process frames in batches
|
|
211
|
+
# Use helper function directly to avoid nested multiprocessing (we're already in a worker process)
|
|
212
|
+
# Larger batch size reduces overhead
|
|
213
|
+
batch_size_gpu = 500
|
|
214
|
+
frm_batches = np.array_split(frm_range, (len(frm_range) + batch_size_gpu - 1) // batch_size_gpu)
|
|
215
|
+
|
|
216
|
+
# Create a mapping from frame_id to index in frm_range for fast lookup
|
|
217
|
+
frm_id_to_idx = {frame_id: idx for idx, frame_id in enumerate(frm_range)}
|
|
218
|
+
|
|
219
|
+
for frm_batch_cnt, frm_ids in enumerate(frm_batches):
|
|
220
|
+
# Read batch of frames directly using helper (no multiprocessing)
|
|
221
|
+
frm_idx_array = np.array(frm_ids)
|
|
222
|
+
frms_dict = _read_img_batch_from_video_helper(
|
|
223
|
+
frm_idx=frm_idx_array,
|
|
224
|
+
video_path=video_path,
|
|
225
|
+
greyscale=False,
|
|
226
|
+
verbose=False,
|
|
227
|
+
black_and_white=False,
|
|
228
|
+
clahe=False
|
|
229
|
+
)
|
|
230
|
+
frms = np.stack([frms_dict[f] for f in frm_ids], axis=0)
|
|
231
|
+
|
|
232
|
+
# Process all frames in batch using optimized CPU cv2.warpAffine with combined matrices
|
|
233
|
+
for i, frame_id in enumerate(frm_ids):
|
|
234
|
+
# Fast dictionary lookup instead of np.where
|
|
235
|
+
frame_idx_in_range = frm_id_to_idx.get(frame_id)
|
|
236
|
+
if frame_idx_in_range is None:
|
|
237
|
+
continue
|
|
238
|
+
|
|
239
|
+
img = frms[i]
|
|
240
|
+
m = m_combined[frame_idx_in_range].astype(np.float32)
|
|
241
|
+
final_frame = cv2.warpAffine(img, m, (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
242
|
+
writer.write(final_frame)
|
|
243
|
+
|
|
244
|
+
if verbose:
|
|
245
|
+
print(f'Creating frame {frame_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
246
|
+
else:
|
|
247
|
+
cap = cv2.VideoCapture(video_path)
|
|
248
|
+
for frm_idx, frm_id in enumerate(frm_range):
|
|
249
|
+
img = read_frm_of_video(video_path=cap, frame_index=frm_id)
|
|
250
|
+
rotated_frame = cv2.warpAffine(img, m_rotates[frm_idx], (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
251
|
+
final_frame = cv2.warpAffine(rotated_frame, m_translations[frm_idx], (video_meta['width'], video_meta['height']), borderValue=fill_clr)
|
|
252
|
+
writer.write(final_frame)
|
|
253
|
+
if verbose:
|
|
254
|
+
print(f'Creating frame {frm_id}/{video_meta["frame_count"]} ({video_name}, CPU core: {batch + 1}).')
|
|
255
|
+
writer.release()
|
|
256
|
+
return batch + 1
|
|
257
|
+
|
|
258
|
+
class EgocentricVideoRotator():
|
|
259
|
+
"""
|
|
260
|
+
Perform egocentric rotation of a video using CPU multiprocessing.
|
|
261
|
+
|
|
262
|
+
.. video:: _static/img/EgocentricalAligner_2.webm
|
|
263
|
+
:width: 800
|
|
264
|
+
:autoplay:
|
|
265
|
+
:loop:
|
|
266
|
+
|
|
267
|
+
.. seealso::
|
|
268
|
+
To perform joint egocentric alignment of both pose and video, or pose only, use :func:`~simba.data_processors.egocentric_aligner.EgocentricalAligner`.
|
|
269
|
+
To produce rotation vectors, use :func:`~simba.utils.data.egocentrically_align_pose_numba` or :func:`~simba.utils.data.egocentrically_align_pose`.
|
|
270
|
+
|
|
271
|
+
:param Union[str, os.PathLike] video_path: Path to a video file.
|
|
272
|
+
: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`.
|
|
273
|
+
: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`.
|
|
274
|
+
:param bool verbose: If True, prints progress. Deafult True.
|
|
275
|
+
:param Tuple[int, int, int] fill_clr: The color of the additional pixels. Deafult black. (0, 0, 0).
|
|
276
|
+
:param int core_cnt: Number of CPU cores to use for video rotation; `-1` uses all available cores.
|
|
277
|
+
: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.
|
|
278
|
+
|
|
279
|
+
:example:
|
|
280
|
+
>>> DATA_PATH = "C:\501_MA142_Gi_Saline_0513.csv"
|
|
281
|
+
>>> VIDEO_PATH = "C:\501_MA142_Gi_Saline_0513.mp4"
|
|
282
|
+
>>> SAVE_PATH = "C:\501_MA142_Gi_Saline_0513_rotated.mp4"
|
|
283
|
+
>>> ANCHOR_LOC = np.array([250, 250])
|
|
284
|
+
|
|
285
|
+
>>> df = read_df(file_path=DATA_PATH, file_type='csv')
|
|
286
|
+
>>> bp_cols = [x for x in df.columns if not x.endswith('_p')]
|
|
287
|
+
>>> data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32)
|
|
288
|
+
>>> _, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=6, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0)
|
|
289
|
+
>>> rotater = EgocentricVideoRotator(video_path=VIDEO_PATH, centers=centers, rotation_vectors=rotation_vectors, anchor_location=ANCHOR_LOC, save_path=SAVE_PATH)
|
|
290
|
+
>>> rotater.run()
|
|
291
|
+
"""
|
|
292
|
+
|
|
293
|
+
def __init__(self,
|
|
294
|
+
video_path: Union[str, os.PathLike],
|
|
295
|
+
centers: np.ndarray,
|
|
296
|
+
rotation_vectors: np.ndarray,
|
|
297
|
+
anchor_location: Tuple[int, int],
|
|
298
|
+
verbose: bool = True,
|
|
299
|
+
fill_clr: Tuple[int, int, int] = (0, 0, 0),
|
|
300
|
+
core_cnt: int = -1,
|
|
301
|
+
save_path: Optional[Union[str, os.PathLike]] = None,
|
|
302
|
+
gpu: Optional[bool] = True):
|
|
303
|
+
|
|
304
|
+
check_file_exist_and_readable(file_path=video_path)
|
|
305
|
+
self.video_meta_data = get_video_meta_data(video_path=video_path)
|
|
306
|
+
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)
|
|
307
|
+
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)
|
|
308
|
+
check_valid_tuple(x=anchor_location, source=f'{self.__class__.__name__} anchor_location', accepted_lengths=(2,), valid_dtypes=(int,))
|
|
309
|
+
for i in anchor_location: check_int(name=f'{self.__class__.__name__} anchor_location', value=i, min_value=1)
|
|
310
|
+
check_valid_boolean(value=[verbose], source=f'{self.__class__.__name__} verbose')
|
|
311
|
+
check_if_valid_rgb_tuple(data=fill_clr)
|
|
312
|
+
check_int(name=f'{self.__class__.__name__} core_cnt', value=core_cnt, min_value=-1, unaccepted_vals=[0])
|
|
313
|
+
if core_cnt > find_core_cnt()[0] or core_cnt == -1:
|
|
314
|
+
self.core_cnt = find_core_cnt()[0]
|
|
315
|
+
else:
|
|
316
|
+
self.core_cnt = core_cnt
|
|
317
|
+
video_dir, self.video_name, _ = get_fn_ext(filepath=video_path)
|
|
318
|
+
if save_path is not None:
|
|
319
|
+
self.save_dir = os.path.dirname(save_path)
|
|
320
|
+
check_if_dir_exists(in_dir=self.save_dir, source=f'{self.__class__.__name__} save_path')
|
|
321
|
+
else:
|
|
322
|
+
self.save_dir = video_dir
|
|
323
|
+
save_path = os.path.join(video_dir, f'{self.video_name}_rotated.mp4')
|
|
324
|
+
self.video_path, self.save_path = video_path, save_path
|
|
325
|
+
self.centers, self.rotation_vectors, self.gpu = centers, rotation_vectors, gpu
|
|
326
|
+
self.verbose, self.fill_clr, self.anchor_loc = verbose, fill_clr, anchor_location
|
|
327
|
+
|
|
328
|
+
def run(self):
|
|
329
|
+
video_timer = SimbaTimer(start=True)
|
|
330
|
+
temp_dir = os.path.join(self.save_dir, 'temp')
|
|
331
|
+
if not os.path.isdir(temp_dir):
|
|
332
|
+
create_directory(paths=temp_dir)
|
|
333
|
+
else:
|
|
334
|
+
remove_a_folder(folder_dir=temp_dir)
|
|
335
|
+
create_directory(paths=temp_dir)
|
|
336
|
+
frm_list = np.arange(0, self.video_meta_data['frame_count'])
|
|
337
|
+
frm_list = np.array_split(frm_list, self.core_cnt)
|
|
338
|
+
frm_list = [(cnt, x) for cnt, x in enumerate(frm_list)]
|
|
339
|
+
if self.verbose:
|
|
340
|
+
print(f"Creating rotated video {self.video_name}, multiprocessing (chunksize: {1}, cores: {self.core_cnt})...")
|
|
341
|
+
with multiprocessing.Pool(self.core_cnt, maxtasksperchild=Defaults.LARGE_MAX_TASK_PER_CHILD.value) as pool:
|
|
342
|
+
constants = functools.partial(egocentric_video_aligner,
|
|
343
|
+
temp_dir=temp_dir,
|
|
344
|
+
video_name=self.video_name,
|
|
345
|
+
video_path=self.video_path,
|
|
346
|
+
centers=self.centers,
|
|
347
|
+
rotation_vectors=self.rotation_vectors,
|
|
348
|
+
target=self.anchor_loc,
|
|
349
|
+
verbose=self.verbose,
|
|
350
|
+
fill_clr=self.fill_clr,
|
|
351
|
+
gpu=self.gpu)
|
|
352
|
+
for cnt, result in enumerate(pool.imap(constants, frm_list, chunksize=1)):
|
|
353
|
+
if self.verbose:
|
|
354
|
+
print(f"Rotate batch {result}/{self.core_cnt} complete...")
|
|
355
|
+
pool.terminate()
|
|
356
|
+
pool.join()
|
|
357
|
+
|
|
358
|
+
concatenate_videos_in_folder(in_folder=temp_dir, save_path=self.save_path, remove_splits=True, gpu=self.gpu, verbose=self.verbose)
|
|
359
|
+
video_timer.stop_timer()
|
|
360
|
+
stdout_success(msg=f"Egocentric rotation video {self.save_path} complete", elapsed_time=video_timer.elapsed_time_str, source=self.__class__.__name__)
|
|
361
|
+
|
|
362
|
+
if __name__ == "__main__":
|
|
363
|
+
DATA_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\data\501_MA142_Gi_Saline_0513.csv"
|
|
364
|
+
VIDEO_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513.mp4"
|
|
365
|
+
SAVE_PATH = r"C:\Users\sroni\OneDrive\Desktop\desktop\rotate_ex\videos\501_MA142_Gi_Saline_0513_rotated.mp4"
|
|
366
|
+
ANCHOR_LOC = np.array([250, 250])
|
|
367
|
+
|
|
368
|
+
df = read_df(file_path=DATA_PATH, file_type='csv')
|
|
369
|
+
bp_cols = [x for x in df.columns if not x.endswith('_p')]
|
|
370
|
+
data = df[bp_cols].values.reshape(len(df), int(len(bp_cols)/2), 2).astype(np.int32)
|
|
371
|
+
|
|
372
|
+
_, centers, rotation_vectors = egocentrically_align_pose(data=data, anchor_1_idx=5, anchor_2_idx=2, anchor_location=ANCHOR_LOC, direction=0)
|
|
373
|
+
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, gpu=True)
|
|
374
|
+
rotater.run()
|