We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
load_state_dict bug
## 🐛 Bug
Introducing `load_state_dict` into the models in order to avoid doing `model.model.load_state_dict` caused a bug when trying to load from a previous checkpoint.

When saving the model the state dict looks like: `model.backbone...` but the new `load_state_dict` expects only `backbone`, without the model part.

More specifically, the following fails:
```python
model.load_state_dict(model.state_dict())
```

Related to #137 


--- 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     def load_state_dict(
64         self, state_dict: Dict[str, Tensor], strict: bool = True,
65     ):
66         return self.model.load_state_dict(state_dict=state_dict, strict=strict)
67 
68     @property
69     def param_groups(self):
70         return self._param_groups
71 
72     @staticmethod
73     def convert_raw_prediction(
74         raw_pred: dict, detection_threshold: float, mask_threshold: float
75     ):
76         preds = MantisFasterRCNN.convert_raw_prediction(
77             raw_pred=raw_pred, detection_threshold=detection_threshold
78         )
79 
80         above_threshold = preds["above_threshold"]
81         masks_probs = raw_pred["masks"][above_threshold]
82         masks_probs = masks_probs.detach().cpu().numpy()
83         # convert probabilities to 0 or 1 based on mask_threshold
84         masks = masks_probs > mask_threshold
85         masks = MaskArray(masks.squeeze(1))
86 
87         return {**preds, "masks": masks}
88 
89     @staticmethod
90     def build_training_sample(
91         imageid: int,
92         img: np.ndarray,
93         label: List[int],
94         bbox: List[BBox],
95         mask: MaskArray,
96         **kwargs,
97     ):
98         x, y = MantisFasterRCNN.build_training_sample(
99             imageid=imageid, img=img, label=label, bbox=bbox,
100         )
101         y["masks"] = tensor(mask.data, dtype=torch.uint8)
102         return x, y
103
```

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     def load_state_dict(
58         self, state_dict: Dict[str, Tensor], strict: bool = True,
59     ):
60         return self.model.load_state_dict(state_dict=state_dict, strict=strict)
61 
62     @property
63     def param_groups(self):
64         return self._param_groups
65 
66     @staticmethod
67     def convert_raw_prediction(raw_pred: dict, detection_threshold: float):
68         above_threshold = raw_pred["scores"] >= detection_threshold
69 
70         labels = raw_pred["labels"][above_threshold]
71         labels = labels.detach().cpu().numpy()
72 
73         scores = raw_pred["scores"][above_threshold]
74         scores = scores.detach().cpu().numpy()
75 
76         boxes = raw_pred["boxes"][above_threshold]
77         bboxes = []
78         for box_tensor in boxes:
79             xyxy = box_tensor.cpu().tolist()
80             bbox = BBox.from_xyxy(*xyxy)
81             bboxes.append(bbox)
82 
83         return {
84             "labels": labels,
85             "scores": scores,
86             "bboxes": bboxes,
87             "above_threshold": above_threshold,
88         }
89 
90     @staticmethod
91     def build_training_sample(
92         imageid: int, img: np.ndarray, label: List[int], bbox: List[BBox], **kwargs,
93     ):
94         x = im2tensor(img)
95         # injected values when annotations are empty are disconsidered
96         # because we mark label as 0 (background)
97         _fake_box = [0, 1, 2, 3]
98         y = {
99             "image_id": tensor(imageid, dtype=torch.int64),
100             "labels": tensor(label or [0], dtype=torch.int64),
101             "boxes": tensor([o.xyxy for o in bbox] or [_fake_box], dtype=torch.float),
102         }
103         return x, y
104
```

Path: `mantisshrimp/models/mantis_module.py`
Content:
```
1 __all__ = ["MantisModule"]
2 
3 from mantisshrimp.imports import *
4 from mantisshrimp.models.device_module_mixin import *
5 from mantisshrimp.models.parameters_splits_module_mixin import *
6 
7 
8 class MantisModule(DeviceModuleMixin, ParametersSplitsModuleMixin, nn.Module, ABC):
9     @abstractmethod
10     def load_state_dict(
11         self, state_dict: Dict[str, Tensor], strict: bool = True,
12     ):
13         """ Load weights
14         """
15 
16     @classmethod
17     @abstractmethod
18     def dataloader(cls, **kwargs) -> DataLoader:
19         """ Returns the specific dataloader for this class
20         """
21
```
--- 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_module.py b/mantisshrimp/models/mantis_module.py
--- a/mantisshrimp/models/mantis_module.py
+++ b/mantisshrimp/models/mantis_module.py
@@ -6,13 +6,6 @@
 
 
 class MantisModule(DeviceModuleMixin, ParametersSplitsModuleMixin, nn.Module, ABC):
-    @abstractmethod
-    def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: bool = True,
-    ):
-        """ Load weights
-        """
-
     @classmethod
     @abstractmethod
     def dataloader(cls, **kwargs) -> DataLoader:
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
@@ -54,11 +54,6 @@
             images=images, convert_raw_prediction=convert_raw_prediction
         )
 
-    def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: bool = True,
-    ):
-        return self.model.load_state_dict(state_dict=state_dict, strict=strict)
-
     @property
     def param_groups(self):
         return self._param_groups
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
@@ -60,11 +60,6 @@
             images=images, convert_raw_prediction=convert_raw_prediction
         )
 
