Update app.py

app.py

@@ -1,21 +1,43 @@
 import gradio as gr
-from rembg import remove
 from PIL import Image
-import io
-
-# 背景透過処理
-def remove_bg(img):
-    # rembgで背景を削除
-    result = remove(img)
-    return result
-
-demo = gr.Interface(
-    fn=remove_bg,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="背景透過テスト",
-    description="アップロードした画像から背景を消して透過PNGを返します。"
-)
+import torch
+from segment_anything import sam_model_registry, SamPredictor
+
+# SAMの準備（モデルは小さい vit_b を利用）
+sam_checkpoint = "sam_vit_b_01ec64.pth"  # モデルファイル
+model_type = "vit_b"
+
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+predictor = SamPredictor(sam)
+
+def segment(img, points):
+    predictor.set_image(img)
+
+    input_points = []
+    for p in points:
+        input_points.append([p["x"], p["y"]])
+
+    input_labels = [1] * len(input_points)  # クリック位置は前景扱い
+    masks, _, _ = predictor.predict(
+        point_coords=torch.tensor(input_points),
+        point_labels=torch.tensor(input_labels),
+        multimask_output=False,
+    )
+
+    mask = masks[0]
+    out = Image.fromarray((mask * 255).astype("uint8"))
+    return out
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(type="pil", label="画像アップロード")
+            pts = gr.Point(label="クリックして分割位置指定")
+            btn = gr.Button("分割実行")
+        with gr.Column():
+            out = gr.Image(type="pil", label="分割結果")
+
+    btn.click(segment, inputs=[inp, pts], outputs=out)
 
 if __name__ == "__main__":
     demo.launch()