✨ [Add] v9-c-segment model, inference WIP

yolo/model/module.py

@@ -130,6 +130,38 @@ class MultiheadDetection(nn.Module):
         return [head(x) for x, head in zip(x_list, self.heads)]
 
 
+class Segmentation(nn.Module):
+    def __init__(self, in_channels: Tuple[int], num_maskes: int):
+        super().__init__()
+        first_neck, in_channels = in_channels
+
+        mask_neck = max(first_neck // 4, num_maskes)
+        self.mask_conv = nn.Sequential(
+            Conv(in_channels, mask_neck, 3), Conv(mask_neck, mask_neck, 3), nn.Conv2d(mask_neck, num_maskes, 1)
+        )
+
+    def forward(self, x: Tensor) -> Tuple[Tensor]:
+        x = self.mask_conv(x)
+        return x
+
+
+class MultiheadSegmentation(nn.Module):
+    """Mutlihead Segmentation module for Dual segment or Triple segment"""
+
+    def __init__(self, in_channels: List[int], num_classes: int, num_maskes: int, **head_kwargs):
+        super().__init__()
+        mask_channels, proto_channels = in_channels[:-1], in_channels[-1]
+
+        self.detect = MultiheadDetection(mask_channels, num_classes, **head_kwargs)
+        self.heads = nn.ModuleList(
+            [Segmentation((in_channels[0], in_channel), num_maskes) for in_channel in mask_channels]
+        )
+        self.heads.append(Conv(proto_channels, num_maskes, 1))
+
+    def forward(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
+        return [head(x) for x, head in zip(x_list, self.heads)]
+
+
 class Anchor2Vec(nn.Module):
     def __init__(self, reg_max: int = 16) -> None:
         super().__init__()

yolo/model/yolo.py

@@ -45,7 +45,7 @@ class YOLO(nn.Module):
                 # Find in channels
                 if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "AConv", "CBLinear"]):
                     layer_args["in_channels"] = output_dim[source]
-                if "Detection" in layer_type:
+                if "Detection" in layer_type or "Segmentation" in layer_type:
                     layer_args["in_channels"] = [output_dim[idx] for idx in source]
                     layer_args["num_classes"] = self.num_classes
                     layer_args["reg_max"] = self.reg_max