55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
|
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}")
|