🔀 [Merge] branch 'MODEL' into TEST

yolo/config/model/v9-c-seg.yaml

@@ -0,0 +1,151 @@
+name: v9-c-seg
+
+anchor:
+  reg_max: 16
+  strides: [8, 16, 32]
+
+model:
+  backbone:
+    - Conv:
+        args: {out_channels: 64, kernel_size: 3, stride: 2}
+        source: 0
+    - Conv:
+        args: {out_channels: 128, kernel_size: 3, stride: 2}
+    - RepNCSPELAN:
+        args: {out_channels: 256, part_channels: 128}
+
+    - ADown:
+        args: {out_channels: 256}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 256}
+        tags: B3
+
+    - ADown:
+        args: {out_channels: 512}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: B4
+
+    - ADown:
+        args: {out_channels: 512}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: B5
+
+  neck:
+    - SPPELAN:
+        args: {out_channels: 512}
+        tags: N3
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B4]
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: N4
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B3]
+
+  head:
+    - RepNCSPELAN:
+        args: {out_channels: 256, part_channels: 256}
+        tags: P3
+
+    - ADown:
+        args: {out_channels: 256}
+    - Concat:
+        source: [-1, N4]
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: P4
+
+    - ADown:
+        args: {out_channels: 512}
+    - Concat:
+        source: [-1, N3]
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: P5
+
+  detection:
+    - RepNCSPELAN:
+        source: P3
+        args: {out_channels: 256, part_channels: 256, csp_args: {repeat_num: 2}}
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Conv:
+        args: {out_channels: 256, kernel_size: 3}
+
+    - MultiheadSegmentation:
+        source: [P3, P4, P5, -1]
+        args: {num_maskes: 32}
+        tags: Main
+        output: True
+
+  auxiliary:
+    - CBLinear:
+        source: B3
+        args: {out_channels: [256]}
+        tags: R3
+    - CBLinear:
+        source: B4
+        args: {out_channels: [256, 512]}
+        tags: R4
+    - CBLinear:
+        source: B5
+        args: {out_channels: [256, 512, 512]}
+        tags: R5
+
+    - Conv:
+        args: {out_channels: 64, kernel_size: 3, stride: 2}
+        source: 0
+    - Conv:
+        args: {out_channels: 128, kernel_size: 3, stride: 2}
+    - RepNCSPELAN:
+        args: {out_channels: 256, part_channels: 128}
+
+    - ADown:
+        args: {out_channels: 256}
+    - CBFuse:
+        source: [R3, R4, R5, -1]
+        args: {index: [0, 0, 0]}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 256}
+        tags: A3
+
+    - ADown:
+        args: {out_channels: 512}
+    - CBFuse:
+        source: [R4, R5, -1]
+        args: {index: [1, 1]}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: A4
+
+    - ADown:
+        args: {out_channels: 512}
+    - CBFuse:
+        source: [R5, -1]
+        args: {index: [2]}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: A5
+
+    - RepNCSPELAN:
+        source: A3
+        args: {out_channels: 512, part_channels: 256, csp_args: {repeat_num: 2}}
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Conv:
+        args: {out_channels: 256, kernel_size: 3}
+
+    - MultiheadSegmentation:
+        source: [A3, A4, A5, -1]
+        args: {num_maskes: 32}
+        tags: AUX
+        output: True

yolo/model/module.py

@@ -81,7 +81,7 @@ class Detection(nn.Module):
         self.anc2vec = Anchor2Vec(reg_max=reg_max)
 
         self.anchor_conv[-1].bias.data.fill_(1.0)
-        self.class_conv[-1].bias.data.fill_(-10)
+        self.class_conv[-1].bias.data.fill_(-10)  # TODO: math.log(5 * 4 ** idx / 80 ** 3)
 
     def forward(self, x: Tensor) -> Tuple[Tensor]:
         anchor_x = self.anchor_conv(x)
@@ -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

yolo/tools/format_converters.py

@@ -83,3 +83,55 @@ def convert_weight_v7(old_state_dict, new_state_dict):
         assert new_shape == old_shape, "Weight Shape Mismatch!! {old_key_name}"
         new_state_dict[new_key_name] = old_state_dict[old_key_name]
     return new_state_dict
+
+
+replace_dict = {"cv": "conv", ".m.": ".bottleneck."}
+
+
+def convert_weight_seg(old_state_dict, new_state_dict):
+    diff = -1
+    for old_weight_name in old_state_dict.keys():
+        old_idx = int(old_weight_name.split(".")[1])
+        if old_idx == 23:
+            diff = 3
+        elif old_idx == 41:
+            diff = -19
+        new_idx = old_idx + diff
+        new_weight_name = old_weight_name.replace(f".{old_idx}.", f".{new_idx}.")
+        for key, val in replace_dict.items():
+            new_weight_name = new_weight_name.replace(key, val)
+
+        if new_weight_name not in new_state_dict.keys():
+            heads = "heads"
+            _, _, conv_name, conv_idx, *details = old_weight_name.split(".")
+            if "proto" in conv_name:
+                conv_idx = "3"
+                new_weight_name = ".".join(["model", str(layer_idx), heads, conv_task, *details])
+                continue
+            if "dfl" in old_weight_name:
+                continue
+            if conv_name == "cv2" or conv_name == "cv3" or conv_name == "cv6":
+                layer_idx = 44
+                heads = "detect.heads"
+            if conv_name == "cv4" or conv_name == "cv5" or conv_name == "cv7":
+                layer_idx = 25
+                heads = "detect.heads"
+
+            if conv_name == "cv2" or conv_name == "cv4":
+                conv_task = "anchor_conv"
+            if conv_name == "cv3" or conv_name == "cv5":
+                conv_task = "class_conv"
+            if conv_name == "cv6" or conv_name == "cv7":
+                conv_task = "mask_conv"
+                heads = "heads"
+
+            new_weight_name = ".".join(["model", str(layer_idx), heads, conv_idx, conv_task, *details])
+
+        if (
+            new_weight_name not in new_state_dict.keys()
+            or new_state_dict[new_weight_name].shape != old_state_dict[old_weight_name].shape
+        ):
+            print(f"new: {new_weight_name}, old: {old_weight_name}")
+            print(f"{new_state_dict[new_weight_name].shape} {old_state_dict[old_weight_name].shape}")
+        new_state_dict[new_weight_name] = old_state_dict[old_weight_name]
+    return new_state_dict