from fastapi import APIRouter, Request, UploadFile, File, Form from fastapi.responses import HTMLResponse, FileResponse, JSONResponse from fastapi.templating import Jinja2Templates from starlette.background import BackgroundTask import shutil import os import uuid from pathlib import Path from typing import Optional import json import base64 from ultralytics import YOLO import cv2 import numpy as np # Templates directory TEMPLATES_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "templates") templates = Jinja2Templates(directory=TEMPLATES_DIR) router = APIRouter() UPLOAD_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static", "uploads") RESULTS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static", "results") os.makedirs(UPLOAD_DIR, exist_ok=True) os.makedirs(RESULTS_DIR, exist_ok=True) ALLOWED_EXTENSIONS = {"jpg", "jpeg", "png", "tiff", "tif"} # Model paths BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) DAMAGE_MODEL_PATH = os.path.join(BASE_DIR, "models", "damage", "weights", "weights", "best.pt") PARTS_MODEL_PATH = os.path.join(BASE_DIR, "models", "parts", "weights", "weights", "best.pt") # Class names for parts PARTS_CLASS_NAMES = ['headlamp', 'front_bumper', 'hood', 'door', 'rear_bumper'] # Helper: Run YOLO inference and return results def run_yolo_inference(model_path, image_path, task='segment'): model = YOLO(model_path) results = model.predict(source=image_path, imgsz=640, conf=0.25, save=False, task=task) return results[0] # Helper: Draw masks and confidence on image def draw_masks_and_conf(image_path, yolo_result, class_names=None): img = cv2.imread(image_path) overlay = img.copy() out_img = img.copy() colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255), (0,255,255)] for i, box in enumerate(yolo_result.boxes): conf = float(box.conf[0]) cls = int(box.cls[0]) color = colors[cls % len(colors)] # Draw bbox x1, y1, x2, y2 = map(int, box.xyxy[0]) cv2.rectangle(overlay, (x1, y1), (x2, y2), color, 2) label = f"{class_names[cls] if class_names else 'damage'}: {conf:.2f}" cv2.putText(overlay, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) # Draw mask if available if hasattr(yolo_result, 'masks') and yolo_result.masks is not None: mask = yolo_result.masks.data[i].cpu().numpy() mask = (mask * 255).astype(np.uint8) mask = cv2.resize(mask, (x2-x1, y2-y1)) roi = overlay[y1:y2, x1:x2] colored_mask = np.zeros_like(roi) colored_mask[mask > 127] = color overlay[y1:y2, x1:x2] = cv2.addWeighted(roi, 0.5, colored_mask, 0.5, 0) out_img = cv2.addWeighted(overlay, 0.7, img, 0.3, 0) return out_img # Helper: Generate JSON output def generate_json_output(filename, damage_result, parts_result): # Damage severity: use max confidence severity_score = float(max([float(box.conf[0]) for box in damage_result.boxes], default=0)) damage_regions = [] for box in damage_result.boxes: x1, y1, x2, y2 = map(float, box.xyxy[0]) conf = float(box.conf[0]) damage_regions.append({"bbox": [x1, y1, x2, y2], "confidence": conf}) # Parts parts = [] for i, box in enumerate(parts_result.boxes): x1, y1, x2, y2 = map(float, box.xyxy[0]) conf = float(box.conf[0]) cls = int(box.cls[0]) # Damage %: use mask area / bbox area if available damage_percentage = None if hasattr(parts_result, 'masks') and parts_result.masks is not None: mask = parts_result.masks.data[i].cpu().numpy() mask_area = np.sum(mask > 0.5) bbox_area = (x2-x1)*(y2-y1) damage_percentage = float(mask_area / bbox_area) if bbox_area > 0 else None parts.append({ "part": PARTS_CLASS_NAMES[cls] if cls < len(PARTS_CLASS_NAMES) else str(cls), "damaged": True, "confidence": conf, "damage_percentage": damage_percentage, "bbox": [x1, y1, x2, y2] }) # Optionally, add base64 masks # (not implemented here for brevity) return { "filename": filename, "damage": { "severity_score": severity_score, "regions": damage_regions }, "parts": parts, "cost_estimate": None } # Dummy login credentials def check_login(username: str, password: str) -> bool: return username == "demo" and password == "demo123" @router.get("/", response_class=HTMLResponse) def home(request: Request): return templates.TemplateResponse("index.html", {"request": request, "result": None}) @router.post("/login", response_class=HTMLResponse) def login(request: Request, username: str = Form(...), password: str = Form(...)): if check_login(username, password): return templates.TemplateResponse("index.html", {"request": request, "result": None, "user": username}) return templates.TemplateResponse("login.html", {"request": request, "error": "Invalid credentials"}) @router.get("/login", response_class=HTMLResponse) def login_page(request: Request): return templates.TemplateResponse("login.html", {"request": request}) @router.post("/upload", response_class=HTMLResponse) def upload_image(request: Request, file: UploadFile = File(...)): ext = file.filename.split(".")[-1].lower() if ext not in ALLOWED_EXTENSIONS: return templates.TemplateResponse("index.html", {"request": request, "error": "Unsupported file type."}) # Save uploaded file session_id = str(uuid.uuid4()) upload_path = os.path.join(UPLOAD_DIR, f"{session_id}.{ext}") with open(upload_path, "wb") as buffer: shutil.copyfileobj(file.file, buffer) # Run both inferences try: damage_result = run_yolo_inference(DAMAGE_MODEL_PATH, upload_path) parts_result = run_yolo_inference(PARTS_MODEL_PATH, upload_path) # Save annotated images damage_img_path = os.path.join(RESULTS_DIR, f"{session_id}_damage.png") parts_img_path = os.path.join(RESULTS_DIR, f"{session_id}_parts.png") json_path = os.path.join(RESULTS_DIR, f"{session_id}_result.json") damage_img_url = f"/static/results/{session_id}_damage.png" parts_img_url = f"/static/results/{session_id}_parts.png" json_url = f"/static/results/{session_id}_result.json" # Defensive: set to None by default damage_img = None parts_img = None json_output = None # Only save and set if inference returns boxes if hasattr(damage_result, 'boxes') and len(damage_result.boxes) > 0: damage_img = draw_masks_and_conf(upload_path, damage_result) cv2.imwrite(damage_img_path, damage_img) if hasattr(parts_result, 'boxes') and len(parts_result.boxes) > 0: parts_img = draw_masks_and_conf(upload_path, parts_result, class_names=PARTS_CLASS_NAMES) cv2.imwrite(parts_img_path, parts_img) if (hasattr(damage_result, 'boxes') and len(damage_result.boxes) > 0) or (hasattr(parts_result, 'boxes') and len(parts_result.boxes) > 0): json_output = generate_json_output(file.filename, damage_result, parts_result) with open(json_path, "w") as jf: json.dump(json_output, jf, indent=2) # Prepare URLs for download (only if files exist) result = { "filename": file.filename, "damage_image": damage_img_url if damage_img is not None else None, "parts_image": parts_img_url if parts_img is not None else None, "json": json_output, "json_download": json_url if json_output is not None else None } # Debug log print("[DEBUG] Result dict:", result) except Exception as e: result = { "filename": file.filename, "error": f"Inference failed: {str(e)}", "damage_image": None, "parts_image": None, "json": None, "json_download": None } print("[ERROR] Inference failed:", e) import threading import time def delayed_cleanup(): time.sleep(300) # 5 minutes try: os.remove(upload_path) except Exception: pass for suffix in ["_damage.png", "_parts.png", "_result.json"]: try: os.remove(os.path.join(RESULTS_DIR, f"{session_id}{suffix}")) except Exception: pass threading.Thread(target=delayed_cleanup, daemon=True).start() return templates.TemplateResponse( "index.html", { "request": request, "result": result, "original_image": f"/static/uploads/{session_id}.{ext}" } )