Spaces:
Sleeping
Sleeping
| 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") | |
| 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 ----------------------------- | |
| 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) | |