Update main.py

main.py

@@ -1,17 +1,18 @@
 import io
-import shutil
+from typing import List
 
 import uvicorn
 import numpy as np
 import uuid
 from datetime import datetime
 from pathlib import Path
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, HTTPException, Form
 from fastapi.responses import JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from PIL import Image
 import cv2
+import yaml
 from src.detection import YOLOv11Detector
 from src.comparison import DamageComparator
 from src.visualization import DamageVisualizer
@@ -32,10 +33,34 @@ app.add_middleware(
 )
 
 # Initialize components
-detector = YOLOv11Detector()
+detector = None
 comparator = DamageComparator()
 visualizer = DamageVisualizer()
 
+# Model paths mapping
+MODEL_PATHS = {
+    0: "models_small/best.pt",
+    1: "models_medium/best.pt",
+    2: "models_large/best.pt"
+}
+
+def load_detector(select_models: int = 2):
+    """Load detector with specified model"""
+    global detector
+    # Update config with selected model path
+    with open('config.yaml', 'r') as f:
+        config = yaml.safe_load(f)
+
+    config['model']['path'] = MODEL_PATHS.get(select_models, MODEL_PATHS[2])
+
+    with open('config.yaml', 'w') as f:
+        yaml.dump(config, f, default_flow_style=False)
+    detector = YOLOv11Detector()
+    # Reload detector with new config
+    return detector
+
+detector = load_detector(2)
+
 # Create necessary directories
 UPLOADS_DIR = Path("uploads")
 RESULTS_DIR = Path("results")
@@ -66,88 +91,200 @@ async def health_check():
     """Health check endpoint"""
     return {"status": "healthy", "model": "YOLOv11"}
 
-
 @app.post("/detect")
-async def detect_single_image(file: UploadFile = File(...)):
+async def detect_single_image(
+        file: UploadFile = File(None),
+        files: List[UploadFile] = File(None),
+        select_models: int = Form(2)
+):
     """
-    Detect damages in a single image
+    - Multi-view detection với deduplication
 
     Args:
-        file: Image file
-
-    Returns:
-        Detection results with bounding boxes and path to visualized image
+        file: Single image (backward compatibility)
+        select_models: Model selection (0=small, 1=medium, 2=large)
     """
     try:
-        # Read and process image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        image_np = np.array(image)
-        image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+        # Validate select_models
+        if select_models not in [0, 1, 2]:
+            raise HTTPException(status_code=400, detail="select_models must be 0, 1, or 2")
+
+        # Load appropriate detector
+        current_detector = load_detector(select_models)
+
+        # Case 1: Single image (backward compatibility)
+        if file is not None:
+            contents = await file.read()
+            image = Image.open(io.BytesIO(contents)).convert("RGB")
+            image_np = np.array(image)
+            image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+
+            # Perform detection
+            detections = current_detector.detect(image_bgr)
+
+            # Create visualization
+            visualized = visualizer.draw_detections(image_bgr, detections, 'new_damage')
+
+            # Save and return 
+            filename = f"detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
+            output_path = UPLOADS_DIR / filename
+            cv2.imwrite(str(output_path), visualized)
+
+            return JSONResponse({
+                "status": "success",
+                "detections": detections,
+                "statistics": {
+                    "total_damages": len(detections['boxes']),
+                    "damage_types": list(set(detections['classes']))
+                },
+                "visualized_image_path": f"uploads/{filename}",
+                "visualized_image_url": f"http://localhost:8000/uploads/{filename}",
+                "filename": filename
+            })
 
-        # Perform detection
-        detections = detector.detect(image_bgr)
+        # Case 2: Multiple images - MULTI-VIEW DETECTION với ReID
+        elif files is not None and len(files) > 0:
+            print(f"\nMulti-view detection with {len(files)} images")
 
-        # Create visualization
-        visualized = visualizer.draw_detections(image_bgr, detections, 'new_damage')
+            images_list = []
+            detections_list = []
 
-        # Generate unique filename
-        filename = f"detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
+            # 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)
 
-        output_path = UPLOADS_DIR / filename
+                images_list.append(image_bgr)
+                detections = current_detector.detect(image_bgr)
+                detections_list.append(detections)
 
