Update app.py

app.py

@@ -1,8 +1,21 @@
+import torch
+import numpy as np
+import cv2
 import gradio as gr
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+from PIL import Image
 
-def greet(name):
-    return f"Hello {name}!"
+# 加载 Segment Anything 模型
+sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth").to("cuda")
+mask_generator = SamAutomaticMaskGenerator(sam)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+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)
 
-demo.launch()
+# Gradio API
+demo = gr.Interface(fn=segment, inputs="image", outputs="image")
+demo.launch()