from transformers import pipeline
import gradio as gr
import numpy as np


def predict(image):
    result = pipe(image)
    return image, [
        (
            np.array(subsection["mask"]) / 255,
            subsection["label"],
        )
        for subsection in result
    ]


pipe = pipeline(
    "image-segmentation",
    model="SatwikKambham/segformer-b0-finetuned-suim",
)
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(
        type="pil",
        height=400,
    ),
    outputs=gr.AnnotatedImage(
        color_map={
            "Background (waterbody)": "#000000",
            "Human divers": "#0000FF",
            "Aquatic plants and sea-grass": "#00FF00",
            "Wrecks and ruins": "#FF0000",
            "Robots (AUVs/ROVs/instruments)": "#FFFF00",
            "Reefs and invertebrates": "#00FFFF",
            "Fish and vertebrates": "#FF00FF",
            "Sea-floor and rocks": "#FFFFFF",
        },
        height=400,
    ),
    examples="examples",
)

demo.launch()