line_net.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768
  1. from typing import Any, Callable, List, Optional, Tuple, Union
  2. import torch
  3. from torch import nn
  4. from torchvision.ops import MultiScaleRoIAlign
  5. from libs.vision_libs.models import MobileNet_V3_Large_Weights, mobilenet_v3_large
  6. from libs.vision_libs.models.detection.anchor_utils import AnchorGenerator
  7. from libs.vision_libs.models.detection.rpn import RPNHead, RegionProposalNetwork
  8. from libs.vision_libs.models.detection.ssdlite import _mobilenet_extractor
  9. from libs.vision_libs.models.detection.transform import GeneralizedRCNNTransform
  10. from libs.vision_libs.ops import misc as misc_nn_ops
  11. from libs.vision_libs.transforms._presets import ObjectDetection
  12. from .line_head import LineRCNNHeads
  13. from .line_predictor import LineRCNNPredictor
  14. from libs.vision_libs.models._api import register_model, Weights, WeightsEnum
  15. from libs.vision_libs.models._meta import _COCO_PERSON_CATEGORIES, _COCO_PERSON_KEYPOINT_NAMES, _COCO_CATEGORIES
  16. from libs.vision_libs.models._utils import _ovewrite_value_param, handle_legacy_interface
  17. from libs.vision_libs.models.resnet import resnet50, ResNet50_Weights
  18. from libs.vision_libs.models.detection._utils import overwrite_eps
  19. from libs.vision_libs.models.detection.backbone_utils import _resnet_fpn_extractor, _validate_trainable_layers
  20. from libs.vision_libs.models.detection.faster_rcnn import FasterRCNN, TwoMLPHead, FastRCNNPredictor
  21. from .roi_heads import RoIHeads
  22. from .trainer import Trainer
  23. from ..base import backbone_factory
  24. from ..base.base_detection_net import BaseDetectionNet
  25. import torch.nn.functional as F
  26. from ..config.config_tool import read_yaml
  27. FEATURE_DIM = 8
  28. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  29. __all__ = [
  30. "LineNet",
  31. "LineNet_ResNet50_FPN_Weights",
  32. "LineNet_ResNet50_FPN_V2_Weights",
  33. "LineNet_MobileNet_V3_Large_FPN_Weights",
  34. "LineNet_MobileNet_V3_Large_320_FPN_Weights",
  35. "linenet_resnet50_fpn",
  36. "linenet_resnet50_fpn_v2",
  37. "linenet_mobilenet_v3_large_fpn",
  38. "linenet_mobilenet_v3_large_320_fpn",
  39. ]
  40. def _default_anchorgen():
  41. anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
  42. aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
  43. return AnchorGenerator(anchor_sizes, aspect_ratios)
  44. class LineNet(BaseDetectionNet):
  45. def __init__(self, cfg, **kwargs):
  46. cfg = read_yaml(cfg)
  47. self.cfg=cfg
  48. backbone = cfg['backbone']
  49. num_classes = cfg['num_classes']
  50. if backbone == 'resnet50_fpn':
  51. backbone=backbone_factory.get_resnet50_fpn()
  52. print(f'out_chanenels:{backbone.out_channels}')
  53. self.__construct__(backbone=backbone, num_classes=num_classes, **kwargs)
  54. def __construct__(
  55. self,
  56. backbone,
  57. num_classes=None,
  58. # transform parameters
  59. min_size=512,
  60. max_size=1333,
  61. image_mean=None,
  62. image_std=None,
  63. # RPN parameters
  64. rpn_anchor_generator=None,
  65. rpn_head=None,
  66. rpn_pre_nms_top_n_train=2000,
  67. rpn_pre_nms_top_n_test=1000,
  68. rpn_post_nms_top_n_train=2000,
  69. rpn_post_nms_top_n_test=1000,
  70. rpn_nms_thresh=0.7,
  71. rpn_fg_iou_thresh=0.7,
  72. rpn_bg_iou_thresh=0.3,
  73. rpn_batch_size_per_image=256,
  74. rpn_positive_fraction=0.5,
  75. rpn_score_thresh=0.0,
  76. # Box parameters
  77. box_roi_pool=None,
  78. box_head=None,
  79. box_predictor=None,
  80. box_score_thresh=0.05,
  81. box_nms_thresh=0.5,
  82. box_detections_per_img=100,
  83. box_fg_iou_thresh=0.5,
  84. box_bg_iou_thresh=0.5,
  85. box_batch_size_per_image=512,
  86. box_positive_fraction=0.25,
  87. bbox_reg_weights=None,
  88. # line parameters
  89. line_head=None,
  90. line_predictor=None,
  91. **kwargs,
  92. ):
  93. if not hasattr(backbone, "out_channels"):
  94. raise ValueError(
  95. "backbone should contain an attribute out_channels "
  96. "specifying the number of output channels (assumed to be the "
  97. "same for all the levels)"
  98. )
  99. if not isinstance(rpn_anchor_generator, (AnchorGenerator, type(None))):
  100. raise TypeError(
  101. f"rpn_anchor_generator should be of type AnchorGenerator or None instead of {type(rpn_anchor_generator)}"
  102. )
  103. if not isinstance(box_roi_pool, (MultiScaleRoIAlign, type(None))):
  104. raise TypeError(
  105. f"box_roi_pool should be of type MultiScaleRoIAlign or None instead of {type(box_roi_pool)}"
  106. )
  107. if num_classes is not None:
  108. if box_predictor is not None:
  109. raise ValueError("num_classes should be None when box_predictor is specified")
  110. else:
  111. if box_predictor is None:
  112. raise ValueError("num_classes should not be None when box_predictor is not specified")
  113. out_channels = backbone.out_channels
  114. if line_head is None:
  115. num_class = 5
  116. line_head = LineRCNNHeads(out_channels, num_class)
  117. if line_predictor is None:
  118. line_predictor = LineRCNNPredictor(self.cfg)
  119. if rpn_anchor_generator is None:
  120. rpn_anchor_generator = _default_anchorgen()
  121. if rpn_head is None:
  122. rpn_head = RPNHead(out_channels, rpn_anchor_generator.num_anchors_per_location()[0])
  123. rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
  124. rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)
  125. rpn = RegionProposalNetwork(
  126. rpn_anchor_generator,
  127. rpn_head,
  128. rpn_fg_iou_thresh,
  129. rpn_bg_iou_thresh,
  130. rpn_batch_size_per_image,
  131. rpn_positive_fraction,
  132. rpn_pre_nms_top_n,
  133. rpn_post_nms_top_n,
  134. rpn_nms_thresh,
  135. score_thresh=rpn_score_thresh,
  136. )
  137. if box_roi_pool is None:
  138. box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2)
  139. if box_head is None:
  140. resolution = box_roi_pool.output_size[0]
  141. representation_size = 1024
  142. box_head = TwoMLPHead(out_channels * resolution ** 2, representation_size)
  143. if box_predictor is None:
  144. representation_size = 1024
  145. box_predictor = BoxPredictor(representation_size, num_classes)
  146. roi_heads = RoIHeads(
  147. # Box
  148. box_roi_pool,
  149. box_head,
  150. box_predictor,
  151. line_head,
  152. line_predictor,
  153. box_fg_iou_thresh,
  154. box_bg_iou_thresh,
  155. box_batch_size_per_image,
  156. box_positive_fraction,
  157. bbox_reg_weights,
  158. box_score_thresh,
  159. box_nms_thresh,
  160. box_detections_per_img,
  161. )
  162. if image_mean is None:
  163. image_mean = [0.485, 0.456, 0.406]
  164. if image_std is None:
  165. image_std = [0.229, 0.224, 0.225]
  166. transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std, **kwargs)
  167. super().__init__(backbone, rpn, roi_heads, transform)
  168. self.roi_heads = roi_heads
  169. self.roi_heads.line_head = line_head
  170. self.roi_heads.line_predictor = line_predictor
  171. def train_by_cfg(self, cfg):
  172. # cfg = read_yaml(cfg)
  173. self.trainer = Trainer()
  174. self.trainer.train_cfg(model=self,cfg=cfg)
  175. class TwoMLPHead(nn.Module):
  176. """
  177. Standard heads for FPN-based models
  178. Args:
  179. in_channels (int): number of input channels
  180. representation_size (int): size of the intermediate representation
  181. """
  182. def __init__(self, in_channels, representation_size):
  183. super().__init__()
  184. self.fc6 = nn.Linear(in_channels, representation_size)
  185. self.fc7 = nn.Linear(representation_size, representation_size)
  186. def forward(self, x):
  187. x = x.flatten(start_dim=1)
  188. x = F.relu(self.fc6(x))
  189. x = F.relu(self.fc7(x))
  190. return x
  191. class LineNetConvFCHead(nn.Sequential):
  192. def __init__(
  193. self,
  194. input_size: Tuple[int, int, int],
  195. conv_layers: List[int],
  196. fc_layers: List[int],
  197. norm_layer: Optional[Callable[..., nn.Module]] = None,
  198. ):
  199. """
  200. Args:
  201. input_size (Tuple[int, int, int]): the input size in CHW format.
  202. conv_layers (list): feature dimensions of each Convolution layer
  203. fc_layers (list): feature dimensions of each FCN layer
  204. norm_layer (callable, optional): Module specifying the normalization layer to use. Default: None
  205. """
  206. in_channels, in_height, in_width = input_size
  207. blocks = []
  208. previous_channels = in_channels
  209. for current_channels in conv_layers:
  210. blocks.append(misc_nn_ops.Conv2dNormActivation(previous_channels, current_channels, norm_layer=norm_layer))
  211. previous_channels = current_channels
  212. blocks.append(nn.Flatten())
  213. previous_channels = previous_channels * in_height * in_width
  214. for current_channels in fc_layers:
  215. blocks.append(nn.Linear(previous_channels, current_channels))
  216. blocks.append(nn.ReLU(inplace=True))
  217. previous_channels = current_channels
  218. super().__init__(*blocks)
  219. for layer in self.modules():
  220. if isinstance(layer, nn.Conv2d):
  221. nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu")
  222. if layer.bias is not None:
  223. nn.init.zeros_(layer.bias)
  224. class BoxPredictor(nn.Module):
  225. """
  226. Standard classification + bounding box regression layers
  227. for Fast R-CNN.
  228. Args:
  229. in_channels (int): number of input channels
  230. num_classes (int): number of output classes (including background)
  231. """
  232. def __init__(self, in_channels, num_classes):
  233. super().__init__()
  234. self.cls_score = nn.Linear(in_channels, num_classes)
  235. self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
  236. def forward(self, x):
  237. if x.dim() == 4:
  238. torch._assert(
  239. list(x.shape[2:]) == [1, 1],
  240. f"x has the wrong shape, expecting the last two dimensions to be [1,1] instead of {list(x.shape[2:])}",
  241. )
  242. x = x.flatten(start_dim=1)
  243. scores = self.cls_score(x)
  244. bbox_deltas = self.bbox_pred(x)
  245. return scores, bbox_deltas
  246. _COMMON_META = {
  247. "categories": _COCO_CATEGORIES,
  248. "min_size": (1, 1),
  249. }
  250. class LineNet_ResNet50_FPN_Weights(WeightsEnum):
  251. COCO_V1 = Weights(
  252. url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
  253. transforms=ObjectDetection,
  254. meta={
  255. **_COMMON_META,
  256. "num_params": 41755286,
  257. "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-resnet-50-fpn",
  258. "_metrics": {
  259. "COCO-val2017": {
  260. "box_map": 37.0,
  261. }
  262. },
  263. "_ops": 134.38,
  264. "_file_size": 159.743,
  265. "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
  266. },
  267. )
  268. DEFAULT = COCO_V1
  269. class LineNet_ResNet50_FPN_V2_Weights(WeightsEnum):
  270. COCO_V1 = Weights(
  271. url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_v2_coco-dd69338a.pth",
  272. transforms=ObjectDetection,
  273. meta={
  274. **_COMMON_META,
  275. "num_params": 43712278,
  276. "recipe": "https://github.com/pytorch/vision/pull/5763",
  277. "_metrics": {
  278. "COCO-val2017": {
  279. "box_map": 46.7,
  280. }
  281. },
  282. "_ops": 280.371,
  283. "_file_size": 167.104,
  284. "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
  285. },
  286. )
  287. DEFAULT = COCO_V1
  288. class LineNet_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
  289. COCO_V1 = Weights(
  290. url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth",
  291. transforms=ObjectDetection,
  292. meta={
  293. **_COMMON_META,
  294. "num_params": 19386354,
  295. "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-fpn",
  296. "_metrics": {
  297. "COCO-val2017": {
  298. "box_map": 32.8,
  299. }
  300. },
  301. "_ops": 4.494,
  302. "_file_size": 74.239,
  303. "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
  304. },
  305. )
  306. DEFAULT = COCO_V1
  307. class LineNet_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
  308. COCO_V1 = Weights(
  309. url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_320_fpn-907ea3f9.pth",
  310. transforms=ObjectDetection,
  311. meta={
  312. **_COMMON_META,
  313. "num_params": 19386354,
  314. "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-320-fpn",
  315. "_metrics": {
  316. "COCO-val2017": {
  317. "box_map": 22.8,
  318. }
  319. },
  320. "_ops": 0.719,
  321. "_file_size": 74.239,
  322. "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
  323. },
  324. )
  325. DEFAULT = COCO_V1
  326. @register_model()
  327. @handle_legacy_interface(
  328. weights=("pretrained", LineNet_ResNet50_FPN_Weights.COCO_V1),
  329. weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1),
  330. )
  331. def linenet_resnet50_fpn(
  332. *,
  333. weights: Optional[LineNet_ResNet50_FPN_Weights] = None,
  334. progress: bool = True,
  335. num_classes: Optional[int] = None,
  336. weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1,
  337. trainable_backbone_layers: Optional[int] = None,
  338. **kwargs: Any,
  339. ) -> LineNet:
  340. """
  341. Faster R-CNN model with a ResNet-50-FPN backbone from the `Faster R-CNN: Towards Real-Time Object
  342. Detection with Region Proposal Networks <https://arxiv.org/abs/1506.01497>`__
  343. paper.
  344. .. betastatus:: detection module
  345. The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each
  346. image, and should be in ``0-1`` range. Different images can have different sizes.
  347. The behavior of the model changes depending on if it is in training or evaluation mode.
  348. During training, the model expects both the input tensors and a targets (list of dictionary),
  349. containing:
  350. - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with
  351. ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
  352. - labels (``Int64Tensor[N]``): the class label for each ground-truth box
  353. The model returns a ``Dict[Tensor]`` during training, containing the classification and regression
  354. losses for both the RPN and the R-CNN.
  355. During inference, the model requires only the input tensors, and returns the post-processed
  356. predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as
  357. follows, where ``N`` is the number of detections:
  358. - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with
  359. ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
  360. - labels (``Int64Tensor[N]``): the predicted labels for each detection
  361. - scores (``Tensor[N]``): the scores of each detection
  362. For more details on the output, you may refer to :ref:`instance_seg_output`.
  363. Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
  364. Example::
  365. >>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
  366. >>> # For training
  367. >>> images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4)
  368. >>> boxes[:, :, 2:4] = boxes[:, :, 0:2] + boxes[:, :, 2:4]
  369. >>> labels = torch.randint(1, 91, (4, 11))
  370. >>> images = list(image for image in images)
  371. >>> targets = []
  372. >>> for i in range(len(images)):
  373. >>> d = {}
  374. >>> d['boxes'] = boxes[i]
  375. >>> d['labels'] = labels[i]
  376. >>> targets.append(d)
  377. >>> output = model(images, targets)
  378. >>> # For inference
  379. >>> model.eval()
  380. >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
  381. >>> predictions = model(x)
  382. >>>
  383. >>> # optionally, if you want to export the model to ONNX:
  384. >>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)
  385. Args:
  386. weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights`, optional): The
  387. pretrained weights to use. See
  388. :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights` below for
  389. more details, and possible values. By default, no pre-trained
  390. weights are used.
  391. progress (bool, optional): If True, displays a progress bar of the
  392. download to stderr. Default is True.
  393. num_classes (int, optional): number of output classes of the model (including the background)
  394. weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The
  395. pretrained weights for the backbone.
  396. trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
  397. final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are
  398. trainable. If ``None`` is passed (the default) this value is set to 3.
  399. **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
  400. base class. Please refer to the `source code
  401. <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
  402. for more details about this class.
  403. .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights
  404. :members:
  405. """
  406. weights = LineNet_ResNet50_FPN_Weights.verify(weights)
  407. weights_backbone = ResNet50_Weights.verify(weights_backbone)
  408. if weights is not None:
  409. weights_backbone = None
  410. num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
  411. elif num_classes is None:
  412. num_classes = 91
  413. is_trained = weights is not None or weights_backbone is not None
  414. trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)
  415. norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d
  416. backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
  417. backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
  418. model = LineNet(backbone, num_classes=num_classes, **kwargs)
  419. if weights is not None:
  420. model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
  421. if weights == LineNet_ResNet50_FPN_Weights.COCO_V1:
  422. overwrite_eps(model, 0.0)
  423. return model
  424. @register_model()
  425. @handle_legacy_interface(
  426. weights=("pretrained", LineNet_ResNet50_FPN_V2_Weights.COCO_V1),
  427. weights_backbone=("pretrained_backbone", ResNet50_Weights.IMAGENET1K_V1),
  428. )
  429. def linenet_resnet50_fpn_v2(
  430. *,
  431. weights: Optional[LineNet_ResNet50_FPN_V2_Weights] = None,
  432. progress: bool = True,
  433. num_classes: Optional[int] = None,
  434. weights_backbone: Optional[ResNet50_Weights] = None,
  435. trainable_backbone_layers: Optional[int] = None,
  436. **kwargs: Any,
  437. ) -> LineNet:
  438. """
  439. Constructs an improved Faster R-CNN model with a ResNet-50-FPN backbone from `Benchmarking Detection
  440. Transfer Learning with Vision Transformers <https://arxiv.org/abs/2111.11429>`__ paper.
  441. .. betastatus:: detection module
  442. It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See
  443. :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more
  444. details.
  445. Args:
  446. weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights`, optional): The
  447. pretrained weights to use. See
  448. :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights` below for
  449. more details, and possible values. By default, no pre-trained
  450. weights are used.
  451. progress (bool, optional): If True, displays a progress bar of the
  452. download to stderr. Default is True.
  453. num_classes (int, optional): number of output classes of the model (including the background)
  454. weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The
  455. pretrained weights for the backbone.
  456. trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
  457. final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are
  458. trainable. If ``None`` is passed (the default) this value is set to 3.
  459. **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
  460. base class. Please refer to the `source code
  461. <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
  462. for more details about this class.
  463. .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_V2_Weights
  464. :members:
  465. """
  466. weights = LineNet_ResNet50_FPN_V2_Weights.verify(weights)
  467. weights_backbone = ResNet50_Weights.verify(weights_backbone)
  468. if weights is not None:
  469. weights_backbone = None
  470. num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
  471. elif num_classes is None:
  472. num_classes = 91
  473. is_trained = weights is not None or weights_backbone is not None
  474. trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)
  475. backbone = resnet50(weights=weights_backbone, progress=progress)
  476. backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers, norm_layer=nn.BatchNorm2d)
  477. rpn_anchor_generator = _default_anchorgen()
  478. rpn_head = RPNHead(backbone.out_channels, rpn_anchor_generator.num_anchors_per_location()[0], conv_depth=2)
  479. box_head = LineNetConvFCHead(
  480. (backbone.out_channels, 7, 7), [256, 256, 256, 256], [1024], norm_layer=nn.BatchNorm2d
  481. )
  482. model = LineNet(
  483. backbone,
  484. num_classes=num_classes,
  485. rpn_anchor_generator=rpn_anchor_generator,
  486. rpn_head=rpn_head,
  487. box_head=box_head,
  488. **kwargs,
  489. )
  490. if weights is not None:
  491. model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
  492. return model
  493. def _linenet_mobilenet_v3_large_fpn(
  494. *,
  495. weights: Optional[Union[LineNet_MobileNet_V3_Large_FPN_Weights, LineNet_MobileNet_V3_Large_320_FPN_Weights]],
  496. progress: bool,
  497. num_classes: Optional[int],
  498. weights_backbone: Optional[MobileNet_V3_Large_Weights],
  499. trainable_backbone_layers: Optional[int],
  500. **kwargs: Any,
  501. ) -> LineNet:
  502. if weights is not None:
  503. weights_backbone = None
  504. num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))
  505. elif num_classes is None:
  506. num_classes = 91
  507. is_trained = weights is not None or weights_backbone is not None
  508. trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 6, 3)
  509. norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d
  510. backbone = mobilenet_v3_large(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
  511. backbone = _mobilenet_extractor(backbone, True, trainable_backbone_layers)
  512. anchor_sizes = (
  513. (
  514. 32,
  515. 64,
  516. 128,
  517. 256,
  518. 512,
  519. ),
  520. ) * 3
  521. aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
  522. model = LineNet(
  523. backbone, num_classes, rpn_anchor_generator=AnchorGenerator(anchor_sizes, aspect_ratios), **kwargs
  524. )
  525. if weights is not None:
  526. model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
  527. return model
  528. @register_model()
  529. @handle_legacy_interface(
  530. weights=("pretrained", LineNet_MobileNet_V3_Large_320_FPN_Weights.COCO_V1),
  531. weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1),
  532. )
  533. def linenet_mobilenet_v3_large_320_fpn(
  534. *,
  535. weights: Optional[LineNet_MobileNet_V3_Large_320_FPN_Weights] = None,
  536. progress: bool = True,
  537. num_classes: Optional[int] = None,
  538. weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1,
  539. trainable_backbone_layers: Optional[int] = None,
  540. **kwargs: Any,
  541. ) -> LineNet:
  542. """
  543. Low resolution Faster R-CNN model with a MobileNetV3-Large backbone tuned for mobile use cases.
  544. .. betastatus:: detection module
  545. It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See
  546. :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more
  547. details.
  548. Example::
  549. >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(weights=FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT)
  550. >>> model.eval()
  551. >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
  552. >>> predictions = model(x)
  553. Args:
  554. weights (:class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights`, optional): The
  555. pretrained weights to use. See
  556. :class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights` below for
  557. more details, and possible values. By default, no pre-trained
  558. weights are used.
  559. progress (bool, optional): If True, displays a progress bar of the
  560. download to stderr. Default is True.
  561. num_classes (int, optional): number of output classes of the model (including the background)
  562. weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The
  563. pretrained weights for the backbone.
  564. trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
  565. final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are
  566. trainable. If ``None`` is passed (the default) this value is set to 3.
  567. **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
  568. base class. Please refer to the `source code
  569. <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
  570. for more details about this class.
  571. .. autoclass:: torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights
  572. :members:
  573. """
  574. weights = LineNet_MobileNet_V3_Large_320_FPN_Weights.verify(weights)
  575. weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone)
  576. defaults = {
  577. "min_size": 320,
  578. "max_size": 640,
  579. "rpn_pre_nms_top_n_test": 150,
  580. "rpn_post_nms_top_n_test": 150,
  581. "rpn_score_thresh": 0.05,
  582. }
  583. kwargs = {**defaults, **kwargs}
  584. return _linenet_mobilenet_v3_large_fpn(
  585. weights=weights,
  586. progress=progress,
  587. num_classes=num_classes,
  588. weights_backbone=weights_backbone,
  589. trainable_backbone_layers=trainable_backbone_layers,
  590. **kwargs,
  591. )
  592. @register_model()
  593. @handle_legacy_interface(
  594. weights=("pretrained", LineNet_MobileNet_V3_Large_FPN_Weights.COCO_V1),
  595. weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1),
  596. )
  597. def linenet_mobilenet_v3_large_fpn(
  598. *,
  599. weights: Optional[LineNet_MobileNet_V3_Large_FPN_Weights] = None,
  600. progress: bool = True,
  601. num_classes: Optional[int] = None,
  602. weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1,
  603. trainable_backbone_layers: Optional[int] = None,
  604. **kwargs: Any,
  605. ) -> LineNet:
  606. """
  607. Constructs a high resolution Faster R-CNN model with a MobileNetV3-Large FPN backbone.
  608. .. betastatus:: detection module
  609. It works similarly to Faster R-CNN with ResNet-50 FPN backbone. See
  610. :func:`~torchvision.models.detection.fasterrcnn_resnet50_fpn` for more
  611. details.
  612. Example::
  613. >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights=FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT)
  614. >>> model.eval()
  615. >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
  616. >>> predictions = model(x)
  617. Args:
  618. weights (:class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights`, optional): The
  619. pretrained weights to use. See
  620. :class:`~torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights` below for
  621. more details, and possible values. By default, no pre-trained
  622. weights are used.
  623. progress (bool, optional): If True, displays a progress bar of the
  624. download to stderr. Default is True.
  625. num_classes (int, optional): number of output classes of the model (including the background)
  626. weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The
  627. pretrained weights for the backbone.
  628. trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from
  629. final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are
  630. trainable. If ``None`` is passed (the default) this value is set to 3.
  631. **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``
  632. base class. Please refer to the `source code
  633. <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_
  634. for more details about this class.
  635. .. autoclass:: torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights
  636. :members:
  637. """
  638. weights = LineNet_MobileNet_V3_Large_FPN_Weights.verify(weights)
  639. weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone)
  640. defaults = {
  641. "rpn_score_thresh": 0.05,
  642. }
  643. kwargs = {**defaults, **kwargs}
  644. return _linenet_mobilenet_v3_large_fpn(
  645. weights=weights,
  646. progress=progress,
  647. num_classes=num_classes,
  648. weights_backbone=weights_backbone,
  649. trainable_backbone_layers=trainable_backbone_layers,
  650. **kwargs,
  651. )