123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839 |
- from typing import Any, Callable, List, Optional, Tuple, Union
- import torch
- from torch import nn
- from torchvision.ops import MultiScaleRoIAlign
- from libs.vision_libs.models import MobileNet_V3_Large_Weights, mobilenet_v3_large
- from libs.vision_libs.models.detection.anchor_utils import AnchorGenerator
- from libs.vision_libs.models.detection.rpn import RPNHead, RegionProposalNetwork
- from libs.vision_libs.models.detection.ssdlite import _mobilenet_extractor
- from libs.vision_libs.models.detection.transform import GeneralizedRCNNTransform
- from libs.vision_libs.ops import misc as misc_nn_ops
- from libs.vision_libs.transforms._presets import ObjectDetection
- from .line_head import LineRCNNHeads
- from .line_predictor import LineRCNNPredictor
- from libs.vision_libs.models._api import register_model, Weights, WeightsEnum
- from libs.vision_libs.models._meta import _COCO_PERSON_CATEGORIES, _COCO_PERSON_KEYPOINT_NAMES, _COCO_CATEGORIES
- from libs.vision_libs.models._utils import _ovewrite_value_param, handle_legacy_interface
- from libs.vision_libs.models.resnet import resnet50, ResNet50_Weights
- from libs.vision_libs.models.detection._utils import overwrite_eps
- from libs.vision_libs.models.detection.backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers
- from libs.vision_libs.models.detection.faster_rcnn import FasterRCNN, TwoMLPHead, FastRCNNPredictor
- from .roi_heads import RoIHeads
- from .trainer import Trainer
- from ..base.base_detection_net import BaseDetectionNet
- import torch.nn.functional as F
- from ..config.config_tool import read_yaml
- FEATURE_DIM = 8
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- __all__ = [
- "LineNet",
- "LineNet_ResNet50_FPN_Weights",
- "LineNet_ResNet50_FPN_V2_Weights",
- "LineNet_MobileNet_V3_Large_FPN_Weights",
- "LineNet_MobileNet_V3_Large_320_FPN_Weights",
- "linenet_resnet50_fpn",
- "linenet_resnet50_fpn_v2",
- "linenet_mobilenet_v3_large_fpn",
- "linenet_mobilenet_v3_large_320_fpn",
- ]
- def _default_anchorgen():
- anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
- aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
- return AnchorGenerator(anchor_sizes, aspect_ratios)
- class LineNet(BaseDetectionNet):
- def __init__(self, cfg, **kwargs):
- cfg = read_yaml(cfg)
- backbone = cfg['backbone']
- num_classes = cfg['num_classes']
- if backbone == 'resnet50_fpn':
- is_trained = False
- trainable_backbone_layers = _validate_trainable_layers(is_trained, None, 5, 3)
- norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d
- backbone = resnet50(weights=None, progress=True, norm_layer=norm_layer)
- backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
- out_channels = backbone.out_channels
- min_size = 512,
- max_size = 1333,
- rpn_pre_nms_top_n_train = 2000,
- rpn_pre_nms_top_n_test = 1000,
- rpn_post_nms_top_n_train = 2000,
- rpn_post_nms_top_n_test = 1000,
- rpn_nms_thresh = 0.7,
- rpn_fg_iou_thresh = 0.7,
- rpn_bg_iou_thresh = 0.3,
- rpn_batch_size_per_image = 256,
- rpn_positive_fraction = 0.5,
- rpn_score_thresh = 0.0,
- box_score_thresh = 0.05,
- box_nms_thresh = 0.5,
- box_detections_per_img = 100,
- box_fg_iou_thresh = 0.5,
- box_bg_iou_thresh = 0.5,
- box_batch_size_per_image = 512,
- box_positive_fraction = 0.25,
- bbox_reg_weights = None,
- line_head = LineRCNNHeads(out_channels, 5)
- line_predictor = LineRCNNPredictor(cfg)
- rpn_anchor_generator = _default_anchorgen()
- rpn_head = RPNHead(out_channels, rpn_anchor_generator.num_anchors_per_location()[0])
- rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
- rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)
- rpn = RegionProposalNetwork(
- rpn_anchor_generator,
- rpn_head,
- rpn_fg_iou_thresh,
- rpn_bg_iou_thresh,
- rpn_batch_size_per_image,
- rpn_positive_fraction,
- rpn_pre_nms_top_n,
- rpn_post_nms_top_n,
- rpn_nms_thresh,
- score_thresh=rpn_score_thresh,
- )
- box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2)
- resolution = box_roi_pool.output_size[0]
- representation_size = 1024
- box_head = TwoMLPHead(out_channels * resolution ** 2, representation_size)
- representation_size = 1024
- box_predictor = BoxPredictor(representation_size, num_classes)
- roi_heads = RoIHeads(
- # Box
- box_roi_pool,
- box_head,
- box_predictor,
- line_head,
- line_predictor,
- box_fg_iou_thresh,
- box_bg_iou_thresh,
- box_batch_size_per_image,
- box_positive_fraction,
- bbox_reg_weights,
- box_score_thresh,
- box_nms_thresh,
- box_detections_per_img,
- )
- image_mean = [0.485, 0.456, 0.406]
- image_std = [0.229, 0.224, 0.225]
- transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs)
- super().__init__(backbone, rpn, roi_heads, transform)
- self.roi_heads = roi_heads
- # def __init__(
- # self,
- # backbone,
- # num_classes=None,
- # # transform parameters
- # min_size=512,
- # max_size=1333,
- # image_mean=None,
- # image_std=None,
- # # RPN parameters
- # rpn_anchor_generator=None,
- # rpn_head=None,
- # rpn_pre_nms_top_n_train=2000,
- # rpn_pre_nms_top_n_test=1000,
- # rpn_post_nms_top_n_train=2000,
- # rpn_post_nms_top_n_test=1000,
- # rpn_nms_thresh=0.7,
- # rpn_fg_iou_thresh=0.7,
- # rpn_bg_iou_thresh=0.3,
- # rpn_batch_size_per_image=256,
- # rpn_positive_fraction=0.5,
- # rpn_score_thresh=0.0,
- # # Box parameters
- # box_roi_pool=None,
- # box_head=None,
- # box_predictor=None,
- # box_score_thresh=0.05,
- # box_nms_thresh=0.5,
- # box_detections_per_img=100,
- # box_fg_iou_thresh=0.5,
- # box_bg_iou_thresh=0.5,
- # box_batch_size_per_image=512,
- # box_positive_fraction=0.25,
- # bbox_reg_weights=None,
- # # line parameters
- # line_head=None,
- # line_predictor=None,
- # **kwargs,
- # ):
- #
- # if not hasattr(backbone, "out_channels"):
- # raise ValueError(
- # "backbone should contain an attribute out_channels "
- # "specifying the number of output channels (assumed to be the "
- # "same for all the levels)"
- # )
- #
- # if not isinstance(rpn_anchor_generator, (AnchorGenerator, type(None))):
- # raise TypeError(
- # f"rpn_anchor_generator should be of type AnchorGenerator or None instead of {type(rpn_anchor_generator)}"
- # )
- # if not isinstance(box_roi_pool, (MultiScaleRoIAlign, type(None))):
- # raise TypeError(
- # f"box_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(box_roi_pool)}"
- # )
- #
- # if num_classes is not None:
- # if box_predictor is not None:
- # raise ValueError("num_classes should be None when box_predictor is specified")
- # else:
- # if box_predictor is None:
- # raise ValueError("num_classes should not be None when box_predictor is not specified")
- #
- # out_channels = backbone.out_channels
- #
- # if line_head is None:
- # num_class = 5
- # line_head = LineRCNNHeads(out_channels, num_class)
- #
- # if line_predictor is None:
- # line_predictor = LineRCNNPredictor()
- #
- # if rpn_anchor_generator is None:
- # rpn_anchor_generator = _default_anchorgen()
- # if rpn_head is None:
- # rpn_head = RPNHead(out_channels, rpn_anchor_generator.num_anchors_per_location()[0])
- #
- # rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
- # rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)
- #
- # rpn = RegionProposalNetwork(
- # rpn_anchor_generator,
- # rpn_head,
- # rpn_fg_iou_thresh,
- # rpn_bg_iou_thresh,
- # rpn_batch_size_per_image,
- # rpn_positive_fraction,
- # rpn_pre_nms_top_n,
- # rpn_post_nms_top_n,
- # rpn_nms_thresh,
- # score_thresh=rpn_score_thresh,
- # )
- #
- # if box_roi_pool is None:
- # box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2)
- #
- # if box_head is None:
- # resolution = box_roi_pool.output_size[0]
- # representation_size = 1024
- # box_head = TwoMLPHead(out_channels * resolution ** 2, representation_size)
- #
- # if box_predictor is None:
- # representation_size = 1024
- # box_predictor = BoxPredictor(representation_size, num_classes)
- #
- # roi_heads = RoIHeads(
- # # Box
- # box_roi_pool,
- # box_head,
- # box_predictor,
- # line_head,
- # line_predictor,
- # box_fg_iou_thresh,
- # box_bg_iou_thresh,
- # box_batch_size_per_image,
- # box_positive_fraction,
- # bbox_reg_weights,
- # box_score_thresh,
- # box_nms_thresh,
- # box_detections_per_img,
- # )
- #
- # if image_mean is None:
- # image_mean = [0.485, 0.456, 0.406]
- # if image_std is None:
- # image_std = [0.229, 0.224, 0.225]
- # transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs)
- #
- # super().__init__(backbone, rpn, roi_heads, transform)
- #
- # self.roi_heads = roi_heads
- # self.roi_heads.line_head = line_head
- # self.roi_heads.line_predictor = line_predictor
- def train_by_cfg(self, cfg):
- # cfg = read_yaml(cfg)
- self.trainer = Trainer()
- self.trainer.train_cfg(model=self,cfg=cfg)
- class TwoMLPHead(nn.Module):
- """
- Standard heads for FPN-based models
- Args:
- in_channels (int): number of input channels
- representation_size (int): size of the intermediate representation
- """
- def __init__(self, in_channels, representation_size):
- super().__init__()
- self.fc6 = nn.Linear(in_channels, representation_size)
- self.fc7 = nn.Linear(representation_size, representation_size)
- def forward(self, x):
- x = x.flatten(start_dim=1)
- x = F.relu(self.fc6(x))
- x = F.relu(self.fc7(x))
- return x
- class LineNetConvFCHead(nn.Sequential):
- def __init__(
- self,
- input_size: Tuple[int, int, int],
- conv_layers: List[int],
- fc_layers: List[int],
- norm_layer: Optional[Callable[..., nn.Module]] = None,
- ):
- """
- Args:
- input_size (Tuple[int, int, int]): the input size in CHW format.
- conv_layers (list): feature dimensions of each Convolution layer
- fc_layers (list): feature dimensions of each FCN layer
- norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None
- """
- in_channels, in_height, in_width = input_size
- blocks = []
- previous_channels = in_channels
- for current_channels in conv_layers:
- blocks.append(misc_nn_ops.Conv2dNormActivation(previous_channels, current_channels, norm_layer=norm_layer))
- previous_channels = current_channels
- blocks.append(nn.Flatten())
- previous_channels = previous_channels * in_height * in_width
- for current_channels in fc_layers:
- blocks.append(nn.Linear(previous_channels, current_channels))
- blocks.append(nn.ReLU(inplace=True))
- previous_channels = current_channels
- super().__init__(*blocks)
- for layer in self.modules():
- if isinstance(layer, nn.Conv2d):
- nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu")
- if layer.bias is not None:
- nn.init.zeros_(layer.bias)
- class BoxPredictor(nn.Module):
- """
- Standard classification + bounding box regression layers
- for Fast R-CNN.
- Args:
- in_channels (int): number of input channels
- num_classes (int): number of output classes (including background)
- """
- def __init__(self, in_channels, num_classes):
- super().__init__()
- self.cls_score = nn.Linear(in_channels, num_classes)
- self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
- def forward(self, x):
- if x.dim() == 4:
- torch._assert(
- list(x.shape[2:]) == [1, 1],
- f"x has the wrong shape, expecting the last two dimensions to be [1,1] instead of {list(x.shape[2:])}",
- )
- x = x.flatten(start_dim=1)
- scores = self.cls_score(x)
- bbox_deltas = self.bbox_pred(x)
- return scores, bbox_deltas
- _COMMON_META = {
- "categories": _COCO_CATEGORIES,
- "min_size": (1, 1),
- }
- class LineNet_ResNet50_FPN_Weights(WeightsEnum):
- COCO_V1 = Weights(
- url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
- transforms=ObjectDetection,
- meta={
- **_COMMON_META,
- "num_params": 41755286,
- "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-resnet-50-fpn",
- "_metrics": {
- "COCO-val2017": {
- "box_map": 37.0,
- }
- },
- "_ops": 134.38,
- "_file_size": 159.743,
- "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
- },
- )
- DEFAULT = COCO_V1
- class LineNet_ResNet50_FPN_V2_Weights(WeightsEnum):
- COCO_V1 = Weights(
- url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_v2_coco-dd69338a.pth",
- transforms=ObjectDetection,
- meta={
- **_COMMON_META,
- "num_params": 43712278,
- "recipe": "https://github.com/pytorch/vision/pull/5763",
- "_metrics": {
- "COCO-val2017": {
- "box_map": 46.7,
- }
- },
- "_ops": 280.371,
- "_file_size": 167.104,
- "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
- },
- )
- DEFAULT = COCO_V1
- class LineNet_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
- COCO_V1 = Weights(
- url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth",
- transforms=ObjectDetection,
- meta={
- **_COMMON_META,
- "num_params": 19386354,
- "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-fpn",
- "_metrics": {
- "COCO-val2017": {
- "box_map": 32.8,
- }
- },
- "_ops": 4.494,
- "_file_size": 74.239,
- "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
- },
- )
- DEFAULT = COCO_V1
- class LineNet_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
- COCO_V1 = Weights(
- url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_320_fpn-907ea3f9.pth",
- transforms=ObjectDetection,
- meta={
- **_COMMON_META,
- "num_params": 19386354,
- "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-320-fpn",
- "_metrics": {
- "COCO-val2017": {
- "box_map": 22.8,
- }
- },
- "_ops": 0.719,
- "_file_size": 74.239,
- "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
- },
- )
- DEFAULT = COCO_V1
- @register_model()
- @handle_legacy_interface(
- weights=("pretrained", LineNet_ResNet50_FPN_Weights.COCO_V1),
- weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1),
- )
- def linenet_resnet50_fpn(
- *,
- weights: Optional[LineNet_ResNet50_FPN_Weights] = None,
- progress: bool = True,
- num_classes: Optional[int] = None,
- weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1,
- trainable_backbone_layers: Optional[int] = None,
- **kwargs: Any,
- ) -> LineNet:
- """
- Faster R-CNN model with a ResNet-50-FPN backbone from the `Faster R-CNN: Towards Real-Time Object
- Detection with Region Proposal Networks <https://arxiv.org/abs/1506.01497>`__
- paper.
- .. betastatus:: detection module
- The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
- image, and should be in ``0-1`` range. Different images can have different sizes.
- The behavior of the model changes depending on if it is in training or evaluation mode.
- During training, the model expects both the input tensors and a targets (list of dictionary),
- containing:
- - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with
- ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
- - labels (``Int64Tensor[N]``): the class label for each ground-truth box
- The model returns a ``Dict[Tensor]`` during training, containing the classification and regression
- losses for both the RPN and the R-CNN.
- During inference, the model requires only the input tensors, and returns the post-processed
- predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as
- follows, where ``N`` is the number of detections:
- - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with
- ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
- - labels (``Int64Tensor[N]``): the predicted labels for each detection
- - scores (``Tensor[N]``): the scores of each detection
- For more details on the output, you may refer to :ref:`instance_seg_output`.
- Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
- Example::
- >>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
- >>> # For training
- >>> images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4)
- >>> boxes[:, :, 2:4] = boxes[:, :, 0:2] + boxes[:, :, 2:4]
- >>> labels = torch.randint(1, 91, (4, 11))
- >>> images = list(image for image in images)
- >>> targets = []
- >>> for i in range(len(images)):
- >>> d = {}
- >>> d['boxes'] = boxes[i]
- >>> d['labels'] = labels[i]
- >>> targets.append(d)
- >>> output = model(images, targets)
- >>> # For inference
- >>> model.eval()
- >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
- >>> predictions = model(x)
- >>>
- >>> # optionally, if you want to export the model to ONNX:
- >>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)
- Args:
- weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights`, optional): The
- pretrained weights to use. See
- :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights` below for
- more details, and possible values. By default, no pre-trained
- weights are used.
- progress (bool, optional): If True, displays a progress bar of the
- download to stderr. Default is True.
- num_classes (int, optional): number of output classes of the model (including the background)
- weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The
- pretrained weights for the backbone.
- trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
- final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are
- trainable. If ``None`` is passed (the default) this value is set to 3.
- **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
- base class. Please refer to the `source code
- <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
- for more details about this class.
- .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights
- :members:
- """
- weights = LineNet_ResNet50_FPN_Weights.verify(weights)
- weights_backbone = ResNet50_Weights.verify(weights_backbone)
- if weights is not None:
- weights_backbone = None
- num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
- elif num_classes is None:
- num_classes = 91
- is_trained = weights is not None or weights_backbone is not None
- trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)
- norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d
- backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
- backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
- model = LineNet(backbone, num_classes=num_classes, **kwargs)
- if weights is not None:
- model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
- if weights == LineNet_ResNet50_FPN_Weights.COCO_V1:
- overwrite_eps(model, 0.0)
- return model
- @register_model()
- @handle_legacy_interface(
- weights=("pretrained", LineNet_ResNet50_FPN_V2_Weights.COCO_V1),
- weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1),
- )
- def linenet_resnet50_fpn_v2(
- *,
- weights: Optional[LineNet_ResNet50_FPN_V2_Weights] = None,
- progress: bool = True,
- num_classes: Optional[int] = None,
- weights_backbone: Optional[ResNet50_Weights] = None,
- trainable_backbone_layers: Optional[int] = None,
- **kwargs: Any,
- ) -> LineNet:
- """
- Constructs an improved Faster R-CNN model with a ResNet-50-FPN backbone from `Benchmarking Detection
- Transfer Learning with Vision Transformers <https://arxiv.org/abs/2111.11429>`__ paper.
- .. betastatus:: detection module
- It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See
- :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more
- details.
- Args:
- weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights`, optional): The
- pretrained weights to use. See
- :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights` below for
- more details, and possible values. By default, no pre-trained
- weights are used.
- progress (bool, optional): If True, displays a progress bar of the
- download to stderr. Default is True.
- num_classes (int, optional): number of output classes of the model (including the background)
- weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The
- pretrained weights for the backbone.
- trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
- final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are
- trainable. If ``None`` is passed (the default) this value is set to 3.
- **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
- base class. Please refer to the `source code
- <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
- for more details about this class.
- .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights
- :members:
- """
- weights = LineNet_ResNet50_FPN_V2_Weights.verify(weights)
- weights_backbone = ResNet50_Weights.verify(weights_backbone)
- if weights is not None:
- weights_backbone = None
- num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
- elif num_classes is None:
- num_classes = 91
- is_trained = weights is not None or weights_backbone is not None
- trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)
- backbone = resnet50(weights=weights_backbone, progress=progress)
- backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers, norm_layer=nn.BatchNorm2d)
- rpn_anchor_generator = _default_anchorgen()
- rpn_head = RPNHead(backbone.out_channels, rpn_anchor_generator.num_anchors_per_location()[0], conv_depth=2)
- box_head = LineNetConvFCHead(
- (backbone.out_channels, 7, 7), [256, 256, 256, 256], [1024], norm_layer=nn.BatchNorm2d
- )
- model = LineNet(
- backbone,
- num_classes=num_classes,
- rpn_anchor_generator=rpn_anchor_generator,
- rpn_head=rpn_head,
- box_head=box_head,
- **kwargs,
- )
- if weights is not None:
- model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
- return model
- def _linenet_mobilenet_v3_large_fpn(
- *,
- weights: Optional[Union[LineNet_MobileNet_V3_Large_FPN_Weights, LineNet_MobileNet_V3_Large_320_FPN_Weights]],
- progress: bool,
- num_classes: Optional[int],
- weights_backbone: Optional[MobileNet_V3_Large_Weights],
- trainable_backbone_layers: Optional[int],
- **kwargs: Any,
- ) -> LineNet:
- if weights is not None:
- weights_backbone = None
- num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
- elif num_classes is None:
- num_classes = 91
- is_trained = weights is not None or weights_backbone is not None
- trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 6, 3)
- norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d
- backbone = mobilenet_v3_large(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
- backbone = _mobilenet_extractor(backbone, True, trainable_backbone_layers)
- anchor_sizes = (
- (
- 32,
- 64,
- 128,
- 256,
- 512,
- ),
- ) * 3
- aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
- model = LineNet(
- backbone, num_classes, rpn_anchor_generator=AnchorGenerator(anchor_sizes, aspect_ratios), **kwargs
- )
- if weights is not None:
- model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
- return model
- @register_model()
- @handle_legacy_interface(
- weights=("pretrained", LineNet_MobileNet_V3_Large_320_FPN_Weights.COCO_V1),
- weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1),
- )
- def linenet_mobilenet_v3_large_320_fpn(
- *,
- weights: Optional[LineNet_MobileNet_V3_Large_320_FPN_Weights] = None,
- progress: bool = True,
- num_classes: Optional[int] = None,
- weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1,
- trainable_backbone_layers: Optional[int] = None,
- **kwargs: Any,
- ) -> LineNet:
- """
- Low resolution Faster R-CNN model with a MobileNetV3-Large backbone tuned for mobile use cases.
- .. betastatus:: detection module
- It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See
- :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more
- details.
- Example::
- >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(weights=FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT)
- >>> model.eval()
- >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
- >>> predictions = model(x)
- Args:
- weights (:class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights`, optional): The
- pretrained weights to use. See
- :class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights` below for
- more details, and possible values. By default, no pre-trained
- weights are used.
- progress (bool, optional): If True, displays a progress bar of the
- download to stderr. Default is True.
- num_classes (int, optional): number of output classes of the model (including the background)
- weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The
- pretrained weights for the backbone.
- trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
- final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are
- trainable. If ``None`` is passed (the default) this value is set to 3.
- **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
- base class. Please refer to the `source code
- <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
- for more details about this class.
- .. autoclass:: torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights
- :members:
- """
- weights = LineNet_MobileNet_V3_Large_320_FPN_Weights.verify(weights)
- weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone)
- defaults = {
- "min_size": 320,
- "max_size": 640,
- "rpn_pre_nms_top_n_test": 150,
- "rpn_post_nms_top_n_test": 150,
- "rpn_score_thresh": 0.05,
- }
- kwargs = {**defaults, **kwargs}
- return _linenet_mobilenet_v3_large_fpn(
- weights=weights,
- progress=progress,
- num_classes=num_classes,
- weights_backbone=weights_backbone,
- trainable_backbone_layers=trainable_backbone_layers,
- **kwargs,
- )
- @register_model()
- @handle_legacy_interface(
- weights=("pretrained", LineNet_MobileNet_V3_Large_FPN_Weights.COCO_V1),
- weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1),
- )
- def linenet_mobilenet_v3_large_fpn(
- *,
- weights: Optional[LineNet_MobileNet_V3_Large_FPN_Weights] = None,
- progress: bool = True,
- num_classes: Optional[int] = None,
- weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1,
- trainable_backbone_layers: Optional[int] = None,
- **kwargs: Any,
- ) -> LineNet:
- """
- Constructs a high resolution Faster R-CNN model with a MobileNetV3-Large FPN backbone.
- .. betastatus:: detection module
- It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See
- :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more
- details.
- Example::
- >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights=FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT)
- >>> model.eval()
- >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
- >>> predictions = model(x)
- Args:
- weights (:class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights`, optional): The
- pretrained weights to use. See
- :class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights` below for
- more details, and possible values. By default, no pre-trained
- weights are used.
- progress (bool, optional): If True, displays a progress bar of the
- download to stderr. Default is True.
- num_classes (int, optional): number of output classes of the model (including the background)
- weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The
- pretrained weights for the backbone.
- trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
- final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are
- trainable. If ``None`` is passed (the default) this value is set to 3.
- **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
- base class. Please refer to the `source code
- <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
- for more details about this class.
- .. autoclass:: torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights
- :members:
- """
- weights = LineNet_MobileNet_V3_Large_FPN_Weights.verify(weights)
- weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone)
- defaults = {
- "rpn_score_thresh": 0.05,
- }
- kwargs = {**defaults, **kwargs}
- return _linenet_mobilenet_v3_large_fpn(
- weights=weights,
- progress=progress,
- num_classes=num_classes,
- weights_backbone=weights_backbone,
- trainable_backbone_layers=trainable_backbone_layers,
- **kwargs,
- )
|