Update app.py

app.py

@@ -4,33 +4,54 @@ from PIL import Image
 import torch
 import numpy as np
 import traceback
+import random
 
-# SAM Modell laden (CPU-kompatibel)
+# SAM Modell laden (Meta)
 model_id = "facebook/sam-vit-base"
 processor = SamProcessor.from_pretrained(model_id)
 model = SamModel.from_pretrained(model_id)
 
+def random_color():
+    """Erzeugt eine 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)
 
-        # Bild vorbereiten
+        # Eingabe vorbereiten
         inputs = processor(images=image, return_tensors="pt").to(device)
 
         with torch.no_grad():
             outputs = model(**inputs)
 
-        # Die alte API erwartet NICHT das Schlüsselwort 'outputs'
+        # Masken verarbeiten (ohne Parameternamen)
         masks = processor.post_process_masks(
             outputs.pred_masks.cpu(),
             inputs["original_sizes"].cpu(),
             inputs["reshaped_input_sizes"].cpu()
         )
 
-        mask = masks[0][0][0].numpy()
-        mask_image = Image.fromarray((mask * 255).astype(np.uint8))
-        return mask_image
+        mask_arrays = masks[0].numpy()
+
+        # Originalbild als Basis
+        img_array = np.array(image)
+        color_overlay = np.zeros_like(img_array)
+
+        # Jede Maske anders einfärben
+        for i in range(len(mask_arrays)):
+            mask = mask_arrays[i][0]
+            color = random_color()
+            for c in range(3):
+                color_overlay[:, :, c] += (mask * color[c]).astype(np.uint8)
+
+        # Farbmischung: 60% Original, 40% Maske
+        blended = Image.fromarray(
+            (0.6 * img_array + 0.4 * color_overlay).astype(np.uint8)
+        )
+
+        return blended
 
     except Exception as e:
         return f"Fehler:\n{traceback.format_exc()}"
@@ -39,8 +60,8 @@ 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 Segment Anything (Meta SAM - CPU Safe)",
-    description="Stable version for Hugging Face CPU runtime. Uses Meta's SAM model."
+    title="FishBoost – Colorful SAM Segmentation (Meta Model)",
+    description="Zeigt farbige Objekt-Segmente basierend auf Meta's SAM-Modell. Jede erkannte Region erhält eine zufällige Farbe."
 )
 
 demo.launch()