Create app3.py

app3.py

@@ -0,0 +1,90 @@
+import os
+import io
+import base64
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.responses import JSONResponse
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultPredictor
+from detectron2 import model_zoo
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog
+
+app = FastAPI(title="Roof Segmentation API")
+
+@app.get("/")
+def home():
+    return {"status": "running"}
+
+MODEL_PATH = "model_final (4).pth"
+
+# ----------------------------- Detectron2 Config -----------------------------
+cfg = get_cfg()
+cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
+cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+cfg.MODEL.WEIGHTS = MODEL_PATH
+cfg.MODEL.DEVICE = "cpu"  # or "cuda" if GPU available
+
+predictor = DefaultPredictor(cfg)
+kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+
+# ----------------------------- Helper -----------------------------
+def encode_image(img: np.ndarray) -> str:
+    """Convert BGR image to base64 PNG string."""
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    pil_img = Image.fromarray(img_rgb)
+    buf = io.BytesIO()
+    pil_img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+# ----------------------------- API Endpoint -----------------------------
+@app.post("/predict")
+async def predict(
+    file: UploadFile = File(...),
+    output_type: str = Form("both")  # 'predicted_output', 'polygon_overlay', or 'both'
+):
+    contents = await file.read()
+    image = np.array(Image.open(io.BytesIO(contents)).convert("RGB"))[:, :, ::-1]  # BGR
+
+    outputs = predictor(image)
+    instances = outputs["instances"].to("cpu")
+
+    # --- Polygon overlay ---
+    polygon_overlay = image.copy()
+    if instances.has("pred_masks"):
+        masks = instances.pred_masks
+        combined_mask = torch.any(masks, dim=0).numpy().astype("uint8") * 255
+        clean_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
+        clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_CLOSE, kernel)
+        smooth_mask = cv2.GaussianBlur(clean_mask, (5, 5), 0)
+        _, smooth_mask = cv2.threshold(smooth_mask, 127, 255, cv2.THRESH_BINARY)
+        contours, _ = cv2.findContours(smooth_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if contours:
+            largest_contour = max(contours, key=cv2.contourArea)
+            polygon = cv2.approxPolyDP(largest_contour, epsilon=10, closed=True)
+            cv2.polylines(polygon_overlay, [polygon], isClosed=True, color=(0, 255, 0), thickness=2)
+
+    # --- Detectron2 predicted output ---
+    v = Visualizer(image[:, :, ::-1], MetadataCatalog.get("toproof_train"), scale=1.2)
+    out = v.draw_instance_predictions(instances)
+    predicted_output = out.get_image()[:, :, ::-1]  # BGR
+
+    # Encode images
+    predicted_b64 = encode_image(predicted_output)
+    polygon_b64 = encode_image(polygon_overlay)
+
+    # Build response based on output_type
+    result = {}
+    if output_type == "predicted_output":
+        result["predicted_output"] = predicted_b64
+    elif output_type == "polygon_overlay":
+        result["polygon_overlay"] = polygon_b64
+    else:  # both
+        result["predicted_output"] = predicted_b64
+        result["polygon_overlay"] = polygon_b64
+
+    return JSONResponse(result)