CoolFace
Apppublic

MonumentDetection/ContinualLearningFastAPI

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
fasterrcnn.py353 linesDownload Raw Back to faster_models
1from torchvision.ops import MultiScaleRoIAlign2from torchvision.models.detection import FasterRCNN3from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights4from torchvision.models.resnet import resnet50, ResNet50_Weights5from torchvision.models._utils import  _ovewrite_value_param6from torchvision.models.detection._utils import  overwrite_eps7from torchvision.ops import misc as misc_nn_ops8from torch import nn9from torchvision.models.detection.backbone_utils import _validate_trainable_layers, _resnet_fpn_extractor10from typing import Any, Optional, TypeVar11import torch12 13V = TypeVar("V")14# _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))15def _ovewrite_value_param(param: str, actual: Optional[V], expected: V) -> V:16    if actual is not None:17        if actual != expected:18            raise ValueError(f"The parameter '{param}' expected value {expected} but got {actual} instead.")19    return expected20 21def fasterrcnn_resnet50_fpn(22    *,23    weights: Optional[FasterRCNN_ResNet50_FPN_Weights] = None,24    progress: bool = True,25    num_classes: Optional[int] = None,26    weights_backbone: Optional[ResNet50_Weights] = ResNet50_Weights.IMAGENET1K_V1,27    trainable_backbone_layers: Optional[int] = None,28    extend =0,29    **kwargs: Any,30) -> FasterRCNN:31    """32    Faster R-CNN model with a ResNet-50-FPN backbone from the `Faster R-CNN: Towards Real-Time Object33    Detection with Region Proposal Networks <https://arxiv.org/abs/1506.01497>`__34    paper.35 36    .. betastatus:: detection module37 38    The input to the model is expected to be a list of tensors, each of shape ``[C, H, W]``, one for each39    image, and should be in ``0-1`` range. Different images can have different sizes.40 41    The behavior of the model changes depending on if it is in training or evaluation mode.42 43    During training, the model expects both the input tensors and a targets (list of dictionary),44    containing:45 46        - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with47          ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.48        - labels (``Int64Tensor[N]``): the class label for each ground-truth box49 50    The model returns a ``Dict[Tensor]`` during training, containing the classification and regression51    losses for both the RPN and the R-CNN.52 53    During inference, the model requires only the input tensors, and returns the post-processed54    predictions as a ``List[Dict[Tensor]]``, one for each input image. The fields of the ``Dict`` are as55    follows, where ``N`` is the number of detections:56 57        - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with58          ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.59        - labels (``Int64Tensor[N]``): the predicted labels for each detection60        - scores (``Tensor[N]``): the scores of each detection61 62    For more details on the output, you may refer to :ref:`instance_seg_output`.63 64    Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.65 66    Example::67 68        >>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)69        >>> # For training70        >>> images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4)71        >>> boxes[:, :, 2:4] = boxes[:, :, 0:2] + boxes[:, :, 2:4]72        >>> labels = torch.randint(1, 91, (4, 11))73        >>> images = list(image for image in images)74        >>> targets = []75        >>> for i in range(len(images)):76        >>>     d = {}77        >>>     d['boxes'] = boxes[i]78        >>>     d['labels'] = labels[i]79        >>>     targets.append(d)80        >>> output = model(images, targets)81        >>> # For inference82        >>> model.eval()83        >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]84        >>> predictions = model(x)85        >>>86        >>> # optionally, if you want to export the model to ONNX:87        >>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)88 89    Args:90        weights (:class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights`, optional): The91            pretrained weights to use. See92            :class:`~torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights` below for93            more details, and possible values. By default, no pre-trained94            weights are used.95        progress (bool, optional): If True, displays a progress bar of the96            download to stderr. Default is True.97        num_classes (int, optional): number of output classes of the model (including the background)98        weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The99            pretrained weights for the backbone.100        trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from101            final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are102            trainable. If ``None`` is passed (the default) this value is set to 3.103        **kwargs: parameters passed to the ``torchvision.models.detection.faster_rcnn.FasterRCNN``104            base class. Please refer to the `source code105            <https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py>`_106            for more details about this class.107 108    .. autoclass:: torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights109        :members:110    """111    weights = FasterRCNN_ResNet50_FPN_Weights.verify(weights)112    weights_backbone = ResNet50_Weights.verify(weights_backbone)113 114    if weights is not None:115        weights_backbone = None116        num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"]))117    elif num_classes is None:118        num_classes = 91119 120    is_trained = weights is not None or weights_backbone is not None121    trainable_backbone_layers = _validate_trainable_layers(is_trained, trainable_backbone_layers, 5, 3)122    norm_layer = misc_nn_ops.FrozenBatchNorm2d if is_trained else nn.BatchNorm2d123 124    backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer)125    backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)126    model = FasterRCNN(backbone, num_classes=num_classes, **kwargs)127 128    if weights is not None:129        model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))130        if weights == FasterRCNN_ResNet50_FPN_Weights.COCO_V1:131            overwrite_eps(model, 0.0)132 133    return model134 135def filter_pred(predicted):136    filtered_predictions = []137 138    for pred in predicted:139        scores = pred['scores']140        indices = scores > 0.5 141        if indices.any():142            filtered_boxes = pred['boxes'][indices]143            filtered_labels = pred['labels'][indices]144            filtered_scores = pred['scores'][indices]145            filtered_pred = {'boxes': filtered_boxes,146                             'labels': filtered_labels,147                             'scores': filtered_scores}148        else:149            filtered_boxes = torch.tensor([[0, 0, 0, 0]])150            filtered_labels = torch.tensor([0])151            filtered_scores =torch.tensor([0])152            filtered_pred = {'boxes': filtered_boxes,153                             'labels': filtered_labels,154                             'scores': filtered_scores}155        filtered_predictions.append(filtered_pred)156    return filtered_predictions157 158CLASSES1 = [159    'bg',160    'Akash Bhairav',161    'Bhadrakali Temple',162    'Jalbinayak',163    'Lumadhi Bhadrakali Temple Sankata',164    'Maitidevi Temple',165    'Patan Dhoka',166    'Sano Pashupati',167    'Swoyambhunath',168    'Tridevi Temple',169    'ashok stupa',170    'birupakshya',171    'chamunda mai',172    'charumati',173    'mahadev temple',174    'taleju bell_KDS',175    'pratappur temple',176    'chakku bakku',177    'Ghantaghar',178    'kumaristhan',179    'uma maheshwor'180]181CLASSES2= [182    'bg',183      'BalNilkantha',184        'Chandeshwori Temple',185          'Dakshin Barahi',186            'Dharahara',187              'Jamachen Monastry',188                'Khumbeshwor mahadev',189                  'Kotilingeshvara',190                    'Mahabauddha Asan',191                      'Pilot Baba',192                        'Ram Mandir',193                          'Ranipokhari',194                            'Fasidega Temple',195                              'Guyeshwori',196                                'Hanuman Idol',197                                 'Jame Masjid',198                                   'Red Gumba',199                                     'Santaneshwor Mahadev',200                                       'Sankha Statue',201                                         'Shantidham',202                                           'Yetkha Bahal']203 204 205CLASSES3 = [206    'bg',207     'Naxal Bhagwati',208     'Basantapur Tower',209     'Bhaktapur Tower',210     'Bhimeleshvara',211     'Degu Tale',212     'Gaddi Durbar',213     'Garud',214     'Kasthamandap',215     'Kavindrapura Sattal',216     'Kirtipur Tower',217     'Kumari Ghar',218     'Lalitpur Tower',219     'Shiva Temple',220     'Simha Sattal',221     'Trailokya Mohan',222     'Hanuman Idol',223     'Panchamukhi Hanuman',224     'Kotilingeshvara',225     'Sano Pashupati',226     'Mahabauddha Asan']227 228 229CLASSES4 = [230    'bg',231     'Bhimsen Temple',232     'Bhupatindra Malla Column',233     'Chayasilin Mandap',234     'Dakshin Barahi',235     'Fasidega Temple',236     'Gopinath Krishna Temple',237     'Nyatapola Temple',238     'Siddhi Lakshmi Temple',239     'Badrinath Temple',240     'Bhairavnath Temple',241     'Golden Gate',242     'Mahadev Temple',243     'Palace of the 55 Windows',244     'Taleju Bell_BDS',245     'Vastala Temple',246     'Balkumari, Bhaktapur',247     'Golden Temple',248     'Jaya Bageshwori',249     'Lokeshwor Temple Bhaktapur',250     'Wakupati Narayan Temple'251     ]252 253 254 255classes=['Akash Bhairav','ashok stupa','Badrinath','Bagbairav',256'Balkumari, Bhaktapur',257'BalNilkantha',258'basantapur tower',259'Bhadrakali Temple',260'bhairavnath temple',261'bhaktapur tower','bhimeleshvara',262'Bhimsen Temple','Bhupatindra Malla Column',263'bhuvana lakshmeshvara',264'birupakshya',265'Buddha Statue',266'chakku bakku',267'chamunda mai',268'Chandeshwori Temple',269'Char Narayan Temple',270'charumati',271'chasin dega',272'Chayasilin Mandap',273'Dakshin Barahi',274'degu tale',275'Dharahara',276'Fasidega Temple',277'Garud Statue',278'garud',279'Ghantaghar',280'golden gate',281'golden temple',282'Gopinath krishna Temple',283'guyeshwori',284'hanuman idol',285'Harishankar Temple',286'indrapura',287'Isckon Temple',288'jagannatha temple',289'Jalbinayak',290'Jamachen Monastry',291'jame masjid',292'jaya bageshwori',293'kala-bhairava',294'kasthamandap',295'kavindrapura sattal',296'Kedamatha Tirtha',297'Khumbeshwor mahadev',298'kiranteshwor mahadev',299'kirtipur tower',300'Kotilingeshvara',301'Krishna mandir PDS',302'Krishna_temple _kobahal',303'Kumari Ghar',304'kumaristhan',305'kumbheshwor mahadev',306'lalitpur tower',307'lokeshwor temple bhaktapur',308'Lumadhi Bhadrakali Temple Sankata',309'Mahabauddha Asan',310'mahadev temple',311'Maipi Temple',312'Maitidevi Temple',313'manamaiju temple',314'nagarmandap shree kriti bihar',315'narayan temple',316'National Gallery',317'Naxal Bhagwati',318'Nyatapola temple',319'Palace of 55 Windows',320'Panchamukhi Hanuman',321'Patan Dhoka',322'Pilot Baba',323'PimBahal Gumba',324'pratap malla column',325'pratappur temple',326'Ram Mandir',327'Ranipokhari',328'red gumba',329'sahid gate',330'Sankha Statue',331'Sano Pashupati',332'Santaneshwor Mahadev',333'shantidham',334'Shiva Temple',335'shveta bhairava',336'Siddhi Lakshmi temple',337'simha sattal',338'Swoyambhunath',339'taleju bell pds',340'taleju bell_BDS',341'taleju bell_KDS',342'taleju temple',343'taleju_temple_south',344'trailokya mohan',345'Tridevi Temple',346'uma maheshwor',347'ume_maheshwara',348'Vastala Temple',349'vishnu temple',350'Wakupati Narayan Temple',351'wishing well budhha statue',352'Yetkha Bahal',353'yog_narendra_malla_statue']