Spaces:
Build error
Build error
File size: 5,895 Bytes
821daa3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import io
import gdown
import base64
from typing import Optional
import cv2
import numpy as np
from PIL import Image
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}
# -------------------------------
# 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 # Only irregular-flat
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("/app/model_irregular_flat.pth")
# -------------------------------
# Post-processing functions
# -------------------------------
def postprocess_rect(mask: np.ndarray, epsilon: float) -> Optional[np.ndarray]:
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)
eps = epsilon * cv2.arcLength(c, True)
approx = cv2.approxPolyDP(c, eps, True)
simp = np.zeros_like(mask_uint8)
cv2.fillPoly(simp, [approx], 255)
return simp
def postprocess_irregular(mask: np.ndarray, epsilon: float) -> Optional[np.ndarray]:
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)
eps = epsilon * cv2.arcLength(c, True)
polygon = cv2.approxPolyDP(c, eps, True)
return polygon.reshape(-1, 2)
def mask_to_polygon(mask: np.ndarray) -> Optional[np.ndarray]:
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return None
largest = max(contours, key=cv2.contourArea)
return largest.reshape(-1, 2)
def im_to_b64_png(im: np.ndarray) -> str:
_, buffer = cv2.imencode(".png", im)
return base64.b64encode(buffer).decode()
def overlay_polygon(im: np.ndarray, polygon: Optional[np.ndarray]) -> np.ndarray:
overlay = im.copy()
if polygon is not None:
cv2.polylines(overlay, [polygon.astype(np.int32)], True, (0,0,255), 2)
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(...), # "rectangular" or "irregular"
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
# Choose predictor and post-processing based on rooftop type
if rooftop_type.lower() == "rectangular":
predictor = predictor_rect
post_fn = lambda mask: postprocess_rect(mask, epsilon)
model_used = "model_rect_final.pth"
elif rooftop_type.lower() == "irregular":
predictor = predictor_irregular_flat
post_fn = lambda mask: postprocess_irregular(mask, epsilon)
model_used = "model_irregular_flat.pth"
else:
return JSONResponse(status_code=400, content={"error": "Invalid rooftop_type. Choose 'rectangular' or 'irregular'."})
# Run prediction
outputs = predictor(im)
instances = outputs["instances"].to("cpu")
if len(instances) == 0:
return {"polygon": None, "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)
# Post-process
result_mask = post_fn(raw_mask)
polygon = mask_to_polygon(result_mask) if rooftop_type.lower() == "rectangular" else result_mask
# Overlay
overlay = overlay_polygon(im, polygon)
img_b64 = im_to_b64_png(overlay)
return {
"polygon": polygon.tolist() if polygon is not None else None,
"image": img_b64,
"model_used": model_used,
"rooftop_type": rooftop_type,
"epsilon": epsilon
} |