Visitar URL original
Adding cropping to pose_estimation_pytorch/apis/analyze_videos.py (#2… · DeepLabCut/DeepLabCut@5f41812 · GitHub
Skip to content

Commit 5f41812

Browse files
authored
Adding cropping to pose_estimation_pytorch/apis/analyze_videos.py (#2737)
* Adding cropping to pose_estimation_pytorch/apis/analyze_videos.py * Adding docstring plus addressing comments
1 parent 597cf6b commit 5f41812

3 files changed

Lines changed: 36 additions & 21 deletions

File tree

‎deeplabcut/compat.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -874,6 +874,7 @@ def analyze_videos(
874874
auto_track=auto_track,
875875
identity_only=identity_only,
876876
overwrite=False,
877+
cropping=cropping,
877878
**torch_kwargs,
878879
)
879880

‎deeplabcut/pose_estimation_pytorch/apis/analyze_videos.py‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,14 @@ class VideoIterator(VideoReader):
4545
"""A class to iterate over videos, with possible added context"""
4646

4747
def __init__(
48-
self, video_path: str | Path, context: list[dict[str, Any]] | None = None
48+
self, video_path: str | Path, context: list[dict[str, Any]] | None = None, cropping: list[int] | None = None
4949
) -> None:
5050
super().__init__(str(video_path))
5151
self._context = context
5252
self._index = 0
53+
self._crop = cropping is not None
54+
if self._crop:
55+
self.set_bbox(*cropping)
5356

5457
def get_context(self) -> list[dict[str, Any]] | None:
5558
if self._context is None:
@@ -68,7 +71,7 @@ def __iter__(self):
6871
return self
6972

7073
def __next__(self) -> np.ndarray | tuple[str, dict[str, Any]]:
71-
frame = self.read_frame()
74+
frame = self.read_frame(crop=self._crop)
7275
if frame is None:
7376
self._index = 0
7477
self.reset()
@@ -94,9 +97,10 @@ def video_inference(
9497
detector_runner: InferenceRunner | None = None,
9598
with_identity: bool = False,
9699
return_video_metadata: bool = False,
100+
cropping: list[int] | None = None,
97101
) -> list[dict[str, np.ndarray]]:
98102
"""Runs inference on a video"""
99-
video = VideoIterator(str(video_path))
103+
video = VideoIterator(str(video_path), cropping=cropping)
100104
n_frames = video.get_n_frames()
101105
vid_w, vid_h = video.dimensions
102106
print(f"Starting to analyze {video_path}")
@@ -167,6 +171,7 @@ def analyze_videos(
167171
auto_track: bool | None = True,
168172
identity_only: bool | None = False,
169173
overwrite: bool = False,
174+
cropping: list[int] | None = None,
170175
) -> str:
171176
"""Makes prediction based on a trained network.
172177
@@ -220,6 +225,11 @@ def analyze_videos(
220225
identity_only: sub-call for auto_track. If ``True`` and animal identity was
221226
learned by the model, assembly and tracking rely exclusively on identity
222227
prediction.
228+
cropping: list or None, optional, default=None
229+
List of cropping coordinates as [x1, x2, y1, y2].
230+
Note that the same cropping parameters will then be used for all videos.
231+
If different video crops are desired, run ``analyze_videos`` on individual
232+
videos with the corresponding cropping coordinates.
223233
224234
Returns:
225235
The scorer used to analyze the videos
@@ -252,6 +262,9 @@ def analyze_videos(
252262
cfg, model_cfg, snapshot_index, detector_snapshot_index,
253263
)
254264

265+
if cropping is None and cfg.get("cropping", False):
266+
cropping = cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]
267+
255268
# Get general project parameters
256269
bodyparts = model_cfg["metadata"]["bodyparts"]
257270
unique_bodyparts = model_cfg["metadata"]["unique_bodyparts"]
@@ -328,6 +341,7 @@ def analyze_videos(
328341
pose_runner=pose_runner,
329342
task=pose_task,
330343
detector_runner=detector_runner,
344+
cropping=cropping,
331345
)
332346
runtime.append(time.time())
333347
metadata = _generate_metadata(

‎deeplabcut/utils/auxfun_videos.py‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,24 @@ def get_bbox(self, relative=False):
159159
y2 = int(self._height * y2)
160160
return x1, x2, y1, y2
161161

162+
def set_bbox(self, x1, x2, y1, y2, relative=False):
163+
if x2 <= x1 or y2 <= y1:
164+
raise ValueError(
165+
f"Coordinates look wrong... " f"Ensure {x1} < {x2} and {y1} < {y2}."
166+
)
167+
if not relative:
168+
x1 /= self._width
169+
x2 /= self._width
170+
y1 /= self._height
171+
y2 /= self._height
172+
bbox = x1, x2, y1, y2
173+
if any(coord > 1 for coord in bbox):
174+
warnings.warn(
175+
"Bounding box larger than the video... " "Clipping to video dimensions."
176+
)
177+
bbox = tuple(map(lambda x: min(x, 1), bbox))
178+
self._bbox = bbox
179+
162180
@property
163181
def fps(self):
164182
return self._fps
@@ -205,24 +223,6 @@ def __init__(self, video_path, codec="h264", dpi=100, fps=None):
205223
if fps:
206224
self.fps = fps
207225

208-
def set_bbox(self, x1, x2, y1, y2, relative=False):
209-
if x2 <= x1 or y2 <= y1:
210-
raise ValueError(
211-
f"Coordinates look wrong... " f"Ensure {x1} < {x2} and {y1} < {y2}."
212-
)
213-
if not relative:
214-
x1 /= self._width
215-
x2 /= self._width
216-
y1 /= self._height
217-
y2 /= self._height
218-
bbox = x1, x2, y1, y2
219-
if any(coord > 1 for coord in bbox):
220-
warnings.warn(
221-
"Bounding box larger than the video... " "Clipping to video dimensions."
222-
)
223-
bbox = tuple(map(lambda x: min(x, 1), bbox))
224-
self._bbox = bbox
225-
226226
def shorten(
227227
self, start, end, suffix="short", dest_folder=None, validate_inputs=True
228228
):

0 commit comments

Comments
 (0)