Update app.py

app.py

@@ -1,58 +1,38 @@
 # app.py
-# app.py
-import io
-import os
-import time
-import base64
+import io, os, time, base64
 from pathlib import Path
 from typing import Optional
-
-import uvicorn
+import cv2
+import numpy as np
 from fastapi import FastAPI, File, UploadFile, Query, HTTPException
 from fastapi.responses import JSONResponse
-import numpy as np
-import cv2
 from ultralytics import YOLO
-import requests
+import uvicorn
 
-# ------------------------
-# Config
-# ------------------------
-# Put your model file next to this script and set MODEL_PATH to the filename
+# ---------- config ----------
 MODEL_PATH = os.getenv("MODEL_PATH", "best_rockfall_model.pt")
 CONF_DEFAULT = float(os.getenv("CONF_DEFAULT", 0.25))
 IMG_SIZE_DEFAULT = int(os.getenv("IMG_SIZE_DEFAULT", 640))
 
-if not Path(MODEL_PATH).exists():
-    print(f"Warning: MODEL_PATH '{MODEL_PATH}' not found. Please place your .pt file next to app.py or set MODEL_PATH env var.")
-
-# ------------------------
-# App + Model loader
-# ------------------------
-app = FastAPI(title="Rockfall Detection API (fixed single-class behavior)")
+app = FastAPI(title="Rockfall API - mask + bbox + probability")
 
-MODEL = {"model": None, "path": None}
-
-def load_model(path: str = MODEL_PATH):
-    # Use YOLO.from_pretrained if using a named model; here we load a local .pt
-    model = YOLO(path)
-    MODEL["model"] = model
-    MODEL["path"] = path
-    print("Model loaded from:", path)
-    return model
+# load model
+MODEL = {"model": None}
+def load_model(path=MODEL_PATH):
+    if not Path(path).exists():
+        print(f"Warning: model file not found at {path}")
+    MODEL["model"] = YOLO(path)
+    print("Model loaded:", path)
 
 @app.on_event("startup")
 def startup():
     try:
         load_model(MODEL_PATH)
     except Exception as e:
-        # don't crash server; let user fix model and call reload endpoint if implemented
         MODEL["model"] = None
-        print("Failed to load model on startup:", e)
+        print("Model load failed:", e)
 
-# ------------------------
-# Helpers
-# ------------------------
+# ---------- helpers ----------
 def read_imagefile(upload_file: UploadFile) -> np.ndarray:
     data = upload_file.file.read()
     upload_file.file.close()
@@ -62,80 +42,104 @@ def read_imagefile(upload_file: UploadFile) -> np.ndarray:
         raise ValueError("Cannot decode image")
     return img
 
-def fetch_image_from_url(url: str, timeout: int = 10) -> np.ndarray:
-    r = requests.get(url, timeout=timeout)
-    if r.status_code != 200:
-        raise ValueError(f"Failed to fetch image: status {r.status_code}")
-    arr = np.frombuffer(r.content, np.uint8)
-    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
-    if img is None:
-        raise ValueError("Invalid image content")
-    return img
+def image_to_datauri_png(img_bgr: np.ndarray) -> str:
+    # encode as PNG
+    _, buf = cv2.imencode(".png", img_bgr)
+    b64 = base64.b64encode(buf).decode("utf-8")
+    return f"data:image/png;base64,{b64}"
 
-def run_inference(img_bgr: np.ndarray, conf: float, imgsz: int):
+def run_predict_and_build_mask(img_bgr: np.ndarray, conf: float, imgsz: int):
     if MODEL["model"] is None:
-        raise RuntimeError("Model is not loaded. Ensure MODEL_PATH is correct and model file exists.")
-    # Ultralytics accepts numpy image (BGR). Use predict() to get full results object.
-    start = time.time()
+        raise RuntimeError("Model not loaded")
+
+    # Run the model (ultralytics). Provide numpy BGR image; returns list
+    t0 = time.time()
     results = MODEL["model"].predict(source=img_bgr, conf=conf, imgsz=imgsz, verbose=False)
-    elapsed_ms = (time.time() - start) * 1000.0
-    res = results[0]  # single image result
-    return res, elapsed_ms
+    elapsed_ms = (time.time() - t0) * 1000.0
+    res = results[0]
+
+    h, w = img_bgr.shape[:2]
+    # Start with empty mask
+    mask = np.zeros((h, w), dtype=np.uint8)
 
