Lib restructure

docs/autogen.py

@@ -72,10 +72,10 @@
         "icevision.models.ross.efficientdet.dataloaders.build_infer_batch",
     ],
     "efficientdet_fastai.md": [
-        "icevision.models.ross.efficientdet.fastai.learner.learner",
+        "icevision.models.ross.efficientdet.fastai.learner",
     ],
     "efficientdet_lightning.md": [
-        "icevision.models.ross.efficientdet.lightning.model_adapter.ModelAdapter",
+        "icevision.models.ross.efficientdet.lightning.ModelAdapter",
     ],
 }
 

icevision/models/fastai/unet/backbones/mobilenet_configs.py

@@ -6,11 +6,16 @@
 from icevision.utils.torch_utils import check_all_model_params_in_groups2
 from torch import nn
 import torchvision
-from icevision.models.fastai.unet.backbones.backbone_config import (
-    TorchvisionUNetBackboneConfig,
-)
+from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig
 
 
+# utils
+class TorchvisionUNetBackboneConfig(TorchvisionBackboneConfig):
+    def __init__(self, **kwargs):
+        super().__init__(model_name="unet", **kwargs)
+
+
+# backbones
 def mobilenet_fn(pretrained: bool = True):
     model = torchvision.models.mobilenet_v2(pretrained=pretrained)
 

icevision/models/fastai/unet/backbones/resnet_configs.py

@@ -1,5 +1,4 @@
 __all__ = [
-    "resnet_param_groups",
     "resnet18",
     "resnet34",
     "resnet50",
@@ -16,9 +15,12 @@
 from torch import nn
 import torchvision
 from icevision.utils.torch_utils import check_all_model_params_in_groups2
-from icevision.models.fastai.unet.backbones.backbone_config import (
-    TorchvisionUNetBackboneConfig,
-)
+from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig
+
+
+class TorchvisionUNetBackboneConfig(TorchvisionBackboneConfig):
+    def __init__(self, **kwargs):
+        super().__init__(model_name="unet", **kwargs)
 
 
 def _resnet_features(model: nn.Module, out_channels: int):

icevision/models/ross/efficientdet/backbones/__init__.py

@@ -0,0 +1 @@
+from icevision.models.ross.efficientdet.backbones.backbones import *

icevision/models/ross/efficientdet/fastai/learner.py → icevision/models/ross/efficientdet/fastai.py

@@ -1,9 +1,30 @@
-__all__ = ["learner"]
+__all__ = ["learner", "EfficientDetCallback"]
 
 from icevision.imports import *
 from icevision.engines.fastai import *
+from icevision.models.ross import efficientdet
 from icevision.models.ross.efficientdet.loss_fn import loss_fn
-from icevision.models.ross.efficientdet.fastai.callbacks import EfficientDetCallback
+
+
+class EfficientDetCallback(fastai.Callback):
+    def before_batch(self):
+        assert len(self.xb) == len(self.yb) == 1, "Only works for single input-output"
+        self.learn.xb = self.xb[0]
+        self.learn.records = self.yb[0]
+        self.learn.yb = ()
+
+    def after_pred(self):
+        self.learn.yb = [self.learn.xb[1]]
+        self.learn.xb = [self.learn.xb[0]]
+
+        if not self.training:
+            preds = efficientdet.convert_raw_predictions(
+                batch=(*self.xb, *self.yb),
+                raw_preds=self.pred["detections"],
+                records=self.learn.records,
+                detection_threshold=0.0,
+            )
+            self.learn.converted_preds = preds
 
 
 def learner(

icevision/models/ross/efficientdet/show_batch.py

@@ -1,5 +1,6 @@
 __all__ = ["show_batch"]
 
+import torch
 from icevision.utils import *
 from icevision.visualize import *
 

icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_configs.py

@@ -10,15 +10,37 @@
     "wide_resnet101_2_fpn",
 ]
 
-from icevision.models.torchvision.faster_rcnn.backbones.resnet_fpn_utils import (
-    patch_param_groups,
-)
+from icevision.imports import *
+from icevision.utils import *
 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
-from icevision.models.torchvision.faster_rcnn.backbones.backbone_config import (
-    TorchvisionFasterRCNNBackboneConfig,
-)
+from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig
+
+
+# utils
+class TorchvisionFasterRCNNBackboneConfig(TorchvisionBackboneConfig):
+    def __init__(self, **kwargs):
+        super().__init__(model_name="faster_rcnn", **kwargs)
+
+
+def param_groups(model: nn.Module) -> List[nn.Parameter]:
+    body = model.body
+
+    layers = []
+    layers += [nn.Sequential(body.conv1, body.bn1)]
+    layers += [getattr(body, l) for l in list(body) if l.startswith("layer")]
+    layers += [model.fpn]
+
+    _param_groups = [list(layer.parameters()) for layer in layers]
+    check_all_model_params_in_groups2(model, _param_groups)
+
+    return _param_groups
+
+
+def patch_param_groups(model: nn.Module) -> None:
+    model.param_groups = MethodType(param_groups, model)
 
 
+# backbones
 def _resnet_fpn(name: str, pretrained: bool = True, **kwargs):
     model = resnet_fpn_backbone(backbone_name=name, pretrained=pretrained, **kwargs)
     patch_param_groups(model)

icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_configs.py

@@ -10,15 +10,38 @@
     "wide_resnet101_2_fpn",
 ]
 
-from icevision.models.torchvision.keypoint_rcnn.backbones.resnet_fpn_utils import (
-    patch_param_groups,
-)
+
+from icevision.imports import *
+from icevision.utils import *
 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
-from icevision.models.torchvision.keypoint_rcnn.backbones.backbone_config import (
-    TorchvisionKeypointRCNNBackboneConfig,
-)
+from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig
+
+
+# utils
+class TorchvisionKeypointRCNNBackboneConfig(TorchvisionBackboneConfig):
+    def __init__(self, **kwargs):
+        super().__init__(model_name="keypoint_rcnn", **kwargs)
+
+
+def param_groups(model: nn.Module) -> List[nn.Parameter]:
+    body = model.body
+
+    layers = []
+    layers += [nn.Sequential(body.conv1, body.bn1)]
+    layers += [getattr(body, l) for l in list(body) if l.startswith("layer")]
+    layers += [model.fpn]
+
+    _param_groups = [list(layer.parameters()) for layer in layers]
+    check_all_model_params_in_groups2(model, _param_groups)
+
+    return _param_groups
+
+
+def patch_param_groups(model: nn.Module) -> None:
+    model.param_groups = MethodType(param_groups, model)
 
 
+# backbones
 def _resnet_fpn(name: str, pretrained: bool = True, **kwargs):
     model = resnet_fpn_backbone(backbone_name=name, pretrained=pretrained, **kwargs)
     patch_param_groups(model)

icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_configs.py

@@ -10,15 +10,37 @@
     "wide_resnet101_2_fpn",
 ]
 
-from icevision.models.torchvision.mask_rcnn.backbones.resnet_fpn_utils import (
-    patch_param_groups,
-)
+from icevision.imports import *
+from icevision.utils import *
 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
-from icevision.models.torchvision.mask_rcnn.backbones.backbone_config import (
-    TorchvisionMaskRCNNBackboneConfig,
-)
+from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig
+
+
+# utils
+class TorchvisionMaskRCNNBackboneConfig(TorchvisionBackboneConfig):
+    def __init__(self, **kwargs):
+        super().__init__(model_name="mask_rcnn", **kwargs)
+
+
+def param_groups(model: nn.Module) -> List[nn.Parameter]:
+    body = model.body
+
+    layers = []
+    layers += [nn.Sequential(body.conv1, body.bn1)]
+    layers += [getattr(body, l) for l in list(body) if l.startswith("layer")]
+    layers += [model.fpn]
+
+    _param_groups = [list(layer.parameters()) for layer in layers]
+    check_all_model_params_in_groups2(model, _param_groups)
+
+    return _param_groups
+
+
+def patch_param_groups(model: nn.Module) -> None:
+    model.param_groups = MethodType(param_groups, model)
 
 
+# backbones
 def _resnet_fpn(name: str, pretrained: bool = True, **kwargs):
     model = resnet_fpn_backbone(backbone_name=name, pretrained=pretrained, **kwargs)
     patch_param_groups(model)