import subprocess
import os
import ultralytics
import onnx
from os import path

model = ultralytics.YOLO("yolov8n.pt")

model.export(format="onnx", imgsz=640, opset=11)

model = onnx.load("/home/vlad/recycle4me/yolov8n.onnx")
print("Inputs:")
for i in model.graph.input:
    print(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim])
print("Outputs:")
for o in model.graph.output:
    print(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim])

current_dir = os.getcwd()
onnx_path = os.path.join(current_dir, "yolov8n.onnx")
val_data_path = os.path.join(current_dir, "val_data")

result = subprocess.run([
    path.expanduser(".venv/bin/hailomz"),
    "compile",
    "yolov8n",
    "--hw-arch", "hailo8l",
    "--ckpt", onnx_path,
    "--calib-path", val_data_path,
    "--start-node-names", "images",
    "--end-node-names", "output0",
    "--classes", "7",
], check=True)