| import gradio as gr |
| import numpy as np |
| import torch |
| import cv2 |
| from PIL import Image |
| import os |
| import urllib.request |
| from segment_anything import sam_model_registry, SamAutomaticMaskGenerator |
|
|
| |
| MODEL_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth" |
| MODEL_PATH = "sam_vit_b_01ec64.pth" |
|
|
| if not os.path.exists(MODEL_PATH): |
| print("Modell wird heruntergeladen...") |
| urllib.request.urlretrieve(MODEL_URL, MODEL_PATH) |
| print("Modell heruntergeladen.") |
|
|
| |
| 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) |
| 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 |
| 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)) |
|
|
| demo = gr.Interface( |
| fn=segment_all_objects, |
| inputs=gr.Image(type="pil", label="Bild hochladen"), |
| outputs=gr.Image(type="pil", label="Segmentiertes Ergebnis"), |
| title="FishBoost SAM (Meta Original)", |
| description="Segmentiert automatisch alle Objekte im Bild mit Metas offiziellem SAM-Modell." |
| ) |
|
|
| demo.launch() |
|
|
