import gradio as gr
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
from PIL import Image
import time # --- PERUBAHAN 1: Impor modul 'time' ---

# Muat model
model = AutoModelForImageSegmentation.from_pretrained(
    "briaai/RMBG-2.0",
    trust_remote_code=True
)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"✅ Menggunakan perangkat: {device}")
model.to(device)
model.eval()

# --- PERUBAHAN 2: Ubah ukuran resize menjadi 512 ---
transform_image = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# --- PERUBAHAN 3: Modifikasi fungsi untuk 'yield' (menampilkan status real-time) ---
def remove_background(image: Image.Image):
    """
    Menghilangkan latar belakang dari gambar yang diunggah dan menampilkan timer.
    """
    # Menampilkan status awal di UI
    yield None, "Memulai proses..."

    # Mulai hitung waktu
    start_time = time.time()
    
    img = image.convert("RGB")
    orig_size = img.size

    # Pre‑process dan pindahkan ke device (CPU/GPU)
    input_images = transform_image(img).unsqueeze(0).to(device)
    with torch.no_grad():
        preds = model(input_images)[-1].sigmoid().cpu()

    # Konversi prediksi ke PIL dan ubah ukurannya ke ukuran asli
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(orig_size)

    # Tambahkan channel alpha (latar belakang transparan)
    img.putalpha(mask)

    # Selesai hitung waktu
    end_time = time.time()
    duration = end_time - start_time
    status_text = f"✅ Selesai dalam: {duration:.2f} detik"

    # Mengirimkan hasil akhir (gambar dan teks status) ke UI
    yield img, status_text

# Membuat antarmuka Gradio
iface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil", label="Unggah gambar"),
    # --- PERUBAHAN 4: Sesuaikan outputs untuk menerima gambar dan teks ---
    outputs=[
        gr.Image(type="pil", label="Gambar tanpa latar"),
        gr.Textbox(label="Status")
    ],
    title="BRIA RMBG‑2.0 – Penghapus Latar Belakang",
    description="Model RMBG‑2.0 dari BRIA AI (non-komersial). Unggah gambar dan dapatkan gambar dengan latar transparan."
)

if __name__ == "__main__":
    iface.launch()