SuperAnimal Model Updates

This pull request refactors the code to use SuperAnimal models on downstream datasets (using transfer learning, fine-tuning and fine-tuning with memory replay). It also adds compatibility with weights released for SuperAnimal-Bird through dlclibrary.

Before merging this pull request, we need:

• merge DLClibrary #29
• make new DLClibrary release (v0.0.7) https://pypi.org/project/dlclibrary/#history

New Features & Updates

SuperAnimal - Separation of Detectors and Pose Models

The modelzoo code was re-written to handle pose models and detectors separately. This allows users to mix-and-match which detectors are used with which pose models.

The ModelZoo page in the GUI was updated so users can pick which detector/pose estimation model combination they want.

SuperAnimal - Bird

This pull request (combined with a dlclibrary release) adds a SuperAnimal-Bird model, using an SSDLite detector and a pose model with a ResNet50 backbone.

WeightInitialization

The WeightInitialization class is updated to directly store the paths to the snapshots from which to initialize models. This makes it much better encapsulated, and makes it agnostic to the SuperAnimal/ModelZoo models, allowing the class to be used to initialize weights from any snapshot desired.

To create a WeightInitialization instance to fine-tune SuperAnimal models, deeplabcut.modelzoo.build_weight_init can be used. This weight_init can then be used to create a new training dataset.

import deeplabcut
from deeplabcut.modelzoo import build_weight_init

config = "/path/to/my/project/config.yaml"
net_type = "hrnet_w32"
detector_type = "fasterrcnn_resnet50_fpn_v2"

weight_init = build_weight_init(
    config, 
    super_animal="superanimal_topviewmouse",
    model_name=net_type,
    detector_name=detector_type,
    with_decoder=True,
    memory_replay=False,
)

deeplabcut.create_training_dataset(
    config,
    Shuffles=[10],
    net_type=net_type,
    detector_type=detector_type,
    weight_init=weight_init,
)

Calling build_weight_init will look for the models using dlclibrary, and download them if they aren't available yet. Adding this weight initialization will add the following element to the shuffle's pytorch_config.yaml:

train_settings:
  ...
  weight_init:
    dataset: superanimal_topviewmouse
    snapshot_path: /path/to/DeepLabCut/deeplabcut/modelzoo/checkpoints/superanimal_topviewmouse_hrnet_w32.pt
    detector_snapshot_path: /path/to/DeepLabCut/deeplabcut/modelzoo/checkpoints/superanimal_topviewmouse_fasterrcnn_resnet50_fpn_v2.pt
    with_decoder: true
    memory_replay: false
    conversion_array:
    - 0
    - 1
    - 2
    - 3
    - 7
    - 8
    - 9
    - 10
    - 13
    - 14
    - 15
    - 25

The conversion array maps the bodyparts labeled in the DeepLabCut project to the indices of SuperAnimal bodyparts.

Bug Fixes

• the test/pose_cfg.yaml file is created correctly for PyTorch shuffles that are fine-tuning SuperAnimal models with memory replay: the bodyparts needed to be updated to contain the SuperAnimal bodyparts
• use the default bottom-up image augmentation when fine-tuning a SuperAnimal model without a detector (i.e. when converting it to a single-animal, bottom-up model)
• create_labeled_video: when using a fine-tuned SuperAnimal model (with memory replay), the bodyparts would not be plot correctly as the project bodyparts were used
• fixes issue RuntimeError in DataLoader: Stack Expects Each Tensor to Be Equal Size During Training #2728
• fixes issue cannot pickle 'memoryview' object on pytorch deeplabcut #2739
• fixes issue PyTorch maDLC on 1 indiv creates saDLC tracks which fail in create_labeled_video #2741

Consistency/Style Updates

• train_network and evaluate_network now accept the project configuration filepath as either a Path or a str
• the train/pose_cfg.yaml file is no longer created for PyTorch shuffles