We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Rescale transforms and torchvision RCNNs
## 🐛 Bug
torchvisions `FasterRCNN` and `MaskRCNN` internally rescales the images via `GeneralizedRCNNTransform`.

This means that any resizing transform previously applied to the image (probably on the `Dataset` stage) will be unexpectedly  overriden.

It becomes really confusing if we should apply a resize transform via `Dataset` or via the model, ideally we want all transform to be applied at the `Dataset` stage. Going one step further, can this even be said about normalization?

## Solution 1
Changes the default of `min_size` and `max_size` in the model to 2 and 9999 respectively, practically making them non effective.

## Solution 2
Monkey patch `model.transform` to a function that returns the same items it receives (being careful with what are the expected output types)

---
Solution 1 seems like a better idea, it's less magical, we can still have `min/max_size` as parameters, we just change their defaults.

---
Should we do the same with normalize? Changing the mean to 0 and std to 1? It's them clear that if we want to normalize we should do it in the `Dataset` stage.

I'm thinking this because models that we integrate in the future might not have this internal transforms, and it will be very confusing to a new user why sometimes he has to explicity define normalization and sometimes not.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py`
Content:
```
1 __all__ = ["MantisMaskRCNN"]
2 
3 from mantisshrimp.imports import *
4 from mantisshrimp.core import *
5 from mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *
6 from mantisshrimp.models.mantis_rcnn.mantis_rcnn import *
7 from mantisshrimp.models.mantis_rcnn.mantis_faster_rcnn import *
8 from mantisshrimp.backbones import *
9 
10 
11 class MantisMaskRCNN(MantisRCNN):
12     @delegates(MaskRCNN.__init__)
13     def __init__(
14         self,
15         num_classes: int,
16         backbone: nn.Module = None,
17         param_groups: List[nn.Module] = None,
18         **kwargs,
19     ):
20         super().__init__()
21         self.num_classes = num_classes
22 
23         if backbone is None:
24             # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset
25             self.model = maskrcnn_resnet50_fpn(pretrained=True, **kwargs)
26             in_features = self.model.roi_heads.box_predictor.cls_score.in_features
27             self.model.roi_heads.box_predictor = FastRCNNPredictor(
28                 in_features, num_classes
29             )
30             in_features_mask = (
31                 self.model.roi_heads.mask_predictor.conv5_mask.in_channels
32             )
33             self.model.roi_heads.mask_predictor = MaskRCNNPredictor(
34                 in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes
35             )
36             param_groups = resnet_fpn_backbone_param_groups(self.model.backbone)
37         else:
38             self.model = MaskRCNN(backbone, num_classes=num_classes, **kwargs)
39             param_groups = param_groups or [backbone]
40 
41         self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]
42         check_all_model_params_in_groups(self.model, self.param_groups)
43 
44     def forward(self, images, targets=None):
45         return self.model(images, targets)
46 
47     def predict(
48         self,
49         images: List[np.ndarray],
50         detection_threshold: float = 0.5,
51         mask_threshold: float = 0.5,
52     ):
53         convert_raw_prediction = partial(
54             self.convert_raw_prediction,
55             detection_threshold=detection_threshold,
56             mask_threshold=mask_threshold,
57         )
58 
59         return self._predict(
60             images=images, convert_raw_prediction=convert_raw_prediction
61         )
62 
63     @property
64     def param_groups(self):
65         return self._param_groups
66 
67     @staticmethod
68     def convert_raw_prediction(
69         raw_pred: dict, detection_threshold: float, mask_threshold: float
70     ):
71         preds = MantisFasterRCNN.convert_raw_prediction(
72             raw_pred=raw_pred, detection_threshold=detection_threshold
73         )
74 
75         above_threshold = preds["above_threshold"]
76         masks_probs = raw_pred["masks"][above_threshold]
77         masks_probs = masks_probs.detach().cpu().numpy()
78         # convert probabilities to 0 or 1 based on mask_threshold
79         masks = masks_probs > mask_threshold
80         masks = MaskArray(masks.squeeze(1))
81 
82         return {**preds, "masks": masks}
83 
84     @staticmethod
85     def build_training_sample(
86         imageid: int,
87         img: np.ndarray,
88         labels: List[int],
89         bboxes: List[BBox],
90         masks: MaskArray,
91         **kwargs,
92     ):
93         x, y = MantisFasterRCNN.build_training_sample(
94             imageid=imageid, img=img, labels=labels, bboxes=bboxes,
95         )
96         y["masks"] = tensor(masks.data, dtype=torch.uint8)
97         return x, y
98
```

