Semantic Nodes Categorization

.github/workflows/tests.yaml

@@ -14,7 +14,6 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest]
-        version: ['3.10', '3.11']
 
     runs-on: ${{ matrix.os }}
 
@@ -27,7 +26,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: ${{ matrix.version }}
+        python-version: '3.10'
         cache: pip
 
     - name: Install dependencies [Ubuntu]

luxonis_train/nodes/__init__.py

@@ -1,37 +1,4 @@
-from .base_node import BaseNode
-from .bisenet_head import BiSeNetHead
-from .classification_head import ClassificationHead
-from .contextspatial import ContextSpatial
-from .efficient_bbox_head import EfficientBBoxHead
-from .efficient_keypoint_bbox_head import EfficientKeypointBBoxHead
-from .efficientnet import EfficientNet
-from .efficientrep import EfficientRep
-from .implicit_keypoint_bbox_head import ImplicitKeypointBBoxHead
-from .micronet import MicroNet
-from .mobilenetv2 import MobileNetV2
-from .mobileone import MobileOne
-from .reppan_neck import RepPANNeck
-from .repvgg import RepVGG
-from .resnet import ResNet
-from .rexnetv1 import ReXNetV1_lite
-from .segmentation_head import SegmentationHead
-
-__all__ = [
-    "BiSeNetHead",
-    "ClassificationHead",
-    "ContextSpatial",
-    "EfficientBBoxHead",
-    "EfficientNet",
-    "EfficientRep",
-    "EfficientKeypointBBoxHead",
-    "ImplicitKeypointBBoxHead",
-    "BaseNode",
-    "MicroNet",
-    "MobileNetV2",
-    "MobileOne",
-    "ReXNetV1_lite",
-    "RepPANNeck",
-    "RepVGG",
-    "ResNet",
-    "SegmentationHead",
-]
+from .backbones import *
+from .base_node import *
+from .heads import *
+from .necks import *

luxonis_train/nodes/backbones/__init__.py

@@ -0,0 +1,21 @@
+from .contextspatial import ContextSpatial
+from .efficientnet import EfficientNet
+from .efficientrep import EfficientRep
+from .micronet import MicroNet
+from .mobilenetv2 import MobileNetV2
+from .mobileone import MobileOne
+from .repvgg import RepVGG
+from .resnet import ResNet
+from .rexnetv1 import ReXNetV1_lite
+
+__all__ = [
+    "ContextSpatial",
+    "EfficientNet",
+    "EfficientRep",
+    "MicroNet",
+    "MobileNetV2",
+    "MobileOne",
+    "ReXNetV1_lite",
+    "RepVGG",
+    "ResNet",
+]

luxonis_train/nodes/contextspatial.py → luxonis_train/nodes/backbones/contextspatial.py

@@ -7,15 +7,14 @@
 from torch import Tensor, nn
 from torch.nn import functional as F
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import (
     AttentionRefinmentBlock,
     ConvModule,
     FeatureFusionBlock,
 )
 from luxonis_train.utils.registry import NODES
 
-from .base_node import BaseNode
-
 
 class ContextSpatial(BaseNode[Tensor, list[Tensor]]):
     def __init__(self, context_backbone: str = "MobileNetV2", **kwargs):

luxonis_train/nodes/efficientnet.py → luxonis_train/nodes/backbones/efficientnet.py

@@ -7,7 +7,7 @@
 import torch
 from torch import Tensor, nn
 
-from .base_node import BaseNode
+from luxonis_train.nodes.base_node import BaseNode
 
 
 class EfficientNet(BaseNode[Tensor, list[Tensor]]):

luxonis_train/nodes/efficientrep.py → luxonis_train/nodes/backbones/efficientrep.py

@@ -9,15 +9,14 @@
 
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import (
     BlockRepeater,
     RepVGGBlock,
     SpatialPyramidPoolingBlock,
 )
 from luxonis_train.utils.general import make_divisible
 
-from .base_node import BaseNode
-
 logger = logging.getLogger(__name__)
 
 

luxonis_train/nodes/micronet.py → luxonis_train/nodes/backbones/micronet.py

@@ -4,10 +4,9 @@
 from torch import Tensor, nn
 
 from luxonis_train.nodes.activations import HSigmoid, HSwish
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
 
-from .base_node import BaseNode
-
 
 class MicroNet(BaseNode[Tensor, list[Tensor]]):
     """

luxonis_train/nodes/mobilenetv2.py → luxonis_train/nodes/backbones/mobilenetv2.py

@@ -6,7 +6,7 @@
 import torchvision
 from torch import Tensor, nn
 
-from .base_node import BaseNode
+from luxonis_train.nodes.base_node import BaseNode
 
 
 class MobileNetV2(BaseNode[Tensor, list[Tensor]]):

luxonis_train/nodes/mobileone.py → luxonis_train/nodes/backbones/mobileone.py

@@ -10,10 +10,9 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule, SqueezeExciteBlock
 
-from .base_node import BaseNode
-
 
 class MobileOne(BaseNode[Tensor, list[Tensor]]):
     """Implementation of MobileOne backbone.

luxonis_train/nodes/repvgg.py → luxonis_train/nodes/backbones/repvgg.py

