Skip to content

deeplabcut.pose_estimation_pytorch.models.detectors.torchvision

Module to adapt torchvision detectors for DeepLabCut.

Classes:

Name Description
TorchvisionDetectorAdaptor

An adaptor for torchvision detectors.

TorchvisionDetectorAdaptor

Bases: BaseDetector

An adaptor for torchvision detectors.

This class is an adaptor for torchvision detectors to DeepLabCut detectors. Some of the models (from fastest to most powerful) available are: - ssdlite320_mobilenet_v3_large - fasterrcnn_mobilenet_v3_large_fpn - fasterrcnn_resnet50_fpn_v2

This class should not be used out-of-the-box. Subclasses (such as FasterRCNN or SSDLite) should be used instead.

The torchvision implementation does not allow to get both predictions and losses with a single forward pass. Therefore, during evaluation only bounding box metrics (mAP, mAR) are available for the test set. See validation loss issue: - https://discuss.pytorch.org/t/compute-validation-loss-for-faster-rcnn/62333/12 - https://stackoverflow.com/a/65347721

Parameters:

Name Type Description Default

model

str

The torchvision model to use (see all options at https://pytorch.org/vision/stable/models.html#object-detection).

 required

weights

str | None

The weights to load for the model. If None, no pre-trained weights are loaded.

 None

num_classes

int | None

Number of classes that the model should output. If None, the number of classes the model is pre-trained on is used.

 2

freeze_bn_stats

bool

Whether to freeze stats for BatchNorm layers.

 False

freeze_bn_weights

bool

Whether to freeze weights for BatchNorm layers.

 False

box_score_thresh

float

during inference, only return proposals with a classification score greater than box_score_thresh

 0.01

Methods:

Name Description
forward

Forward pass of the torchvision detector.
get_target

Returns target in a format a torchvision detector can handle.
Source code in deeplabcut/pose_estimation_pytorch/models/detectors/torchvision.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
class TorchvisionDetectorAdaptor(BaseDetector):
    """An adaptor for torchvision detectors.

    This class is an adaptor for torchvision detectors to DeepLabCut detectors. Some of
    the models (from fastest to most powerful) available are:
      - ssdlite320_mobilenet_v3_large
      - fasterrcnn_mobilenet_v3_large_fpn
      - fasterrcnn_resnet50_fpn_v2

    This class should not be used out-of-the-box. Subclasses (such as FasterRCNN or
    SSDLite) should be used instead.

    The torchvision implementation does not allow to get both predictions and losses
    with a single forward pass. Therefore, during evaluation only bounding box metrics
    (mAP, mAR) are available for the test set. See validation loss issue:
    - https://discuss.pytorch.org/t/compute-validation-loss-for-faster-rcnn/62333/12
    - https://stackoverflow.com/a/65347721

    Args:
        model: The torchvision model to use (see all options at
            https://pytorch.org/vision/stable/models.html#object-detection).
        weights: The weights to load for the model. If None, no pre-trained weights are
            loaded.
        num_classes: Number of classes that the model should output. If None, the number
            of classes the model is pre-trained on is used.
        freeze_bn_stats: Whether to freeze stats for BatchNorm layers.
        freeze_bn_weights: Whether to freeze weights for BatchNorm layers.
        box_score_thresh: during inference, only return proposals with a classification
            score greater than box_score_thresh
    """

    def __init__(
        self,
        model: str,
        weights: str | None = None,
        num_classes: int | None = 2,
        freeze_bn_stats: bool = False,
        freeze_bn_weights: bool = False,
        box_score_thresh: float = 0.01,
        model_kwargs: dict | None = None,
    ) -> None:
        super().__init__(
            freeze_bn_stats=freeze_bn_stats,
            freeze_bn_weights=freeze_bn_weights,
            pretrained=weights is not None,
        )

        # Load the model
        model_fn = getattr(detection, model)
        if model_kwargs is None:
            model_kwargs = {}

        self.model = model_fn(
            weights=weights,
            box_score_thresh=box_score_thresh,
            num_classes=num_classes,
            **model_kwargs,
        )

        # See source:  https://stackoverflow.com/a/65347721
        self.model.eager_outputs = lambda losses, detections: (losses, detections)

    def forward(
        self, x: torch.Tensor, targets: list[dict[str, torch.Tensor]] | None = None
    ) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]]]:
        """Forward pass of the torchvision detector.

        Args:
            x: images to be processed, of shape (b, c, h, w)
            targets: ground-truth boxes present in the images

        Returns:
            losses: {'loss_name': loss_value}
            detections: for each of the b images, {"boxes": bounding_boxes}
        """
        return self.model(x, targets)

    def get_target(self, labels: dict) -> list[dict[str, torch.Tensor]]:
        """Returns target in a format a torchvision detector can handle.

        Args:
            labels: dict of annotations, must contain the keys:
                area: tensor containing area information for each annotation
                labels: tensor containing class labels for each annotation
                is_crowd: tensor indicating if each annotation is a crowd (1) or not (0)
                image_id: tensor containing image ids for each annotation
                boxes: tensor containing bounding box information for each annotation

        Returns:
            res: list of dictionaries, each representing target information for a single
                annotation. Each dictionary contains the following keys:
                    'area'
                    'labels'
                    'is_crowd'
                    'boxes'

        Examples:
            input:
                annotations = {
                    "area": torch.Tensor([100, 200]),
                    "labels": torch.Tensor([1, 2]),
                    "is_crowd": torch.Tensor([0, 1]),
                    "boxes": torch.Tensor([[10, 20, 30, 40], [50, 60, 70, 80]])
                }
            output:
                res =  [
                    {
                        'area': tensor([100.]),
                        'labels': tensor([1]),
                        'image_id': tensor([1]),
                        'is_crowd': tensor([0]),
                        'boxes': tensor([[10., 20., 40., 60.]])
                    },
                    {
                        'area': tensor([200.]),
                        'labels': tensor([2]),
                        'image_id': tensor([1]),
                        'is_crowd': tensor([1]),
                        'boxes': tensor([[50., 60., 70., 80.]])
                    }
                ]
        """
        res = []
        for i, box_ann in enumerate(labels["boxes"]):
            mask = (box_ann[:, 2] > 0.0) & (box_ann[:, 3] > 0.0)
            box_ann = box_ann[mask]
            # bbox format conversion (x, y, w, h) -> (x1, y1, x2, y2)
            box_ann[:, 2] += box_ann[:, 0]
            box_ann[:, 3] += box_ann[:, 1]
            res.append(
                {
                    "area": labels["area"][i][mask],
                    "labels": labels["labels"][i][mask].long(),
                    "is_crowd": labels["is_crowd"][i][mask].long(),
                    "boxes": box_ann,
                }
            )

        return res

