Update main.py

main.py

@@ -1,283 +1,285 @@
-import io
-import shutil
-
-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.responses import JSONResponse, FileResponse
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
-from PIL import Image
-import cv2
-from src.detection import YOLOv11Detector
-from src.comparison import DamageComparator
-from src.visualization import DamageVisualizer
-
-
-app = FastAPI(
-    title="Car Damage Detection API",
-    description="YOLOv11-based car damage detection and comparison system",
-    version="1.0.0"
-)
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Initialize components
-detector = YOLOv11Detector()
-comparator = DamageComparator()
-visualizer = DamageVisualizer()
-
-# Create necessary directories
-Path("uploads").mkdir(exist_ok=True)
-Path("results").mkdir(exist_ok=True)
-
-# Mount static files directory
-app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
-
-@app.get("/")
-async def root():
-    """Root endpoint"""
-    return {
-        "message": "Car Damage Detection API with YOLOv11",
-        "endpoints": {
-            "/docs": "API documentation",
-            "/detect": "Single image detection",
-            "/compare": "Compare before/after images (6 pairs)",
-            "/uploads/{filename}": "Access saved visualization images",
-            "/health": "Health check"
-        }
-    }
-
-
-def save_temp_file(upload_file: UploadFile) -> str:
-    """Save an uploaded file into /tmp and return the temp file path"""
-    tmp_dir = Path("/tmp")
-    tmp_dir.mkdir(exist_ok=True)
-
-    temp_path = tmp_dir / upload_file.filename
-
-    with open(temp_path, "wb") as buffer:
-        shutil.copyfileobj(upload_file.file, buffer)
-
-    return str(temp_path)
-@app.get("/health")
-async def health_check():
-    """Health check endpoint"""
-    return {"status": "healthy", "model": "YOLOv11"}
-
-@app.post("/detect")
-async def detect_single_image(file: UploadFile = File(...)):
-    """
-    Detect damages in a single image
-
-    Args:
-        file: Image file
-
-    Returns:
-        Detection results with bounding boxes and path to visualized image
-    """
-    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)
-
-        # Perform detection
-        detections = detector.detect(image_bgr)
-
-        # Create visualization
-        visualized = visualizer.draw_detections(image_bgr, detections, 'new_damage')
-
-
-        filename = f"detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
-        output_path = Path("/tmp") / 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"/tmp/{filename}",
-        })
-
-    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"),
-    # 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"),
-):
-    """
-    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
-    """
-    try:
-        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 = []
-
-        # 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()
-
-            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 image pairs for grid visualization
-            image_pairs.append((before_bgr, after_bgr))
-
-            # Detect damages
-            before_detections = detector.detect(before_bgr)
-            after_detections = detector.detect(after_bgr)
-
-            # Compare damages
-            comparison = comparator.analyze_damage_status(before_detections, after_detections)
-
-            # Update overall 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
-            vis_img = visualizer.create_comparison_visualization(
-                before_bgr, after_bgr,
-                before_detections, after_detections,
-                comparison
-            )
-
-            # Save visualization image with unique filename
-            vis_filename = f"comparison_{timestamp_str}_{session_id}_pos{i+1}.jpg"
-            vis_path = Path("/tmp") / 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
-            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'],
-                    "visualization_path": f"/tmp/{vis_filename}",
-                    "visualization_url": vis_url
-                }
-            })
-
-        # Determine overall case
-        overall_case = "CASE_3_SUCCESS"
-        overall_message = "Successful delivery - No damage detected"
-
-        if total_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_case = "CASE_1_EXISTING"
-            overall_message = "Error from the beginning, not during the delivery process -> Delivery completed"
-
-        # Create summary grid visualization
-        grid_results = [res[f"position_{i+1}"] for i, res in enumerate(position_results)]
-        grid_img = visualizer.create_summary_grid(grid_results, image_pairs)
-
-        # Save grid summary image
-        grid_filename = f"summary_grid_{timestamp_str}_{session_id}.jpg"
-        grid_path = Path("uploads") / 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,
-            "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
-                }
-            },
-            "position_results": position_results,
-            "summary_visualization_path": f"/uploads/{grid_filename}",
-            "summary_visualization_url": grid_url,
-            "recommendations": {
-                "action_required": total_new_damages > 0,
-                "suggested_action": "Investigate delivery process" if total_new_damages > 0 else "Proceed with delivery completion"
-            }
-        })
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Comparison failed: {str(e)}")
-
-
-
-if __name__ == "__main__":
-    uvicorn.run(
-        "main:app",
-        host="0.0.0.0",
-        port=8000,
-        reload=True,
-        log_level="info"
-    )
+import io
+import shutil
+
+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.responses import JSONResponse, FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from PIL import Image
+import cv2
+from src.detection import YOLOv11Detector
+from src.comparison import DamageComparator
+from src.visualization import DamageVisualizer
+
+
+app = FastAPI(
+    title="Car Damage Detection API",
+    description="YOLOv11-based car damage detection and comparison system",
+    version="1.0.0"
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize components
+detector = YOLOv11Detector()
+comparator = DamageComparator()
+visualizer = DamageVisualizer()
+
+# Create necessary directories
+Path("uploads").mkdir(exist_ok=True)
+Path("results").mkdir(exist_ok=True)
+
+# Mount static files directory
+app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
+
+@app.get("/")
+async def root():
+    """Root endpoint"""
+    return {
+        "message": "Car Damage Detection API with YOLOv11",
+        "endpoints": {
+            "/docs": "API documentation",
+            "/detect": "Single image detection",
+            "/compare": "Compare before/after images (6 pairs)",
+            "/uploads/{filename}": "Access saved visualization images",
+            "/health": "Health check"
+        }
+    }
+
+
+def save_temp_file(upload_file: UploadFile) -> str:
+    """Save an uploaded file into /tmp and return the temp file path"""
+    tmp_dir = Path("/tmp")
+    tmp_dir.mkdir(exist_ok=True)
+
+    temp_path = tmp_dir / upload_file.filename
+
+    with open(temp_path, "wb") as buffer:
+        shutil.copyfileobj(upload_file.file, buffer)
+    return str(temp_path)
+
+    
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {"status": "healthy", "model": "YOLOv11"}
+
+@app.post("/detect")
+async def detect_single_image(file: UploadFile = File(...)):
+    """
+    Detect damages in a single image
+
+    Args:
+        file: Image file
+
+    Returns:
+        Detection results with bounding boxes and path to visualized image
+    """
+    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)
+
+        # Perform detection
+        detections = detector.detect(image_bgr)
+
+        # Create visualization
+        visualized = visualizer.draw_detections(image_bgr, detections, 'new_damage')
+
+
+        filename = f"detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jpg"
+        output_path = Path("/tmp") / 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"/tmp/{filename}",
+        })
+
+    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"),
+    # 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"),
+):
+    """
+    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
+    """
+    try:
+        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 = []
+
+        # 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()
+
+            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 image pairs for grid visualization
+            image_pairs.append((before_bgr, after_bgr))
+
+            # Detect damages
+            before_detections = detector.detect(before_bgr)
+            after_detections = detector.detect(after_bgr)
+
+            # Compare damages
+            comparison = comparator.analyze_damage_status(before_detections, after_detections)
+
+            # Update overall 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
+            vis_img = visualizer.create_comparison_visualization(
+                before_bgr, after_bgr,
+                before_detections, after_detections,
+                comparison
+            )
+
+            # Save visualization image with unique filename
+            vis_filename = f"comparison_{timestamp_str}_{session_id}_pos{i+1}.jpg"
+            vis_path = Path("/tmp") / 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
+            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'],
+                    "visualization_path": f"/tmp/{vis_filename}",
+                    "visualization_url": vis_url
+                }
+            })
+
+        # Determine overall case
+        overall_case = "CASE_3_SUCCESS"
+        overall_message = "Successful delivery - No damage detected"
+
+        if total_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_case = "CASE_1_EXISTING"
+            overall_message = "Error from the beginning, not during the delivery process -> Delivery completed"
+
+        # Create summary grid visualization
+        grid_results = [res[f"position_{i+1}"] for i, res in enumerate(position_results)]
+        grid_img = visualizer.create_summary_grid(grid_results, image_pairs)
+
+        # Save grid summary image
+        grid_filename = f"summary_grid_{timestamp_str}_{session_id}.jpg"
+        grid_path = Path("uploads") / 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,
+            "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
+                }
+            },
+            "position_results": position_results,
+            "summary_visualization_path": f"/uploads/{grid_filename}",
+            "summary_visualization_url": grid_url,
+            "recommendations": {
+                "action_required": total_new_damages > 0,
+                "suggested_action": "Investigate delivery process" if total_new_damages > 0 else "Proceed with delivery completion"
+            }
+        })
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Comparison failed: {str(e)}")
+
+
+
+if __name__ == "__main__":
+    import os
+    uvicorn.run(
+        "main:app",
+        host="0.0.0.0",
+        port=int(os.environ.get("PORT", 7860)),  # Lấy port do HF set
+        reload=False,  # Production không cần reload
+        log_level="info"
+    )