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()