forward 

forward(
    x: Tensor, targets: list[dict[str, Tensor]] | None = None
) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]]]

Forward pass of the torchvision detector.

Parameters:

Name Type Description Default

x

Tensor

images to be processed, of shape (b, c, h, w)

 required

targets

list[dict[str, Tensor]] | None

ground-truth boxes present in the images

 None

Returns:

Name Type Description
losses tuple[dict[str, Tensor], list[dict[str, Tensor]]]

{'loss_name': loss_value} detections: for each of the b images, {"boxes": bounding_boxes}
Source code in deeplabcut/pose_estimation_pytorch/models/detectors/torchvision.py 
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def forward(
    self, x: torch.Tensor, targets: list[dict[str, torch.Tensor]] | None = None
) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]]]:
    """Forward pass of the torchvision detector.

    Args:
        x: images to be processed, of shape (b, c, h, w)
        targets: ground-truth boxes present in the images

    Returns:
        losses: {'loss_name': loss_value}
        detections: for each of the b images, {"boxes": bounding_boxes}
    """
    return self.model(x, targets)

get_target 

get_target(labels: dict) -> list[dict[str, torch.Tensor]]

Returns target in a format a torchvision detector can handle.

Parameters:

Name Type Description Default

labels

dict

dict of annotations, must contain the keys: area: tensor containing area information for each annotation labels: tensor containing class labels for each annotation is_crowd: tensor indicating if each annotation is a crowd (1) or not (0) image_id: tensor containing image ids for each annotation boxes: tensor containing bounding box information for each annotation

 required

Returns:

Name Type Description
res list[dict[str, Tensor]]

list of dictionaries, each representing target information for a single annotation. Each dictionary contains the following keys: 'area' 'labels' 'is_crowd' 'boxes'

Examples:

input: annotations = { "area": torch.Tensor([100, 200]), "labels": torch.Tensor([1, 2]), "is_crowd": torch.Tensor([0, 1]), "boxes": torch.Tensor([[10, 20, 30, 40], [50, 60, 70, 80]]) } output: res = [ { 'area': tensor([100.]), 'labels': tensor([1]), 'image_id': tensor([1]), 'is_crowd': tensor([0]), 'boxes': tensor([[10., 20., 40., 60.]]) }, { 'area': tensor([200.]), 'labels': tensor([2]), 'image_id': tensor([1]), 'is_crowd': tensor([1]), 'boxes': tensor([[50., 60., 70., 80.]]) } ]

Source code in deeplabcut/pose_estimation_pytorch/models/detectors/torchvision.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def get_target(self, labels: dict) -> list[dict[str, torch.Tensor]]:
    """Returns target in a format a torchvision detector can handle.

    Args:
        labels: dict of annotations, must contain the keys:
            area: tensor containing area information for each annotation
            labels: tensor containing class labels for each annotation
            is_crowd: tensor indicating if each annotation is a crowd (1) or not (0)
            image_id: tensor containing image ids for each annotation
            boxes: tensor containing bounding box information for each annotation

    Returns:
        res: list of dictionaries, each representing target information for a single
            annotation. Each dictionary contains the following keys:
                'area'
                'labels'
                'is_crowd'
                'boxes'

    Examples:
        input:
            annotations = {
                "area": torch.Tensor([100, 200]),
                "labels": torch.Tensor([1, 2]),
                "is_crowd": torch.Tensor([0, 1]),
                "boxes": torch.Tensor([[10, 20, 30, 40], [50, 60, 70, 80]])
            }
        output:
            res =  [
                {
                    'area': tensor([100.]),
                    'labels': tensor([1]),
                    'image_id': tensor([1]),
                    'is_crowd': tensor([0]),
                    'boxes': tensor([[10., 20., 40., 60.]])
                },
                {
                    'area': tensor([200.]),
                    'labels': tensor([2]),
                    'image_id': tensor([1]),
                    'is_crowd': tensor([1]),
                    'boxes': tensor([[50., 60., 70., 80.]])
                }
            ]
    """
    res = []
    for i, box_ann in enumerate(labels["boxes"]):
        mask = (box_ann[:, 2] > 0.0) & (box_ann[:, 3] > 0.0)
        box_ann = box_ann[mask]
        # bbox format conversion (x, y, w, h) -> (x1, y1, x2, y2)
        box_ann[:, 2] += box_ann[:, 0]
        box_ann[:, 3] += box_ann[:, 1]
        res.append(
            {
                "area": labels["area"][i][mask],
                "labels": labels["labels"][i][mask].long(),
                "is_crowd": labels["is_crowd"][i][mask].long(),
                "boxes": box_ann,
            }
        )

    return res