Update main.py

main.py

@@ -1,11 +1,11 @@
 import io
-from typing import List, Dict
+from typing import List
 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
+from fastapi.responses import JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from PIL import Image
@@ -15,7 +15,6 @@ 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",
@@ -44,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
@@ -59,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
 }
 
@@ -81,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(
@@ -90,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
@@ -113,12 +110,10 @@ 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
@@ -148,14 +143,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"
@@ -298,7 +293,6 @@ async def detect_single_image(
 ):
     """
     Multi-view detection with ONNX Runtime optimizations
-
     Args:
         file: Single image (backward compatibility)
         files: Multiple images for multi-view detection
@@ -309,7 +303,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)")
 
@@ -349,39 +343,52 @@ async def detect_single_image(
                 },
                 "visualized_image_path": f"uploads/{filename}",
                 "visualized_image_url": f"http://localhost:8000/uploads/{filename}",
-                "filename": filename
+                "filename": filename,
+                "performance_note": "Using ONNX optimizations" if model_type == "ONNX" else "Consider using ONNX models (5-8) for better performance"
             })
 
-        # Case 2: Multiple images
-        elif files is not None:
+        # 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 = []
             detections_list = []
-            images = []
-            unique_damages = {}
-            for idx, f in enumerate(files):
-                contents = await f.read()
+
+            # Process all images
+            for idx, img_file in enumerate(files):
+                contents = await img_file.read()
                 image = Image.open(io.BytesIO(contents)).convert("RGB")
                 image_np = np.array(image)
                 image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
-                images.append(image_bgr)
+
+                images_list.append(image_bgr)
                 detections = current_detector.detect(image_bgr)
                 detections_list.append(detections)
 
-            # Deduplicate across views
-            unique_damages = comparator.deduplicate_detections_across_views(detections_list, images)
+                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
+            )
 
             # Create combined visualization
-            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
+            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
 
-            for img_idx, image in enumerate(images):
+            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 != max_height:
-                    image = cv2.resize(image, (w, max_height))
-                detections = detections_list[img_idx]
+                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
                 combined_img[:, x_offset:x_offset + w] = image
 
                 # Draw detections with unique IDs
@@ -456,73 +463,6 @@ 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)
@@ -545,7 +485,6 @@ async def compare_vehicle_damages(
 ):
     """
     Enhanced comparison with ONNX Runtime optimizations and ReID
-
     Args:
         before_1-6: Before delivery images from 6 positions
         after_1-6: After delivery images from 6 positions
@@ -554,7 +493,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)")
 
@@ -564,13 +503,15 @@ 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]
 
-        # Đọ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]
+        position_results = []
+        all_visualizations = []
+        image_pairs = []
 
-        # 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")
+        # Collect all before/after images and detections
+        all_before_images = []
+        all_after_images = []
+        all_before_detections = []
+        all_after_detections = []
 
         # Overall statistics
         total_new_damages = 0
@@ -580,51 +521,73 @@ async def compare_vehicle_damages(
         session_id = str(uuid.uuid4())[:8]
         timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
 
-        position_results = []
-        all_visualizations = []
-        image_pairs = []
-        all_before_images = []
-        all_after_images = []
-        all_before_detections = []
-        all_after_detections = []
+        # Process each position pair
+        for i in range(6):
+            before_contents = await before_images[i].read()
+            after_contents = await after_images[i].read()
 
-        # 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)
-            ]
+            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)
 
-            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"])
+            image_pairs.append((before_bgr, after_bgr))
 
-                # Preprocess images for image_pairs (since before_bgr and after_bgr are returned)
-                image_pairs.append((result["before_bgr"], result["after_bgr"]))
+            # Detect damages
+            before_detections = current_detector.detect(before_bgr)
+            after_detections = current_detector.detect(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"])
+            all_before_detections.append(before_detections)
+            all_after_detections.append(after_detections)
 
-                # 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"])
+            # Enhanced comparison with ReID
+            comparison = comparator.analyze_damage_status(
+                before_detections, after_detections,
+                before_bgr, after_bgr
+            )
 
-        # Sort position_results by position number
-        position_results.sort(key=lambda x: int(list(x.keys())[0].split('_')[1]))
+            # Update statistics
+            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
+                }
+            })
 
         # Deduplicate BEFORE damages across all 6 views
         unique_before = comparator.deduplicate_detections_across_views(
@@ -653,7 +616,7 @@ async def compare_vehicle_damages(
             overall_message = "Existing damages from beginning → Delivery completed"
 
         # Create summary grid
-        grid_results = [res[list(res.keys())[0]] for res in position_results]
+        grid_results = [res[f"position_{i + 1}"] for i, res in enumerate(position_results)]
         grid_img = visualizer.create_summary_grid(grid_results, image_pairs)
 
         grid_filename = f"summary_grid_{timestamp_str}_{session_id}.jpg"
@@ -711,6 +674,7 @@ async def compare_vehicle_damages(
         raise HTTPException(status_code=500, detail=f"Comparison failed: {str(e)}")
 
 
+
 if __name__ == "__main__":
     import os
     uvicorn.run(