Update app.py

app.py

@@ -3,16 +3,16 @@ from transformers import SamProcessor, SamModel
 from PIL import Image
 import torch
 import numpy as np
-import traceback
 import random
+import traceback
 
-# SAM Modell laden (Meta)
+# Modell laden
 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"""
+    """Zufällige RGB-Farbe"""
     return [random.randint(0, 255) for _ in range(3)]
 
 def segment_image(image):
@@ -20,13 +20,11 @@ def segment_image(image):
         device = torch.device("cpu")
         model.to(device)
 
-        # Eingabe vorbereiten
         inputs = processor(images=image, return_tensors="pt").to(device)
 
         with torch.no_grad():
             outputs = model(**inputs)
 
-        # Masken verarbeiten (ohne Parameternamen)
         masks = processor.post_process_masks(
             outputs.pred_masks.cpu(),
             inputs["original_sizes"].cpu(),
@@ -34,34 +32,32 @@ def segment_image(image):
         )
 
         mask_arrays = masks[0].numpy()
-
-        # Originalbild als Basis
         img_array = np.array(image)
-        color_overlay = np.zeros_like(img_array)
+        overlay = np.zeros_like(img_array, dtype=np.uint8)
 
-        # Jede Maske anders einfärben
-        for i in range(len(mask_arrays)):
-            mask = mask_arrays[i][0]
+        # Jede Maske farbig einfärben
+        for mask in mask_arrays:
+            mask = mask[0]
             color = random_color()
             for c in range(3):
-                color_overlay[:, :, c] += (mask * color[c]).astype(np.uint8)
+                overlay[:, :, c] = np.where(mask > 0.5, color[c], overlay[:, :, c])
 
-        # Farbmischung: 60% Original, 40% Maske
+        # Stärkere Farbmischung (80 % Maske / 20 % Original)
         blended = Image.fromarray(
-            (0.6 * img_array + 0.4 * color_overlay).astype(np.uint8)
+            (0.2 * img_array + 0.8 * overlay).astype(np.uint8)
         )
 
         return blended
 
-    except Exception as e:
+    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 (Meta Model)",
-    description="Zeigt farbige Objekt-Segmente basierend auf Meta's SAM-Modell. Jede erkannte Region erhält eine zufällige Farbe."
+    title="FishBoost – Colorful SAM Segmentation (Enhanced Colors)",
+    description="Erzeugt kräftige, farbige Masken mit Meta SAM (CPU-Version)."
 )
 
 demo.launch()