Upload app.py

app.py

@@ -0,0 +1,202 @@
+import os
+import uuid
+import shutil
+import logging
+from contextlib import asynccontextmanager
+
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, HTTPException, Query
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from dotenv import load_dotenv
+
+from scripts.gradcam import get_resnet_gradcam
+from scripts.yolo_predict import get_yolo_damage_boxes
+from scripts.load_models import initialize_models
+from scripts.prediction_helper import ResnetONNXPredictor, FusionONNXPredictor
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+UPLOAD_DIR = "static/uploads"
+RESULT_DIR = "static/results"
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(RESULT_DIR, exist_ok=True)
+
+CLASS_MAP = {
+    0: "Front Breakage",
+    1: "Front Crushed",
+    2: "Front Normal",
+    3: "Rear Breakage",
+    4: "Rear Crushed",
+    5: "Rear Normal"
+}
+
+# Unified dictionary to hold all loaded models
+ml_models = {}
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logger.info("Loading models at startup...")
+    try:
+        models = initialize_models(CLASS_MAP)
+        ml_models["resnet_onnx"] = ResnetONNXPredictor(models["resnet_onnx"], CLASS_MAP)
+        ml_models["fusion_onnx"] = FusionONNXPredictor(models["fusion_onnx"], CLASS_MAP)
+        ml_models["resnet_pt"] = models["resnet_pt"]
+        ml_models["yolo_onnx"] = models["yolo_onnx"]
+        
+        logger.info("All models loaded successfully.")
+    except Exception as e:
+        logger.exception("Model loading failed.")
+        raise RuntimeError(str(e))
+    
+    yield
+    ml_models.clear()
+    logger.info("Application shutdown.")
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+def validate_image(upload_file: UploadFile):
+    if not upload_file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")
+
+def save_upload(upload_file: UploadFile):
+    unique_id = str(uuid.uuid4())
+    filename = f"{unique_id}_input.jpg"
+    file_path = os.path.join(UPLOAD_DIR, filename)
+
+    with open(file_path, "wb") as buffer:
+        shutil.copyfileobj(upload_file.file, buffer)
+
+    return unique_id, filename, file_path
+
+# Endpoints
+@app.get("/")
+def api_status():
+    return {"status": "API is running"}
+
+@app.post("/predict/resnet")
+async def resnet_prediction(file: UploadFile = File(...)):
+    validate_image(file)
+    try:
+        pil_image = Image.open(file.file).convert("RGB")
+        result = ml_models["resnet_onnx"].predict(pil_image)
+        return {"status": "success", "prediction": result}
+    except Exception as e:
+        logger.exception("ResNet ONNX prediction failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/predict/fusion")
+async def fusion_prediction(file: UploadFile = File(...)):
+    validate_image(file)
+    try:
+        pil_image = Image.open(file.file).convert("RGB")
+        result = ml_models["fusion_onnx"].predict(pil_image)
+        return {"status": "success", "prediction": result}
+    except Exception as e:
+        logger.exception("Fusion ONNX prediction failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/predict/yolo")
+async def yolo_detection(file: UploadFile = File(...)):
+    validate_image(file)
+    try:
+        unique_id, input_filename, input_path = save_upload(file)
+        output_name = f"{unique_id}_yolo.jpg"
+        output_path = os.path.join(RESULT_DIR, output_name)
+
+        result = get_yolo_damage_boxes(input_path, ml_models["yolo_onnx"], output_path)
+
+        return {
+            "status": "success",
+            "original_image": f"/static/uploads/{input_filename}",
+            "yolo_image": f"/static/results/{output_name}",
+            "detections": result["detections"],
+            "total_detections": result["total_detections"],
+            "message": result["message"]
+        }
+    except Exception as e:
+        logger.exception("YOLO ONNX detection failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/predict/gradcam")
+async def gradcam_generation(file: UploadFile = File(...)):
+    validate_image(file)
+    try:
+        unique_id, input_filename, input_path = save_upload(file)
+        output_name = f"{unique_id}_gradcam.jpg"
+        output_path = os.path.join(RESULT_DIR, output_name)
+
+        get_resnet_gradcam(input_path, ml_models["resnet_pt"], output_path)
+
+        return {
+            "status": "success",
+            "original_image": f"/static/uploads/{input_filename}",
+            "gradcam_image": f"/static/results/{output_name}"
+        }
+    except Exception as e:
+        logger.exception("Grad-CAM generation failed.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/predict/comprehensive")
+async def comprehensive_prediction(
+    file: UploadFile = File(...),
+    mode: str = Query("fusion", description="Classification model to use: 'resnet' or 'fusion'")
+):
+    validate_image(file)
+    mode = mode.lower()
+
+    if mode not in {"resnet", "fusion"}:
+        raise HTTPException(status_code=400, detail="mode must be 'resnet' or 'fusion'")
+
+    try:
+        unique_id, input_filename, input_path = save_upload(file)
+        pil_image = Image.open(input_path).convert("RGB")
+
+        # 1. Classification
+        if mode == "resnet":
+            classification_result = ml_models["resnet_onnx"].predict(pil_image)
+        else:
+            classification_result = ml_models["fusion_onnx"].predict(pil_image)
+
+        # 2. YOLO Bounding Boxes
+        yolo_output_name = f"{unique_id}_yolo.jpg"
+        yolo_output_path = os.path.join(RESULT_DIR, yolo_output_name)
+        yolo_result = get_yolo_damage_boxes(input_path, ml_models["yolo_onnx"], yolo_output_path)
+
+        # 3. Grad-CAM
+        gradcam_output_name = f"{unique_id}_gradcam.jpg"
+        gradcam_output_path = os.path.join(RESULT_DIR, gradcam_output_name)
+        get_resnet_gradcam(input_path, ml_models["resnet_pt"], gradcam_output_path)
+
+        return {
+            "status": "success",
+            "mode": mode,
+            "original_image": f"/static/uploads/{input_filename}",
+            "classification": classification_result,
+            "yolo": {
+                "image": f"/static/results/{yolo_output_name}",
+                "detections": yolo_result["detections"],
+                "total_detections": yolo_result["total_detections"]
+            },
+            "gradcam": f"/static/results/{gradcam_output_name}"
+        }
+
+    except Exception as e:
+        logger.exception("Comprehensive prediction pipeline failed.")
+        raise HTTPException(status_code=500, detail=str(e))