-        # Save visualization image
-        cv2.imwrite(str(output_path), visualized)
+                print(f"   View {idx + 1}: {len(detections['boxes'])} detections")
 
-        # Generate URL for accessing the image
-        image_url = f"http://localhost:8000/uploads/{filename}"
+            # DEDUPLICATION using ReID
+            print("\nPerforming cross-view deduplication...")
+            unique_damages = comparator.deduplicate_detections_across_views(
+                detections_list, images_list
+            )
 
-        return JSONResponse({
-            "status": "success",
-            "detections": detections,
-            "statistics": {
-                "total_damages": len(detections['boxes']),
-                "damage_types": list(set(detections['classes']))
-            },
-            "visualized_image_path": f"uploads/{filename}",
-            "visualized_image_url": image_url,
-            "filename": filename
-        })
+            # 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
+
+            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
+                combined_img[:, x_offset:x_offset + w] = image
+
+                # Draw detections with unique IDs
+                for det_idx, bbox in enumerate(detections['boxes']):
+                    # Find unique damage ID for this detection
+                    damage_id = None
+                    for uid, damage_info in unique_damages.items():
+                        for d in damage_info['detections']:
+                            if d['view_idx'] == img_idx and d['bbox'] == bbox:
+                                damage_id = uid
+                                break
+
+                    # Draw with unique ID
+                    x1, y1, x2, y2 = bbox
+                    x1 += x_offset
+                    x2 += x_offset
+
+                    # Color based on unique ID
+                    if damage_id:
+                        color_hash = int(damage_id[-6:], 16)
+                        color = ((color_hash >> 16) & 255, (color_hash >> 8) & 255, color_hash & 255)
+                    else:
+                        color = (0, 0, 255)
+
+                    cv2.rectangle(combined_img, (x1, y1), (x2, y2), color, 2)
+
+                    # Label
+                    label = f"{damage_id[:8] if damage_id else 'Unknown'}"
+                    cv2.putText(combined_img, label, (x1, y1 - 5),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+                x_offset += w
+
+            # Save combined visualization
+            filename = f"multiview_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
+            output_path = UPLOADS_DIR / filename
+            cv2.imwrite(str(output_path), combined_img)
+
+            # Return results
+            total_detections = sum(len(d['boxes']) for d in detections_list)
+
+            return JSONResponse({
+                "status": "success",
+                "mode": "multi_view_with_reid",
+                "total_detections_all_views": total_detections,
+                "unique_damages_count": len(unique_damages),
+                "unique_damages": {
+                    damage_id: {
+                        "appears_in_views": info['views'],
+                        "class": info['class'],
+                        "avg_confidence": float(info['avg_confidence']),
+                        "detection_count": len(info['detections'])
+                    }
+                    for damage_id, info in unique_damages.items()
+                },
+                "reduction_rate": f"{(1 - len(unique_damages) / total_detections) * 100:.1f}%" if total_detections > 0 else "0%",
+                "visualized_image_path": f"uploads/{filename}",
+                "visualized_image_url": f"http://localhost:8000/uploads/{filename}",
+                "message": f"Detected {total_detections} damages across {len(files)} views, "
+                           f"identified {len(unique_damages)} unique damages using ReID"
+            })
+
+        else:
+            raise HTTPException(status_code=400, detail="No image provided")
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}")
 
 
+
+
 @app.post("/compare")
 async def compare_vehicle_damages(
         # Before delivery images (6 positions)
-        before_1: UploadFile = File(..., description="Before - Position 1"),
-        before_2: UploadFile = File(..., description="Before - Position 2"),
-        before_3: UploadFile = File(..., description="Before - Position 3"),
-        before_4: UploadFile = File(..., description="Before - Position 4"),
-        before_5: UploadFile = File(..., description="Before - Position 5"),
-        before_6: UploadFile = File(..., description="Before - Position 6"),
+        before_1: UploadFile = File(...),
+        before_2: UploadFile = File(...),
+        before_3: UploadFile = File(...),
+        before_4: UploadFile = File(...),
+        before_5: UploadFile = File(...),
+        before_6: UploadFile = File(...),
         # After delivery images (6 positions)
-        after_1: UploadFile = File(..., description="After - Position 1"),
-        after_2: UploadFile = File(..., description="After - Position 2"),
-        after_3: UploadFile = File(..., description="After - Position 3"),
-        after_4: UploadFile = File(..., description="After - Position 4"),
-        after_5: UploadFile = File(..., description="After - Position 5"),
-        after_6: UploadFile = File(..., description="After - Position 6"),
+        after_1: UploadFile = File(...),
+        after_2: UploadFile = File(...),
+        after_3: UploadFile = File(...),
+        after_4: UploadFile = File(...),
+        after_5: UploadFile = File(...),
+        after_6: UploadFile = File(...),
+        # Model selection
+        select_models: int = Form(2),
 ):
     """
-    Compare vehicle damages before and after delivery
-
-    Analyzes 6 pairs of images (before/after) from different positions
-    and determines the damage status according to 3 cases:
-    - Case 1: Existing damages (from before) -> Delivery completed
-    - Case 2: New damages detected -> Error during delivery
-    - Case 3: No damages -> Successful delivery
-
-    Returns:
-        Detailed comparison results for each position and overall status
+    Enhanced comparison với ReID
     """
     try:
