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()