yjh0410 преди 2 години
родител
ревизия
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