-def boxes_from_result(res) -> list:
-    """
-    Convert ultralytics result to a unified list of detections.
-    Force single-class mapping: class_id=0, class_name='rockfall' for every detection.
-    """
     detections = []
-    if hasattr(res, "boxes") and len(res.boxes) > 0:
-        xyxy = res.boxes.xyxy.cpu().numpy()  # x1,y1,x2,y2
-        confs = res.boxes.conf.cpu().numpy()
-        # cls = res.boxes.cls.cpu().numpy().astype(int)  # original classes (ignored)
-        for b, score in zip(xyxy, confs):
-            x1, y1, x2, y2 = map(int, b)
-            detections.append({
-                "class_id": 0,                # force to single-class 0
-                "class_name": "rockfall",     # canonical class name
-                "confidence": float(score),
-                "bbox": [x1, y1, x2, y2]
-            })
-    return detections
-
-def annotate_image_with_boxes(img_bgr: np.ndarray, detections: list) -> np.ndarray:
-    """
-    Draw boxes (green) and labels on a copy of the input image (BGR).
-    """
+    max_conf = 0.0
+    first_bbox = None
+
+    # If model provides segmentation masks (res.masks), prefer them
+    if hasattr(res, "masks") and res.masks is not None and len(res.masks.xy) > 0:
+        # res.masks.xy is list of polygons; res.masks.data contains bitmasks depending on version
+        # Try to produce a binary mask by rasterizing the polygons
+        for poly in res.masks.xy:
+            # poly is list of (x,y) floats; draw fill polygon
+            pts = np.array(poly, dtype=np.int32).reshape(-1, 2)
+            cv2.fillPoly(mask, [pts], 255)
+        # get confidences / boxes if available
+        if hasattr(res, "boxes") and len(res.boxes) > 0:
+            xyxy = res.boxes.xyxy.cpu().numpy()
+            confs = res.boxes.conf.cpu().numpy()
+            max_conf = float(confs.max()) if len(confs)>0 else 0.0
+            first_bbox = list(map(int, xyxy[0]))
+    else:
+        # No segmentation output -> build mask using bounding boxes
+        if hasattr(res, "boxes") and len(res.boxes) > 0:
+            xyxy = res.boxes.xyxy.cpu().numpy()   # x1,y1,x2,y2
+            confs = res.boxes.conf.cpu().numpy()
+            max_conf = float(confs.max()) if len(confs)>0 else 0.0
+            for i, (b, c) in enumerate(zip(xyxy, confs)):
+                x1, y1, x2, y2 = map(int, b)
+                # clip
+                x1, y1 = max(0,x1), max(0,y1)
+                x2, y2 = min(w-1,x2), min(h-1,y2)
+                if x2<=x1 or y2<=y1:
+                    continue
+                # fill the rectangle in mask (255 for rockfall)
+                cv2.rectangle(mask, (x1,y1), (x2,y2), 255, thickness=-1)
+                detections.append({"bbox":[x1,y1,x2,y2], "confidence": float(c)})
+            # pick first bbox (if exists)
+            if detections:
+                first_bbox = detections[0]["bbox"]
+        else:
+            # no boxes -> empty mask, prob 0
+            max_conf = 0.0
+
+    # compute mask stats
+    mask_area = int((mask>0).sum())
+    mask_fraction = float(mask_area) / (w*h)
+
+    return {
+        "mask": mask,
+        "mask_area": mask_area,
+        "mask_fraction": mask_fraction,
+        "probability": float(max_conf),
+        "bbox": first_bbox,
+        "detections": detections,
+        "inference_ms": elapsed_ms
+    }
+
+def annotate_image(img_bgr: np.ndarray, bbox: Optional[list], probability: float):
     out = img_bgr.copy()
