yjh0410 2 жил өмнө
parent
commit
884bc45f2c
1 өөрчлөгдсөн 2 нэмэгдсэн , 2 устгасан
  1. 2 2
      config/__init__.py

+ 2 - 2
config/__init__.py

@@ -28,8 +28,8 @@ def build_model_config(args):
     elif args.model == 'yolov5':
         cfg = yolov5_cfg
     # YOLOv7
-    elif args.model == 'yolov7':
-        cfg = yolov7_cfg
+    elif args.model in ['yolov7_nano' 'yolov7_tiny', 'yolov7_large', 'yolov7_huge']:
+        cfg = yolov7_cfg[args.model]
     # YOLOv8
     elif args.model == 'yolov8':
         cfg = yolov8_cfg