@@ -6,7 +6,7 @@
 
 from luxonis_train.nodes.blocks import RepVGGBlock
 
-from .base_node import BaseNode
+from ..base_node import BaseNode
 
 logger = logging.getLogger(__name__)
 

luxonis_train/nodes/resnet.py → luxonis_train/nodes/backbones/resnet.py

@@ -8,7 +8,7 @@
 import torchvision
 from torch import Tensor, nn
 
-from .base_node import BaseNode
+from ..base_node import BaseNode
 
 
 class ResNet(BaseNode[Tensor, list[Tensor]]):

luxonis_train/nodes/rexnetv1.py → luxonis_train/nodes/backbones/rexnetv1.py

@@ -7,11 +7,10 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
 from luxonis_train.utils.general import make_divisible
 
-from .base_node import BaseNode
-
 
 class ReXNetV1_lite(BaseNode[Tensor, list[Tensor]]):
     def __init__(

luxonis_train/nodes/base_node.py

@@ -19,6 +19,8 @@
 ForwardOutputT = TypeVar("ForwardOutputT")
 ForwardInputT = TypeVar("ForwardInputT")
 
+__all__ = ["BaseNode"]
+
 
 class BaseNode(
     nn.Module,

luxonis_train/nodes/heads/__init__.py

@@ -0,0 +1,15 @@
+from .bisenet_head import BiSeNetHead
+from .classification_head import ClassificationHead
+from .efficient_bbox_head import EfficientBBoxHead
+from .efficient_keypoint_bbox_head import EfficientKeypointBBoxHead
+from .implicit_keypoint_bbox_head import ImplicitKeypointBBoxHead
+from .segmentation_head import SegmentationHead
+
+__all__ = [
+    "BiSeNetHead",
+    "ClassificationHead",
+    "EfficientBBoxHead",
+    "EfficientKeypointBBoxHead",
+    "ImplicitKeypointBBoxHead",
+    "SegmentationHead",
+]

luxonis_train/nodes/bisenet_head.py → luxonis_train/nodes/heads/bisenet_head.py

@@ -7,12 +7,11 @@
 
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
 from luxonis_train.utils.general import infer_upscale_factor
 from luxonis_train.utils.types import LabelType, Packet
 
-from .base_node import BaseNode
-
 
 class BiSeNetHead(BaseNode[Tensor, Tensor]):
     in_height: int

luxonis_train/nodes/classification_head.py → luxonis_train/nodes/heads/classification_head.py

@@ -1,9 +1,8 @@
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.utils.types import LabelType
 
-from .base_node import BaseNode
-
 
 class ClassificationHead(BaseNode[Tensor, Tensor]):
     in_channels: int

luxonis_train/nodes/efficient_bbox_head.py → luxonis_train/nodes/heads/efficient_bbox_head.py

@@ -9,6 +9,7 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import EfficientDecoupledBlock
 from luxonis_train.utils.boxutils import (
     anchors_for_fpn_features,
@@ -17,8 +18,6 @@
 )
 from luxonis_train.utils.types import LabelType, Packet
 
-from .base_node import BaseNode
-
 
 class EfficientBBoxHead(
     BaseNode[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]]

luxonis_train/nodes/implicit_keypoint_bbox_head.py → luxonis_train/nodes/heads/implicit_keypoint_bbox_head.py

@@ -5,6 +5,7 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import KeypointBlock, LearnableMulAddConv
 from luxonis_train.utils.boxutils import (
     non_max_suppression,
@@ -13,8 +14,6 @@
 )
 from luxonis_train.utils.types import LabelType, Packet
 
-from .base_node import BaseNode
-
 logger = logging.getLogger(__name__)
 
 

luxonis_train/nodes/segmentation_head.py → luxonis_train/nodes/heads/segmentation_head.py

@@ -7,12 +7,11 @@
 import torch.nn as nn
 from torch import Tensor
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import UpBlock
 from luxonis_train.utils.general import infer_upscale_factor
 from luxonis_train.utils.types import LabelType
 
-from .base_node import BaseNode
-
 
 class SegmentationHead(BaseNode[Tensor, Tensor]):
     in_height: int

luxonis_train/nodes/necks/__init__.py

@@ -0,0 +1,3 @@
+from .reppan_neck import RepPANNeck
+
+__all__ = ["RepPANNeck"]

luxonis_train/nodes/reppan_neck.py → luxonis_train/nodes/necks/reppan_neck.py

@@ -10,11 +10,10 @@
 
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import RepDownBlock, RepUpBlock
 from luxonis_train.utils.general import make_divisible
 
-from .base_node import BaseNode
-
 
 class RepPANNeck(BaseNode[list[Tensor], list[Tensor]]):
     def __init__(

tests/integration/test_sanity.py

@@ -35,9 +35,17 @@ def clear_output():
 
 
 @pytest.mark.parametrize(
-    "config_file", [str(path) for path in Path("configs").glob("*model*")]
+    "config_file",
+    [
+        "classification_model",
+        "segmentation_model",
+        "detection_model",
+        "keypoint_bbox_model",
+        "resnet_model",
+    ],
 )
 def test_simple_models(config_file: str):
+    config_file = f"configs/{config_file}.yaml"
     model = LuxonisModel(config_file, opts=OPTS)
     model.train()
     model.test()