import gradio as gr
import cv2
import numpy as np
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from PIL import Image

# --- Inicialização do modelo 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.7
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
    "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
)
cfg.MODEL.DEVICE = "cpu"

predictor = DefaultPredictor(cfg)

# --- Função de inferência ---
def segment_image(img: Image.Image) -> Image.Image:
    image = np.array(img.convert("RGB"))[:, :, ::-1]  # RGB → BGR
    outputs = predictor(image)
    v = Visualizer(image[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.0)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    result = Image.fromarray(out.get_image()[:, :, ::-1])
    return result

# --- Interface Gradio ---
iface = gr.Interface(
    fn=segment_image,
    inputs=gr.Image(type="pil", label="Imagem de entrada"),
    outputs=gr.Image(type="pil", label="Segmentação Detectron2"),
    title="🎯 Segmentador Detectron2",
    description="Faça upload de uma imagem e veja a segmentação em tempo real."
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)