File size: 954 Bytes
af15e39 08e7a4a af15e39 6792216 9cb3e33 08e7a4a 6792216 af15e39 08e7a4a af15e39 bc13c3b af15e39 9cb3e33 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import torch
import numpy as np
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
# 下载模型文件
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 模型（使用 CPU）
sam = sam_model_registry["vit_b"](checkpoint=chkpt_path).to("cpu") # 改为 CPU
mask_generator = SamAutomaticMaskGenerator(sam)
def segment(image):
image = np.array(image)
masks = mask_generator.generate(image)
largest_mask = max(masks, key=lambda x: x['area'])['segmentation']
binary_mask = np.where(largest_mask, 255, 0).astype(np.uint8)
return Image.fromarray(binary_mask)
# Gradio API
demo = gr.Interface(fn=segment, inputs="image", outputs="image")
demo.launch() |