from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.model_zoo import model_zoo
from detectron2.engine import DefaultPredictor
import numpy as np
import torch
import gradio as gr

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = "artifacts/models/model_final.pth"

cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)

# Set the model to inference mode
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # Confidence threshold
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # Update this for custom models
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize the predictor
predictor = DefaultPredictor(cfg)

def predictor_ctrl(img):
    # Perform inference using the predictor
    pred = predictor(img)

    # Visualize predictions
    v = Visualizer(img[:, :, ::-1], scale=0.8)
    out = v.draw_instance_predictions(pred["instances"].to("cpu"))

    # Return the image with predictions (converted back to RGB)
    return out.get_image()[:, :, ::-1]

# Gradio interface
gr.Interface(
    fn=predictor_ctrl,
    inputs=gr.Image(width=720,height=520),  # Gradio provides a NumPy array
    outputs=gr.Image(width=720, height=520),  # Return NumPy array
    examples=[
        "artifacts/car_back.jpg",
        "artifacts/front.png",
        "artifacts/side.webp"
    ]
).launch()