Update app.py

app.py

@@ -4,9 +4,16 @@ import cv2
 import gradio as gr
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 from PIL import Image
+from huggingface_hub import hf_hub_download  # 导入 Hugging Face Hub 下载工具
+
+# 下载模型文件
+chkpt_path = hf_hub_download(
+    repo_id="ybelkada/segment-anything",  # 使用 ybelkada 的仓库
+    filename="checkpoints/sam_vit_b_01ec64.pth"  # 下载 sam_vit_b_01ec64.pth 模型
+)
 
 # 加载 Segment Anything 模型
-sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth").to("cuda")
+sam = sam_model_registry["vit_b"](checkpoint=chkpt_path).to("cuda")  # 注意模型类型改为 vit_b
 mask_generator = SamAutomaticMaskGenerator(sam)
 
 def segment(image):
@@ -18,4 +25,4 @@ def segment(image):
 
 # Gradio API
 demo = gr.Interface(fn=segment, inputs="image", outputs="image")
-demo.launch()
+demo.launch()