import gradio as gr
import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt
from celldetection import fetch_model, to_tensor

# ✅ Load the model
device = 'cpu'
model = fetch_model('ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c').to(device).eval()

# ✅ Inference function
def segment(image):
    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
    x = to_tensor(img_rgb, transpose=True, device=device, dtype=torch.float32)[None]

    with torch.no_grad():
        output = model(x)

    contours = output['contours'][0]
    original = (img_rgb * 255).astype(np.uint8).copy()
    segmented = original.copy()

    for contour in contours:
        contour = np.array(contour.cpu(), dtype=np.int32)
        cv2.drawContours(segmented, [contour], -1, (255, 0, 0), 2)

    h, w, c = original.shape
    gap = 60
    canvas = np.zeros((h, w * 2 + gap, c), dtype=np.uint8)
    canvas[:, :w, :] = original
    canvas[:, w + gap:, :] = segmented

    return cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)

# ✅ Example images list
examples = [
    ["1.png"],
    ["2.png"],
    ["3.png"]
]

# ✅ Launch the Gradio interface
gr.Interface(
    fn=segment,
    inputs=gr.Image(type="numpy"),
    outputs="image",
    title="Cell Segmentation Demo (FZJ-INM1)",
    description="Upload a microscopy image to see side-by-side segmentation.",
    examples=examples
).launch()