from ultralytics import YOLO
import torch
import os
import sys

# All weights live in models/ (see models/modelREADME.md).
MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")

def main():
    # Usage: python export_model.py [model.pt]  (default: karung-dimuat seg model)
    default_model = os.path.join(MODELS_DIR, "karung-dimuat-detection-di-feedmill-yolo26n-seg-200e.pt")
    model_path = sys.argv[1] if len(sys.argv) > 1 else default_model
    if not os.path.exists(model_path):
        print(f"Error: {model_path} tidak ditemukan!")
        print(f"Usage: python export_model.py [path-ke-model.pt]")
        return

    print("=" * 60)
    print("--- PROSES EKSPOR MODEL KE TENSORRT (.engine) ---")
    print("=" * 60)
    print(f"CUDA Terdeteksi: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU Device Name: {torch.cuda.get_device_name(0)}")
    
    print("\nLoading model ke memori...")
    model = YOLO(model_path)
    
    print("\nMengekspor model ke format TensorRT (FP16)...")
    try:
        # Eksport ke TensorRT (.engine)
        # half=True mengaktifkan kuantisasi FP16 (sangat cepat di GPU Jetson)
        engine_path = model.export(format="engine", device=0, half=True)
        print(f"\n[SUKSES] Model berhasil diekspor ke: {engine_path}")
        print("\nSelanjutnya:")
        print("1. Ganti MODEL_FILE di predict.py menjadi nama file .engine yang baru dibuat.")
        print("2. Jalankan kembali predict.py untuk performa GPU maksimal!")
    except Exception as e:
        print(f"\n[Gagal ekspor langsung ke Engine]: {e}")
        print("\nMencoba metode alternatif: Ekspor ke ONNX terlebih dahulu...")
        try:
            onnx_path = model.export(format="onnx", half=True, dynamic=False, opset=12)
            print(f"[SUKSES] Model berhasil diekspor ke ONNX: {onnx_path}")
            
            onnx_file = os.path.basename(onnx_path)
            engine_file = onnx_file.replace(".onnx", ".engine")
            print("\nAnda bisa mengompilasi file ONNX tersebut ke Engine secara manual di Jetson dengan menjalankan:")
            print(f"  /usr/src/tensorrt/bin/trtexec --onnx={onnx_file} --saveEngine={engine_file} --fp16")
            print("\nSetelah kompilasi manual selesai, ganti MODEL_FILE di predict.py dengan file .engine hasil kompilasi tersebut.")
        except Exception as ex:
            print(f"[Gagal ekspor ke ONNX]: {ex}")

if __name__ == "__main__":
    main()