main.py

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 (flat rooftop)
# -------------------------------
MODEL_PATH_FLAT = "/tmp/model_flat.pth"
DRIVE_FILE_ID = "1GO_Ut-N2e2we8t9mnsysb0P1qMsBF8FW"

def download_flat_model():
    if not os.path.exists(MODEL_PATH_FLAT):
        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 flat rooftop Detectron2 model...")
        gdown.download(url, MODEL_PATH_FLAT, quiet=False, fuzzy=True, use_cookies=False)
        print("Download complete.")
    else:
        print("Flat model already exists, skipping download.")

download_flat_model()

if os.path.exists(MODEL_PATH_FLAT):
    print("Flat rooftop model is ready at", MODEL_PATH_FLAT)
else:
    print("Flat rooftop model NOT found! Something went wrong!")

# -------------------------------
# Detectron2 model setup
# -------------------------------
def setup_model_flat(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 = 1
    cfg.MODEL.WEIGHTS = weights_path
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
    cfg.MODEL.DEVICE = "cpu"
    return DefaultPredictor(cfg)

# Load flat rooftop model
predictor_flat = setup_model_flat(MODEL_PATH_FLAT)

# -------------------------------
# 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)
            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() != "flat":
        return JSONResponse(status_code=400, content={"error": "Invalid rooftop_type. Only 'flat' is supported."})

    predictor = predictor_flat
    model_used = "model_flat.pth"

    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
    }