import gradio as gr
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_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()