zedwone commited on
Commit
6792216
·
verified ·
1 Parent(s): 9cb3e33

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -4,7 +4,7 @@ import cv2
4
  import gradio as gr
5
  from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
6
  from PIL import Image
7
- from huggingface_hub import hf_hub_download # 导入 Hugging Face Hub 下载工具
8
 
9
  # 下载模型文件
10
  chkpt_path = hf_hub_download(
@@ -12,8 +12,8 @@ chkpt_path = hf_hub_download(
12
  filename="checkpoints/sam_vit_b_01ec64.pth" # 下载 sam_vit_b_01ec64.pth 模型
13
  )
14
 
15
- # 加载 Segment Anything 模型
16
- sam = sam_model_registry["vit_b"](checkpoint=chkpt_path).to("cuda") # 注意模型类型改为 vit_b
17
  mask_generator = SamAutomaticMaskGenerator(sam)
18
 
19
  def segment(image):
 
4
  import gradio as gr
5
  from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
6
  from PIL import Image
7
+ from huggingface_hub import hf_hub_download
8
 
9
  # 下载模型文件
10
  chkpt_path = hf_hub_download(
 
12
  filename="checkpoints/sam_vit_b_01ec64.pth" # 下载 sam_vit_b_01ec64.pth 模型
13
  )
14
 
15
+ # 加载 Segment Anything 模型(使用 CPU)
16
+ sam = sam_model_registry["vit_b"](checkpoint=chkpt_path).to("cpu") # 改为 CPU
17
  mask_generator = SamAutomaticMaskGenerator(sam)
18
 
19
  def segment(image):