Path: `mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py`
Content:
```
1 __all__ = ["MantisFasterRCNN"]
2 
3 from mantisshrimp.imports import *
4 from mantisshrimp.core import *
5 from mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *
6 from mantisshrimp.models.mantis_rcnn.mantis_rcnn import *
7 from mantisshrimp.backbones import *
8 
9 
10 class MantisFasterRCNN(MantisRCNN):
11     """
12     Creates a flexible Faster RCNN implementation based on torchvision library.
13     Args: 
14     n_class (int) : number of classes. Do not have class_id "0" it is reserved as background.
15                     n_class = number of classes to label + 1 for background.
16     """
17 
18     @delegates(FasterRCNN.__init__)
19     def __init__(
20         self,
21         num_classes: int,
22         backbone: nn.Module = None,
23         param_groups: List[nn.Module] = None,
24         metrics=None,
25         **kwargs,
26     ):
27         super().__init__(metrics=metrics)
28         self.num_classes = num_classes
29         self.backbone = backbone
30         if backbone is None:
31             # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset
32             self.model = fasterrcnn_resnet50_fpn(pretrained=True, **kwargs)
33             in_features = self.model.roi_heads.box_predictor.cls_score.in_features
34             self.model.roi_heads.box_predictor = FastRCNNPredictor(
35                 in_features, num_classes
36             )
37             param_groups = resnet_fpn_backbone_param_groups(self.model.backbone)
38         else:
39             self.model = FasterRCNN(backbone, num_classes=num_classes, **kwargs)
40             param_groups = param_groups or [backbone]
41 
42         self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]
43         check_all_model_params_in_groups(self.model, self._param_groups)
44 
45     def forward(self, images, targets=None):
46         return self.model(images, targets)
47 
48     def predict(self, images: List[np.ndarray], detection_threshold: float = 0.5):
49         convert_raw_prediction = partial(
50             self.convert_raw_prediction, detection_threshold=detection_threshold,
51         )
52 
53         return self._predict(
54             images=images, convert_raw_prediction=convert_raw_prediction
55         )
56 
57     @property
58     def param_groups(self):
59         return self._param_groups
60 
61     @staticmethod
62     def convert_raw_prediction(raw_pred: dict, detection_threshold: float):
63         above_threshold = raw_pred["scores"] >= detection_threshold
64 
65         labels = raw_pred["labels"][above_threshold]
66         labels = labels.detach().cpu().numpy()
67 
68         scores = raw_pred["scores"][above_threshold]
69         scores = scores.detach().cpu().numpy()
70 
71         boxes = raw_pred["boxes"][above_threshold]
72         bboxes = []
73         for box_tensor in boxes:
74             xyxy = box_tensor.cpu().tolist()
75             bbox = BBox.from_xyxy(*xyxy)
76             bboxes.append(bbox)
77 
78         return {
79             "labels": labels,
80             "scores": scores,
81             "bboxes": bboxes,
82             "above_threshold": above_threshold,
83         }
84 
85     @staticmethod
86     def build_training_sample(
87         imageid: int, img: np.ndarray, labels: List[int], bboxes: List[BBox], **kwargs,
88     ):
89         x = im2tensor(img)
90         # injected values when annotations are empty are disconsidered
91         # because we mark label as 0 (background)
92         _fake_box = [0, 1, 2, 3]
93         y = {
94             "image_id": tensor(imageid, dtype=torch.int64),
95             "labels": tensor(labels or [0], dtype=torch.int64),
96             "boxes": tensor([o.xyxy for o in bboxes] or [_fake_box], dtype=torch.float),
97         }
98         return x, y
99
```

