Create main3.py

main3.py

@@ -0,0 +1,191 @@
+import io
+import os
+import gdown
+import base64
+import cv2
+import numpy as np
+from PIL import Image
+from typing import Optional
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.projects.point_rend import add_pointrend_config
+
+# -------------------------------
+# FastAPI setup
+# -------------------------------
+app = FastAPI(title="Rooftop Segmentation API")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# -------------------------------
+# Available epsilons
+# -------------------------------
+EPSILONS = [0.01, 0.005, 0.004, 0.003, 0.001]
+
+@app.get("/epsilons")
+def get_epsilons():
+    return {"epsilons": EPSILONS}
+
+# -------------------------------
+# Google Drive model download (irregular-flat)
+# -------------------------------
+MODEL_PATH_IRREGULAR = "/tmp/model_irregular_flat.pth"
+DRIVE_FILE_ID = "15vi4zPhCs3aBnGepVnXFOqQjxdK1jpnA"
+
+def download_irregular_model():
+    if not os.path.exists(MODEL_PATH_IRREGULAR):
+        url = f"https://drive.google.com/uc?id={DRIVE_FILE_ID}"
+        tmp_dir = "/tmp/gdown"
+        os.makedirs(tmp_dir, exist_ok=True)
+        os.environ["GDOWN_CACHE_DIR"] = tmp_dir
+        print("Downloading irregular-flat Detectron2 model...")
+        gdown.download(url, MODEL_PATH_IRREGULAR, quiet=False, fuzzy=True, use_cookies=False)
+        print("Download complete.")
+    else:
+        print("Irregular-flat model already exists, skipping download.")
+
+download_irregular_model()
+
+if os.path.exists(MODEL_PATH_IRREGULAR):
+    print("Irregular-flat model is ready at", MODEL_PATH_IRREGULAR)
+else:
+    print("Irregular-flat model NOT found! Something went wrong!")
+
+# -------------------------------
+# Detectron2 model setup
+# -------------------------------
+def setup_model_rect(weights_path: str):
+    cfg = get_cfg()
+    add_pointrend_config(cfg)
+    cfg_path = "detectron2/projects/PointRend/configs/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco.yaml"
+    cfg.merge_from_file(cfg_path)
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
+    cfg.MODEL.POINT_HEAD.NUM_CLASSES = cfg.MODEL.ROI_HEADS.NUM_CLASSES
+    cfg.MODEL.WEIGHTS = weights_path
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+    cfg.MODEL.DEVICE = "cpu"
+    return DefaultPredictor(cfg)
+
+def setup_model_irregular(weights_path: str):
+    cfg = get_cfg()
+    add_pointrend_config(cfg)
+    cfg_path = "detectron2/projects/PointRend/configs/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco.yaml"
+    cfg.merge_from_file(cfg_path)
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+    cfg.MODEL.POINT_HEAD.NUM_CLASSES = cfg.MODEL.ROI_HEADS.NUM_CLASSES
+    cfg.MODEL.WEIGHTS = weights_path
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+    cfg.MODEL.DEVICE = "cpu"
+    return DefaultPredictor(cfg)
+
+# Load models
+predictor_rect = setup_model_rect("/app/model_rect_final.pth")
+predictor_irregular_flat = setup_model_irregular(MODEL_PATH_IRREGULAR)
+
+# -------------------------------
+# Utility functions
+# -------------------------------
+def im_to_b64_png(im: np.ndarray) -> str:
+    _, buffer = cv2.imencode(".png", im)
+    return base64.b64encode(buffer).decode()
+
+def extract_polygon(mask: np.ndarray, epsilon_ratio: float = 0.004):
+    mask_uint8 = (mask * 255).astype(np.uint8)
+    contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return None
+    c = max(contours, key=cv2.contourArea)
+    epsilon = epsilon_ratio * cv2.arcLength(c, True)
+    polygon = cv2.approxPolyDP(c, epsilon, True)
+    return polygon.reshape(-1, 2)
+
+def overlay_polygon(im: np.ndarray, polygon: Optional[np.ndarray], vertex_color=(0,0,255), line_color=(0,255,0)):
+    overlay = im.copy()
+    if polygon is not None:
+        # Draw polygon outline (thin)
+        cv2.polylines(overlay, [polygon.astype(np.int32)], True, line_color, thickness=2)
+
+        # Draw vertices
+        for i, (x, y) in enumerate(polygon):
+            cv2.circle(overlay, (int(x), int(y)), 4, vertex_color, -1)
+            # Draw vertex index (black number)
+            cv2.putText(overlay, str(i+1), (int(x)+5, int(y)-5),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (20,20,20), 1, cv2.LINE_AA)
+
+        # Display vertex count on top
+        vertex_count = len(polygon)
+        cv2.putText(overlay, f"num_vertices = {vertex_count}", (20, 35),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (20,20,20), 2, cv2.LINE_AA)
+
+    return overlay
+
+# -------------------------------
+# API endpoints
+# -------------------------------
+@app.get("/")
+def root():
+    return {"message": "Rooftop Segmentation API is running!"}
+
+@app.post("/predict")
+async def predict(
+    file: UploadFile = File(...),
+    rooftop_type: str = Form(...),
+    epsilon: float = Form(0.004)
+):
+    contents = await file.read()
+    try:
+        im_pil = Image.open(io.BytesIO(contents)).convert("RGB")
+    except Exception as e:
+        return JSONResponse(status_code=400, content={"error": "Invalid image", "detail": str(e)})
+
+    im = np.array(im_pil)[:, :, ::-1].copy()  # RGB -> BGR
+
+    if rooftop_type.lower() == "rectangular":
+        predictor = predictor_rect
+        model_used = "model_rect_final.pth"
+    elif rooftop_type.lower() == "irregular":
+        predictor = predictor_irregular_flat
+        model_used = "model_irregular_flat.pth"
+    else:
+        return JSONResponse(status_code=400, content={"error": "Invalid rooftop_type. Choose 'rectangular' or 'irregular'."})
+
+    outputs = predictor(im)
+    instances = outputs["instances"].to("cpu")
+
+    if len(instances) == 0:
+        return {
+            "polygon": None,
+            "vertices": None,
+            "vertex_count": 0,
+            "image": None,
+            "model_used": model_used,
+            "rooftop_type": rooftop_type,
+            "epsilon": epsilon
+        }
+
+    idx = int(instances.scores.argmax().item())
+    raw_mask = instances.pred_masks[idx].numpy().astype(np.uint8)
+
+    polygon = extract_polygon(raw_mask, epsilon_ratio=epsilon)
+    vertex_count = len(polygon) if polygon is not None else 0
+
+    overlay = overlay_polygon(im, polygon)
+    img_b64 = im_to_b64_png(overlay)
+
+    return {
+        "polygon": polygon.tolist() if polygon is not None else None,
+        "vertex_count": vertex_count,
+        "image": img_b64,
+        "model_used": model_used,
+        "rooftop_type": rooftop_type,
+        "epsilon": epsilon
+    }