Visitar URL original
New Top-Down Pose Estimation Model - RTMPose (#2795) · DeepLabCut/DeepLabCut@1d26325 · GitHub
Skip to content

Commit 1d26325

Browse files
n-poulsenMMathisLab
n-poulsen
and
MMathisLab
authored
New Top-Down Pose Estimation Model - RTMPose (#2795)
Co-authored-by: Mackenzie Mathis <mathis@rowland.harvard.edu>
1 parent b7f1c3d commit 1d26325

48 files changed

Lines changed: 2783 additions & 180 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ snapshot-*
2121
# Modelzoo checkpoints
2222
deeplabcut/modelzoo/checkpoints/
2323

24+
# PyTorch backbone weights
25+
deeplabcut/pose_estimation_pytorch/models/backbones/pretrained_weights/
26+
2427
# Wandb files
2528
wandb/
2629

‎deeplabcut/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
"DLC loaded in light mode; you cannot use any GUI (labeling, relabeling and standalone GUI)"
3131
)
3232

33+
from deeplabcut.core.engine import Engine
3334
from deeplabcut.create_project import (
3435
create_new_project,
3536
create_new_project_3d,

‎deeplabcut/pose_estimation_pytorch/README.md‎

Lines changed: 62 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -420,16 +420,20 @@ train(
420420

421421
### Running Video Analysis outside a DeepLabCut Project
422422

423-
DeepLabCut provides high-level APIs (via the GUI or the python package) to analyze your data. The usage of this API assumes the existance of a DLC project (with `config.yaml` file, etc.).
423+
DeepLabCut provides high-level APIs (via the GUI or the python package) to analyze your
424+
data. The usage of this API assumes the existance of a DLC project (with `config.yaml`
425+
file, etc.).
424426

425-
Sometimes it might be more convenient to just run a model on your data via a low-level API. We also use this API under the hood, in particular for the Model Zoo. Check out the example below:
427+
Sometimes it might be more convenient to just run a model on your data via a low-level
428+
API. We also use this API under the hood, in particular for the Model Zoo. Check out the
429+
example below:
426430

427431
```python
428432
from pathlib import Path
429433

434+
from deeplabcut.pose_estimation_pytorch import Task
430435
from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import video_inference
431436
from deeplabcut.pose_estimation_pytorch.config import read_config_as_dict
432-
from deeplabcut.pose_estimation_pytorch.task import Task
433437
from deeplabcut.pose_estimation_pytorch.apis.utils import get_inference_runners
434438

435439
train_dir = Path("/Users/Jaylen/my-dlc-models/train")
@@ -447,30 +451,76 @@ detector_batch_size = 8
447451

448452
# read model configuration
449453
model_cfg = read_config_as_dict(pytorch_config_path)
450-
bodyparts = model_cfg["metadata"]["bodyparts"]
451-
unique_bodyparts = model_cfg["metadata"]["unique_bodyparts"]
452-
with_identity = model_cfg["metadata"].get("with_identity", False)
453454

454455
pose_task = Task(model_cfg["method"])
455456
pose_runner, detector_runner = get_inference_runners(
456457
model_config=model_cfg,
457458
snapshot_path=snapshot_path,
458459
max_individuals=max_num_animals,
459-
num_bodyparts=len(bodyparts),
460-
num_unique_bodyparts=len(unique_bodyparts),
461460
batch_size=batch_size,
462-
with_identity=with_identity,
463-
transform=None,
464461
detector_batch_size=detector_batch_size,
465462
detector_path=detector_snapshot_path,
466-
detector_transform=None,
467463
)
468464

469465
predictions = video_inference(
470466
video=video_path,
471467
task=pose_task,
472468
pose_runner=pose_runner,
473469
detector_runner=detector_runner,
474-
with_identity=False,
475470
)
476471
```
472+
473+
474+
### Running Top-Down Video Analysis with Existing Bounding Boxes
475+
476+
When `deeplabcut.pose_estimation_pytorch.apis.analyze_videos.video_inference` is called
477+
with a top-down model, it is assumed that a detector snapshot is given as well to obtain
478+
bounding boxes with which to run pose estimation. It's possible that you've already
479+
obtained bounding boxes for your video (with another object detector or through some
480+
other means), and you want to re-use those bounding boxes instead of running an object
481+
detector again.
482+
483+
You can easily do so by writing a bit of custom code, as shown in the example below:
484+
485+
```python
486+
from pathlib import Path
487+
488+
import numpy as np
489+
from deeplabcut.pose_estimation_pytorch import get_inference_runners
490+
from deeplabcut.pose_estimation_pytorch.apis import VideoIterator
491+
from deeplabcut.pose_estimation_pytorch.config import read_config_as_dict
492+
from tqdm import tqdm
493+
494+
# create an iterator for your video
495+
video = VideoIterator("/Users/Jayson/my-cool-video.mp4")
496+
497+
# dummy bboxes - you can load yours from a file or in another way
498+
# the bboxes should be in `xywh` format, i.e. (x_top_left, y_top_left, width, height)
499+
bounding_boxes = [
500+
dict( # frame 0 bounding boxes
501+
bboxes=np.array([[12, 37, 120, 78]]),
502+
),
503+
dict( # frame 1 bounding boxes
504+
bboxes=np.array([[17, 45, 128, 73], [532, 34, 117, 87]]),
505+
),
506+
# ...
507+
dict( # frame N bboxes -> must be equal to the number of frames in the video!
508+
bboxes=np.array([[17, 45, 128, 73], [532, 34, 117, 87]]),
509+
),
510+
]
511+
video.set_context(bounding_boxes)
512+
max_individuals = np.max([len(context["bboxes"]) for context in bounding_boxes])
513+
514+
# run inference!
515+
model_cfg = read_config_as_dict("/Users/Jayson/pytorch_config.yaml")
516+
pose_runner, _ = get_inference_runners(
517+
model_config=model_cfg,
518+
snapshot_path=Path("/Users/Jayson/model-snapshot.pt"),
519+
max_individuals=max_individuals,
520+
batch_size=32,
521+
)
522+
523+
# your predictions will be a list, containing the predictions made for each frame
524+
# as a dict (with keys for "bodyparts" but also "bboxes")!
525+
predictions = pose_runner.inference(images=tqdm(video))
526+
```

‎deeplabcut/pose_estimation_pytorch/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
analyze_images,
1313
analyze_videos,
1414
convert_detections2tracklets,
15+
get_inference_runners,
1516
evaluate_network,
1617
extract_maps,
1718
extract_save_all_maps,

‎deeplabcut/pose_estimation_pytorch/apis/__init__.py‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,32 @@
1010
#
1111

1212
from deeplabcut.pose_estimation_pytorch.apis.analyze_images import (
13+
analyze_image_folder,
1314
analyze_images,
1415
superanimal_analyze_images,
1516
)
16-
from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import analyze_videos
17+
from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import (
18+
analyze_videos,
19+
video_inference,
20+
VideoIterator,
21+
)
1722
from deeplabcut.pose_estimation_pytorch.apis.convert_detections_to_tracklets import (
1823
convert_detections2tracklets,
1924
)
20-
from deeplabcut.pose_estimation_pytorch.apis.evaluate import evaluate_network
25+
from deeplabcut.pose_estimation_pytorch.apis.evaluate import (
26+
evaluate,
27+
evaluate_network,
28+
)
2129
from deeplabcut.pose_estimation_pytorch.apis.export import export_model
22-
from deeplabcut.pose_estimation_pytorch.apis.train import train_network
30+
from deeplabcut.pose_estimation_pytorch.apis.train import (
31+
train,
32+
train_network,
33+
)
34+
from deeplabcut.pose_estimation_pytorch.apis.utils import (
35+
get_detector_inference_runner,
36+
get_inference_runners,
37+
get_pose_inference_runner,
38+
)
2339
from deeplabcut.pose_estimation_pytorch.apis.visualization import (
2440
extract_maps,
2541
extract_save_all_maps,

‎deeplabcut/pose_estimation_pytorch/apis/utils.py‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -435,9 +435,9 @@ def build_bboxes_dict_for_dataframe(
435435
def get_inference_runners(
436436
model_config: dict,
437437
snapshot_path: str | Path,
438-
max_individuals: int,
439-
num_bodyparts: int,
440-
num_unique_bodyparts: int,
438+
max_individuals: int | None = None,
439+
num_bodyparts: int | None = None,
440+
num_unique_bodyparts: int | None = None,
441441
batch_size: int = 1,
442442
device: str | None = None,
443443
with_identity: bool = False,
@@ -452,9 +452,12 @@ def get_inference_runners(
452452
Args:
453453
model_config: the pytorch configuration file
454454
snapshot_path: the path of the snapshot from which to load the weights
455-
max_individuals: the maximum number of individuals per image
456-
num_bodyparts: the number of bodyparts predicted by the model
457-
num_unique_bodyparts: the number of unique_bodyparts predicted by the model
455+
max_individuals: the maximum number of individuals per image (if None, uses the
456+
individuals defined in the model_config metadata)
457+
num_bodyparts: the number of bodyparts predicted by the model (if None, uses the
458+
bodyparts defined in the model_config metadata)
459+
num_unique_bodyparts: the number of unique_bodyparts predicted by the model (if
460+
None, uses the unique bodyparts defined in the model_config metadata)
458461
batch_size: the batch size to use for the pose model.
459462
with_identity: whether the pose model has an identity head
460463
device: if defined, overwrites the device selection from the model config
@@ -474,6 +477,13 @@ def get_inference_runners(
474477
a runner for pose estimation
475478
a runner for detection, if detector_path is not None
476479
"""
480+
if max_individuals is None:
481+
max_individuals = len(model_config["metadata"]["individuals"])
482+
if num_bodyparts is None:
483+
num_bodyparts = len(model_config["metadata"]["bodyparts"])
484+
if num_unique_bodyparts is None:
485+
num_unique_bodyparts = len(model_config["metadata"]["unique_bodyparts"])
486+
477487
pose_task = Task(model_config["method"])
478488
if device is None:
479489
device = resolve_device(model_config)
@@ -499,10 +509,15 @@ def get_inference_runners(
499509
if device == "mps":
500510
detector_device = "cpu"
501511

512+
crop_cfg = model_config["data"]["inference"].get("top_down_crop", {})
513+
width, height = crop_cfg.get("width", 256), crop_cfg.get("height", 256)
514+
margin = crop_cfg.get("margin", 0)
515+
502516
pose_preprocessor = build_top_down_preprocessor(
503517
color_mode=model_config["data"]["colormode"],
504518
transform=transform,
505-
cropped_image_size=(256, 256),
519+
top_down_crop_size=(width, height),
520+
top_down_crop_margin=margin,
506521
)
507522
pose_postprocessor = build_top_down_postprocessor(
508523
max_individuals=max_individuals,
@@ -659,10 +674,15 @@ def get_pose_inference_runner(
659674
with_identity=with_identity,
660675
)
661676
else:
677+
crop_cfg = model_config["data"]["inference"].get("top_down_crop", {})
678+
width, height = crop_cfg.get("width", 256), crop_cfg.get("height", 256)
679+
margin = crop_cfg.get("margin", 0)
680+
662681
pose_preprocessor = build_top_down_preprocessor(
663682
color_mode=model_config["data"]["colormode"],
664683
transform=transform,
665-
cropped_image_size=(256, 256),
684+
top_down_crop_size=(width, height),
685+
top_down_crop_margin=margin,
666686
)
667687
pose_postprocessor = build_top_down_postprocessor(
668688
max_individuals=max_individuals,

‎deeplabcut/pose_estimation_pytorch/config/backbones/cspnext_m.yaml‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
model:
2+
backbone:
3+
type: CSPNeXt
4+
model_name: cspnext_m
5+
freeze_bn_stats: false
6+
freeze_bn_weights: false
7+
deepen_factor: 0.67
8+
widen_factor: 0.75
9+
backbone_output_channels: 768
10+
runner:
11+
optimizer:
12+
type: AdamW
13+
params:
14+
lr: 0.0005
15+
scheduler:
16+
type: LRListScheduler
17+
params:
18+
lr_list: [ [ 1e-4 ], [ 1e-5 ] ]
19+
milestones: [ 90, 190 ]

‎deeplabcut/pose_estimation_pytorch/config/backbones/cspnext_s.yaml‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
model:
2+
backbone:
3+
type: CSPNeXt
4+
model_name: cspnext_s
5+
freeze_bn_stats: false
6+
freeze_bn_weights: false
7+
deepen_factor: 0.33
8+
widen_factor: 0.5
9+
backbone_output_channels: 512
10+
runner:
11+
optimizer:
12+
type: AdamW
13+
params:
14+
lr: 0.0005
15+
scheduler:
16+
type: LRListScheduler
17+
params:
18+
lr_list: [ [ 1e-4 ], [ 1e-5 ] ]
19+
milestones: [ 90, 190 ]

‎deeplabcut/pose_estimation_pytorch/config/backbones/cspnext_x.yaml‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
model:
2+
backbone:
3+
type: CSPNeXt
4+
model_name: cspnext_x
5+
freeze_bn_stats: false
6+
freeze_bn_weights: false
7+
deepen_factor: 1.33
8+
widen_factor: 1.25
9+
backbone_output_channels: 1280
10+
runner:
11+
optimizer:
12+
type: AdamW
13+
params:
14+
lr: 0.0005
15+
scheduler:
16+
type: LRListScheduler
17+
params:
18+
lr_list: [ [ 1e-4 ], [ 1e-5 ] ]
19+
milestones: [ 90, 190 ]

‎deeplabcut/pose_estimation_pytorch/config/base/aug_top_down.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
colormode: RGB
22
inference:
33
normalize_images: true
4+
top_down_crop:
5+
width: 256
6+
height: 256
47
train:
58
affine:
69
p: 0.5
@@ -13,3 +16,6 @@ train:
1316
hist_eq: false
1417
motion_blur: false
1518
normalize_images: true
19+
top_down_crop:
20+
width: 256
21+
height: 256

0 commit comments

Comments
 (0)