import os
from ultralytics import YOLO

def main():
    print("==================================================")
    print("🚀 YOLO26 Real Evaluation Metrics Calculator 🚀")
    print("==================================================\n")
    
    # Minta pengguna memasukkan path file dataset.yaml
    yaml_path = input("Masukkan path lengkap ke file dataset Anda (contoh: dataset.yaml): ").strip()
    
    if not os.path.exists(yaml_path):
        print(f"\n❌ Error: File '{yaml_path}' tidak ditemukan.")
        print("Pastikan Anda memiliki file .yaml yang berisi konfigurasi dataset YOLO (train/val/test).")
        return

    base_dir = os.path.dirname(os.path.abspath(__file__))
    inst_model_path = os.path.join(base_dir, "yolo26x-seg.pt")
    sem_model_path = os.path.join(base_dir, "yolo26x-sem.pt")
    
    print("\n[1/2] Mengevaluasi YOLO26 Instance Model...")
    if os.path.exists(inst_model_path):
        try:
            inst_model = YOLO(inst_model_path)
            # Menjalankan validasi. Menggunakan split='test' atau 'val' tergantung yaml
            inst_metrics = inst_model.val(data=yaml_path, split='test') # Ganti ke 'val' jika tidak ada test
            
            print("\n✅ HASIL INSTANCE MODEL (yolo26x-seg.pt):")
            print(f"  - mAP@0.5 (Box): {inst_metrics.box.map50:.4f}")
            print(f"  - mAP@0.5:0.95 (Box): {inst_metrics.box.map:.4f}")
            if hasattr(inst_metrics, 'seg'):
                print(f"  - mAP@0.5 (Mask/Seg): {inst_metrics.seg.map50:.4f}")
                print(f"  - mAP@0.5:0.95 (Mask/Seg): {inst_metrics.seg.map:.4f}")
        except Exception as e:
            print(f"Gagal mengevaluasi Instance Model: {e}")
            print("Coba ganti split='test' menjadi split='val' di dalam script jika error dataset.")
    else:
        print(f"❌ Model {inst_model_path} tidak ditemukan.")

    print("\n--------------------------------------------------\n")
    
    print("[2/2] Mengevaluasi YOLO26 Semantic Model...")
    if os.path.exists(sem_model_path):
        try:
            sem_model = YOLO(sem_model_path)
            sem_metrics = sem_model.val(data=yaml_path, split='test') # Ganti ke 'val' jika tidak ada test
            
            print("\n✅ HASIL SEMANTIC MODEL (yolo26x-sem.pt):")
            print(f"  - mAP@0.5 (Box): {sem_metrics.box.map50:.4f}")
            print(f"  - mAP@0.5:0.95 (Box): {sem_metrics.box.map:.4f}")
            if hasattr(sem_metrics, 'seg'):
                print(f"  - mAP@0.5 (Mask/Seg): {sem_metrics.seg.map50:.4f}")
                print(f"  - mAP@0.5:0.95 (Mask/Seg): {sem_metrics.seg.map:.4f}")
        except Exception as e:
            print(f"Gagal mengevaluasi Semantic Model: {e}")
    else:
        print(f"❌ Model {sem_model_path} tidak ditemukan.")

    print("\n==================================================")
    print("Selesai! Silakan gunakan angka di atas untuk mengganti")
    print("angka mockup di draf artikel paper Anda.")
    print("==================================================")

if __name__ == "__main__":
    main()