-    for det in detections:
-        x1, y1, x2, y2 = det["bbox"]
-        conf = det["confidence"]
-        label = f"{det['class_name']} {conf:.2f}"
-        # draw rectangle
-        cv2.rectangle(out, (x1, y1), (x2, y2), (0, 255, 0), thickness=2)
-        # label background
+    if bbox:
+        x1,y1,x2,y2 = bbox
+        cv2.rectangle(out, (x1,y1), (x2,y2), (0,255,0), 2)
+        label = f"rockfall {probability:.2f}"
         (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
-        cv2.rectangle(out, (x1, y1 - th - 6), (x1 + tw + 6, y1), (0, 255, 0), -1)
-        cv2.putText(out, label, (x1 + 3, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1, cv2.LINE_AA)
+        cv2.rectangle(out, (x1, y1 - th - 6), (x1 + tw + 6, y1), (0,255,0), -1)
+        cv2.putText(out, label, (x1+3, y1-4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,0), 1, cv2.LINE_AA)
     return out
 
-def image_to_base64(img_bgr: np.ndarray) -> str:
-    _, buf = cv2.imencode(".jpg", img_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
-    b64 = base64.b64encode(buf.tobytes()).decode("utf-8")
-    return f"data:image/jpeg;base64,{b64}"
-
-# ------------------------
-# Endpoints
-# ------------------------
+# ---------- endpoints ----------
 @app.get("/health")
 def health():
-    return {"status": "ok", "model_loaded": MODEL["model"] is not None, "model_path": MODEL["path"]}
+    return {"status":"ok", "model_loaded": MODEL["model"] is not None}
 
 @app.post("/predict/image")
 async def predict_image(file: UploadFile = File(...), conf: Optional[float] = Query(CONF_DEFAULT), imgsz: Optional[int] = Query(IMG_SIZE_DEFAULT)):
     """
-    Upload one image file for prediction.
-    Returns JSON with detections and annotated image (base64).
+    Returns:
+      - mask_base64: binary mask PNG as data URI (rockfall=255)
+      - annotated_image_base64: annotated image PNG as data URI (bbox + label)
+      - mask_area, mask_fraction, probability, bbox (first)
     """
     try:
         img = read_imagefile(file)
@@ -143,56 +147,26 @@ async def predict_image(file: UploadFile = File(...), conf: Optional[float] = Qu
         raise HTTPException(status_code=400, detail=f"Invalid image: {e}")
 
     try:
-        res, t_ms = run_inference(img, conf=float(conf), imgsz=int(imgsz))
+        out = run_predict_and_build_mask(img, float(conf), int(imgsz))
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Inference error: {e}")
 
-    detections = boxes_from_result(res)
-    annotated = annotate_image_with_boxes(img, detections)
-    return {
-        "filename": file.filename,
-        "num_detections": len(detections),
-        "predictions": detections,
-        "annotated_image_base64": image_to_base64(annotated),
-        "inference_time_ms": t_ms
-    }
-
-@app.post("/predict/url")
-async def predict_url(image_url: str, conf: Optional[float] = Query(CONF_DEFAULT), imgsz: Optional[int] = Query(IMG_SIZE_DEFAULT)):
-    """
-    Predict on an image fetched from a URL.
-    """
-    try:
-        img = fetch_image_from_url(image_url)
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-    try:
-        res, t_ms = run_inference(img, conf=float(conf), imgsz=int(imgsz))
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Inference error: {e}")
+    mask_png_uri = image_to_datauri_png(out["mask"])
+    annotated = annotate_image(img, out["bbox"], out["probability"])
+    annotated_png_uri = image_to_datauri_png(annotated)
 
-    detections = boxes_from_result(res)
-    annotated = annotate_image_with_boxes(img, detections)
     return {
-        "image_url": image_url,
-        "num_detections": len(detections),
-        "predictions": detections,
-        "annotated_image_base64": image_to_base64(annotated),
-        "inference_time_ms": t_ms
+        "filename": file.filename,
+        "num_detections": len(out["detections"]),
+        "probability": out["probability"],
+        "bbox": out["bbox"],               # null or [x1,y1,x2,y2]
+        "mask_area_pixels": out["mask_area"],
+        "mask_area_fraction": out["mask_fraction"],
+        "mask_base64": mask_png_uri,
+        "annotated_image_base64": annotated_png_uri,
+        "inference_time_ms": out["inference_ms"]
     }
 
-# Optional: admin reload endpoint
-@app.get("/model/reload")
-def reload_model():
-    try:
-        load_model(MODEL_PATH)
-        return {"status": "reloaded", "model_path": MODEL["path"]}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-# ------------------------
-# Run (dev)
-# ------------------------
+# run
 if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
+    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)