fix models main

mindyolo/models/heads/yolov5_head.py

@@ -1,3 +1,4 @@
+import math
 import numpy as np
 
 import mindspore as ms
@@ -48,6 +49,16 @@ def construct(self, x):
         # return outs
         return outs if self.training else (ops.concat(z, 1), outs)
 
+    def initialize_biases(self, cf=None):  # initialize biases into Detect(), cf is class frequency
+        # https://arxiv.org/abs/1708.02002 section 3.3
+        m = self
+        for mi, s in zip(m.m, m.stride):  # from
+            s = s.asnumpy()
+            b = mi.bias.view(m.na, -1).asnumpy()  # conv.bias(255) to (3,85)
+            b[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
+            b[:, 5:] += math.log(0.6 / (m.nc - 0.999999)) if cf is None else np.log(cf / cf.sum())  # cls
+            mi.bias = ops.assign(mi.bias, Tensor(b, ms.float32).view(-1))
+
     @staticmethod
     def _make_grid(nx=20, ny=20, dtype=ms.float32):
         # FIXME: Not supported on a specific model of machine

mindyolo/models/yolov3.py

@@ -3,9 +3,9 @@
 import mindspore as ms
 from mindspore import Tensor, nn
 
-from .heads.yolov3_head import YOLOv3Head
-from .model_factory import build_model_from_cfg
-from .registry import register_model
+from mindyolo.models.heads.yolov3_head import YOLOv3Head
+from mindyolo.models.model_factory import build_model_from_cfg
+from mindyolo.models.registry import register_model
 
 __all__ = ["YOLOv3", "yolov3"]
 
@@ -50,16 +50,17 @@ def yolov3(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv3:
 
 if __name__ == "__main__":
     from mindyolo.models.model_factory import create_model
-    from mindyolo.utils.config import parse_config
-
-    opt = parse_config()
-    model = create_model(
-        model_name="yolov3",
-        model_cfg=opt.net,
-        num_classes=opt.data.nc,
-        sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False,
+    from mindyolo.utils.config import load_config, Config
+
+    cfg, _, _ = load_config('../../configs/yolov3/yolov3.yaml')
+    cfg = Config(cfg)
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False,
     )
     x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32)
-    out = model(x)
+    out = network(x)
     out = out[0] if isinstance(out, (list, tuple)) else out
     print(f"Output shape is {[o.shape for o in out]}")

mindyolo/models/yolov4.py

@@ -1,7 +1,10 @@
-import mindspore.nn as nn
+import numpy as np
 
-from .model_factory import build_model_from_cfg
-from .registry import register_model
+import mindspore as ms
+from mindspore import Tensor, nn
+
+from mindyolo.models.model_factory import build_model_from_cfg
+from mindyolo.models.registry import register_model
 
 __all__ = ["YOLOv4", "yolov4"]
 
@@ -32,3 +35,21 @@ def yolov4(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv4:
     """Get yolov4 model."""
     model = YOLOv4(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
     return model
+
+
+if __name__ == "__main__":
+    from mindyolo.models.model_factory import create_model
+    from mindyolo.utils.config import load_config, Config
+
+    cfg, _, _ = load_config('../../configs/yolov4/yolov4.yaml')
+    cfg = Config(cfg)
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False,
+    )
+    x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32)
+    out = network(x)
+    out = out[0] if isinstance(out, (list, tuple)) else out
+    print(f"Output shape is {[o.shape for o in out]}")

mindyolo/models/yolov5.py

@@ -3,9 +3,9 @@
 import mindspore as ms
 from mindspore import Tensor, nn
 
-from .heads.yolov7_head import YOLOv7AuxHead, YOLOv7Head
-from .model_factory import build_model_from_cfg
-from .registry import register_model
+from mindyolo.models.heads.yolov5_head import YOLOv5Head
+from mindyolo.models.model_factory import build_model_from_cfg
+from mindyolo.models.registry import register_model
 
 __all__ = ["YOLOv5", "yolov5"]
 
@@ -37,10 +37,8 @@ def construct(self, x):
     def initialize_weights(self):
         # reset parameter for Detect Head
         m = self.model.model[-1]
-        if isinstance(m, YOLOv7Head):
+        if isinstance(m, YOLOv5Head):
             m.initialize_biases()
-        if isinstance(m, YOLOv7AuxHead):
-            m.initialize_aux_biases()
 
 
 @register_model
@@ -55,16 +53,17 @@ def yolov5(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv5:
 
 if __name__ == "__main__":
     from mindyolo.models.model_factory import create_model
-    from mindyolo.utils.config import parse_config
-
-    opt = parse_config()
-    model = create_model(
-        model_name="yolov5",
-        model_cfg=opt.net,
-        num_classes=opt.data.nc,
-        sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False,
+    from mindyolo.utils.config import load_config, Config
+
+    cfg, _, _ = load_config('../../configs/yolov5/yolov5s.yaml')
+    cfg = Config(cfg)
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False,
     )
     x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32)
-    out = model(x)
+    out = network(x)
     out = out[0] if isinstance(out, (list, tuple)) else out
     print(f"Output shape is {[o.shape for o in out]}")

mindyolo/models/yolov7.py

@@ -3,8 +3,8 @@
 import mindspore as ms
 from mindspore import Tensor, nn
 
-from .heads.yolov7_head import YOLOv7AuxHead, YOLOv7Head
-from .model_factory import build_model_from_cfg
+from mindyolo.models.heads.yolov7_head import YOLOv7AuxHead, YOLOv7Head
+from mindyolo.models.model_factory import build_model_from_cfg
 from .registry import register_model
 
 __all__ = ["YOLOv7", "yolov7"]
@@ -55,16 +55,17 @@ def yolov7(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv7:
 
 if __name__ == "__main__":
     from mindyolo.models.model_factory import create_model
-    from mindyolo.utils.config import parse_config
-
-    opt = parse_config()
-    model = create_model(
-        model_name="yolov7",
-        model_cfg=opt.net,
-        num_classes=opt.data.nc,
-        sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False,
+    from mindyolo.utils.config import load_config, Config
+
+    cfg, _, _ = load_config('../../configs/yolov7/yolov7-tiny.yaml')
+    cfg = Config(cfg)
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False,
     )
     x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32)
-    out = model(x)
+    out = network(x)
     out = out[0] if isinstance(out, (list, tuple)) else out
     print(f"Output shape is {[o.shape for o in out]}")

mindyolo/models/yolov8.py

@@ -3,9 +3,9 @@
 import mindspore as ms
 from mindspore import Tensor, nn
 
-from .heads.yolov8_head import YOLOv8Head
-from .model_factory import build_model_from_cfg
-from .registry import register_model
+from mindyolo.models.heads.yolov8_head import YOLOv8Head
+from mindyolo.models.model_factory import build_model_from_cfg
+from mindyolo.models.registry import register_model
 
 __all__ = ["YOLOv8", "yolov8"]
 
@@ -54,16 +54,17 @@ def yolov8(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv8:
 
 if __name__ == "__main__":
     from mindyolo.models.model_factory import create_model
-    from mindyolo.utils.config import parse_config
-
-    opt = parse_config()
-    model = create_model(
-        model_name="yolov8",
-        model_cfg=opt.net,
-        num_classes=opt.data.nc,
-        sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False,
+    from mindyolo.utils.config import load_config, Config
+
+    cfg, _, _ = load_config('../../configs/yolov8/yolov8s.yaml')
+    cfg = Config(cfg)
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False,
     )
     x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32)
-    out = model(x)
+    out = network(x)
     out = out[0] if isinstance(out, (list, tuple)) else out
     print(f"Output shape is {[o.shape for o in out]}")

mindyolo/models/yolox.py

@@ -4,8 +4,8 @@
 from mindspore import Tensor, nn
 
 from mindyolo.models.registry import register_model
-from .heads import YOLOXHead
-from .model_factory import build_model_from_cfg
+from mindyolo.models.heads import YOLOXHead
+from mindyolo.models.model_factory import build_model_from_cfg
 
 __all__ = ["YOLOX", "yolox"]
 
@@ -45,3 +45,21 @@ def yolox(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOX:
     """Get yolox model."""
     model = YOLOX(cfg, in_channels=in_channels, num_classes=num_classes, **kwargs)
     return model
+
+
+if __name__ == "__main__":
+    from mindyolo.models.model_factory import create_model
+    from mindyolo.utils.config import load_config, Config
+
+    cfg, _, _ = load_config('../../configs/yolox/yolox-s.yaml')
+    cfg = Config(cfg)
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False,
+    )
+    x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32)
+    out = network(x)
+    out = out[0] if isinstance(out, (list, tuple)) else out
+    print(f"Output shape is {[o.shape for o in out]}")