Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. | |
| """ | |
| Implements the Generalized R-CNN framework | |
| """ | |
| import torch | |
| from torch import nn | |
| from maskrcnn_benchmark.structures.image_list import to_image_list | |
| from ..backbone import build_backbone | |
| from ..rpn.rpn import build_rpn | |
| from ..segmentation.segmentation import build_segmentation | |
| from ..roi_heads.roi_heads import build_roi_heads | |
| import time | |
| class GeneralizedRCNN(nn.Module): | |
| """ | |
| Main class for Generalized R-CNN. Currently supports boxes and masks. | |
| It consists of three main parts: | |
| - backbone | |
| = rpn | |
| - heads: takes the features + the proposals from the RPN and computes | |
| detections / masks from it. | |
| """ | |
| def __init__(self, cfg): | |
| super(GeneralizedRCNN, self).__init__() | |
| self.cfg = cfg | |
| self.backbone = build_backbone(cfg) | |
| if cfg.MODEL.SEG_ON: | |
| self.proposal = build_segmentation(cfg) | |
| else: | |
| self.proposal = build_rpn(cfg) | |
| if cfg.MODEL.TRAIN_DETECTION_ONLY: | |
| self.roi_heads = None | |
| else: | |
| self.roi_heads = build_roi_heads(cfg) | |
| def forward(self, images, targets=None): | |
| """ | |
| Arguments: | |
| images (list[Tensor] or ImageList): images to be processed | |
| targets (list[BoxList]): ground-truth boxes present in the image (optional) | |
| Returns: | |
| result (list[BoxList] or dict[Tensor]): the output from the model. | |
| During training, it returns a dict[Tensor] which contains the losses. | |
| During testing, it returns list[BoxList] contains additional fields | |
| like `scores`, `labels` and `mask` (for Mask R-CNN models). | |
| """ | |
| if self.training and targets is None: | |
| raise ValueError("In training mode, targets should be passed") | |
| # torch.cuda.synchronize() | |
| # start_time = time.time() | |
| images = to_image_list(images) | |
| # torch.cuda.synchronize() | |
| # end_time = time.time() | |
| # print('image load time:', end_time - start_time) | |
| # torch.cuda.synchronize() | |
| # start_time = time.time() | |
| features = self.backbone(images.tensors) | |
| # torch.cuda.synchronize() | |
| # end_time = time.time() | |
| # print('backbone time:', end_time - start_time) | |
| if self.cfg.MODEL.SEG_ON and not self.training: | |
| # torch.cuda.synchronize() | |
| # start_time = time.time() | |
| (proposals, seg_results), fuse_feature = self.proposal(images, features, targets) | |
| # torch.cuda.synchronize() | |
| # end_time = time.time() | |
| # print('seg time:', end_time - start_time) | |
| else: | |
| if self.cfg.MODEL.SEG_ON: | |
| (proposals, proposal_losses), fuse_feature = self.proposal(images, features, targets) | |
| else: | |
| proposals, proposal_losses = self.proposal(images, features, targets) | |
| if self.roi_heads is not None: | |
| if self.cfg.MODEL.SEG_ON and self.cfg.MODEL.SEG.USE_FUSE_FEATURE: | |
| x, result, detector_losses = self.roi_heads(fuse_feature, proposals, targets) | |
| else: | |
| x, result, detector_losses = self.roi_heads(features, proposals, targets) | |
| else: | |
| # RPN-only models don't have roi_heads | |
| # x = features | |
| result = proposals | |
| detector_losses = {} | |
| if self.training: | |
| losses = {} | |
| if self.roi_heads is not None: | |
| losses.update(detector_losses) | |
| losses.update(proposal_losses) | |
| return losses | |
| else: | |
| if self.cfg.MODEL.SEG_ON: | |
| return result, proposals, seg_results | |
| else: | |
| return result | |
| # return result | |