Update meaw.py

meaw.py

@@ -1,116 +1,88 @@
 import os
+import io
+import base64
 import cv2
 import torch
-import gdown
-import base64
+import numpy as np
+from PIL import Image
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 from detectron2.config import get_cfg
-from detectron2 import model_zoo
 from detectron2.engine import DefaultPredictor
-from detectron2.utils.visualizer import Visualizer, ColorMode
+from detectron2 import model_zoo
+from detectron2.utils.visualizer import Visualizer
 from detectron2.data import MetadataCatalog
-import io
-# === Model Setup ===
-MODEL_PATH = "/tmp/model_final.pth"
-DRIVE_FILE_ID = "1bazIVYG0CYMubDLoMu5pgH6ArC1sayzg"
 
-def download_model():
-    if not os.path.exists(MODEL_PATH):
-        url = f"https://drive.google.com/uc?id={DRIVE_FILE_ID}"
+app = FastAPI(title="Roof Segmentation API")
 
-        # Create a writable temporary directory for gdown
-        tmp_dir = "/tmp/gdown"
-        os.makedirs(tmp_dir, exist_ok=True)
+@app.get("/")
+def home():
+    return {"status": "running"}
 
-        # Set GDOWN cache path to a safe location
-        os.environ["GDOWN_CACHE_DIR"] = tmp_dir
+MODEL_PATH = "model_final (4).pth"
 
-        print("Downloading Detectron2 model...")
-        gdown.download(
-            url,
-            MODEL_PATH,
-            quiet=False,
-            fuzzy=True,
-            use_cookies=False  # Important for Hugging Face
-        )
-        print("Download complete.")
-
-download_model()
-
-
-# Register dummy metadata (for visualization only)
-MetadataCatalog.get(metadata_name).set(thing_classes=class_names)
-
-# === Config Setup ===
+# -----------------------------
+# 1. Configure Detectron2
+# -----------------------------
 cfg = get_cfg()
-cfg.merge_from_file(
-    model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
-)
-cfg.MODEL.WEIGHTS = MODEL_PATH
+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.DEVICE = "cpu"  # or "cuda" if GPU is enabled
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+cfg.MODEL.WEIGHTS = MODEL_PATH
+cfg.MODEL.DEVICE = "cpu"   # Use "cuda" if you have GPU
 
-# === Predictor ===
 predictor = DefaultPredictor(cfg)
+
+# Morphology kernel for polygon smoothing
+kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+
 # -----------------------------
-# 3. Helper: Encode mask to Base64
+# Helper: Encode image to Base64
 # -----------------------------
-def encode_mask(mask: np.ndarray) -> str:
-    """Convert mask numpy array to base64 PNG string."""
-    mask_img = Image.fromarray(mask.astype(np.uint8))
+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()
-    mask_img.save(buf, format="PNG")
+    pil_img.save(buf, format="PNG")
     return base64.b64encode(buf.getvalue()).decode("utf-8")
-# === Predict Function ===
-def predict(image_path: str):
-    im = cv2.imread(image_path)
-    outputs = predictor(im)
-    instances = outputs["instances"].to("cpu")
 
-    # Visualization
-    v = Visualizer(
-        im[:, :, ::-1],
-        metadata=metadata,
-        scale=1.2,
-        instance_mode=ColorMode.IMAGE
-    )
-    out = v.draw_instance_predictions(instances)
-    result_img = out.get_image()[:, :, ::-1]
-
- # -----------------------------
-# 4. API Endpointt
+# -----------------------------
+# API Endpoint
 # -----------------------------
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # Read image
     contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-    image = np.array(image)[:, :, ::-1]  # to BGR for OpenCV/Detectron2
+    image = np.array(Image.open(io.BytesIO(contents)).convert("RGB"))[:, :, ::-1]  # BGR
 
-    # Run inference
     outputs = predictor(image)
     instances = outputs["instances"].to("cpu")
 
-    results = []
-    mask_b64 = None
-
+    # --- Polygon overlay ---
+    polygon_overlay = image.copy()
     if instances.has("pred_masks"):
-        masks = instances.pred_masks.numpy()
-        boxes = instances.pred_boxes.tensor.numpy()
-        scores = instances.scores.numpy()
-
-        # Combine masks into one
-        combined_mask = np.any(masks, axis=0).astype(np.uint8) * 255
-        mask_b64 = encode_mask(combined_mask)
+        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]  # back to BGR
 
-        for i in range(len(masks)):
-            results.append({
-                "box": boxes[i].tolist(),
-                "score": float(scores[i])
-            })
+    # Encode both images as Base64
+    predicted_b64 = encode_image(predicted_output)
+    polygon_b64 = encode_image(polygon_overlay)
 
     return JSONResponse({
-        "predictions": results,
-        "mask": mask_b64  # base64 string (PNG)
+        "predicted_output": predicted_b64,
+        "polygon_overlay": polygon_b64
     })