Visitar URL original
Improve `deeplabcut.pose_estimation_pytorch` imports (#2833) · DeepLabCut/DeepLabCut@d03da3c · GitHub
Skip to content

Commit d03da3c

Browse files
n-poulsenn-poulsen
authored
Improve deeplabcut.pose_estimation_pytorch imports (#2833)
1 parent c9ab9c2 commit d03da3c

14 files changed

Lines changed: 257 additions & 101 deletions

File tree

‎deeplabcut/pose_estimation_pytorch/README.md‎

Lines changed: 45 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,17 @@ disk.
6464
```python
6565
from pathlib import Path
6666

67-
from deeplabcut.pose_estimation_pytorch.config import (
68-
make_pytorch_pose_config,
69-
write_config,
70-
)
67+
import deeplabcut.pose_estimation_pytorch as dlc_torch
7168

7269
project_cfg = { "Task": "mice", ... } # the configuration for your DLC project
7370
pose_config_path = Path("/path/to/my/config/pytorch_cfg.yaml")
74-
model_cfg = make_pytorch_pose_config(
71+
model_cfg = dlc_torch.config.make_pytorch_pose_config(
7572
project_config=project_cfg,
7673
pose_config_path=pose_config_path,
7774
net_type="hrnet_w32",
7875
top_down=True,
76+
save=True,
7977
)
80-
write_config(pose_config_path, model_cfg)
8178
```
8279

8380
#### Adding Models
@@ -238,9 +235,9 @@ dataset creation and test/train splitting. The `DLCLoader` class is used to load
238235
labeled data for a specific shuffle.
239236

240237
```python3
241-
from deeplabcut.pose_estimation_pytorch.data import DLCLoader
238+
import deeplabcut.pose_estimation_pytorch as dlc_torch
242239

243-
loader = DLCLoader(
240+
loader = dlc_torch.DLCLoader(
244241
config="/path/to/my/project/config.yaml",
245242
trainset_index=0,
246243
shuffle=1,
@@ -265,23 +262,20 @@ images and keypoints to a tensor dataset for training and evaluation. You can ge
265262
an instance of training/test dataset with your `DLCLoader`:
266263

267264
```python3
268-
from deeplabcut.pose_estimation_pytorch.data import (
269-
build_transforms,
270-
DLCLoader,
271-
)
265+
import deeplabcut.pose_estimation_pytorch as dlc_torch
272266

273-
loader = DLCLoader(
267+
loader = dlc_torch.DLCLoader(
274268
config="/path/to/my/project/config.yaml",
275269
trainset_index=0,
276270
shuffle=1,
277271
)
278272
train_dataset = loader.create_dataset(
279-
transform=build_transforms(loader.model_cfg["data"]["train"]),
273+
transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
280274
mode="train",
281275
task=loader.pose_task,
282276
)
283277
valid_dataset = loader.create_dataset(
284-
transform=build_transforms(loader.model_cfg["data"]["train"]),
278+
transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
285279
mode="test",
286280
task=loader.pose_task,
287281
)
@@ -320,58 +314,48 @@ configuration file (as described in [#model_configuration_files]).
320314
```python3
321315
from pathlib import Path
322316

323-
from deeplabcut.pose_estimation_pytorch.config import (
324-
make_pytorch_pose_config,
325-
write_config,
326-
)
327-
from deeplabcut.pose_estimation_pytorch.data import (
328-
build_transforms,
329-
COCOLoader,
330-
)
331-
from deeplabcut.pose_estimation_pytorch.task import Task
317+
import deeplabcut.pose_estimation_pytorch as dlc_torch
332318

333319
# Specify project paths
334320
project_root = Path("/path/to/my/COCOProject")
335321
train_json_filename = "train.json"
336322
test_json_filename = "test.json"
337323

338324
# Parse information about the project
339-
train_dict = COCOLoader.load_json(project_root, filename=train_json_filename)
340-
max_num_individuals, bodyparts = COCOLoader.get_project_parameters(train_dict)
325+
train_dict = dlc_torch.COCOLoader.load_json(project_root, filename=train_json_filename)
326+
max_num_individuals, bodyparts = dlc_torch.COCOLoader.get_project_parameters(train_dict)
341327

342328
# Generate a configuration file for your PyTorch model
343329
# In this case, it's for a Top-Down HRNet_w32
344330
experiment_path = project_root / "experiments" / "hrnet_w32"
345331
model_cfg_path = experiment_path / "train" / "pytorch_cfg.yaml"
346-
model_cfg = make_pytorch_pose_config(
347-
project_config={
348-
"project_path": str(project_root.resolve()),
349-
"multianimalproject": max_num_individuals > 1,
350-
"bodyparts": bodyparts,
351-
"multianimalbodyparts": bodyparts,
352-
"uniquebodyparts": [],
353-
"individuals": [f"idv{i}" for i in range(max_num_individuals)],
354-
},
332+
model_cfg = dlc_torch.config.make_pytorch_pose_config(
333+
project_config=dlc_torch.config.make_basic_project_config(
334+
dataset_path=str(project_root.resolve()),
335+
bodyparts=bodyparts,
336+
max_individuals=max_num_individuals,
337+
multi_animal=True,
338+
),
355339
pose_config_path=experiment_path,
356340
net_type="hrnet_w32",
357341
top_down=True,
342+
save=True,
358343
)
359-
write_config(config_path=model_cfg_path, config=model_cfg)
360344

361345
# Create the loader for the COCO dataset
362-
loader = COCOLoader(
346+
loader = dlc_torch.COCOLoader(
363347
project_root=project_root,
364348
model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
365349
train_json_filename=train_json_filename,
366350
test_json_filename=test_json_filename,
367351
)
368352
train_dataset = loader.create_dataset(
369-
transform=build_transforms(loader.model_cfg["data"]["train"]),
353+
transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
370354
mode="train",
371355
task=loader.pose_task,
372356
)
373357
valid_dataset = loader.create_dataset(
374-
transform=build_transforms(loader.model_cfg["data"]["train"]),
358+
transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
375359
mode="test",
376360
task=loader.pose_task,
377361
)
@@ -389,16 +373,15 @@ pretrained weights, and either train them or run inference with them.
389373
```python
390374
from pathlib import Path
391375

376+
import deeplabcut.pose_estimation_pytorch as dlc_torch
392377
from deeplabcut.pose_estimation_pytorch.apis.train import train
393-
from deeplabcut.pose_estimation_pytorch.data import COCOLoader
394-
from deeplabcut.pose_estimation_pytorch.task import Task
395378

396379
# Specify project paths
397380
project_root = Path("/path/to/my/COCOProject")
398381
train_json_filename = "train.json"
399382
test_json_filename = "test.json"
400383

401-
loader = COCOLoader(
384+
loader = dlc_torch.COCOLoader(
402385
project_root=project_root,
403386
model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
404387
train_json_filename=train_json_filename,
@@ -407,7 +390,7 @@ loader = COCOLoader(
407390
train(
408391
loader=loader,
409392
run_config=loader.model_cfg,
410-
task=Task(loader.model_cfg["method"]),
393+
task=dlc_torch.Task(loader.model_cfg["method"]),
411394
device="cuda:2",
412395
logger_config=dict(
413396
type="WandbLogger",
@@ -421,7 +404,7 @@ train(
421404
### Running Video Analysis outside a DeepLabCut Project
422405

423406
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`
407+
data. The usage of this API assumes the existence of a DLC project (with `config.yaml`
425408
file, etc.).
426409

427410
Sometimes it might be more convenient to just run a model on your data via a low-level
@@ -431,10 +414,7 @@ example below:
431414
```python
432415
from pathlib import Path
433416

434-
from deeplabcut.pose_estimation_pytorch import Task
435-
from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import video_inference
436-
from deeplabcut.pose_estimation_pytorch.config import read_config_as_dict
437-
from deeplabcut.pose_estimation_pytorch.apis.utils import get_inference_runners
417+
import deeplabcut.pose_estimation_pytorch as dlc_torch
438418

439419
train_dir = Path("/Users/Jaylen/my-dlc-models/train")
440420
pytorch_config_path = train_dir / "pytorch_config.yaml"
@@ -450,21 +430,26 @@ batch_size = 16
450430
detector_batch_size = 8
451431

452432
# read model configuration
453-
model_cfg = read_config_as_dict(pytorch_config_path)
454-
455-
pose_task = Task(model_cfg["method"])
456-
pose_runner, detector_runner = get_inference_runners(
433+
model_cfg = dlc_torch.config.read_config_as_dict(pytorch_config_path)
434+
pose_task = dlc_torch.Task(model_cfg["method"])
435+
pose_runner = dlc_torch.get_pose_inference_runner(
457436
model_config=model_cfg,
458437
snapshot_path=snapshot_path,
459438
max_individuals=max_num_animals,
460439
batch_size=batch_size,
461-
detector_batch_size=detector_batch_size,
462-
detector_path=detector_snapshot_path,
463440
)
464441

465-
predictions = video_inference(
442+
detector_runner = None
443+
if pose_task == dlc_torch.Task.TOP_DOWN:
444+
detector_runner = dlc_torch.get_detector_inference_runner(
445+
model_config=model_cfg,
446+
snapshot_path=detector_snapshot_path,
447+
max_individuals=max_num_animals,
448+
batch_size=detector_batch_size,
449+
)
450+
451+
predictions = dlc_torch.video_inference(
466452
video=video_path,
467-
task=pose_task,
468453
pose_runner=pose_runner,
469454
detector_runner=detector_runner,
470455
)
@@ -486,13 +471,11 @@ You can easily do so by writing a bit of custom code, as shown in the example be
486471
from pathlib import Path
487472

488473
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
474+
import deeplabcut.pose_estimation_pytorch as dlc_torch
492475
from tqdm import tqdm
493476

494477
# create an iterator for your video
495-
video = VideoIterator("/Users/Jayson/my-cool-video.mp4")
478+
video = dlc_torch.VideoIterator("/Users/Jayson/my-cool-video.mp4")
496479

497480
# dummy bboxes - you can load yours from a file or in another way
498481
# the bboxes should be in `xywh` format, i.e. (x_top_left, y_top_left, width, height)
@@ -512,8 +495,8 @@ video.set_context(bounding_boxes)
512495
max_individuals = np.max([len(context["bboxes"]) for context in bounding_boxes])
513496

514497
# run inference!
515-
model_cfg = read_config_as_dict("/Users/Jayson/pytorch_config.yaml")
516-
pose_runner, _ = get_inference_runners(
498+
model_cfg = dlc_torch.config.read_config_as_dict("/Users/Jayson/pytorch_config.yaml")
499+
pose_runner = dlc_torch.get_pose_inference_runner(
517500
model_config=model_cfg,
518501
snapshot_path=Path("/Users/Jayson/model-snapshot.pt"),
519502
max_individuals=max_individuals,

‎deeplabcut/pose_estimation_pytorch/__init__.py‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,31 +8,53 @@
88
#
99
# Licensed under GNU Lesser General Public License v3.0
1010
#
11+
import deeplabcut.pose_estimation_pytorch.config as config
1112
from deeplabcut.pose_estimation_pytorch.apis import (
13+
analyze_image_folder,
1214
analyze_images,
1315
analyze_videos,
16+
build_predictions_dataframe,
17+
create_labeled_images,
1418
convert_detections2tracklets,
15-
get_inference_runners,
19+
evaluate,
1620
evaluate_network,
1721
extract_maps,
1822
extract_save_all_maps,
23+
get_detector_inference_runner,
24+
get_pose_inference_runner,
25+
predict,
26+
superanimal_analyze_images,
27+
train,
1928
train_network,
29+
video_inference,
30+
VideoIterator,
31+
visualize_predictions,
2032
)
2133
from deeplabcut.pose_estimation_pytorch.config import (
2234
available_detectors,
2335
available_models,
2436
)
25-
from deeplabcut.pose_estimation_pytorch.data.base import Loader
26-
from deeplabcut.pose_estimation_pytorch.data.cocoloader import COCOLoader
27-
from deeplabcut.pose_estimation_pytorch.data.dataset import (
37+
from deeplabcut.pose_estimation_pytorch.data import (
38+
build_transforms,
39+
COCOLoader,
40+
COLLATE_FUNCTIONS,
41+
DLCLoader,
42+
Loader,
2843
PoseDataset,
2944
PoseDatasetParameters,
3045
)
31-
from deeplabcut.pose_estimation_pytorch.data.dlcloader import DLCLoader
32-
from deeplabcut.pose_estimation_pytorch.runners.base import (
46+
from deeplabcut.pose_estimation_pytorch.runners import (
47+
build_inference_runner,
48+
build_training_runner,
49+
DetectorInferenceRunner,
50+
DetectorTrainingRunner,
3351
get_load_weights_only,
52+
InferenceRunner,
53+
PoseInferenceRunner,
54+
PoseTrainingRunner,
3455
set_load_weights_only,
56+
TorchSnapshotManager,
57+
TrainingRunner,
3558
)
36-
from deeplabcut.pose_estimation_pytorch.runners.snapshots import TorchSnapshotManager
3759
from deeplabcut.pose_estimation_pytorch.task import Task
3860
from deeplabcut.pose_estimation_pytorch.utils import fix_seeds

‎deeplabcut/pose_estimation_pytorch/apis/__init__.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from deeplabcut.pose_estimation_pytorch.apis.analyze_images import (
1313
analyze_image_folder,
1414
analyze_images,
15+
analyze_image_folder,
1516
superanimal_analyze_images,
1617
)
1718
from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import (
@@ -23,8 +24,10 @@
2324
convert_detections2tracklets,
2425
)
2526
from deeplabcut.pose_estimation_pytorch.apis.evaluate import (
27+
predict,
2628
evaluate,
2729
evaluate_network,
30+
visualize_predictions,
2831
)
2932
from deeplabcut.pose_estimation_pytorch.apis.export import export_model
3033
from deeplabcut.pose_estimation_pytorch.apis.train import (
@@ -37,6 +40,12 @@
3740
get_pose_inference_runner,
3841
)
3942
from deeplabcut.pose_estimation_pytorch.apis.visualization import (
43+
create_labeled_images,
4044
extract_maps,
4145
extract_save_all_maps,
4246
)
47+
from deeplabcut.pose_estimation_pytorch.apis.utils import (
48+
build_predictions_dataframe,
49+
get_detector_inference_runner,
50+
get_pose_inference_runner,
51+
)

0 commit comments

Comments
 (0)