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)
    })