Update main.py

main.py

@@ -1,11 +1,11 @@
 import io
-from typing import List
+from typing import List, Dict
 import uvicorn
 import numpy as np
 import uuid
 from datetime import datetime
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
-from fastapi.responses import JSONResponse, FileResponse
+from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from PIL import Image
@@ -15,6 +15,7 @@ from src.detection import YOLOv11Detector
 from src.comparison import DamageComparator
 from src.visualization import DamageVisualizer
 from pathlib import Path
+from concurrent.futures import ProcessPoolExecutor, as_completed
 
 app = FastAPI(
     title="Car Damage Detection API",
@@ -43,12 +44,13 @@ 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)
-    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
+    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
 }
 
 # Config paths - ONNX uses same config as PT version
@@ -57,18 +59,19 @@ 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",  # Small v1 ONNX
-    5: "config_version2.yaml",  # Small v2 ONNX
-    6: "config.yaml",  # Medium v1 ONNX
-    7: "config_version2.yaml"  # Medium v2 ONNX
+    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
 }
 
 # Mapping from PT index to ONNX index
 PT_TO_ONNX_MAPPING = {
-    0: 4,  # Small v1 -> ONNX
-    1: 5,  # Small v2 -> ONNX
-    2: 6,  # Medium v1 -> ONNX
-    3: 7,  # Medium v2 -> ONNX
+    0: 5,  # Small v1 -> ONNX
+    1: 6,  # Small v2 -> ONNX
+    2: 7,  # Medium v1 -> ONNX
+    3: 8,  # Medium v2 -> ONNX
     4: None  # Large has no ONNX
 }
 
@@ -78,7 +81,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 (4, 5, 6, 7):
+    if select_models in (5, 6, 7, 8):
         onnx_path = Path(MODEL_PATHS.get(select_models, ""))
         if not onnx_path.exists():
             raise FileNotFoundError(
@@ -87,7 +90,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):
+    if select_models not in (0, 1, 2, 3, 4):
         select_models = 2  # default to medium v1
 
     # PT preferred for 0..4
@@ -110,11 +113,12 @@ 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
@@ -144,14 +148,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", 
+        "Small v1", "Small v2", "Medium v1", "Medium v2", "Large",
         "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-7")
+        raise ValueError(f"select_models={select_models} must be 0-8")
 
     # Enhanced logging for optimization status
     optimization_status = "🚀 MAXIMUM OPTIMIZATIONS" if model_type == "ONNX" else "📦 Standard PyTorch"
