Move yolox to object_detection module

docs/docs/api-reference/object_detection-yolox.md

@@ -0,0 +1,3 @@
+# zamba.object_detection.yolox
+
+::: zamba.object_detection.yolox

docs/docs/index.md

@@ -5,7 +5,10 @@
 [![codecov](https://codecov.io/gh/drivendataorg/zamba/branch/master/graph/badge.svg)](https://codecov.io/gh/drivendataorg/zamba)
 <!-- [![PyPI](https://img.shields.io/pypi/v/zamba.svg)](https://pypi.org/project/zamba/) -->
 
- <div class="embed-responsive embed-responsive-16by9" width=500>     <iframe width=600 height=340 class="embed-responsive-item" src="https://s3.amazonaws.com/drivendata-public-assets/monkey-vid.mp4" frameborder="0" allowfullscreen=""></iframe></div>
+
+<div class="embed-responsive embed-responsive-16by9" width=500> 
+    <iframe width=600 height=340 class="embed-responsive-item" src="https://s3.amazonaws.com/drivendata-public-assets/monkey-vid.mp4" 
+frameborder="0" allowfullscreen=""></iframe></div>
 
 > *Zamba* means "forest" in Lingala, a Bantu language spoken throughout the Democratic Republic of the Congo and the Republic of the Congo.
 

docs/mkdocs.yml

@@ -42,11 +42,11 @@ nav:
             - zamba.models.densepose.config: "api-reference/densepose_config.md"
             - zamba.models.densepose.densepose_manager: "api-reference/densepose_manager.md"
           - zamba.models.efficientnet_models: "api-reference/models-efficientnet_models.md"
-          - zamba.models.megadetector_lite_yolox: "api-reference/models-megadetector_lite_yolox.md"
           - zamba.models.model_manager: "api-reference/models-model_manager.md"
           - zamba.models.slowfast_models: "api-reference/models-slowfast_models.md"
           - zamba.models.utils: "api-reference/models-utils.md"
-          - zamba.models.yolox_models: "api-reference/models-yolox_models.md"
+      - zamba.object_detection:
+          - zamba.object_detection.yolox: "api-reference/object_detection-yolox.md"
       - zamba.pytorch:
           - zamba.pytorch.dataloaders: "api-reference/pytorch-dataloaders.md"
           - zamba.pytorch.finetuning: "api-reference/pytorch-finetuning.md"

tests/conftest.py

@@ -13,13 +13,11 @@
 
 from zamba.data.video import VideoLoaderConfig
 from zamba.models.config import PredictConfig, TrainConfig
-from zamba.models.megadetector_lite_yolox import MegadetectorLiteYoloX
 from zamba.models.model_manager import MODEL_MAPPING, train_model
+from zamba.models.registry import register_model
+from zamba.object_detection.yolox.megadetector_lite_yolox import MegadetectorLiteYoloX
 from zamba.pytorch.transforms import zamba_image_model_transforms
-from zamba.pytorch_lightning.utils import (
-    register_model,
-    ZambaVideoClassificationLightningModule,
-)
+from zamba.pytorch_lightning.utils import ZambaVideoClassificationLightningModule
 
 
 ASSETS_DIR = Path(__file__).parent / "assets"

tests/test_filter_frames.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from zamba.models.megadetector_lite_yolox import MegadetectorLiteYoloXConfig
+from zamba.object_detection.yolox.megadetector_lite_yolox import MegadetectorLiteYoloXConfig
 
 n_frames = 100
 rng = np.random.RandomState(68891)

tests/test_megadetector_lite_yolox.py

@@ -3,8 +3,8 @@
 import pytest
 import torch
 
-from zamba.models.yolox_models import YoloXNano
-from zamba.models.megadetector_lite_yolox import (
+from zamba.object_detection import YoloXNano
+from zamba.object_detection.yolox.megadetector_lite_yolox import (
     MegadetectorLiteYoloX,
     MegadetectorLiteYoloXConfig,
 )

zamba/__init__.py

@@ -4,8 +4,6 @@
 
 from loguru import logger
 
-from zamba.models.efficientnet_models import TimeDistributedEfficientNet  # noqa: F401
-from zamba.models.slowfast_models import SlowFast  # noqa: F401
 from zamba.version import __version__
 
 __version__

zamba/data/video.py

@@ -19,7 +19,7 @@
 from pydantic import BaseModel, root_validator, validator
 
 from zamba.exceptions import ZambaFfmpegException
-from zamba.models.megadetector_lite_yolox import (
+from zamba.object_detection.yolox.megadetector_lite_yolox import (
     MegadetectorLiteYoloX,
     MegadetectorLiteYoloXConfig,
 )

zamba/models/__init__.py

@@ -0,0 +1,2 @@
+from zamba.models.efficientnet_models import TimeDistributedEfficientNet  # noqa: F401
+from zamba.models.slowfast_models import SlowFast  # noqa: F401

zamba/models/efficientnet_models.py

@@ -5,11 +5,9 @@
 import torch
 from torch import nn
 
+from zamba.models.registry import register_model
 from zamba.pytorch.layers import TimeDistributed
-from zamba.pytorch_lightning.utils import (
-    register_model,
-    ZambaVideoClassificationLightningModule,
-)
+from zamba.pytorch_lightning.utils import ZambaVideoClassificationLightningModule
 
 
 @register_model

zamba/models/model_manager.py

@@ -26,13 +26,10 @@
     PredictConfig,
     RegionEnum,
 )
+from zamba.models.registry import available_models
 from zamba.models.utils import download_weights
 from zamba.pytorch.finetuning import BackboneFinetuning
-from zamba.pytorch_lightning.utils import (
-    available_models,
-    ZambaDataModule,
-    ZambaVideoClassificationLightningModule,
-)
+from zamba.pytorch_lightning.utils import ZambaDataModule, ZambaVideoClassificationLightningModule
 
 
 def instantiate_model(

zamba/models/registry.py

@@ -0,0 +1,15 @@
+from zamba.pytorch_lightning.utils import ZambaVideoClassificationLightningModule
+
+available_models = {}
+
+
+def register_model(cls):
+    """Used to decorate subclasses of ZambaVideoClassificationLightningModule so that they are
+    included in available_models."""
+    if not issubclass(cls, ZambaVideoClassificationLightningModule):
+        raise TypeError(
+            "Cannot register object that is not a subclass of "
+            "ZambaVideoClassificationLightningModule."
+        )
+    available_models[cls.__name__] = cls
+    return cls

zamba/models/slowfast_models.py

@@ -3,11 +3,9 @@
 from pytorchvideo.models.head import ResNetBasicHead
 import torch
 
+from zamba.models.registry import register_model
 from zamba.pytorch.utils import build_multilayer_perceptron
-from zamba.pytorch_lightning.utils import (
-    register_model,
-    ZambaVideoClassificationLightningModule,
-)
+from zamba.pytorch_lightning.utils import ZambaVideoClassificationLightningModule
 
 
 @register_model

zamba/object_detection/__init__.py

@@ -0,0 +1,4 @@
+from zamba.object_detection.yolox.yolox_base import YoloXBase
+from zamba.object_detection.yolox.yolox_nano import YoloXNano
+
+__all__ = ["YoloXBase", "YoloXNano"]

zamba/models/megadetector_lite_yolox.py → zamba/object_detection/yolox/megadetector_lite_yolox.py

@@ -11,9 +11,9 @@
 from tqdm import tqdm
 from yolox.utils.boxes import postprocess
 
-from zamba.models.yolox_models import YoloXNano
+from zamba.object_detection import YoloXNano
 
-LOCAL_MD_LITE_MODEL = Path(__file__).parent / "yolox_models" / "assets" / "yolox_nano_20210901.pth"
+LOCAL_MD_LITE_MODEL = Path(__file__).parent / "assets" / "yolox_nano_20210901.pth"
 
 
 class FillModeEnum(str, Enum):

zamba/models/yolox_models/yolox_nano.py → zamba/object_detection/yolox/yolox_nano.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import torch
-from zamba.models.yolox_models.yolox_base import YoloXBase
+from zamba.object_detection.yolox.yolox_base import YoloXBase
 
 
 class YoloXNano(YoloXBase):

zamba/pytorch_lightning/utils.py

@@ -19,21 +19,6 @@
 from zamba.pytorch.transforms import ConvertTHWCtoCTHW
 
 
-available_models = {}
-
-
-def register_model(cls):
-    """Used to decorate subclasses of ZambaVideoClassificationLightningModule so that they are
-    included in available_models."""
-    if not issubclass(cls, ZambaVideoClassificationLightningModule):
-        raise TypeError(
-            "Cannot register object that is not a subclass of "
-            "ZambaVideoClassificationLightningModule."
-        )
-    available_models[cls.__name__] = cls
-    return cls
-
-
 default_transform = transforms.Compose(
     [
         ConvertTHWCtoCTHW(),