fabrix13/aivala
0
1from __future__ import annotations2 3import os4import base645import io6import json7import logging8import time9import uuid10from pathlib import Path11from typing import Any12 13from pydantic import BaseModel14 15logging.basicConfig(level=logging.INFO)16logger = logging.getLogger("aivala-inference")17 18try:19 from fastapi import FastAPI, File, UploadFile, Request20 from fastapi.middleware.cors import CORSMiddleware21 from fastapi.responses import JSONResponse22 from fastapi.staticfiles import StaticFiles23except Exception as exc: # pragma: no cover24 raise RuntimeError(25 "FastAPI is required to run this server. Install fastapi and uvicorn."26 ) from exc27 28 29ROOT = Path(__file__).resolve().parent30MODEL_PATH = ROOT / "best.pt"31ULTRALYTICS_DIR = ROOT / ".ultralytics"32ULTRALYTICS_DIR.mkdir(exist_ok=True)33os.environ.setdefault("YOLO_CONFIG_DIR", str(ULTRALYTICS_DIR))34os.environ.setdefault("ULTRALYTICS_SETTINGS_DIR", str(ULTRALYTICS_DIR))35 36STATIC_DIR = ROOT / "static"37STATIC_DIR.mkdir(exist_ok=True)38COLAB_SEVERITY_URL = os.getenv("COLAB_SEVERITY_URL", "").strip()39print("COLAB_SEVERITY_URL =", COLAB_SEVERITY_URL)40SEVERITY_LEVELS = {"none", "minor", "moderate", "severe"}41 42app = FastAPI(title="YOLO Damage Inference API")43app.add_middleware(44 CORSMiddleware,45 allow_origins=["*"],46 allow_credentials=True,47 allow_methods=["*"],48 allow_headers=["*"],49)50 51# Mount the static directory to serve static assets52app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")53 54MODEL = None55 56 57class Detection(BaseModel):58 class_id: int59 label: str60 confidence: float61 bbox: list[float]62 severity: str | None = None63 severity_note: str | None = None64 severity_source: str | None = None65 66 67class FaceVerifyRequest(BaseModel):68 frame: str69 70 71def _normalize_severity(value: Any) -> str:72 severity = str(value or "").strip().lower()73 if any(word in severity for word in ["none", "nope", "low clarity", "unclear", "no damage", "false", "nothing", "invalid", "not visible"]):74 return "none"75 if severity in SEVERITY_LEVELS:76 return severity77 return "moderate"78 79 80def _heuristic_severity(label: str, confidence: float) -> tuple[str, str]:81 label_lower = label.lower()82 if any(term in label_lower for term in ["frame", "windshield", "window", "headlight"]):83 return "severe", "Major component damage detected; replacement may be required."84 if any(term in label_lower for term in ["crack", "puncture"]) or confidence >= 0.82:85 return "moderate", "Visible damage detected; repair is likely required."86 return "minor", "Small visible damage area; likely cosmetic repair."87 88 89async def _request_colab_severity(items: list[dict[str, Any]]) -> list[dict[str, str]] | None:90 if not COLAB_SEVERITY_URL:91 logger.warning("Colab severity skipped: COLAB_SEVERITY_URL is not configured.")92 return None93 94 if not items:95 logger.warning("Colab severity skipped: no detection crops were prepared.")96 return None97 try:98 import httpx99 except Exception as exc:100 logger.warning("Colab severity skipped: httpx import failed: %s", exc)101 return None102 payload = {103 "model": "Qwen-3.5-VL-4B",104 "severity_schema": {105 "none": ["No damage visible", "False detection"],106 "minor": ["Cosmetic damage only", "Small area affected"],107 "moderate": ["Visible deformation", "Repair likely required"],108 "severe": ["Major deformation", "Part replacement likely"],109 },110 "detections": items,111 }112 try:113 logger.info("Calling Colab severity endpoint for %d detection crop(s).", len(items))114 115 import socket116 117 logger.info(118 "DNS lookup: %s",119 socket.gethostbyname("curfew-stump-tripod.ngrok-free.dev")120 ) 121 async with httpx.AsyncClient(122 timeout=60.0,123 verify=False,124 ) as client:125 health_url = COLAB_SEVERITY_URL.replace("/severity", "/health")126 health = await client.get(health_url)127 logger.info("HEALTH STATUS: %s", health.status_code)128 response = await client.post(COLAB_SEVERITY_URL, json=payload)129 response.raise_for_status()130 data = response.json()131 except Exception as exc:132 import traceback133 134 logger.error("COLAB REQUEST EXCEPTION TYPE: %s", type(exc).__name__)135 logger.error("COLAB REQUEST EXCEPTION REPR: %r", exc)136 logger.exception("Full traceback")137 138 return None139 raw_results = data.get("results") if isinstance(data, dict) else data140 if not isinstance(raw_results, list):141 logger.warning(142 "Colab severity returned unexpected response shape; falling back to heuristic severity. Response type: %s",143 type(data).__name__,144 )145 return None146 results: list[dict[str, str]] = []147 for index, result in enumerate(raw_results):148 if not isinstance(result, dict):149 continue150 results.append({151 "detection_id": str(result.get("detection_id", items[index].get("detection_id", index))),152 "severity": _normalize_severity(result.get("severity")),153 "note": str(result.get("note") or result.get("severity_note") or "Severity estimated from the detection frame."),154 })155 if not results:156 logger.warning("Colab severity returned no usable result items; falling back to heuristic severity.")157 return None158 159 logger.info("Colab severity succeeded for %d detection crop(s).", len(results))160 return results161 162 163def _encode_detection_crop(frame: Any, bbox: list[float]) -> str | None:164 import cv2165 166 h, w = frame.shape[:2]167 x1, y1, x2, y2 = [int(round(v)) for v in bbox]168 pad_x = max(8, int((x2 - x1) * 0.18))169 pad_y = max(8, int((y2 - y1) * 0.18))170 x1 = max(0, x1 - pad_x)171 y1 = max(0, y1 - pad_y)172 x2 = min(w, x2 + pad_x)173 y2 = min(h, y2 + pad_y)174 if x2 <= x1 or y2 <= y1:175 return None176 crop = frame[y1:y2, x1:x2]177 ch, cw = crop.shape[:2]178 if cw > 768:179 scale = 768.0 / cw180 crop = cv2.resize(crop, (768, int(ch * scale)), interpolation=cv2.INTER_AREA)181 ok, buffer = cv2.imencode(".jpg", crop, [int(cv2.IMWRITE_JPEG_QUALITY), 82])182 if not ok:183 return None184 return "data:image/jpeg;base64," + base64.b64encode(buffer.tobytes()).decode("utf-8")185 186 187async def _annotate_and_predict(video_path: Path, base_url: str) -> dict[str, Any]:188 try:189 from ultralytics import YOLO190 import cv2191 except Exception as exc: # pragma: no cover192 raise RuntimeError(193 "ultralytics and opencv-python are required for inference."194 ) from exc195 196 if not MODEL_PATH.exists():197 raise FileNotFoundError(f"Missing model file: {MODEL_PATH}")198 199 global MODEL200 model = MODEL or YOLO(str(MODEL_PATH))201 capture = cv2.VideoCapture(str(video_path))202 if not capture.isOpened():203 raise RuntimeError("Could not open uploaded video.")204 205 # Retrieve video dimensions and properties206 width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))207 height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))208 fps = float(capture.get(cv2.CAP_PROP_FPS))209 210 logger.info(211 "Video properties: width=%s height=%s fps=%s",212 width,213 height,214 fps,215 )216 217 if width <= 0 or height <= 0:218 capture.release()219 raise RuntimeError(f"Invalid video dimensions: {width}x{height}")220 221 if fps <= 0 or fps > 240:222 logger.warning("Invalid FPS reported (%s). Falling back to 30 FPS.", fps)223 fps = 30.0224 225 sample_interval = max(1, int(fps / 10.0))226 227 # Video generation removed; returning images only.228 229 frame_index = 0230 all_detections: list[dict[str, Any]] = []231 severity_items_by_key: dict[str, dict[str, Any]] = {}232 best_frame = None233 max_conf_sum = -1.0234 best_frame_index = 0235 236 while True:237 ok, frame = capture.read()238 if not ok:239 break240 frame_index += 1241 242 # Only process frames that match the sample interval243 if sample_interval > 1 and frame_index % sample_interval != 0:244 continue245 246 # Run inference on a resized frame for speed247 frame_small = frame248 fh, fw = frame.shape[:2]249 scale = min(1280 / max(fh, fw), 1.0)250 251 if scale < 1.0:252 frame_small = cv2.resize(253 frame,254 (int(fw * scale), int(fh * scale)),255 interpolation=cv2.INTER_AREA,256 )257 258 results = model.predict(frame_small, conf=0.25, verbose=False)259 annotated = frame.copy()260 261 scale_x = fw / frame_small.shape[1]262 scale_y = fh / frame_small.shape[0]263 264 if results and len(results[0].boxes) > 0:265 result = results[0]266 annotated = frame.copy()267 names = result.names268 conf_sum = 0.0269 270 for box in result.boxes:271 cls_id = int(box.cls[0])272 confidence = float(box.conf[0])273 label = names.get(cls_id, str(cls_id))274 x1, y1, x2, y2 = box.xyxy[0].tolist()275 276 xyxy = [277 float(x1 * scale_x),278 float(y1 * scale_y),279 float(x2 * scale_x),280 float(y2 * scale_y),281 ]282 283 x1i, y1i, x2i, y2i = map(int, xyxy)284 285 cv2.rectangle(286 annotated,287 (x1i, y1i),288 (x2i, y2i),289 (0, 255, 0),290 2,291 )292 293 cv2.putText(294 annotated,295 f"{label} {confidence:.2f}",296 (x1i, max(25, y1i - 10)),297 cv2.FONT_HERSHEY_SIMPLEX,298 0.7,299 (0, 255, 0),300 2,301 cv2.LINE_AA,302 )303 304 conf_sum += confidence305 306 all_detections.append({307 "class_id": cls_id,308 "label": label,309 "confidence": confidence,310 "bbox": xyxy311 })312 313 crop = _encode_detection_crop(frame, xyxy)314 if crop and (label not in severity_items_by_key or confidence > severity_items_by_key[label]["confidence"]):315 severity_items_by_key[label] = {316 "detection_id": label,317 "label": label,318 "confidence": confidence,319 "bbox": [float(v) for v in xyxy],320 "image": crop,321 }322 323 # Select the frame with the highest confidence sum as the best frame (the thumbnail image)324 if conf_sum > max_conf_sum:325 max_conf_sum = conf_sum326 best_frame = annotated.copy()327 best_frame_index = frame_index328 329 330 # Fallback to the first frame if no detections occurred331 if best_frame is None:332 best_frame = annotated.copy()333 best_frame_index = frame_index334 335 capture.release()336 337 if best_frame is None:338 raise RuntimeError("No frames could be read or written.")339 340 # Create a mobile-friendly preview image.341 preview = best_frame342 ph, pw = preview.shape[:2]343 344 max_dim = 1280345 scale = min(max_dim / max(ph, pw), 1.0)346 347 if scale < 1.0:348 preview = cv2.resize(349 preview,350 (int(pw * scale), int(ph * scale)),351 interpolation=cv2.INTER_AREA,352 )353 354 ok, buffer = cv2.imencode(355 ".jpg",356 preview,357 [cv2.IMWRITE_JPEG_QUALITY, 80],358 )359 if not ok:360 raise RuntimeError("Could not encode annotated frame.")361 362 annotated_image = base64.b64encode(buffer.tobytes()).decode("utf-8")363 364 logger.info(365 "annotated_image length=%s bytes",366 len(annotated_image),367 )368 369 # Deduplicate detections: group by label and keep the highest confidence detection370 best_detections: dict[str, dict[str, Any]] = {}371 for det in all_detections:372 label = det["label"]373 if label not in best_detections or det["confidence"] > best_detections[label]["confidence"]:374 best_detections[label] = det375 376 detections = list(best_detections.values())377 378 severity_by_id: dict[str, dict[str, str]] = {}379 colab_results = await _request_colab_severity(list(severity_items_by_key.values()))380 if colab_results:381 severity_by_id = {item["detection_id"]: item for item in colab_results}382 elif detections:383 logger.warning(384 "Using heuristic severity fallback for %d detection(s). Check COLAB_SEVERITY_URL and Colab server logs.",385 len(detections),386 )387 388 for det in detections:389 result = severity_by_id.get(det["label"])390 if result:391 det["severity"] = result["severity"]392 det["severity_note"] = result["note"]393 det["severity_source"] = "Qwen-3.5-VL-4B"394 else:395 severity, note = _heuristic_severity(det["label"], det["confidence"])396 det["severity"] = severity397 det["severity_note"] = note398 det["severity_source"] = "fallback"399 400 # Build summary label401 if detections:402 summary_parts = []403 for det in detections:404 summary_parts.append(f"{det['label']} ({int(det['confidence']*100)}%)")405 summary = "Detected: " + ", ".join(summary_parts)406 else:407 summary = "No damages detected"408 409 return {410 "annotated_image": f"data:image/jpeg;base64,{annotated_image}",411 "detection_frames": list(severity_items_by_key.values()),412 "detections": detections,413 "source_frame": best_frame_index,414 "summary": summary,415 }416 417 418@app.post("/analyze-video")419async def analyze_video(request: Request, file: UploadFile = File(...)):420 suffix = Path(file.filename).suffix.lower() or ".mp4"421 temp_path = ROOT / f"upload_{uuid.uuid4().hex}{suffix}"422 content = await file.read()423 temp_path.write_bytes(content)424 try:425 base_url = str(request.base_url).rstrip('/')426 result = await _annotate_and_predict(temp_path, base_url)427 return JSONResponse(result)428 finally:429 if temp_path.exists():430 temp_path.unlink()431 432 433@app.post("/verify-face")434async def verify_face(payload: FaceVerifyRequest):435 import cv2436 import numpy as np437 438 try:439 raw = payload.frame.split(",")[-1]440 img_bytes = base64.b64decode(raw)441 nparr = np.frombuffer(img_bytes, np.uint8)442 frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)443 if frame is None:444 return JSONResponse({"verified": False, "face_count": 0, "confidence": 0.0, "face_crop": None})445 446 gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)447 gray = cv2.equalizeHist(gray)448 449 frontal_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")450 profile_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_profileface.xml")451 452 faces = frontal_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60))453 if len(faces) == 0:454 faces = profile_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=4, minSize=(60, 60))455 456 verified = len(faces) > 0457 face_count = len(faces)458 confidence = 0.0459 face_crop_b64 = None460 461 if verified:462 x, y, w, h = max(faces, key=lambda f: f[2] * f[3])463 pad = int(min(w, h) * 0.20)464 x1 = max(0, x - pad)465 y1 = max(0, y - pad)466 x2 = min(frame.shape[1], x + w + pad)467 y2 = min(frame.shape[0], y + h + pad)468 face_crop = frame[y1:y2, x1:x2]469 470 ok, buf = cv2.imencode(".jpg", face_crop, [cv2.IMWRITE_JPEG_QUALITY, 85])471 if ok:472 face_crop_b64 = "data:image/jpeg;base64," + base64.b64encode(buf.tobytes()).decode()473 474 confidence = round((w * h) / (frame.shape[0] * frame.shape[1]) * 10, 2)475 confidence = min(confidence, 0.99)476 477 return JSONResponse({478 "verified": verified,479 "face_count": face_count,480 "confidence": confidence,481 "face_crop": face_crop_b64482 })483 except Exception as e:484 return JSONResponse({"verified": False, "face_count": 0, "confidence": 0.0, "face_crop": None, "error": str(e)}, status_code=500)485 486 487@app.on_event("startup")488async def startup_load_model():489 global MODEL490 try:491 from ultralytics import YOLO492 if MODEL is None and MODEL_PATH.exists():493 MODEL = YOLO(str(MODEL_PATH))494 logger.info("YOLO model loaded once at startup.")495 except Exception:496 logger.exception("Failed loading YOLO model at startup")497 498 499@app.get("/health")500def health():501 return {"status": "ok", "model": str(MODEL_PATH)}502 503 504if __name__ == "__main__": # pragma: no cover505 import uvicorn506 507 port = int(os.getenv("AI_SERVER_PORT", "8001"))508 uvicorn.run(app, host="0.0.0.0", port=port)509