Mjolnir65/FasterRCNN

1"""2Implements the Generalized R-CNN framework3"""4 5import warnings6from collections import OrderedDict7from typing import Dict, List, Optional, Tuple, Union8 9import torch10from torch import nn, Tensor11 12from torchvision.utils import _log_api_usage_once13 14 15class GeneralizedRCNN(nn.Module):16    """17    Main class for Generalized R-CNN.18 19    Args:20        backbone (nn.Module):21        rpn (nn.Module):22        roi_heads (nn.Module): takes the features + the proposals from the RPN and computes23            detections / masks from it.24        transform (nn.Module): performs the data transformation from the inputs to feed into25            the model26    """27 28    def __init__(self, backbone: nn.Module, rpn: nn.Module, roi_heads: nn.Module, transform: nn.Module) -> None:29        super().__init__()30        _log_api_usage_once(self)31        self.transform = transform32        self.backbone = backbone33        self.rpn = rpn34        self.roi_heads = roi_heads35        # used only on torchscript mode36        self._has_warned = False37 38    @torch.jit.unused39    def eager_outputs(self, losses, detections):40        # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]]41        if self.training:42            return losses43 44        return detections45 46    def forward(self, images, targets=None):47        # type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]48        """49        Args:50            images (list[Tensor]): images to be processed51            targets (list[Dict[str, Tensor]]): ground-truth boxes present in the image (optional)52 53        Returns:54            result (list[BoxList] or dict[Tensor]): the output from the model.55                During training, it returns a dict[Tensor] which contains the losses.56                During testing, it returns list[BoxList] contains additional fields57                like `scores`, `labels` and `mask` (for Mask R-CNN models).58 59        """60        if self.training:61            if targets is None:62                torch._assert(False, "targets should not be none when in training mode")63            else:64                for target in targets:65                    boxes = target["boxes"]66                    if isinstance(boxes, torch.Tensor):67                        torch._assert(68                            len(boxes.shape) == 2 and boxes.shape[-1] == 4,69                            f"Expected target boxes to be a tensor of shape [N, 4], got {boxes.shape}.",70                        )71                    else:72                        torch._assert(False, f"Expected target boxes to be of type Tensor, got {type(boxes)}.")73 74        original_image_sizes: List[Tuple[int, int]] = []75        for img in images:76            val = img.shape[-2:]77            torch._assert(78                len(val) == 2,79                f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}",80            )81            original_image_sizes.append((val[0], val[1]))82 83        images, targets = self.transform(images, targets)84 85        # Check for degenerate boxes86        # TODO: Move this to a function87        if targets is not None:88            for target_idx, target in enumerate(targets):89                boxes = target["boxes"]90                degenerate_boxes = boxes[:, 2:4] <= boxes[:, :2]91                if degenerate_boxes.any():92                    # print the first degenerate box93                    bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0]94                    degen_bb: List[float] = boxes[bb_idx].tolist()95                    torch._assert(96                        False,97                        "All bounding boxes should have positive height and width."98                        f" Found invalid box {degen_bb} for target at index {target_idx}.",99                    )100 101        features = self.backbone(images.tensors)102        if isinstance(features, torch.Tensor):103            features = OrderedDict([("0", features)])104        105        # modify targets to remove theta for rpn106        # print(f"{len(targets)=}")107        # print(f"{targets[0]=}")108        # targets_rpn = []109        # for target in targets:110        #     target_rpn = target.copy()111        #     target_rpn['boxes'] = target_rpn['boxes'][:, :-1]112        #     targets_rpn.append(target_rpn)113        # print(f"{targets_rpn[0]=}")114        proposals, proposal_losses = self.rpn(images, features, targets)115        detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)116        detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)  # type: ignore[operator]117 118        losses = {}119        losses.update(detector_losses)120        losses.update(proposal_losses)121 122        if torch.jit.is_scripting():123            if not self._has_warned:124                warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting")125                self._has_warned = True126            return losses, detections127        else:128            return self.eager_outputs(losses, detections)