54 lines
2.4 KiB
Python
54 lines
2.4 KiB
Python
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()
|