app.py

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

# Use a pipeline as a high-level helper
pipe = pipeline("mask-generation", model="lightmedsam")

def predict(image_path):
    # Perform image segmentation
    predictions = pipe(image_path)

    # Access the segmented mask or any other relevant information in predictions
    segmented_mask = predictions["segmentation_mask"]

    # Convert the segmentation mask to an RGB image
    segmented_image = colorize_mask(segmented_mask)

    return segmented_image

def colorize_mask(mask):
    # Assuming `mask` is a single-channel segmentation mask (grayscale)
    # You may need to adjust this function based on the specifics of your model's output

    # Convert single-channel mask to 3-channel (RGB) mask
    colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
    colored_mask[:, :, 0] = mask
    colored_mask[:, :, 1] = mask
    colored_mask[:, :, 2] = mask

    return colored_mask

gr.Interface(
    predict,
    inputs=gr.Image(label="Upload medical image", type="filepath"),
    outputs=gr.Image(label="Segmented image"),
    title="Segmented medical image",
    allow_flagging="manual"
).launch()