Update main.py

main.py

@@ -43,13 +43,12 @@ MODEL_PATHS = {
     1: "models_small_version_2/best.pt",  # Small v2 PT
     2: "models_medium/best.pt",  # Medium v1 PT
     3: "models_medium_version_2/best.pt",  # Medium v2 PT
-    4: "models_large/best.pt",  # Large PT (no ONNX for large)
 
     # ONNX models (optimized with v1.19 + opset 21)
-    5: "models_small/best.onnx",  # Small v1 ONNX
-    6: "models_small_version_2/best.onnx",  # Small v2 ONNX
-    7: "models_medium/best.onnx",  # Medium v1 ONNX
-    8: "models_medium_version_2/best.onnx"  # Medium v2 ONNX
+    4: "models_small/best.onnx",  # Small v1 ONNX
+    5: "models_small_version_2/best.onnx",  # Small v2 ONNX
+    6: "models_medium/best.onnx",  # Medium v1 ONNX
+    7: "models_medium_version_2/best.onnx"  # Medium v2 ONNX
 }
 
 # Config paths - ONNX uses same config as PT version
@@ -58,19 +57,18 @@ CONFIG_PATHS = {
     1: "config_version2.yaml",  # Small v2 PT
     2: "config.yaml",  # Medium v1 PT
     3: "config_version2.yaml",  # Medium v2 PT
-    4: "config.yaml",  # Large PT
-    5: "config.yaml",  # Small v1 ONNX
-    6: "config_version2.yaml",  # Small v2 ONNX
-    7: "config.yaml",  # Medium v1 ONNX
-    8: "config_version2.yaml"  # Medium v2 ONNX
+    4: "config.yaml",  # Small v1 ONNX
+    5: "config_version2.yaml",  # Small v2 ONNX
+    6: "config.yaml",  # Medium v1 ONNX
+    7: "config_version2.yaml"  # Medium v2 ONNX
 }
 
 # Mapping from PT index to ONNX index
 PT_TO_ONNX_MAPPING = {
-    0: 5,  # Small v1 -> ONNX
-    1: 6,  # Small v2 -> ONNX
-    2: 7,  # Medium v1 -> ONNX
-    3: 8,  # Medium v2 -> ONNX
+    0: 4,  # Small v1 -> ONNX
+    1: 5,  # Small v2 -> ONNX
+    2: 6,  # Medium v1 -> ONNX
+    3: 7,  # Medium v2 -> ONNX
     4: None  # Large has no ONNX
 }
 
@@ -80,7 +78,7 @@ def get_optimal_model_index(select_models: int, prefer_onnx: bool = True) -> int
     Enhanced model selection with performance optimization info
     """
     # If user explicitly selects ONNX index (5..8) => use that ONNX with optimizations
-    if select_models in (5, 6, 7, 8):
+    if select_models in (4, 5, 6, 7):
         onnx_path = Path(MODEL_PATHS.get(select_models, ""))
         if not onnx_path.exists():
             raise FileNotFoundError(
@@ -89,7 +87,7 @@ def get_optimal_model_index(select_models: int, prefer_onnx: bool = True) -> int
         return select_models
 
     # Normalize to valid PT indices
-    if select_models not in (0, 1, 2, 3, 4):
+    if select_models not in (0, 1, 2, 3):
         select_models = 2  # default to medium v1
 
     # PT preferred for 0..4
@@ -112,12 +110,11 @@ def get_optimal_model_index(select_models: int, prefer_onnx: bool = True) -> int
 
 def load_detector(select_models: int = 2, prefer_onnx: bool = True):
     """
-    Load detector with optimized ONNX Runtime v1.19 support
+   
 
     Args:
         select_models: Model selection
-            - 0-4: PyTorch models (original logic)
-            - 5-8: ONNX models (with maximum optimizations)
+        
         prefer_onnx: Whether to prefer ONNX format for fallback
     """
     global detector, comparator, visualizer
@@ -147,14 +144,14 @@ def load_detector(select_models: int = 2, prefer_onnx: bool = True):
     # Log model info with optimization status
     model_type = "ONNX" if MODEL_PATHS[actual_model_index].endswith('.onnx') else "PyTorch"
     model_labels = [
-        "Small v1", "Small v2", "Medium v1", "Medium v2", "Large",
+        "Small v1", "Small v2", "Medium v1", "Medium v2", 
         "Small v1 ONNX", "Small v2 ONNX", "Medium v1 ONNX", "Medium v2 ONNX"
     ]
 
     if 0 <= select_models < len(model_labels):
         model_size = model_labels[select_models]
     else:
-        raise ValueError(f"select_models={select_models} must be 0-8")
+        raise ValueError(f"select_models={select_models} must be 0-7")
 
     # Enhanced logging for optimization status
     optimization_status = "🚀 MAXIMUM OPTIMIZATIONS" if model_type == "ONNX" else "📦 Standard PyTorch"
@@ -308,7 +305,7 @@ async def detect_single_image(
     """
     try:
         # Validate select_models
-        if select_models not in list(range(0, 9)):
+        if select_models not in list(range(0, 8)):
             raise HTTPException(status_code=400,
                                 detail="select_models must be 0-8 (0-4=PyTorch, 5-8=ONNX optimized)")
 
@@ -499,7 +496,7 @@ async def compare_vehicle_damages(
     """
     try:
         # Validate select_models
-        if select_models not in list(range(0, 9)):
+        if select_models not in list(range(0, 8)):
             raise HTTPException(status_code=400,
                                 detail="select_models must be 0-8 (0-4=PyTorch, 5-8=ONNX optimized)")