+        # Validate select_models
+        if select_models not in [0, 1, 2]:
+            raise HTTPException(status_code=400, detail="select_models must be 0, 1, or 2")
+
+        # Load appropriate detector
+        current_detector = load_detector(select_models)
+
         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]
 
@@ -155,18 +292,22 @@ async def compare_vehicle_damages(
         all_visualizations = []
         image_pairs = []
 
+        # 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
         total_existing_damages = 0
         total_matched_damages = 0
 
-        # Generate unique session ID for this comparison
         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):
-            # Read images
             before_contents = await before_images[i].read()
             after_contents = await after_images[i].read()
 
@@ -179,22 +320,31 @@ async def compare_vehicle_damages(
             before_bgr = cv2.cvtColor(before_np, cv2.COLOR_RGB2BGR)
             after_bgr = cv2.cvtColor(after_np, cv2.COLOR_RGB2BGR)
 
-            # Store image pairs for grid visualization
+            # Store for multi-view analysis
+            all_before_images.append(before_bgr)
+            all_after_images.append(after_bgr)
+
             image_pairs.append((before_bgr, after_bgr))
 
             # Detect damages
-            before_detections = detector.detect(before_bgr)
-            after_detections = detector.detect(after_bgr)
+            before_detections = current_detector.detect(before_bgr)
+            after_detections = current_detector.detect(after_bgr)
 
-            # Compare damages
-            comparison = comparator.analyze_damage_status(before_detections, after_detections)
+            all_before_detections.append(before_detections)
+            all_after_detections.append(after_detections)
 
-            # Update overall statistics
+            # Enhanced comparison with ReID (pass images for feature extraction)
+            comparison = comparator.analyze_damage_status(
+                before_detections, after_detections,
+                before_bgr, after_bgr
+            )
+
+            # 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 for this position
+            # Create visualization
             vis_img = visualizer.create_comparison_visualization(
                 before_bgr, after_bgr,
                 before_detections, after_detections,
@@ -202,20 +352,13 @@ async def compare_vehicle_damages(
             )
 
             vis_filename = f"comparison_{timestamp_str}_{session_id}_pos{i + 1}.jpg"
-            vis_path = UPLOADS_DIR / vis_filename  # Use consistent relative path
-
-            # Save the image
-            success = cv2.imwrite(str(vis_path), vis_img)
-            if not success:
-                print(f"⚠️ Warning: Failed to save visualization for position {i + 1}")
-            else:
-                print(f" Saved: {vis_path}")
+            vis_path = UPLOADS_DIR / vis_filename
+            cv2.imwrite(str(vis_path), vis_img)
 
-            # Generate URL for accessing the image
             vis_url = f"http://localhost:8000/uploads/{vis_filename}"
             all_visualizations.append(vis_url)
 
-            # Store position result
+            # Store position result with ReID info
             position_results.append({
                 f"position_{i + 1}": {
                     "case": comparison['case'],
@@ -224,62 +367,89 @@ async def compare_vehicle_damages(
                     "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
                 }
             })
 
-        # Determine overall case
+
+        # Deduplicate BEFORE damages across all 6 views
+        unique_before = comparator.deduplicate_detections_across_views(
+            all_before_detections, all_before_images
+        )
+
+        # Deduplicate AFTER damages across all 6 views
+        unique_after = comparator.deduplicate_detections_across_views(
+            all_after_detections, all_after_images
+        )
+
+        print(
+            f"Before: {sum(len(d['boxes']) for d in all_before_detections)} detections → {len(unique_before)} unique")
+        print(f"After: {sum(len(d['boxes']) for d in all_after_detections)} detections → {len(unique_after)} unique")
+
+        # Determine overall case with deduplication
+        actual_new_damages = len(unique_after) - len(unique_before)
+
         overall_case = "CASE_3_SUCCESS"
         overall_message = "Successful delivery - No damage detected"
 
-        if total_new_damages > 0:
+        if actual_new_damages > 0:
             overall_case = "CASE_2_NEW_DAMAGE"
-            overall_message = f"Error during vehicle delivery - Detection {total_new_damages} new damage"
-        elif total_matched_damages > 0 and total_new_damages == 0:
+            overall_message = f"Error during delivery - {actual_new_damages} new unique damage(s) detected"
+        elif len(unique_before) > 0 and actual_new_damages <= 0:
             overall_case = "CASE_1_EXISTING"
-            overall_message = "Error from the beginning, not during the delivery process -> Delivery completed"
+            overall_message = "Existing damages from beginning -> Delivery completed"
 
-        # Create summary grid visualization
+        # Create summary grid
         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"
-        grid_path = UPLOADS_DIR / grid_filename  # Already correct in original
-
-        # Save the grid image
-        grid_success = cv2.imwrite(str(grid_path), grid_img)
-        if not grid_success:
-            print(f" Warning: Failed to save grid summary")
-        else:
-            print(f" Saved grid: {grid_path}")
+        grid_path = UPLOADS_DIR / grid_filename
+        cv2.imwrite(str(grid_path), grid_img)
 
         grid_url = f"http://localhost:8000/uploads/{grid_filename}"
 
-        # Generate timestamp for tracking
         timestamp = datetime.now().isoformat()
 
         return JSONResponse({
             "status": "success",
             "session_id": session_id,
             "timestamp": timestamp,
+            "reid_enabled": True,
             "overall_result": {
                 "case": overall_case,
                 "message": overall_message,
                 "statistics": {
-                    "total_new_damages": total_new_damages,
-                    "total_matched_damages": total_matched_damages,
-                    "total_repaired_damages": total_existing_damages
+                    "total_new_damages": int(total_new_damages),
+                    "total_matched_damages": int(total_matched_damages),
+                    "total_repaired_damages": int(total_existing_damages),
+                    "unique_damages_before": int(len(unique_before)),
+                    "unique_damages_after": int(len(unique_after)),
+                    "actual_new_unique_damages": int(max(0, len(unique_after) - len(unique_before)))
+
                 }
             },
+            "deduplication_info": {
+                "before_total_detections": int(sum(len(d['boxes']) for d in all_before_detections)),
+
+                "before_unique_damages": int(len(unique_before)),
+                "after_total_detections": int(sum(len(d['boxes']) for d in all_after_detections)),
+
+                "after_unique_damages": int(len(unique_after)),
+                "duplicate_reduction_rate": f"{(1 - len(unique_after) / sum(len(d['boxes']) for d in all_after_detections)) * 100:.1f}%"
+                if sum(len(d['boxes']) for d in all_after_detections) > 0 else "0%"
+            },
             "position_results": position_results,
             "summary_visualization_path": f"uploads/{grid_filename}",
             "summary_visualization_url": grid_url,
-            "all_visualizations": all_visualizations,  # Added for convenience
+            "all_visualizations": all_visualizations,
             "recommendations": {
-                "action_required": total_new_damages > 0,
-                "suggested_action": "Investigate delivery process" if total_new_damages > 0 else "Proceed with delivery completion"
+                "action_required": bool(actual_new_damages > 0),
+                "suggested_action": "Investigate delivery process" if actual_new_damages > 0
+                else "Proceed with delivery completion"
             }
         })
 
@@ -287,7 +457,6 @@ async def compare_vehicle_damages(
         raise HTTPException(status_code=500, detail=f"Comparison failed: {str(e)}")
 
 
-
 if __name__ == "__main__":
     import os
     uvicorn.run(