import gradio as gr
from transformers import AutoProcessor, AutoModelForMaskGeneration
from PIL import Image
import numpy as np

# 加载预训练的processor和模型
processor = AutoProcessor.from_pretrained("facebook/sam-vit-base")
model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-base")

def segment_image(image):
    # 将图片预处理
    inputs = processor(images=image, return_tensors="pt")
    
    # 获取模型输出
    outputs = model(**inputs)
    
    # 提取分割掩码, 假设输出是 logits，您可能需要根据实际模型输出进行调整
    masks = outputs.logits.sigmoid().detach().cpu().numpy()
    
    # 提取第一个掩码作为示例
    mask = masks[0][0]
    
    # 转换掩码为图像
    mask_image = Image.fromarray((mask * 255).astype(np.uint8))
    
    return mask_image

# 创建 Gradio 接口
demo = gr.Interface(
    fn=segment_image, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Image(type="pil"),
    title="Image Segmentation with Huggingface",
    description="使用Huggingface的SAM-ViT-base模型进行图像分割"
)

# 启动接口
demo.launch()