Update app.py

app.py

@@ -1,63 +1,52 @@
 import gradio as gr
-from transformers import SamProcessor, SamModel
-from PIL import Image
-import torch
 import numpy as np
-import random
-import traceback
-
-# Modell laden
-model_id = "facebook/sam-vit-base"
-processor = SamProcessor.from_pretrained(model_id)
-model = SamModel.from_pretrained(model_id)
-
-def random_color():
-    """Zufällige RGB-Farbe"""
-    return [random.randint(0, 255) for _ in range(3)]
-
-def segment_image(image):
-    try:
-        device = torch.device("cpu")
-        model.to(device)
-
-        inputs = processor(images=image, return_tensors="pt").to(device)
-
-        with torch.no_grad():
-            outputs = model(**inputs)
-
-        masks = processor.post_process_masks(
-            outputs.pred_masks.cpu(),
-            inputs["original_sizes"].cpu(),
-            inputs["reshaped_input_sizes"].cpu()
-        )
-
-        mask_arrays = masks[0].numpy()
-        img_array = np.array(image)
-        overlay = np.zeros_like(img_array, dtype=np.uint8)
-
-        # Jede Maske farbig einfärben
-        for mask in mask_arrays:
-            mask = mask[0]
-            color = random_color()
-            for c in range(3):
-                overlay[:, :, c] = np.where(mask > 0.5, color[c], overlay[:, :, c])
-
-        # Stärkere Farbmischung (80 % Maske / 20 % Original)
-        blended = Image.fromarray(
-            (0.2 * img_array + 0.8 * overlay).astype(np.uint8)
-        )
-
-        return blended
-
-    except Exception:
-        return f"Fehler:\n{traceback.format_exc()}"
-
-demo = gr.Interface(
-    fn=segment_image,
-    inputs=gr.Image(type="pil", label="Upload your fish image"),
-    outputs=gr.Image(type="pil", label="Segmented Output"),
-    title="FishBoost – Colorful SAM Segmentation (Enhanced Colors)",
-    description="Erzeugt kräftige, farbige Masken mit Meta SAM (CPU-Version)."
-)
-
-demo.launch()
+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_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)
+
+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()