Skip to content

Conversation

@n-poulsen
Copy link
Contributor

@n-poulsen n-poulsen commented Oct 10, 2024

SuperAnimal Model Updates

This pull requests adds the first SuperAnimal-Bird models to dlclibrary, and refactors the DeepLabCut 3.0 SuperAnimal models so that object detectors and pose models can be downloaded independently.

Things to do before the pull request is merged

  • Add the SuperAnimal-Bird models to HuggingFace
  • Update the SuperAnimal-Quadruped models
  • Update the SuperAnimal-TopViewMouse models

SuperBird Models

An SSDLite object detector and ResNet50 pose model are added for the SuperAnimal-bird dataset.

Model URL Refactor

PyTorch models are moved to a modelzoo_urls_pytorch.yaml file, where snapshots are organized by dataset and task (detectors/ pose_models). This makes it easier to track which models are available with which engine.

New API Methods

Two new API methods are added to make it easier for users to know which detectors and pose models are available for a SuperAnimal:

def get_available_detectors(dataset: str) -> list[str]:
    """ Only for PyTorch models.
    Returns:
        The detectors available for the dataset.
    """
    ...

def get_available_models(dataset: str) -> list[str]:
    """ Only for PyTorch models.
    Returns:
        The pose models available for the dataset.
    """
    ...

Example use:

>>> import dlclibrary
>>> dlclibrary.get_available_detectors("superanimal_bird")
['fasterrcnn_mobilenet_v3_large_fpn', 'ssdlite']

>>> dlclibrary.get_available_models("superanimal_bird")
['resnet_50']

@n-poulsen n-poulsen changed the title Niels/sa detectors SuperAnimal Model Updates Oct 10, 2024
@n-poulsen n-poulsen requested a review from AlexEMG October 11, 2024 07:47
Copy link
Member

@AlexEMG AlexEMG left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great, maybe put your comment also in the readme?

@AlexEMG
Copy link
Member

AlexEMG commented Oct 11, 2024

@n-poulsen n-poulsen merged commit 5a22726 into main Oct 11, 2024
@n-poulsen n-poulsen deleted the niels/sa_detectors branch October 11, 2024 16:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants