diff --git a/docs/autogen.py b/docs/autogen.py index cce93aca6..87efd6bb4 100644 --- a/docs/autogen.py +++ b/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", ], } diff --git a/docs/templates/images/mmsegmentation-camvid.png b/docs/templates/images/mmsegmentation-camvid.png new file mode 100644 index 000000000..e06a43563 Binary files /dev/null and b/docs/templates/images/mmsegmentation-camvid.png differ diff --git a/icevision/models/fastai/unet/backbones/backbone_config.py b/icevision/models/fastai/unet/backbones/backbone_config.py deleted file mode 100644 index d02b64a85..000000000 --- a/icevision/models/fastai/unet/backbones/backbone_config.py +++ /dev/null @@ -1,6 +0,0 @@ -from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig - - -class TorchvisionUNetBackboneConfig(TorchvisionBackboneConfig): - def __init__(self, **kwargs): - super().__init__(model_name="unet", **kwargs) diff --git a/icevision/models/fastai/unet/backbones/mobilenet_configs.py b/icevision/models/fastai/unet/backbones/mobilenet_configs.py index 783f5e830..7f9afca01 100644 --- a/icevision/models/fastai/unet/backbones/mobilenet_configs.py +++ b/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) diff --git a/icevision/models/fastai/unet/backbones/resnet_configs.py b/icevision/models/fastai/unet/backbones/resnet_configs.py index 09ea24d4a..956bb45a4 100644 --- a/icevision/models/fastai/unet/backbones/resnet_configs.py +++ b/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): diff --git a/icevision/models/ross/efficientdet/backbones/__init__.py b/icevision/models/ross/efficientdet/backbones/__init__.py new file mode 100644 index 000000000..feac8ab3f --- /dev/null +++ b/icevision/models/ross/efficientdet/backbones/__init__.py @@ -0,0 +1 @@ +from icevision.models.ross.efficientdet.backbones.backbones import * diff --git a/icevision/models/ross/efficientdet/backbones.py b/icevision/models/ross/efficientdet/backbones/backbones.py similarity index 100% rename from icevision/models/ross/efficientdet/backbones.py rename to icevision/models/ross/efficientdet/backbones/backbones.py diff --git a/icevision/models/ross/efficientdet/fastai/learner.py b/icevision/models/ross/efficientdet/fastai.py similarity index 62% rename from icevision/models/ross/efficientdet/fastai/learner.py rename to icevision/models/ross/efficientdet/fastai.py index 0e3e6f916..640e1b6b7 100644 --- a/icevision/models/ross/efficientdet/fastai/learner.py +++ b/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( diff --git a/icevision/models/ross/efficientdet/fastai/__init__.py b/icevision/models/ross/efficientdet/fastai/__init__.py deleted file mode 100644 index f7e8ac0e0..000000000 --- a/icevision/models/ross/efficientdet/fastai/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from icevision.models.ross.efficientdet.fastai.callbacks import * -from icevision.models.ross.efficientdet.fastai.learner import * diff --git a/icevision/models/ross/efficientdet/fastai/callbacks.py b/icevision/models/ross/efficientdet/fastai/callbacks.py deleted file mode 100644 index 17b53b7b9..000000000 --- a/icevision/models/ross/efficientdet/fastai/callbacks.py +++ /dev/null @@ -1,25 +0,0 @@ -__all__ = ["EfficientDetCallback"] - -from icevision.models.ross import efficientdet -from icevision.engines.fastai import * - - -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 diff --git a/icevision/models/ross/efficientdet/lightning/model_adapter.py b/icevision/models/ross/efficientdet/lightning.py similarity index 100% rename from icevision/models/ross/efficientdet/lightning/model_adapter.py rename to icevision/models/ross/efficientdet/lightning.py diff --git a/icevision/models/ross/efficientdet/lightning/__init__.py b/icevision/models/ross/efficientdet/lightning/__init__.py deleted file mode 100644 index c96fa5702..000000000 --- a/icevision/models/ross/efficientdet/lightning/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from icevision.models.ross.efficientdet.lightning.model_adapter import * diff --git a/icevision/models/ross/efficientdet/show_batch.py b/icevision/models/ross/efficientdet/show_batch.py index 76ac1285f..da3de1d08 100644 --- a/icevision/models/ross/efficientdet/show_batch.py +++ b/icevision/models/ross/efficientdet/show_batch.py @@ -1,5 +1,6 @@ __all__ = ["show_batch"] +import torch from icevision.utils import * from icevision.visualize import * diff --git a/icevision/models/torchvision/faster_rcnn/backbones/backbone_config.py b/icevision/models/torchvision/faster_rcnn/backbones/backbone_config.py deleted file mode 100644 index ae86bfdd5..000000000 --- a/icevision/models/torchvision/faster_rcnn/backbones/backbone_config.py +++ /dev/null @@ -1,6 +0,0 @@ -from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig - - -class TorchvisionFasterRCNNBackboneConfig(TorchvisionBackboneConfig): - def __init__(self, **kwargs): - super().__init__(model_name="faster_rcnn", **kwargs) diff --git a/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_configs.py b/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_configs.py index 5bb531c1d..de793fca0 100644 --- a/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_configs.py +++ b/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) diff --git a/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_utils.py b/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_utils.py deleted file mode 100644 index d14dc6f8e..000000000 --- a/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_utils.py +++ /dev/null @@ -1,25 +0,0 @@ -__all__ = [ - "param_groups", - "patch_param_groups", -] - -from icevision.imports import * -from icevision.utils import * - - -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) diff --git a/icevision/models/torchvision/keypoint_rcnn/backbones/backbone_config.py b/icevision/models/torchvision/keypoint_rcnn/backbones/backbone_config.py deleted file mode 100644 index 0eed6b65d..000000000 --- a/icevision/models/torchvision/keypoint_rcnn/backbones/backbone_config.py +++ /dev/null @@ -1,6 +0,0 @@ -from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig - - -class TorchvisionKeypointRCNNBackboneConfig(TorchvisionBackboneConfig): - def __init__(self, **kwargs): - super().__init__(model_name="keypoint_rcnn", **kwargs) diff --git a/icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_configs.py b/icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_configs.py index 2298cf5cd..2655bddee 100644 --- a/icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_configs.py +++ b/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) diff --git a/icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_utils.py b/icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_utils.py deleted file mode 100644 index d14dc6f8e..000000000 --- a/icevision/models/torchvision/keypoint_rcnn/backbones/resnet_fpn_utils.py +++ /dev/null @@ -1,25 +0,0 @@ -__all__ = [ - "param_groups", - "patch_param_groups", -] - -from icevision.imports import * -from icevision.utils import * - - -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) diff --git a/icevision/models/torchvision/mask_rcnn/backbones/backbone_config.py b/icevision/models/torchvision/mask_rcnn/backbones/backbone_config.py deleted file mode 100644 index 41272946f..000000000 --- a/icevision/models/torchvision/mask_rcnn/backbones/backbone_config.py +++ /dev/null @@ -1,6 +0,0 @@ -from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig - - -class TorchvisionMaskRCNNBackboneConfig(TorchvisionBackboneConfig): - def __init__(self, **kwargs): - super().__init__(model_name="mask_rcnn", **kwargs) diff --git a/icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_configs.py b/icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_configs.py index b3c716a3e..875a3b9e4 100644 --- a/icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_configs.py +++ b/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) diff --git a/icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_utils.py b/icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_utils.py deleted file mode 100644 index d14dc6f8e..000000000 --- a/icevision/models/torchvision/mask_rcnn/backbones/resnet_fpn_utils.py +++ /dev/null @@ -1,25 +0,0 @@ -__all__ = [ - "param_groups", - "patch_param_groups", -] - -from icevision.imports import * -from icevision.utils import * - - -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) diff --git a/icevision/models/torchvision/retinanet/backbones/backbone_config.py b/icevision/models/torchvision/retinanet/backbones/backbone_config.py deleted file mode 100644 index 560346b83..000000000 --- a/icevision/models/torchvision/retinanet/backbones/backbone_config.py +++ /dev/null @@ -1,6 +0,0 @@ -from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig - - -class TorchvisionRetinanetBackboneConfig(TorchvisionBackboneConfig): - def __init__(self, **kwargs): - super().__init__(model_name="retinanet", **kwargs) diff --git a/icevision/models/torchvision/retinanet/backbones/resnet_fpn_configs.py b/icevision/models/torchvision/retinanet/backbones/resnet_fpn_configs.py index 73e8293be..132763c4d 100644 --- a/icevision/models/torchvision/retinanet/backbones/resnet_fpn_configs.py +++ b/icevision/models/torchvision/retinanet/backbones/resnet_fpn_configs.py @@ -10,16 +10,37 @@ "wide_resnet101_2_fpn", ] -from icevision.models.torchvision.retinanet.backbones.resnet_fpn_utils import ( - patch_param_groups, -) +from icevision.utils import * +from icevision.imports import * from torchvision.models.detection.backbone_utils import resnet_fpn_backbone +from icevision.models.torchvision.backbone_config import TorchvisionBackboneConfig -from icevision.models.torchvision.retinanet.backbones.backbone_config import ( - TorchvisionRetinanetBackboneConfig, -) + +# utils +class TorchvisionRetinanetBackboneConfig(TorchvisionBackboneConfig): + def __init__(self, **kwargs): + super().__init__(model_name="retinanet", **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) diff --git a/icevision/models/torchvision/retinanet/backbones/resnet_fpn_utils.py b/icevision/models/torchvision/retinanet/backbones/resnet_fpn_utils.py deleted file mode 100644 index d14dc6f8e..000000000 --- a/icevision/models/torchvision/retinanet/backbones/resnet_fpn_utils.py +++ /dev/null @@ -1,25 +0,0 @@ -__all__ = [ - "param_groups", - "patch_param_groups", -] - -from icevision.imports import * -from icevision.utils import * - - -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) diff --git a/icevision/models/ultralytics/yolov5/backbones/__init__.py b/icevision/models/ultralytics/yolov5/backbones/__init__.py new file mode 100644 index 000000000..f973d028f --- /dev/null +++ b/icevision/models/ultralytics/yolov5/backbones/__init__.py @@ -0,0 +1 @@ +from icevision.models.ultralytics.yolov5.backbones.backbones import * diff --git a/icevision/models/ultralytics/yolov5/backbones.py b/icevision/models/ultralytics/yolov5/backbones/backbones.py similarity index 100% rename from icevision/models/ultralytics/yolov5/backbones.py rename to icevision/models/ultralytics/yolov5/backbones/backbones.py diff --git a/icevision/models/ultralytics/yolov5/fastai/learner.py b/icevision/models/ultralytics/yolov5/fastai.py similarity index 63% rename from icevision/models/ultralytics/yolov5/fastai/learner.py rename to icevision/models/ultralytics/yolov5/fastai.py index b3f2957d3..4d1e69edc 100644 --- a/icevision/models/ultralytics/yolov5/fastai/learner.py +++ b/icevision/models/ultralytics/yolov5/fastai.py @@ -2,8 +2,31 @@ from icevision.imports import * from icevision.engines.fastai import * -from icevision.models.ultralytics.yolov5.fastai.callbacks import Yolov5Callback from yolov5.utils.loss import ComputeLoss +from icevision.models.ultralytics import yolov5 + + +class Yolov5Callback(fastai.Callback): + def before_batch(self): + assert len(self.xb) == len(self.yb) == 1, "Only works for single input-output" + x, y, records = self.xb[0][0], self.xb[0][1], self.yb + self.learn.xb = [x] + self.learn.yb = [y] + self.learn.records = records[0] + + def after_pred(self): + if not self.training: + inference_out, training_out = self.pred[0], self.pred[1] + self.learn.pred = training_out + batch = (*self.learn.xb, *self.learn.yb) + preds = yolov5.convert_raw_predictions( + batch=batch, + raw_preds=inference_out, + records=self.learn.records, + detection_threshold=0.001, + nms_iou_threshold=0.6, + ) + self.learn.converted_preds = preds def learner( diff --git a/icevision/models/ultralytics/yolov5/fastai/__init__.py b/icevision/models/ultralytics/yolov5/fastai/__init__.py deleted file mode 100644 index 66e9d82c7..000000000 --- a/icevision/models/ultralytics/yolov5/fastai/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from icevision.models.ultralytics.yolov5.fastai.callbacks import * -from icevision.models.ultralytics.yolov5.fastai.learner import * diff --git a/icevision/models/ultralytics/yolov5/fastai/callbacks.py b/icevision/models/ultralytics/yolov5/fastai/callbacks.py deleted file mode 100644 index 7d83adf4d..000000000 --- a/icevision/models/ultralytics/yolov5/fastai/callbacks.py +++ /dev/null @@ -1,27 +0,0 @@ -__all__ = ["Yolov5Callback"] - -from icevision.engines.fastai import * -from icevision.models.ultralytics import yolov5 - - -class Yolov5Callback(fastai.Callback): - def before_batch(self): - assert len(self.xb) == len(self.yb) == 1, "Only works for single input-output" - x, y, records = self.xb[0][0], self.xb[0][1], self.yb - self.learn.xb = [x] - self.learn.yb = [y] - self.learn.records = records[0] - - def after_pred(self): - if not self.training: - inference_out, training_out = self.pred[0], self.pred[1] - self.learn.pred = training_out - batch = (*self.learn.xb, *self.learn.yb) - preds = yolov5.convert_raw_predictions( - batch=batch, - raw_preds=inference_out, - records=self.learn.records, - detection_threshold=0.001, - nms_iou_threshold=0.6, - ) - self.learn.converted_preds = preds diff --git a/icevision/models/ultralytics/yolov5/lightning/model_adapter.py b/icevision/models/ultralytics/yolov5/lightning.py similarity index 100% rename from icevision/models/ultralytics/yolov5/lightning/model_adapter.py rename to icevision/models/ultralytics/yolov5/lightning.py diff --git a/icevision/models/ultralytics/yolov5/lightning/__init__.py b/icevision/models/ultralytics/yolov5/lightning/__init__.py deleted file mode 100644 index cb3a68a8f..000000000 --- a/icevision/models/ultralytics/yolov5/lightning/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from icevision.models.ultralytics.yolov5.lightning.model_adapter import * diff --git a/icevision/utils/model_creation.py b/icevision/utils/model_creation.py new file mode 100644 index 000000000..b958d38b1 --- /dev/null +++ b/icevision/utils/model_creation.py @@ -0,0 +1,222 @@ +import os +import importlib +from typing import types +from typing import Optional, Union, List +import pathlib + +from icevision import models + + +MODEL_CREATION_FILEPATH = pathlib.Path(__file__).parent.resolve() +BASE_PATH = os.path.join(MODEL_CREATION_FILEPATH, "../models") + + +def get_module_element_form_module( + module: types.ModuleType, *args: str +) -> Union[types.ModuleType, callable, object]: + """Loads a submodule for a given module. Args are the elements to navigate the submodules, the order of the args needs to be the order the submodules appear in (sub)module. + + Parameters + ---------- + module: types.ModuleType + Module to load the element from + args: str + Strings with the names of the modules/functions/objects in order to retrieve the element + + Example + ------- + >>> import icevision as iv + >>> # load the efficientdet model type inside the ross module + >>> model_type = get_submodule_form_module(iv.models, 'ross', 'efficientdet') + + Returns + ------- + element_to_return: Union[Module, function, object] + element that corrosponds to the last element in the args + """ + element_to_return = module + for arg in args: + element_to_return = getattr(element_to_return, arg) + return element_to_return + + +def get_backend_libs(base_path: str = BASE_PATH) -> List[str]: + """Returns all backend_libs available in icevision + + Parameters + ---------- + base_path: str + Path from where to read out the available models from + + Returns + ------- + backend_libs: List[str] + List of available backend libs + """ + # FIXME remove the active_backend_libs when the lib restructure is done + active_backend_libs = ["ross", "torchvision", "fastai", "ultralytics"] + backend_libs = [ + i + for i in os.listdir(base_path) + if os.path.isdir(os.path.join(base_path, i)) + and not i.startswith("__") + and i in active_backend_libs + ] + return backend_libs + + +def get_model_types_for_backend_lib( + lib_name: str, base_path: str = BASE_PATH +) -> List[str]: + """Returns all available model_types for a given backend lib + + Parameters + ---------- + lib_name: str + Name of the backend_lib to get the model types of + base_path: str + Path from where to read out the available models from + + Returns + ------- + model_types: List[str] + List of available model types + """ + model_types = [ + i + for i in os.listdir(os.path.join(base_path, lib_name)) + if os.path.isdir(os.path.join(base_path, lib_name, i)) + and not i.startswith("__") + ] + return model_types + + +def get_backbone_names( + lib_name: str, model_name: str, base_path: str = BASE_PATH +) -> List[str]: + """Returns all available backbones for a given backend lib + + Parameters + ---------- + lib_name: str + Name of the backend_lib to get the backbones for + model_name: str + Name of the model type to get the backbones for + base_path: str + Path from where to read out the available models from + + Returns + ------- + backbones: List[str] + List of available backbones + """ + backbone_files = [ + i + for i in os.listdir(os.path.join(base_path, lib_name, model_name, "backbones")) + if os.path.join(base_path, lib_name, model_name, "backbones", i).endswith(".py") + and not i.startswith("__") + ] + backbone_list = [] + for backbone_file in backbone_files: + backbones = importlib.import_module( + "icevision.models." + + ".".join([lib_name, model_name]) + + ".backbones." + + backbone_file.replace(".py", "") + ) + backbone_list += getattr(backbones, "__all__", []) + return backbone_list + + +def load_model_components( + lib_name: str, + model_name: str, + backbone_name: str, + model_base_path: str = BASE_PATH, +): + """Build a model form given backend lib, model_type and backbone. + + Parameters + ---------- + lib_name: str + Name of the backend_lib to use + model_name: str + Name of the model type to use + backbone_name: str + Name of the backbone to use + + Returns + ------- + model_type: Module + model_type to be used to create the dataloaders and so on + backbone: object + backbone for the model + """ + # load the backend lib + try: + backend_lib = get_module_element_form_module(models, lib_name) + except AttributeError: + raise AttributeError( + f"Backend lib: {lib_name} not found. Possible options are: {get_backend_libs(model_base_path)}" + ) + # load the model_type + try: + model_type = get_module_element_form_module(models, lib_name, model_name) + except AttributeError: + raise AttributeError( + f"Model type: {model_name} not found. Possible options are: {get_model_types_for_backend_lib(model_base_path, lib_name)}" + ) + # load backbone + try: + backbone_module = get_module_element_form_module( + models, lib_name, model_name, "backbones" + ) + backbone = getattr(backbone_module, backbone_name) + except AttributeError: + raise AttributeError( + f"Backbone: {backbone_name} not found. Possible options are: {get_backbone_names(lib_name, model_name)}" + ) + + return model_type, backbone + + +def build_model( + lib_name: str, + model_name: str, + backbone_name: str, + num_classes: int, + backbone_config: Optional[dict] = None, + model_config: Optional[dict] = None, +): + """Build a model form given backend lib, model_type and backbone. + + Parameters + ---------- + lib_name: str + Name of the backend_lib to use + model_name: str + Name of the model type to use + backbone_name: str + Name of the backbone to use + num_classes: int + Number of classes to initialize the model for + backbone_config: Optional[dict] + Configuration for the backbone. If None the backbone_config will be set to {"pretrained": True} + model_config: Optional[dict] + Configuration for the model + + + Returns + ------- + model_type: model_type to be used to create the dataloaders and so on + model: model to train + """ + if backbone_config is None: + backbone_config = {"pretrained": True} + if model_config is None: + model_config = {} + model_type, backbone = load_model_components(lib_name, model_name, backbone_name) + model = model_type.model( + backbone=backbone(**backbone_config), num_classes=num_classes, **model_config + ) + return model_type, model diff --git a/tests/models/efficient_det/__init__.py b/tests/models/ross/__init__.py similarity index 100% rename from tests/models/efficient_det/__init__.py rename to tests/models/ross/__init__.py diff --git a/tests/models/efficient_det/fastai/__init__.py b/tests/models/ross/efficient_det/__init__.py similarity index 100% rename from tests/models/efficient_det/fastai/__init__.py rename to tests/models/ross/efficient_det/__init__.py diff --git a/tests/models/efficient_det/conftest.py b/tests/models/ross/efficient_det/conftest.py similarity index 100% rename from tests/models/efficient_det/conftest.py rename to tests/models/ross/efficient_det/conftest.py diff --git a/tests/models/efficient_det/lightning/__init__.py b/tests/models/ross/efficient_det/fastai/__init__.py similarity index 100% rename from tests/models/efficient_det/lightning/__init__.py rename to tests/models/ross/efficient_det/fastai/__init__.py diff --git a/tests/models/efficient_det/fastai/test_train.py b/tests/models/ross/efficient_det/fastai/test_train.py similarity index 100% rename from tests/models/efficient_det/fastai/test_train.py rename to tests/models/ross/efficient_det/fastai/test_train.py diff --git a/tests/models/ross/efficient_det/lightning/__init__.py b/tests/models/ross/efficient_det/lightning/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/models/efficient_det/lightning/test_train.py b/tests/models/ross/efficient_det/lightning/test_train.py similarity index 100% rename from tests/models/efficient_det/lightning/test_train.py rename to tests/models/ross/efficient_det/lightning/test_train.py diff --git a/tests/models/efficient_det/test_dataloaders.py b/tests/models/ross/efficient_det/test_dataloaders.py similarity index 100% rename from tests/models/efficient_det/test_dataloaders.py rename to tests/models/ross/efficient_det/test_dataloaders.py diff --git a/tests/models/efficient_det/test_metrics.py b/tests/models/ross/efficient_det/test_metrics.py similarity index 100% rename from tests/models/efficient_det/test_metrics.py rename to tests/models/ross/efficient_det/test_metrics.py diff --git a/tests/models/efficient_det/test_model.py b/tests/models/ross/efficient_det/test_model.py similarity index 100% rename from tests/models/efficient_det/test_model.py rename to tests/models/ross/efficient_det/test_model.py diff --git a/tests/models/efficient_det/test_prediction.py b/tests/models/ross/efficient_det/test_prediction.py similarity index 100% rename from tests/models/efficient_det/test_prediction.py rename to tests/models/ross/efficient_det/test_prediction.py diff --git a/tests/models/efficient_det/test_show_results.py b/tests/models/ross/efficient_det/test_show_results.py similarity index 95% rename from tests/models/efficient_det/test_show_results.py rename to tests/models/ross/efficient_det/test_show_results.py index d6cd241fd..17211021d 100644 --- a/tests/models/efficient_det/test_show_results.py +++ b/tests/models/ross/efficient_det/test_show_results.py @@ -22,3 +22,7 @@ def test_plot_losses(fridge_efficientdet_model, fridge_ds, monkeypatch): model=model, dataset=ds, sort_by=by, n_samples=2 ) assert len(samples_plus_losses) == len(ds) == len(preds) + + +def test_show_batch(): + pass diff --git a/tests/utils/test_model_creation.py b/tests/utils/test_model_creation.py new file mode 100644 index 000000000..514a4c21a --- /dev/null +++ b/tests/utils/test_model_creation.py @@ -0,0 +1,96 @@ +import pytest +import torchvision + +from icevision.utils.model_creation import * +from icevision import models + + +def test_get_module_element_form_module(): + eff_det_loaded = get_module_element_form_module(models, "torchvision", "retinanet") + assert eff_det_loaded == models.torchvision.retinanet + + +def test_get_backend_libs(): + backend_libs = get_backend_libs() + assert sorted(backend_libs) == sorted( + ["ross", "torchvision", "fastai", "ultralytics"] + ) + + +def test_get_model_types_for_backend_lib(): + torchvision_model_types = get_model_types_for_backend_lib("torchvision") + assert sorted(torchvision_model_types) == sorted( + [ + "retinanet", + "keypoint_rcnn", + "faster_rcnn", + "mask_rcnn", + ] + ) + + +def test_get_backbone_names(): + backbone_names = get_backbone_names("torchvision", "retinanet") + assert sorted(backbone_names) == sorted( + [ + "resnet18_fpn", + "resnet34_fpn", + "resnet50_fpn", + "resnet101_fpn", + "resnet152_fpn", + "resnext50_32x4d_fpn", + "resnext101_32x8d_fpn", + "wide_resnet50_2_fpn", + "wide_resnet101_2_fpn", + ] + ) + + +def test_load_model_components(): + model_type, backbone = load_model_components( + "torchvision", "retinanet", "resnet18_fpn" + ) + assert model_type == models.torchvision.retinanet + assert backbone == models.torchvision.retinanet.backbones.resnet18_fpn + + +def test_load_model_components_throws_error_when_backend_lib_does_not_exist(): + with pytest.raises(AttributeError): + model_type, backbone = load_model_components( + "error", "retinanet", "resnet18_fpn" + ) + + +def test_load_model_components_throws_error_when_model_type_does_not_exist(): + with pytest.raises(AttributeError): + model_type, backbone = load_model_components( + "torchvision", "error", "resnet18_fpn" + ) + + +def test_load_model_components_throws_error_when_backbone_does_not_exist(): + with pytest.raises(AttributeError): + model_type, backbone = load_model_components( + "torchvision", "retinanet", "error" + ) + + +def test_build_model(): + model_type, model = build_model( + "torchvision", + "retinanet", + "resnet18_fpn", + num_classes=1, + backbone_config={"pretrained": False}, + model_config=None, + ) + assert model_type == models.torchvision.retinanet + assert isinstance(model, torchvision.models.detection.RetinaNet) + + +def test_build_model_without_give_backbone_config(): + model_type, model = build_model( + "torchvision", "retinanet", "resnet18_fpn", num_classes=1 + ) + assert model_type == models.torchvision.retinanet + assert isinstance(model, torchvision.models.detection.RetinaNet)