-    def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: bool = True,
-    ):
-        return self.model.load_state_dict(state_dict=state_dict, strict=strict)
-
     @property
     def param_groups(self):
         return self._param_groups

{"golden_diff": "diff --git a/mantisshrimp/models/mantis_module.py b/mantisshrimp/models/mantis_module.py\n--- a/mantisshrimp/models/mantis_module.py\n+++ b/mantisshrimp/models/mantis_module.py\n@@ -6,13 +6,6 @@\n \n \n class MantisModule(DeviceModuleMixin, ParametersSplitsModuleMixin, nn.Module, ABC):\n-    @abstractmethod\n-    def load_state_dict(\n-        self, state_dict: Dict[str, Tensor], strict: bool = True,\n-    ):\n-        \"\"\" Load weights\n-        \"\"\"\n-\n     @classmethod\n     @abstractmethod\n     def dataloader(cls, **kwargs) -> DataLoader:\ndiff --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@@ -54,11 +54,6 @@\n             images=images, convert_raw_prediction=convert_raw_prediction\n         )\n \n-    def load_state_dict(\n-        self, state_dict: Dict[str, Tensor], strict: bool = True,\n-    ):\n-        return self.model.load_state_dict(state_dict=state_dict, strict=strict)\n-\n     @property\n     def param_groups(self):\n         return self._param_groups\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@@ -60,11 +60,6 @@\n             images=images, convert_raw_prediction=convert_raw_prediction\n         )\n \n-    def load_state_dict(\n-        self, state_dict: Dict[str, Tensor], strict: bool = True,\n-    ):\n-        return self.model.load_state_dict(state_dict=state_dict, strict=strict)\n-\n     @property\n     def param_groups(self):\n         return self._param_groups\n", "issue": "load_state_dict bug\n## \ud83d\udc1b Bug\r\nIntroducing `load_state_dict` into the models in order to avoid doing `model.model.load_state_dict` caused a bug when trying to load from a previous checkpoint.\r\n\r\nWhen saving the model the state dict looks like: `model.backbone...` but the new `load_state_dict` expects only `backbone`, without the model part.\r\n\r\nMore specifically, the following fails:\r\n```python\r\nmodel.load_state_dict(model.state_dict())\r\n```\r\n\r\nRelated to #137 \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    def load_state_dict(\n        self, state_dict: Dict[str, Tensor], strict: bool = True,\n    ):\n        return self.model.load_state_dict(state_dict=state_dict, strict=strict)\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        label: List[int],\n        bbox: List[BBox],\n        mask: MaskArray,\n        **kwargs,\n    ):\n        x, y = MantisFasterRCNN.build_training_sample(\n            imageid=imageid, img=img, label=label, bbox=bbox,\n        )\n        y[\"masks\"] = tensor(mask.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    def load_state_dict(\n        self, state_dict: Dict[str, Tensor], strict: bool = True,\n    ):\n        return self.model.load_state_dict(state_dict=state_dict, strict=strict)\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, label: List[int], bbox: 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(label or [0], dtype=torch.int64),\n            \"boxes\": tensor([o.xyxy for o in bbox] or [_fake_box], dtype=torch.float),\n        }\n        return x, y\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py"}, {"content": "__all__ = [\"MantisModule\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.models.device_module_mixin import *\nfrom mantisshrimp.models.parameters_splits_module_mixin import *\n\n\nclass MantisModule(DeviceModuleMixin, ParametersSplitsModuleMixin, nn.Module, ABC):\n    @abstractmethod\n    def load_state_dict(\n        self, state_dict: Dict[str, Tensor], strict: bool = True,\n    ):\n        \"\"\" Load weights\n        \"\"\"\n\n    @classmethod\n    @abstractmethod\n    def dataloader(cls, **kwargs) -> DataLoader:\n        \"\"\" Returns the specific dataloader for this class\n        \"\"\"\n", "path": "mantisshrimp/models/mantis_module.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 *\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        label: List[int],\n        bbox: List[BBox],\n        mask: MaskArray,\n        **kwargs,\n    ):\n        x, y = MantisFasterRCNN.build_training_sample(\n            imageid=imageid, img=img, label=label, bbox=bbox,\n        )\n        y[\"masks\"] = tensor(mask.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, label: List[int], bbox: 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(label or [0], dtype=torch.int64),\n            \"boxes\": tensor([o.xyxy for o in bbox] or [_fake_box], dtype=torch.float),\n        }\n        return x, y\n", "path": "mantisshrimp/models/mantis_rcnn/mantis_faster_rcnn.py"}, {"content": "__all__ = [\"MantisModule\"]\n\nfrom mantisshrimp.imports import *\nfrom mantisshrimp.models.device_module_mixin import *\nfrom mantisshrimp.models.parameters_splits_module_mixin import *\n\n\nclass MantisModule(DeviceModuleMixin, ParametersSplitsModuleMixin, nn.Module, ABC):\n    @classmethod\n    @abstractmethod\n    def dataloader(cls, **kwargs) -> DataLoader:\n        \"\"\" Returns the specific dataloader for this class\n        \"\"\"\n", "path": "mantisshrimp/models/mantis_module.py"}]}