parse architecture yolo11

supervisely-ecosystem · Nov 28, 2024 · 38246ae · 38246ae
1 parent 1364d31
commit 38246ae
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/serve/src/yolov8.py b/serve/src/yolov8.py
@@ -691,7 +691,10 @@ def parse_model_name(checkpoint_name: str):
 
 def get_arch_from_model_name(model_name: str):
     # yolov8n-det
-    p = r"yolov(\d+)"
+    if "11" in model_name:
+        p = r"yolo(\d+)"
+    else:
+        p = r"yolov(\d+)"
     match = re.match(p, model_name.lower())
     if match:
         return f"YOLOv{match.group(1)}"