line_net.py 31 KB

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