Path: `mantisshrimp/models/mantis_rcnn/mantis_rcnn.py`
Content:
```
1 __all__ = ["MantisRCNN"]
2 
3 from mantisshrimp.imports import *
4 from mantisshrimp.utils import *
5 from mantisshrimp.models.mantis_module import *
6 from mantisshrimp.backbones import *
7 
8 
9 class MantisRCNN(MantisModule, ABC):
10     def __init__(self, metrics=None):
11         super().__init__()
12         self.metrics = metrics or []
13 
14     @staticmethod
15     @abstractmethod
16     def build_training_sample(*args, **kwargs):
17         """
18         Converts a record to a format understood by the model.
19         """
20 
21     def _predict(
22         self, images: List[np.ndarray], convert_raw_prediction,
23     ):
24         self.eval()
25         tensor_images = [im2tensor(img).to(self.device) for img in images]
26         raw_preds = self(tensor_images)
27 
28         return [convert_raw_prediction(raw_pred) for raw_pred in raw_preds]
29 
30     @staticmethod
31     def loss(preds, targs) -> Tensor:
32         return sum(preds.values())
33 
34     @classmethod
35     def collate_fn(cls, data):
36         ts = [cls.build_training_sample(**o) for o in data]
37         xb, yb = zip(*ts)
38         return xb, list(yb)
39 
40     @classmethod
41     def dataloader(cls, dataset, **kwargs) -> DataLoader:
42         return DataLoader(dataset=dataset, collate_fn=cls.collate_fn, **kwargs)
43 
44     @staticmethod
45     def get_backbone_by_name(
46         name: str, fpn: bool = True, pretrained: bool = True
47     ) -> nn.Module:
48         """
49         Args:
50             backbone (str): If none creates a default resnet50_fpn model trained on MS COCO 2017
51                 Supported backones are: "resnet18", "resnet34","resnet50", "resnet101", "resnet152",
52                  "resnext50_32x4d", "resnext101_32x8d", "wide_resnet50_2", "wide_resnet101_2", as resnets with fpn backbones.
53                 Without fpn backbones supported are: "resnet18", "resnet34", "resnet50","resnet101",
54                  "resnet152", "resnext101_32x8d", "mobilenet", "vgg11", "vgg13", "vgg16", "vgg19",
55             pretrained (bool): Creates a pretrained backbone with imagenet weights.
56         """
57         # Giving string as a backbone, which is either supported resnet or backbone
58         if fpn:
59             # Creates a torchvision resnet model with fpn added
60             # It returns BackboneWithFPN model
61             backbone = resnet_fpn_backbone(name, pretrained=pretrained)
62         else:
63             # This does not create fpn backbone, it is supported for all models
64             backbone = create_torchvision_backbone(name, pretrained=pretrained)
65         return backbone
66
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py
--- a/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py
+++ b/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py
@@ -4,7 +4,6 @@
 from mantisshrimp.core import *
 from mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *
 from mantisshrimp.models.mantis_rcnn.mantis_rcnn import *
-from mantisshrimp.backbones import *
 
 
 class MantisFasterRCNN(MantisRCNN):
@@ -22,6 +21,7 @@
         backbone: nn.Module = None,
         param_groups: List[nn.Module] = None,
         metrics=None,
+        remove_internal_transforms=True,
         **kwargs,
     ):
         super().__init__(metrics=metrics)
@@ -42,6 +42,9 @@
         self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]
         check_all_model_params_in_groups(self.model, self._param_groups)
 
+        if remove_internal_transforms:
+            self._remove_transforms_from_model(self.model)
+
     def forward(self, images, targets=None):
         return self.model(images, targets)
 
diff --git a/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py
--- a/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py
+++ b/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py
@@ -5,7 +5,6 @@
 from mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *
 from mantisshrimp.models.mantis_rcnn.mantis_rcnn import *
 from mantisshrimp.models.mantis_rcnn.mantis_faster_rcnn import *
-from mantisshrimp.backbones import *
 
 
 class MantisMaskRCNN(MantisRCNN):
@@ -15,6 +14,7 @@
         num_classes: int,
         backbone: nn.Module = None,
         param_groups: List[nn.Module] = None,
+        remove_internal_transforms: bool = True,
         **kwargs,
     ):
         super().__init__()
@@ -41,6 +41,9 @@
         self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]
         check_all_model_params_in_groups(self.model, self.param_groups)
 
+        if remove_internal_transforms:
+            self._remove_transforms_from_model(self.model)
+
     def forward(self, images, targets=None):
         return self.model(images, targets)
 
diff --git a/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py
--- a/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py
+++ b/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py
@@ -27,6 +27,16 @@
 
         return [convert_raw_prediction(raw_pred) for raw_pred in raw_preds]
 
+    def _remove_transforms_from_model(self, model: GeneralizedRCNN):
+        def noop_normalize(image):
+            return image
+
+        def noop_resize(image, target):
+            return image, target
+
+        model.transform.normalize = noop_normalize
+        model.transform.resize = noop_resize
+
     @staticmethod
     def loss(preds, targs) -> Tensor:
         return sum(preds.values())

{"golden_diff": "diff --git a/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py\n--- a/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py\n+++ b/mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py\n@@ -4,7 +4,6 @@\n from mantisshrimp.core import *\n from mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *\n from mantisshrimp.models.mantis_rcnn.mantis_rcnn import *\n-from mantisshrimp.backbones import *\n \n \n class MantisFasterRCNN(MantisRCNN):\n@@ -22,6 +21,7 @@\n         backbone: nn.Module = None,\n         param_groups: List[nn.Module] = None,\n         metrics=None,\n+        remove_internal_transforms=True,\n         **kwargs,\n     ):\n         super().__init__(metrics=metrics)\n@@ -42,6 +42,9 @@\n         self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]\n         check_all_model_params_in_groups(self.model, self._param_groups)\n \n+        if remove_internal_transforms:\n+            self._remove_transforms_from_model(self.model)\n+\n     def forward(self, images, targets=None):\n         return self.model(images, targets)\n \ndiff --git a/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py\n--- a/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py\n+++ b/mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py\n@@ -5,7 +5,6 @@\n from mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *\n from mantisshrimp.models.mantis_rcnn.mantis_rcnn import *\n from mantisshrimp.models.mantis_rcnn.mantis_faster_rcnn import *\n-from mantisshrimp.backbones import *\n \n \n class MantisMaskRCNN(MantisRCNN):\n@@ -15,6 +14,7 @@\n         num_classes: int,\n         backbone: nn.Module = None,\n         param_groups: List[nn.Module] = None,\n+        remove_internal_transforms: bool = True,\n         **kwargs,\n     ):\n         super().__init__()\n@@ -41,6 +41,9 @@\n         self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]\n         check_all_model_params_in_groups(self.model, self.param_groups)\n \n+        if remove_internal_transforms:\n+            self._remove_transforms_from_model(self.model)\n+\n     def forward(self, images, targets=None):\n         return self.model(images, targets)\n \ndiff --git a/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py b/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py\n--- a/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py\n+++ b/mantisshrimp/models/mantis_rcnn/mantis_rcnn.py\n@@ -27,6 +27,16 @@\n \n         return [convert_raw_prediction(raw_pred) for raw_pred in raw_preds]\n \n+    def _remove_transforms_from_model(self, model: GeneralizedRCNN):\n+        def noop_normalize(image):\n+            return image\n+\n+        def noop_resize(image, target):\n+            return image, target\n+\n+        model.transform.normalize = noop_normalize\n+        model.transform.resize = noop_resize\n+\n     @staticmethod\n     def loss(preds, targs) -> Tensor:\n         return sum(preds.values())\n", "issue": "Rescale transforms and torchvision RCNNs\n## \ud83d\udc1b Bug\r\ntorchvisions `FasterRCNN` and `MaskRCNN` internally rescales the images via `GeneralizedRCNNTransform`.\r\n\r\nThis means that any resizing transform previously applied to the image (probably on the `Dataset` stage) will be unexpectedly  overriden.\r\n\r\nIt becomes really confusing if we should apply a resize transform via `Dataset` or via the model, ideally we want all transform to be applied at the `Dataset` stage. Going one step further, can this even be said about normalization?\r\n\r\n## Solution 1\r\nChanges the default of `min_size` and `max_size` in the model to 2 and 9999 respectively, practically making them non effective.\r\n\r\n## Solution 2\r\nMonkey patch `model.transform` to a function that returns the same items it receives (being careful with what are the expected output types)\r\n\r\n---\r\nSolution 1 seems like a better idea, it's less magical, we can still have `min/max_size` as parameters, we just change their defaults.\r\n\r\n---\r\nShould we do the same with normalize? Changing the mean to 0 and std to 1? It's them clear that if we want to normalize we should do it in the `Dataset` stage.\r\n\r\nI'm thinking this because models that we integrate in the future might not have this internal transforms, and it will be very confusing to a new user why sometimes he has to explicity define normalization and sometimes not.\r\n\n", "before_files": [{"content": "__all__ = [\"MantisMaskRCNN\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.core import *\nfrom mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *\nfrom mantisshrimp.models.mantis_rcnn.mantis_rcnn import *\nfrom mantisshrimp.models.mantis_rcnn.mantis_faster_rcnn import *\nfrom mantisshrimp.backbones import *\n\n\nclass MantisMaskRCNN(MantisRCNN):\n    @delegates(MaskRCNN.__init__)\n    def __init__(\n        self,\n        num_classes: int,\n        backbone: nn.Module = None,\n        param_groups: List[nn.Module] = None,\n        **kwargs,\n    ):\n        super().__init__()\n        self.num_classes = num_classes\n\n        if backbone is None:\n            # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset\n            self.model = maskrcnn_resnet50_fpn(pretrained=True, **kwargs)\n            in_features = self.model.roi_heads.box_predictor.cls_score.in_features\n            self.model.roi_heads.box_predictor = FastRCNNPredictor(\n                in_features, num_classes\n            )\n            in_features_mask = (\n                self.model.roi_heads.mask_predictor.conv5_mask.in_channels\n            )\n            self.model.roi_heads.mask_predictor = MaskRCNNPredictor(\n                in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes\n            )\n            param_groups = resnet_fpn_backbone_param_groups(self.model.backbone)\n        else:\n            self.model = MaskRCNN(backbone, num_classes=num_classes, **kwargs)\n            param_groups = param_groups or [backbone]\n\n        self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]\n        check_all_model_params_in_groups(self.model, self.param_groups)\n\n    def forward(self, images, targets=None):\n        return self.model(images, targets)\n\n    def predict(\n        self,\n        images: List[np.ndarray],\n        detection_threshold: float = 0.5,\n        mask_threshold: float = 0.5,\n    ):\n        convert_raw_prediction = partial(\n            self.convert_raw_prediction,\n            detection_threshold=detection_threshold,\n            mask_threshold=mask_threshold,\n        )\n\n        return self._predict(\n            images=images, convert_raw_prediction=convert_raw_prediction\n        )\n\n    @property\n    def param_groups(self):\n        return self._param_groups\n\n    @staticmethod\n    def convert_raw_prediction(\n        raw_pred: dict, detection_threshold: float, mask_threshold: float\n    ):\n        preds = MantisFasterRCNN.convert_raw_prediction(\n            raw_pred=raw_pred, detection_threshold=detection_threshold\n        )\n\n        above_threshold = preds[\"above_threshold\"]\n        masks_probs = raw_pred[\"masks\"][above_threshold]\n        masks_probs = masks_probs.detach().cpu().numpy()\n        # convert probabilities to 0 or 1 based on mask_threshold\n        masks = masks_probs > mask_threshold\n        masks = MaskArray(masks.squeeze(1))\n\n        return {**preds, \"masks\": masks}\n\n    @staticmethod\n    def build_training_sample(\n        imageid: int,\n        img: np.ndarray,\n        labels: List[int],\n        bboxes: List[BBox],\n        masks: MaskArray,\n        **kwargs,\n    ):\n        x, y = MantisFasterRCNN.build_training_sample(\n            imageid=imageid, img=img, labels=labels, bboxes=bboxes,\n        )\n        y[\"masks\"] = tensor(masks.data, dtype=torch.uint8)\n        return x, y\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py"}, {"content": "__all__ = [\"MantisFasterRCNN\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.core import *\nfrom mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *\nfrom mantisshrimp.models.mantis_rcnn.mantis_rcnn import *\nfrom mantisshrimp.backbones import *\n\n\nclass MantisFasterRCNN(MantisRCNN):\n    \"\"\"\n    Creates a flexible Faster RCNN implementation based on torchvision library.\n    Args: \n    n_class (int) : number of classes. Do not have class_id \"0\" it is reserved as background.\n                    n_class = number of classes to label + 1 for background.\n    \"\"\"\n\n    @delegates(FasterRCNN.__init__)\n    def __init__(\n        self,\n        num_classes: int,\n        backbone: nn.Module = None,\n        param_groups: List[nn.Module] = None,\n        metrics=None,\n        **kwargs,\n    ):\n        super().__init__(metrics=metrics)\n        self.num_classes = num_classes\n        self.backbone = backbone\n        if backbone is None:\n            # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset\n            self.model = fasterrcnn_resnet50_fpn(pretrained=True, **kwargs)\n            in_features = self.model.roi_heads.box_predictor.cls_score.in_features\n            self.model.roi_heads.box_predictor = FastRCNNPredictor(\n                in_features, num_classes\n            )\n            param_groups = resnet_fpn_backbone_param_groups(self.model.backbone)\n        else:\n            self.model = FasterRCNN(backbone, num_classes=num_classes, **kwargs)\n            param_groups = param_groups or [backbone]\n\n        self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]\n        check_all_model_params_in_groups(self.model, self._param_groups)\n\n    def forward(self, images, targets=None):\n        return self.model(images, targets)\n\n    def predict(self, images: List[np.ndarray], detection_threshold: float = 0.5):\n        convert_raw_prediction = partial(\n            self.convert_raw_prediction, detection_threshold=detection_threshold,\n        )\n\n        return self._predict(\n            images=images, convert_raw_prediction=convert_raw_prediction\n        )\n\n    @property\n    def param_groups(self):\n        return self._param_groups\n\n    @staticmethod\n    def convert_raw_prediction(raw_pred: dict, detection_threshold: float):\n        above_threshold = raw_pred[\"scores\"] >= detection_threshold\n\n        labels = raw_pred[\"labels\"][above_threshold]\n        labels = labels.detach().cpu().numpy()\n\n        scores = raw_pred[\"scores\"][above_threshold]\n        scores = scores.detach().cpu().numpy()\n\n        boxes = raw_pred[\"boxes\"][above_threshold]\n        bboxes = []\n        for box_tensor in boxes:\n            xyxy = box_tensor.cpu().tolist()\n            bbox = BBox.from_xyxy(*xyxy)\n            bboxes.append(bbox)\n\n        return {\n            \"labels\": labels,\n            \"scores\": scores,\n            \"bboxes\": bboxes,\n            \"above_threshold\": above_threshold,\n        }\n\n    @staticmethod\n    def build_training_sample(\n        imageid: int, img: np.ndarray, labels: List[int], bboxes: List[BBox], **kwargs,\n    ):\n        x = im2tensor(img)\n        # injected values when annotations are empty are disconsidered\n        # because we mark label as 0 (background)\n        _fake_box = [0, 1, 2, 3]\n        y = {\n            \"image_id\": tensor(imageid, dtype=torch.int64),\n            \"labels\": tensor(labels or [0], dtype=torch.int64),\n            \"boxes\": tensor([o.xyxy for o in bboxes] or [_fake_box], dtype=torch.float),\n        }\n        return x, y\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py"}, {"content": "__all__ = [\"MantisRCNN\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.utils import *\nfrom mantisshrimp.models.mantis_module import *\nfrom mantisshrimp.backbones import *\n\n\nclass MantisRCNN(MantisModule, ABC):\n    def __init__(self, metrics=None):\n        super().__init__()\n        self.metrics = metrics or []\n\n    @staticmethod\n    @abstractmethod\n    def build_training_sample(*args, **kwargs):\n        \"\"\"\n        Converts a record to a format understood by the model.\n        \"\"\"\n\n    def _predict(\n        self, images: List[np.ndarray], convert_raw_prediction,\n    ):\n        self.eval()\n        tensor_images = [im2tensor(img).to(self.device) for img in images]\n        raw_preds = self(tensor_images)\n\n        return [convert_raw_prediction(raw_pred) for raw_pred in raw_preds]\n\n    @staticmethod\n    def loss(preds, targs) -> Tensor:\n        return sum(preds.values())\n\n    @classmethod\n    def collate_fn(cls, data):\n        ts = [cls.build_training_sample(**o) for o in data]\n        xb, yb = zip(*ts)\n        return xb, list(yb)\n\n    @classmethod\n    def dataloader(cls, dataset, **kwargs) -> DataLoader:\n        return DataLoader(dataset=dataset, collate_fn=cls.collate_fn, **kwargs)\n\n    @staticmethod\n    def get_backbone_by_name(\n        name: str, fpn: bool = True, pretrained: bool = True\n    ) -> nn.Module:\n        \"\"\"\n        Args:\n            backbone (str): If none creates a default resnet50_fpn model trained on MS COCO 2017\n                Supported backones are: \"resnet18\", \"resnet34\",\"resnet50\", \"resnet101\", \"resnet152\",\n                 \"resnext50_32x4d\", \"resnext101_32x8d\", \"wide_resnet50_2\", \"wide_resnet101_2\", as resnets with fpn backbones.\n                Without fpn backbones supported are: \"resnet18\", \"resnet34\", \"resnet50\",\"resnet101\",\n                 \"resnet152\", \"resnext101_32x8d\", \"mobilenet\", \"vgg11\", \"vgg13\", \"vgg16\", \"vgg19\",\n            pretrained (bool): Creates a pretrained backbone with imagenet weights.\n        \"\"\"\n        # Giving string as a backbone, which is either supported resnet or backbone\n        if fpn:\n            # Creates a torchvision resnet model with fpn added\n            # It returns BackboneWithFPN model\n            backbone = resnet_fpn_backbone(name, pretrained=pretrained)\n        else:\n            # This does not create fpn backbone, it is supported for all models\n            backbone = create_torchvision_backbone(name, pretrained=pretrained)\n        return backbone\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_rcnn.py"}], "after_files": [{"content": "__all__ = [\"MantisMaskRCNN\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.core import *\nfrom mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *\nfrom mantisshrimp.models.mantis_rcnn.mantis_rcnn import *\nfrom mantisshrimp.models.mantis_rcnn.mantis_faster_rcnn import *\n\n\nclass MantisMaskRCNN(MantisRCNN):\n    @delegates(MaskRCNN.__init__)\n    def __init__(\n        self,\n        num_classes: int,\n        backbone: nn.Module = None,\n        param_groups: List[nn.Module] = None,\n        remove_internal_transforms: bool = True,\n        **kwargs,\n    ):\n        super().__init__()\n        self.num_classes = num_classes\n\n        if backbone is None:\n            # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset\n            self.model = maskrcnn_resnet50_fpn(pretrained=True, **kwargs)\n            in_features = self.model.roi_heads.box_predictor.cls_score.in_features\n            self.model.roi_heads.box_predictor = FastRCNNPredictor(\n                in_features, num_classes\n            )\n            in_features_mask = (\n                self.model.roi_heads.mask_predictor.conv5_mask.in_channels\n            )\n            self.model.roi_heads.mask_predictor = MaskRCNNPredictor(\n                in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes\n            )\n            param_groups = resnet_fpn_backbone_param_groups(self.model.backbone)\n        else:\n            self.model = MaskRCNN(backbone, num_classes=num_classes, **kwargs)\n            param_groups = param_groups or [backbone]\n\n        self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]\n        check_all_model_params_in_groups(self.model, self.param_groups)\n\n        if remove_internal_transforms:\n            self._remove_transforms_from_model(self.model)\n\n    def forward(self, images, targets=None):\n        return self.model(images, targets)\n\n    def predict(\n        self,\n        images: List[np.ndarray],\n        detection_threshold: float = 0.5,\n        mask_threshold: float = 0.5,\n    ):\n        convert_raw_prediction = partial(\n            self.convert_raw_prediction,\n            detection_threshold=detection_threshold,\n            mask_threshold=mask_threshold,\n        )\n\n        return self._predict(\n            images=images, convert_raw_prediction=convert_raw_prediction\n        )\n\n    @property\n    def param_groups(self):\n        return self._param_groups\n\n    @staticmethod\n    def convert_raw_prediction(\n        raw_pred: dict, detection_threshold: float, mask_threshold: float\n    ):\n        preds = MantisFasterRCNN.convert_raw_prediction(\n            raw_pred=raw_pred, detection_threshold=detection_threshold\n        )\n\n        above_threshold = preds[\"above_threshold\"]\n        masks_probs = raw_pred[\"masks\"][above_threshold]\n        masks_probs = masks_probs.detach().cpu().numpy()\n        # convert probabilities to 0 or 1 based on mask_threshold\n        masks = masks_probs > mask_threshold\n        masks = MaskArray(masks.squeeze(1))\n\n        return {**preds, \"masks\": masks}\n\n    @staticmethod\n    def build_training_sample(\n        imageid: int,\n        img: np.ndarray,\n        labels: List[int],\n        bboxes: List[BBox],\n        masks: MaskArray,\n        **kwargs,\n    ):\n        x, y = MantisFasterRCNN.build_training_sample(\n            imageid=imageid, img=img, labels=labels, bboxes=bboxes,\n        )\n        y[\"masks\"] = tensor(masks.data, dtype=torch.uint8)\n        return x, y\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_mask_rcnn.py"}, {"content": "__all__ = [\"MantisFasterRCNN\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.core import *\nfrom mantisshrimp.models.mantis_rcnn.rcnn_param_groups import *\nfrom mantisshrimp.models.mantis_rcnn.mantis_rcnn import *\n\n\nclass MantisFasterRCNN(MantisRCNN):\n    \"\"\"\n    Creates a flexible Faster RCNN implementation based on torchvision library.\n    Args: \n    n_class (int) : number of classes. Do not have class_id \"0\" it is reserved as background.\n                    n_class = number of classes to label + 1 for background.\n    \"\"\"\n\n    @delegates(FasterRCNN.__init__)\n    def __init__(\n        self,\n        num_classes: int,\n        backbone: nn.Module = None,\n        param_groups: List[nn.Module] = None,\n        metrics=None,\n        remove_internal_transforms=True,\n        **kwargs,\n    ):\n        super().__init__(metrics=metrics)\n        self.num_classes = num_classes\n        self.backbone = backbone\n        if backbone is None:\n            # Creates the default fasterrcnn as given in pytorch. Trained on COCO dataset\n            self.model = fasterrcnn_resnet50_fpn(pretrained=True, **kwargs)\n            in_features = self.model.roi_heads.box_predictor.cls_score.in_features\n            self.model.roi_heads.box_predictor = FastRCNNPredictor(\n                in_features, num_classes\n            )\n            param_groups = resnet_fpn_backbone_param_groups(self.model.backbone)\n        else:\n            self.model = FasterRCNN(backbone, num_classes=num_classes, **kwargs)\n            param_groups = param_groups or [backbone]\n\n        self._param_groups = param_groups + [self.model.rpn, self.model.roi_heads]\n        check_all_model_params_in_groups(self.model, self._param_groups)\n\n        if remove_internal_transforms:\n            self._remove_transforms_from_model(self.model)\n\n    def forward(self, images, targets=None):\n        return self.model(images, targets)\n\n    def predict(self, images: List[np.ndarray], detection_threshold: float = 0.5):\n        convert_raw_prediction = partial(\n            self.convert_raw_prediction, detection_threshold=detection_threshold,\n        )\n\n        return self._predict(\n            images=images, convert_raw_prediction=convert_raw_prediction\n        )\n\n    @property\n    def param_groups(self):\n        return self._param_groups\n\n    @staticmethod\n    def convert_raw_prediction(raw_pred: dict, detection_threshold: float):\n        above_threshold = raw_pred[\"scores\"] >= detection_threshold\n\n        labels = raw_pred[\"labels\"][above_threshold]\n        labels = labels.detach().cpu().numpy()\n\n        scores = raw_pred[\"scores\"][above_threshold]\n        scores = scores.detach().cpu().numpy()\n\n        boxes = raw_pred[\"boxes\"][above_threshold]\n        bboxes = []\n        for box_tensor in boxes:\n            xyxy = box_tensor.cpu().tolist()\n            bbox = BBox.from_xyxy(*xyxy)\n            bboxes.append(bbox)\n\n        return {\n            \"labels\": labels,\n            \"scores\": scores,\n            \"bboxes\": bboxes,\n            \"above_threshold\": above_threshold,\n        }\n\n    @staticmethod\n    def build_training_sample(\n        imageid: int, img: np.ndarray, labels: List[int], bboxes: List[BBox], **kwargs,\n    ):\n        x = im2tensor(img)\n        # injected values when annotations are empty are disconsidered\n        # because we mark label as 0 (background)\n        _fake_box = [0, 1, 2, 3]\n        y = {\n            \"image_id\": tensor(imageid, dtype=torch.int64),\n            \"labels\": tensor(labels or [0], dtype=torch.int64),\n            \"boxes\": tensor([o.xyxy for o in bboxes] or [_fake_box], dtype=torch.float),\n        }\n        return x, y\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py"}, {"content": "__all__ = [\"MantisRCNN\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.utils import *\nfrom mantisshrimp.models.mantis_module import *\nfrom mantisshrimp.backbones import *\n\n\nclass MantisRCNN(MantisModule, ABC):\n    def __init__(self, metrics=None):\n        super().__init__()\n        self.metrics = metrics or []\n\n    @staticmethod\n    @abstractmethod\n    def build_training_sample(*args, **kwargs):\n        \"\"\"\n        Converts a record to a format understood by the model.\n        \"\"\"\n\n    def _predict(\n        self, images: List[np.ndarray], convert_raw_prediction,\n    ):\n        self.eval()\n        tensor_images = [im2tensor(img).to(self.device) for img in images]\n        raw_preds = self(tensor_images)\n\n        return [convert_raw_prediction(raw_pred) for raw_pred in raw_preds]\n\n    def _remove_transforms_from_model(self, model: GeneralizedRCNN):\n        def noop_normalize(image):\n            return image\n\n        def noop_resize(image, target):\n            return image, target\n\n        model.transform.normalize = noop_normalize\n        model.transform.resize = noop_resize\n\n    @staticmethod\n    def loss(preds, targs) -> Tensor:\n        return sum(preds.values())\n\n    @classmethod\n    def collate_fn(cls, data):\n        ts = [cls.build_training_sample(**o) for o in data]\n        xb, yb = zip(*ts)\n        return xb, list(yb)\n\n    @classmethod\n    def dataloader(cls, dataset, **kwargs) -> DataLoader:\n        return DataLoader(dataset=dataset, collate_fn=cls.collate_fn, **kwargs)\n\n    @staticmethod\n    def get_backbone_by_name(\n        name: str, fpn: bool = True, pretrained: bool = True\n    ) -> nn.Module:\n        \"\"\"\n        Args:\n            backbone (str): If none creates a default resnet50_fpn model trained on MS COCO 2017\n                Supported backones are: \"resnet18\", \"resnet34\",\"resnet50\", \"resnet101\", \"resnet152\",\n                 \"resnext50_32x4d\", \"resnext101_32x8d\", \"wide_resnet50_2\", \"wide_resnet101_2\", as resnets with fpn backbones.\n                Without fpn backbones supported are: \"resnet18\", \"resnet34\", \"resnet50\",\"resnet101\",\n                 \"resnet152\", \"resnext101_32x8d\", \"mobilenet\", \"vgg11\", \"vgg13\", \"vgg16\", \"vgg19\",\n            pretrained (bool): Creates a pretrained backbone with imagenet weights.\n        \"\"\"\n        # Giving string as a backbone, which is either supported resnet or backbone\n        if fpn:\n            # Creates a torchvision resnet model with fpn added\n            # It returns BackboneWithFPN model\n            backbone = resnet_fpn_backbone(name, pretrained=pretrained)\n        else:\n            # This does not create fpn backbone, it is supported for all models\n            backbone = create_torchvision_backbone(name, pretrained=pretrained)\n        return backbone\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_rcnn.py"}]}