import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import numpy as np
import os

# === Load Model ===
try:
    model_path = "best.pt"
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"❌ File {model_path} tidak ditemukan. Upload 'best.pt' ke root.")
    model = YOLO(model_path)
    print("✅ Model loaded successfully!")
except Exception as e:
    print("❌ Gagal load model:", e)
    model = None

# === Peta Label dan Warna ===
label_map = {
    0: "coral or rock",
    1: "pipeline",
    2: "ripple marks",
    3: "shipwreck"
}

color_map = {
    0: (0, 255, 0),      # Hijau
    1: (255, 0, 0),      # Merah
    2: (255, 165, 0),    # Oranye
    3: (0, 0, 255)       # Biru
}

# === Fungsi Prediksi dan Visualisasi ===
def predict_segmentation_with_legend(image):
    try:
        if model is None:
            print("❌ Model belum dimuat.")
            return image

        image = image.convert("RGB")
        print("📥 Gambar diterima:", image.size)

        results = model.predict(image, conf=0.0, iou=0.5)
        result = results[0]

        masks = result.masks.data.cpu().numpy() if result.masks else []
        boxes = result.boxes.xyxy.cpu().numpy() if result.boxes else []
        scores = result.boxes.conf.cpu().numpy() if result.boxes else []
        class_ids = result.boxes.cls.cpu().numpy().astype(int) if result.boxes else []

        image_np = np.array(image).copy()

        try:
            font = ImageFont.truetype("DejaVuSans.ttf", 18)
        except:
            font = ImageFont.load_default()

        CONFIDENCE_THRESHOLD = 0.5

        if len(boxes) == 0:
            draw = ImageDraw.Draw(image)
            draw.text((10, 10), "No detections", fill="red", font=font)
            return image

        # Ambil index dengan skor tertinggi atau di atas threshold
        indices_to_use = [i for i, s in enumerate(scores) if s >= CONFIDENCE_THRESHOLD]
        if len(indices_to_use) == 0:
            indices_to_use = [int(np.argmax(scores))]

        # === MASKING ===
        for i in indices_to_use:
            mask = masks[i]
            class_id = class_ids[i]
            color = color_map.get(class_id, (255, 255, 0))
            mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
            mask_resized = mask_pil.resize((image_np.shape[1], image_np.shape[0]), resample=Image.BILINEAR)
            mask_resized = np.array(mask_resized) / 255.0

            color_mask = np.zeros_like(image_np)
            for c in range(3):
                color_mask[:, :, c] = mask_resized * color[c]

            image_np = np.where(mask_resized[..., None] > 0.5,
                                image_np * 0.5 + color_mask * 0.5, image_np)

        final_image = Image.fromarray(image_np.astype(np.uint8)).convert("RGBA")
        overlay = Image.new("RGBA", final_image.size, (255, 255, 255, 0))
        draw_overlay = ImageDraw.Draw(overlay)

        # === Gambar box dan label ===
        for i in indices_to_use:
            box = boxes[i].astype(int).tolist()
            score = scores[i]
            class_id = class_ids[i]
            label = label_map.get(class_id, str(class_id))
            color = color_map.get(class_id, (255, 255, 0))
            text_color = "white"

            draw_overlay.rectangle(box, outline=color + (255,), width=2)

            # Teks dan background transparan
            label_text = label
            score_text = f"Conf: {score:.2f}"
            label_size = draw_overlay.textbbox((0, 0), label_text, font=font)
            score_size = draw_overlay.textbbox((0, 0), score_text, font=font)

            text_x = max(box[0], 0)
            text_y = max(box[1] - (label_size[3] + score_size[3] + 8), 0)

            draw_overlay.rectangle(
                [text_x - 2, text_y - 2,
                 text_x + max(label_size[2], score_size[2]) + 4,
                 text_y + label_size[3] + score_size[3] + 6],
                fill=(0, 0, 0, 160)
            )

            draw_overlay.text((text_x, text_y), label_text, fill=text_color, font=font)
            draw_overlay.text((text_x, text_y + label_size[3] + 2), score_text, fill=text_color, font=font)

        final_image = Image.alpha_composite(final_image, overlay).convert("RGB")

        # === Buat legenda ===
        legend = Image.new("RGB", (500, 50), (255, 255, 255))
        draw_legend = ImageDraw.Draw(legend)
        x = 10
        for cid, label in label_map.items():
            draw_legend.rectangle([x, 10, x + 20, 30], fill=color_map[cid])
            draw_legend.text((x + 25, 10), label, fill="black", font=font)
            x += 130

        # === Gabungkan gambar & legenda ===
        combined = Image.new("RGB", (final_image.width, final_image.height + 50), (255, 255, 255))
        combined.paste(final_image, (0, 0))
        combined.paste(legend, (10, final_image.height))

        return combined

    except Exception as e:
        print("❌ Error saat segmentasi:", e)
        return image

iface = gr.Interface(
    fn=predict_segmentation_with_legend,
    inputs=gr.Image(type="pil", label="Upload Citra Side Scan Sonar"),
    outputs=gr.Image(type="pil", label="Hasil Segmentasi"),
    title="YOLOv8/YOLOv11 Segmentasi Citra Sonar",
    description="Upload citra sonar dan dapatkan hasil segmentasi lengkap dengan bounding box, mask, label, confidence, dan legenda warna.",
    allow_flagging="never"
)

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