TangYiJay commited on
Commit
bd0e4ed
·
verified ·
1 Parent(s): 7d51e08
Files changed (1) hide show
  1. app.py +60 -0
app.py CHANGED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import requests
3
+ import base64
4
+ import numpy as np
5
+ import cv2
6
+
7
+ # Hugging Face SAM API 配置
8
+ HF_API_URL = "https://api-inference.huggingface.co/models/facebook/sam-vit-base"
9
+ HF_TOKEN = "YOUR_HUGGINGFACE_TOKEN" # 替换成你自己的 token
10
+ headers = {"Authorization": f"Bearer {HF_TOKEN}"}
11
+
12
+
13
+ def segment_and_crop(image):
14
+ # 将输入图像转为 PNG bytes
15
+ _, buffer = cv2.imencode(".png", cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
16
+ img_bytes = buffer.tobytes()
17
+
18
+ # 调用 Hugging Face SAM API
19
+ response = requests.post(HF_API_URL, headers=headers, files={"image": img_bytes})
20
+ if response.status_code != 200:
21
+ return None, f"API error: {response.status_code}\n{response.text}"
22
+
23
+ result = response.json()
24
+
25
+ # SAM 输出 mask(有的模型返回 'masks' base64)
26
+ masks = result.get("masks", [])
27
+ if not masks:
28
+ return None, "No masks returned"
29
+
30
+ cropped_images = []
31
+ for idx, mask_b64 in enumerate(masks):
32
+ mask_bytes = base64.b64decode(mask_b64)
33
+ mask_array = cv2.imdecode(np.frombuffer(mask_bytes, np.uint8), cv2.IMREAD_GRAYSCALE)
34
+ ys, xs = np.where(mask_array > 127)
35
+ if len(xs) == 0 or len(ys) == 0:
36
+ continue
37
+
38
+ # 取边界框并裁剪
39
+ x1, x2 = xs.min(), xs.max()
40
+ y1, y2 = ys.min(), ys.max()
41
+ cropped = image[y1:y2, x1:x2]
42
+ cropped_images.append(cropped)
43
+
44
+ if not cropped_images:
45
+ return None, "No valid crops"
46
+
47
+ # 输出多张裁剪图
48
+ return cropped_images, f"Detected {len(cropped_images)} object(s)."
49
+
50
+
51
+ demo = gr.Interface(
52
+ fn=segment_and_crop,
53
+ inputs=gr.Image(type="numpy", label="Upload or Capture an Image"),
54
+ outputs=[gr.Gallery(label="Cropped Objects"), gr.Textbox(label="Status")],
55
+ title="Smart Object Cropper with SAM",
56
+ description="Upload or capture an image. The app will use the Segment Anything Model (SAM) on Hugging Face to detect and crop main objects.",
57
+ )
58
+
59
+ if __name__ == "__main__":
60
+ demo.launch()