Ejemplo n.º 1
0
def model(num_classes: int,
          backbone: Optional[TorchvisionBackboneConfig] = None,
          remove_internal_transforms: bool = True,
          pretrained: bool = True,
          **retinanet_kwargs) -> nn.Module:
    if backbone is None:
        model = retinanet_resnet50_fpn(pretrained=pretrained,
                                       pretrained_backbone=pretrained,
                                       **retinanet_kwargs)
        model.head = RetinaNetHead(
            in_channels=model.backbone.out_channels,
            num_anchors=model.head.classification_head.num_anchors,
            num_classes=num_classes,
        )
        resnet_fpn.patch_param_groups(model.backbone)
    else:
        model = RetinaNet(backbone=backbone.backbone,
                          num_classes=num_classes,
                          **retinanet_kwargs)

    patch_retinanet_param_groups(model)

    if remove_internal_transforms:
        remove_internal_model_transforms(model)

    return model
Ejemplo n.º 2
0
def create_retinanet(
    num_classes: int = 91,
    backbone: nn.Module = None,
    **kwargs,
):
    """
    Creates RetinaNet implementation based on torchvision library.
    Args:
    num_classes (int) : number of classes.
    Do not have class_id "0" it is reserved as background.
    num_classes = number of classes to label + 1 for background.
    """
    if backbone is None:
        model = retinanet_resnet50_fpn(
            pretrained=True,
            num_classes=91,
            **kwargs,
        )
        model.head = RetinaNetHead(
            in_channels=model.backbone.out_channels,
            num_anchors=model.head.classification_head.num_anchors,
            num_classes=num_classes,
        )
    else:
        model = RetinaNet(backbone, num_classes=num_classes, **kwargs)

    return model
Ejemplo n.º 3
0
    def __init__(self, learning_rate: float = 0.0001, num_classes: int = 91,
                 backbone: str = None, fpn: bool = True,
                 pretrained_backbone: str = None, trainable_backbone_layers: int = 3,
                 **kwargs, ):
        """
        Args:
            learning_rate: the learning rate
            num_classes: number of detection classes (including background)
            pretrained: if true, returns a model pre-trained on COCO train2017
            pretrained_backbone (str): if "imagenet", returns a model with backbone pre-trained on Imagenet
            trainable_backbone_layers: number of trainable resnet layers starting from final block
        """
        super().__init__()
        self.learning_rate = learning_rate
        self.num_classes = num_classes
        self.backbone = backbone
        if backbone is None:
            self.model = retinanet_resnet50_fpn(pretrained=True, **kwargs)

            self.model.head = RetinaNetHead(in_channels=self.model.backbone.out_channels,
                                            num_anchors=self.model.head.classification_head.num_anchors,
                                            num_classes=num_classes, **kwargs)

        else:
            backbone_model = create_retinanet_backbone(self.backbone, fpn, pretrained_backbone,
                                                       trainable_backbone_layers, **kwargs)
            self.model = RetinaNet(backbone_model, num_classes=num_classes, **kwargs)
Ejemplo n.º 4
0
    def get_model(
        model_name,
        num_classes,
        backbone,
        fpn,
        pretrained,
        pretrained_backbone,
        trainable_backbone_layers,
        anchor_generator,
        **kwargs,
    ):
        if backbone is None:
            # Constructs a model with a ResNet-50-FPN backbone when no backbone is specified.
            if model_name == "fasterrcnn":
                model = _models[model_name](
                    pretrained=pretrained,
                    pretrained_backbone=pretrained_backbone,
                    trainable_backbone_layers=trainable_backbone_layers,
                )
                in_features = model.roi_heads.box_predictor.cls_score.in_features
                head = FastRCNNPredictor(in_features, num_classes)
                model.roi_heads.box_predictor = head
            else:
                model = _models[model_name](
                    pretrained=pretrained,
                    pretrained_backbone=pretrained_backbone)
                model.head = RetinaNetHead(
                    in_channels=model.backbone.out_channels,
                    num_anchors=model.head.classification_head.num_anchors,
                    num_classes=num_classes,
                    **kwargs)
        else:
            backbone_model, num_features = ObjectDetector.backbones.get(
                backbone)(
                    pretrained=pretrained_backbone,
                    trainable_layers=trainable_backbone_layers,
                    **kwargs,
                )
            backbone_model.out_channels = num_features
            if anchor_generator is None:
                anchor_generator = AnchorGenerator(
                    sizes=((32, 64, 128, 256, 512), ),
                    aspect_ratios=((0.5, 1.0, 2.0), )) if not hasattr(
                        backbone_model, "fpn") else None

            if model_name == "fasterrcnn":
                model = FasterRCNN(backbone_model,
                                   num_classes=num_classes,
                                   rpn_anchor_generator=anchor_generator)
            else:
                model = RetinaNet(backbone_model,
                                  num_classes=num_classes,
                                  anchor_generator=anchor_generator)
        return model
Ejemplo n.º 5
0
    def fine_tune_model(self, pretrained):
        num_classes = 2

        model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=pretrained)

        # get number of input features and anchor boxed for the classifier
        in_features = model.head.classification_head.conv[0].in_channels
        num_anchors = model.head.classification_head.num_anchors

        # replace the pre-trained head with a new one
        model.head = RetinaNetHead(in_features, num_anchors, num_classes)
        return model
Ejemplo n.º 6
0
    def __init__(
        self,
        learning_rate: float = 0.0001,
        num_classes: int = 91,
        backbone: Optional[str] = None,
        fpn: bool = True,
        pretrained: bool = False,
        pretrained_backbone: bool = True,
        trainable_backbone_layers: int = 3,
        **kwargs: Any,
    ):
        """
        Args:
            learning_rate: the learning rate
            num_classes: number of detection classes (including background)
            backbone: Pretained backbone CNN architecture.
            fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs.
            pretrained: if true, returns a model pre-trained on COCO train2017
            pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet
            trainable_backbone_layers: number of trainable resnet layers starting from final block
        """
        super().__init__()
        self.learning_rate = learning_rate
        self.num_classes = num_classes
        self.backbone = backbone
        if backbone is None:
            self.model = retinanet_resnet50_fpn(pretrained=pretrained,
                                                **kwargs)

            self.model.head = RetinaNetHead(
                in_channels=self.model.backbone.out_channels,
                num_anchors=self.model.head.classification_head.num_anchors,
                num_classes=num_classes,
                **kwargs,
            )

        else:
            backbone_model = create_retinanet_backbone(
                self.backbone, fpn, pretrained_backbone,
                trainable_backbone_layers, **kwargs)
            self.model = torchvision_RetinaNet(backbone_model,
                                               num_classes=num_classes,
                                               **kwargs)