import gradio as gr
from transformers import AutoImageProcessor, SapiensForImageSegmentation
from PIL import Image


model_name = "facebook/sapiens"
processor = AutoImageProcessor.from_pretrained(model_name)
model = SapiensForImageSegmentation.from_pretrained(model_name)

def segment_image(image):
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    segmentation = outputs.logits.argmax(dim=1).detach().cpu().numpy()[0]
    return Image.fromarray(segmentation)

interface = gr.Interface(
    fn=segment_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Sapiens Body Part Segmentation"
)

interface.launch()