Spaces:
Sleeping
Sleeping
File size: 2,333 Bytes
2797b6f 7040cbd 2797b6f 93b9840 9435820 2797b6f 7040cbd 2797b6f 7040cbd 2797b6f 7040cbd 2797b6f 7040cbd 2797b6f 7040cbd |
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 |
import os
import requests
import cv2
import numpy as np
import torch
import base64
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2 import model_zoo
from PIL import Image
import io
app = FastAPI(title="Roof Segmentation API")
@app.get("/")
def home():
return {"status": "running"}
MODEL_PATH = "model_final (4).pth"
# -----------------------------
# 2. 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.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.MODEL.WEIGHTS = MODEL_PATH
cfg.MODEL.DEVICE = "cpu" # Hugging Face Spaces default (no GPU)
predictor = DefaultPredictor(cfg)
# -----------------------------
# 3. Helper: Encode mask 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))
buf = io.BytesIO()
mask_img.save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode("utf-8")
# -----------------------------
# 4. 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
# Run inference
outputs = predictor(image)
instances = outputs["instances"].to("cpu")
results = []
mask_b64 = None
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)
for i in range(len(masks)):
results.append({
"box": boxes[i].tolist(),
"score": float(scores[i])
})
return JSONResponse({
"predictions": results,
"mask": mask_b64 # base64 string (PNG)
})
|