Add application file

app.py

@@ -1,52 +1,40 @@
 import gradio as gr
-import torch
-from segment_anything import sam_model_registry, SamPredictor
 import numpy as np
+import torch
 import cv2
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 from PIL import Image
 
-import os
-import urllib.request
-
-MODEL_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
-MODEL_PATH = "sam_vit_b.pth"
-
-# Eğer model yoksa indir
-if not os.path.exists(MODEL_PATH):
-    print("Model indiriliyor...")
-    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
-    print("Model indirildi.")
-
-# Model yükle 
+# Model yükle
+sam_checkpoint = "sam_vit_b.pth"
 model_type = "vit_b"
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
-sam = sam_model_registry[model_type](checkpoint=MODEL_PATH)
-sam.to(device=device)
-predictor = SamPredictor(sam)
 
-def segment(image, x, y):
-    image = np.array(image)
-    predictor.set_image(image)
-    input_point = np.array([[x, y]])
-    input_label = np.array([1])
-    masks, _, _ = predictor.predict(
-        point_coords=input_point,
-        point_labels=input_label,
-        multimask_output=False,
-    )
-    mask = masks[0]
-    masked_image = image.copy()
-    masked_image[~mask] = 0
-    return Image.fromarray(masked_image)
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        image_input = gr.Image(type="pil")
-        x = gr.Number(label="X")
-        y = gr.Number(label="Y")
-        btn = gr.Button("Segment")
-    output = gr.Image()
-    btn.click(fn=segment, inputs=[image_input, x, y], outputs=output)
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+sam.to(device=device)
 
-demo.launch(share=True)
+mask_generator = SamAutomaticMaskGenerator(sam)
+
+def segment_all_objects(image):
+    image_np = np.array(image)
+    masks = mask_generator.generate(image_np)
+
+    # Maske üzerine çiz
+    overlay = image_np.copy()
+    for i, mask in enumerate(masks):
+        m = mask["segmentation"]
+        color = np.random.randint(0, 255, size=(3,))
+        overlay[m] = overlay[m] * 0.3 + color * 0.7
+        # Maske üstüne label yaz
+        y, x = np.where(m)
+        if len(x) > 0 and len(y) > 0:
+            cx, cy = int(np.mean(x)), int(np.mean(y))
+            cv2.putText(overlay, f"Obj {i+1}", (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
+
+    return Image.fromarray(overlay.astype(np.uint8))
+
+gr.Interface(
+    fn=segment_all_objects,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image()
+).launch()

requirements.txt

@@ -1,7 +1,7 @@
 gradio
 torch
+torchvision
 opencv-python
 numpy
 Pillow
 git+https://github.com/facebookresearch/segment-anything.git
-torchvision