yjh0410 hai 9 meses
pai
achega
d3fb1cb75f
Modificáronse 3 ficheiros con 8 adicións e 8 borrados
  1. 3 1
      yolo/config/__init__.py
  2. 3 5
      yolo/models/__init__.py
  3. 2 2
      yolo/models/yolov10/yolov10_pafpn.py

+ 3 - 1
yolo/config/__init__.py

@@ -21,7 +21,9 @@ def build_config(args):
     print('==============================')
     print('Model: {} ...'.format(args.model.upper()))
     # ----------- YOLO series -----------
-    if   'yolov1' in args.model:
+    if   'yolov10' in args.model:
+        cfg = build_yolov10_config(args)
+    elif 'yolov1' in args.model:
         cfg = build_yolov1_config(args)
     elif 'yolov2' in args.model:
         cfg = build_yolov2_config(args)

+ 3 - 5
yolo/models/__init__.py

@@ -24,7 +24,9 @@ from .rtdetr.build import build_rtdetr
 def build_model(args, cfg, is_val=False):
     # ------------ build object detector ------------
     ## Modified YOLOv1
-    if   'yolov1' in args.model:
+    if   'yolov10' in args.model:
+        model, criterion = build_yolov10(cfg, is_val)
+    elif 'yolov1' in args.model:
         model, criterion = build_yolov1(cfg, is_val)
     ## Modified YOLOv2
     elif 'yolov2' in args.model:
@@ -53,10 +55,6 @@ def build_model(args, cfg, is_val=False):
     ## GElan
     elif 'yolov9' in args.model:
         model, criterion = build_gelan(cfg, is_val)
-    ## YOLOv10
-    elif 'yolov10' in args.model:
-        model, criterion = build_yolov10(cfg, is_val)
-
     ## YOLO11
     elif 'yolo11' in args.model:
         model, criterion = build_yolo11(cfg, is_val)

+ 2 - 2
yolo/models/yolov10/yolov10_pafpn.py

@@ -25,7 +25,7 @@ class Yolov10PaFPN(nn.Module):
                                          expansion  = 0.5,
                                          num_blocks = round(3 * cfg.depth),
                                          shortcut   = False,
-                                         use_cib    = True if self.model_scale in "mlx" else False
+                                         use_cib    = True if self.model_scale in "lx" else False
                                          )
         ## P4 -> P3
         self.top_down_layer_2 = C2fBlock(in_dim     = self.in_dims[2] + round(512*cfg.width),
@@ -52,7 +52,7 @@ class Yolov10PaFPN(nn.Module):
                                           expansion  = 0.5,
                                           num_blocks = round(3 * cfg.depth),
                                           shortcut   = False,
-                                          use_cib    = True if self.model_scale in "mlx" else False
+                                          use_cib    = True
                                           )
 
         self.init_weights()