import os
import torch
from fast_scnn import FastSCNN
from utils import carregar_labelmap_completo

# ⚙️ Configurações
MODELO = "oak-1"
MODEL_NAME = "ervas_full"
RESOLUCAO = (384, 384)
model_path = os.path.join(MODELO, "backup", "fast_scnn", MODEL_NAME)
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
model_name = MODEL_NAME + "_best"

dummy_input = torch.randn(1, 3, RESOLUCAO[0], RESOLUCAO[1])  # (batch, channels, height, width)

_, _, classes, _ = carregar_labelmap_completo(labelmap_path)
NUM_CLASSES = len(classes)

model = FastSCNN(num_classes=NUM_CLASSES)  # ajuste num_classes conforme seu labelmap
model.load_state_dict(torch.load(os.path.join(model_path, model_name + ".pth")))
model.eval()

torch.onnx.export(
    model,
    dummy_input,
    os.path.join(model_path, model_name + ".onnx"),
    input_names=["input"],
    output_names=["output"],
    opset_version=11,
    dynamic_axes=None
)
print(f"Modelo exportado para {os.path.join(model_path, model_name + '.onnx')} com sucesso!")

from openvino.tools.mo import convert_model
from openvino.runtime import serialize
ov_model = convert_model(
    input_model=os.path.join(model_path, model_name + ".onnx"),
    input_shape=[1, 3, RESOLUCAO[0], RESOLUCAO[1]],
)
serialize(
    ov_model,
    os.path.join(model_path, model_name + ".xml"),
    os.path.join(model_path, model_name + ".bin")
)
print("Conversão para IR concluída e arquivos salvos!")

import blobconverter
blob_path = blobconverter.from_openvino(
    xml=os.path.join(model_path, model_name + ".xml"),
    bin=os.path.join(model_path, model_name + ".bin"),
    data_type="FP16",
    shaves=6,
    output_dir=model_path
)
print(f"Blob salvo em: {blob_path}")