@@ -305,7 +309,7 @@ async def detect_single_image(
     """
     try:
         # Validate select_models
-        if select_models not in list(range(0, 8)):
+        if select_models not in list(range(0, 9)):
             raise HTTPException(status_code=400,
                                 detail="select_models must be 0-8 (0-4=PyTorch, 5-8=ONNX optimized)")
 
@@ -345,52 +349,39 @@ async def detect_single_image(
                 },
                 "visualized_image_path": f"uploads/{filename}",
                 "visualized_image_url": f"http://localhost:8000/uploads/{filename}",
-                "filename": filename,
-                "performance_note": "Using ONNX optimizations" if model_type == "ONNX" else "Consider using ONNX models (5-8) for better performance"
+                "filename": filename
             })
 
-        # Case 2: Multiple images - MULTI-VIEW DETECTION with ReID
-        elif files is not None and len(files) > 0:
-            print(f"\nMulti-view detection with {len(files)} images")
-
-            images_list = []
+        # Case 2: Multiple images
+        elif files is not None:
             detections_list = []
-
-            # Process all images
-            for idx, img_file in enumerate(files):
-                contents = await img_file.read()
+            images = []
+            unique_damages = {}
+            for idx, f in enumerate(files):
+                contents = await f.read()
                 image = Image.open(io.BytesIO(contents)).convert("RGB")
                 image_np = np.array(image)
                 image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
-
-                images_list.append(image_bgr)
+                images.append(image_bgr)
                 detections = current_detector.detect(image_bgr)
                 detections_list.append(detections)
 
-                print(f"   View {idx + 1}: {len(detections['boxes'])} detections")
-
-            # DEDUPLICATION using ReID
-            print("\nPerforming cross-view deduplication...")
-            unique_damages = comparator.deduplicate_detections_across_views(
-                detections_list, images_list
-            )
+            # Deduplicate across views
+            unique_damages = comparator.deduplicate_detections_across_views(detections_list, images)
 
             # Create combined visualization
-            combined_height = max(img.shape[0] for img in images_list)
-            combined_width = sum(img.shape[1] for img in images_list)
-            combined_img = np.ones((combined_height, combined_width, 3), dtype=np.uint8) * 255
-
+            heights = [img.shape[0] for img in images]
+            widths = [img.shape[1] for img in images]
+            max_height = max(heights)
+            total_width = sum(widths)
+            combined_img = np.zeros((max_height, total_width, 3), dtype=np.uint8)
             x_offset = 0
-            for img_idx, (image, detections) in enumerate(zip(images_list, detections_list)):
-                # Resize if needed
-                h, w = image.shape[:2]
-                if h != combined_height:
-                    scale = combined_height / h
-                    new_w = int(w * scale)
-                    image = cv2.resize(image, (new_w, combined_height))
-                    w = new_w
 
-                # Draw on combined image
+            for img_idx, image in enumerate(images):
+                h, w = image.shape[:2]
+                if h != max_height:
+                    image = cv2.resize(image, (w, max_height))
+                detections = detections_list[img_idx]
                 combined_img[:, x_offset:x_offset + w] = image
 
                 # Draw detections with unique IDs
@@ -465,6 +456,73 @@ async def detect_single_image(
         raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}")
 
 
+def process_single_position(
+        i: int,
+        before_contents: bytes,
+        after_contents: bytes,
+        config_path: str,
+        timestamp_str: str,
+        session_id: str
+) -> Dict:
+    # Load components trong process (giữ nguyên logic cũ)
+    from src.detection import YOLOv11Detector
+    from src.comparison import DamageComparator
+    from src.visualization import DamageVisualizer
+
+    detector = YOLOv11Detector(config_path=config_path)
+    comparator = DamageComparator(config_path=config_path)
+    visualizer = DamageVisualizer(config_path=config_path)
+
+    # Logic cũ: Preprocess images
+    before_img = Image.open(io.BytesIO(before_contents)).convert("RGB")
+    after_img = Image.open(io.BytesIO(after_contents)).convert("RGB")
+    before_np = np.array(before_img)
+    after_np = np.array(after_img)
+    before_bgr = cv2.cvtColor(before_np, cv2.COLOR_RGB2BGR)
+    after_bgr = cv2.cvtColor(after_np, cv2.COLOR_RGB2BGR)
+
+    # Detect (giữ nguyên)
+    before_detections = detector.detect(before_bgr)
+    after_detections = detector.detect(after_bgr)
+
+    # Compare (giữ nguyên)
+    comparison = comparator.analyze_damage_status(
+        before_detections, after_detections,
+        before_bgr, after_bgr
+    )
+
+    # Visualize và save (giữ nguyên)
+    vis_img = visualizer.create_comparison_visualization(
+        before_bgr, after_bgr,
+        before_detections, after_detections,
+        comparison
+    )
+    vis_filename = f"comparison_{timestamp_str}_{session_id}_pos{i + 1}.jpg"
+    vis_path = UPLOADS_DIR / vis_filename
+    cv2.imwrite(str(vis_path), vis_img)
+    vis_url = f"http://localhost:8000/uploads/{vis_filename}"
+
+    # Return result cho tổng hợp (giữ nguyên structure)
+    return {
+        f"position_{i + 1}": {
+            "case": comparison['case'],
+            "message": comparison['message'],
+            "statistics": comparison['statistics'],
+            "new_damages": comparison['new_damages'],
+            "matched_damages": comparison['matched_damages'],
+            "repaired_damages": comparison['repaired_damages'],
+            "using_reid": comparison['statistics'].get('using_reid', True),
+            "visualization_path": f"uploads/{vis_filename}",
+            "visualization_url": vis_url,
+            "filename": vis_filename
+        },
+        "before_bgr": before_bgr,
+        "after_bgr": after_bgr,
+        "before_detections": before_detections,
+        "after_detections": after_detections
+    }
+
+
 @app.post("/compare")
 async def compare_vehicle_damages(
         # Before delivery images (6 positions)
@@ -496,7 +554,7 @@ async def compare_vehicle_damages(
     """
     try:
         # Validate select_models
-        if select_models not in list(range(0, 8)):
+        if select_models not in list(range(0, 9)):
             raise HTTPException(status_code=400,
                                 detail="select_models must be 0-8 (0-4=PyTorch, 5-8=ONNX optimized)")
 
@@ -506,15 +564,13 @@ async def compare_vehicle_damages(
         before_images = [before_1, before_2, before_3, before_4, before_5, before_6]
         after_images = [after_1, after_2, after_3, after_4, after_5, after_6]
 
-        position_results = []
-        all_visualizations = []
-        image_pairs = []
+        # Đọc contents trước để pass vào processes (vì UploadFile không pickleable)
+        before_contents_list = [await img.read() for img in before_images]
+        after_contents_list = [await img.read() for img in after_images]
 
-        # Collect all before/after images and detections
-        all_before_images = []
-        all_after_images = []
-        all_before_detections = []
-        all_after_detections = []
+        # Config để load trong processes
+        actual_model_index = get_optimal_model_index(select_models, prefer_onnx)
+        config_path = CONFIG_PATHS.get(actual_model_index, "config.yaml")
 
         # Overall statistics
         total_new_damages = 0
@@ -524,73 +580,51 @@ async def compare_vehicle_damages(
         session_id = str(uuid.uuid4())[:8]
         timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
 
-        # Process each position pair
-        for i in range(6):
-            before_contents = await before_images[i].read()
-            after_contents = await after_images[i].read()
-
-            before_img = Image.open(io.BytesIO(before_contents)).convert("RGB")
-            after_img = Image.open(io.BytesIO(after_contents)).convert("RGB")
-
-            before_np = np.array(before_img)
-            after_np = np.array(after_img)
-
-            before_bgr = cv2.cvtColor(before_np, cv2.COLOR_RGB2BGR)
-            after_bgr = cv2.cvtColor(after_np, cv2.COLOR_RGB2BGR)
-
-            # Store for multi-view analysis
-            all_before_images.append(before_bgr)
-            all_after_images.append(after_bgr)
+        position_results = []
+        all_visualizations = []
+        image_pairs = []
+        all_before_images = []
+        all_after_images = []
+        all_before_detections = []
+        all_after_detections = []
 
-            image_pairs.append((before_bgr, after_bgr))
+        # Parallelize
+        with ProcessPoolExecutor(max_workers=6) as executor:
+            futures = [
+                executor.submit(
+                    process_single_position,
+                    i,
+                    before_contents_list[i],
+                    after_contents_list[i],
+                    config_path,
+                    timestamp_str,
+                    session_id
+                )
+                for i in range(6)
+            ]
 
-            # Detect damages
-            before_detections = current_detector.detect(before_bgr)
-            after_detections = current_detector.detect(after_bgr)
+            for future in as_completed(futures):
+                result = future.result()
+                pos_key = list(result.keys())[0]  # e.g., 'position_1'
+                position_results.append(result)
+                all_visualizations.append(result[pos_key]["visualization_url"])
 
-            all_before_detections.append(before_detections)
-            all_after_detections.append(after_detections)
+                # Preprocess images for image_pairs (since before_bgr and after_bgr are returned)
+                image_pairs.append((result["before_bgr"], result["after_bgr"]))
 
-            # Enhanced comparison with ReID
-            comparison = comparator.analyze_damage_status(
-                before_detections, after_detections,
-                before_bgr, after_bgr
-            )
+                all_before_images.append(result["before_bgr"])
+                all_after_images.append(result["after_bgr"])
+                all_before_detections.append(result["before_detections"])
+                all_after_detections.append(result["after_detections"])
 
-            # Update statistics
-            total_new_damages += len(comparison['new_damages'])
-            total_existing_damages += len(comparison['repaired_damages'])
-            total_matched_damages += len(comparison['matched_damages'])
+                # Update statistics from comparison
+                comparison = result[pos_key]
+                total_new_damages += len(comparison["new_damages"])
+                total_existing_damages += len(comparison["repaired_damages"])
+                total_matched_damages += len(comparison["matched_damages"])
 
-            # Create visualization
-            vis_img = visualizer.create_comparison_visualization(
-                before_bgr, after_bgr,
-                before_detections, after_detections,
-                comparison
-            )
-
-            vis_filename = f"comparison_{timestamp_str}_{session_id}_pos{i + 1}.jpg"
-            vis_path = UPLOADS_DIR / vis_filename
-            cv2.imwrite(str(vis_path), vis_img)
-
-            vis_url = f"http://localhost:8000/uploads/{vis_filename}"
-            all_visualizations.append(vis_url)
-
-            # Store position result with ReID info
-            position_results.append({
-                f"position_{i + 1}": {
-                    "case": comparison['case'],
-                    "message": comparison['message'],
-                    "statistics": comparison['statistics'],
-                    "new_damages": comparison['new_damages'],
-                    "matched_damages": comparison['matched_damages'],
-                    "repaired_damages": comparison['repaired_damages'],
-                    "using_reid": comparison['statistics'].get('using_reid', True),
-                    "visualization_path": f"uploads/{vis_filename}",
-                    "visualization_url": vis_url,
-                    "filename": vis_filename
-                }
-            })
+        # Sort position_results by position number
+        position_results.sort(key=lambda x: int(list(x.keys())[0].split('_')[1]))
 
         # Deduplicate BEFORE damages across all 6 views
         unique_before = comparator.deduplicate_detections_across_views(
@@ -619,7 +653,7 @@ async def compare_vehicle_damages(
             overall_message = "Existing damages from beginning → Delivery completed"
 
         # Create summary grid
-        grid_results = [res[f"position_{i + 1}"] for i, res in enumerate(position_results)]
+        grid_results = [res[list(res.keys())[0]] for res in position_results]
         grid_img = visualizer.create_summary_grid(grid_results, image_pairs)
 
         grid_filename = f"summary_grid_{timestamp_str}_{session_id}.jpg"
@@ -677,7 +711,6 @@ async def compare_vehicle_damages(
         raise HTTPException(status_code=500, detail=f"Comparison failed: {str(e)}")
 
 
-
 if __name__ == "__main__":
     import os
     uvicorn.run(