import torch
import numpy as np
import cv2
import gradio as gr
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
from PIL import Image

# 加载 Segment Anything 模型
sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth").to("cuda")
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()