from ultralytics import YOLO
import os
import time

# 1. Load your high-accuracy PC model
model_path = 'best.pt'
if not os.path.exists(model_path):
    print(f"Error: {model_path} not found.")
else:
    model = YOLO(model_path) 

    # 2. Export to TFLite with NMS and Quantization
    # We use a try-except block to handle potential file locking issues during cleanup
    try:
        model.export(
            format='tflite', 
            int8=True, 
            nms=True, 
            imgsz=640,
            data='unified_dataset/data.yaml'
        )
        print("Export command completed.")
    except Exception as e:
        print(f"Export encountered an issue but might have finished: {e}")

    # Check if the model actually exists
    target_path = 'best_saved_model/best_full_integer_quant.tflite'
    if os.path.exists(target_path):
        print(f"Success: {target_path} generated.")
    else:
        # Check other possible names
        tflite_files = []
        if os.path.exists('best_saved_model'):
            for f in os.listdir('best_saved_model'):
                if f.endswith('.tflite'):
                    tflite_files.append(f)
        
        if tflite_files:
            print(f"TFLite files found: {tflite_files}")
        else:
            print("Error: No TFLite model found in best_saved_model/")