line_net-checkpoint.py 31 KB

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