Update app.py

app.py

@@ -1,12 +1,38 @@
 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 greet(name):
-    return "Hello " + name + "!!"
+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.inputs.Image(type="pil"), 
+    outputs="image",
+    title="Image Segmentation with Huggingface",
+    description="使用Huggingface的SAM-ViT-base模型进行图像分割"
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# 启动接口
+demo.launch()