import cv2
import numpy as np
from ultralytics import YOLO
import sys

def main():
    model_path = "/home/jetson/karung/model_karung_truk.engine"
    print("Loading model...")
    model = YOLO(model_path)
    
    # Warm up
    print("Warming up model...")
    dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
    results = model(dummy, imgsz=640, device="cuda", verbose=False)
    print("Warm up complete!")
    
    # Test tracking
    print("Testing track on 5 frames...")
    for i in range(5):
        print(f"Tracking frame {i+1}...")
        results = model.track(dummy, persist=True, verbose=False)
        print(f"Frame {i+1} track complete! Detections count: {len(results[0])}")
        
    print("All tests passed successfully!")

if __name__ == '__main__':
    main()