67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
import json
|
|
import os
|
|
import torch
|
|
from fast_scnn import FastSCNN, FastSCNNWithNorm
|
|
from utils import carregar_labelmap_completo
|
|
|
|
# ⚙️ Configurações
|
|
with open("config.json", "r") as f:
|
|
config = json.load(f)
|
|
MODELO = config["camera"]
|
|
MODEL_NAME = config["model_name"]
|
|
RESOLUCAO = config["resolucao"]
|
|
MAIN_CLASS_NAME = config["main_class_name"]
|
|
use_main_class = config["use_main_class"]
|
|
model_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
|
|
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
|
|
model_name = f"{MODEL_NAME}_best{f'_f1_{MAIN_CLASS_NAME}' if use_main_class else ''}"
|
|
|
|
dummy_input = torch.randn(1, 3, RESOLUCAO[1], RESOLUCAO[0]) # (batch, channels, height, width)
|
|
|
|
_, _, classes, _ = carregar_labelmap_completo(labelmap_path)
|
|
NUM_CLASSES = len(classes)
|
|
|
|
base = FastSCNNWithNorm(num_classes=NUM_CLASSES, to_rgb=True) # ajuste num_classes conforme seu labelmap
|
|
base.backbone.load_state_dict(torch.load(os.path.join(model_path, f"{MODEL_NAME}_best.pth"), map_location="cpu"))
|
|
base.eval()
|
|
|
|
torch.onnx.export(
|
|
base,
|
|
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[1], RESOLUCAO[0]],
|
|
layout="NCHW",
|
|
)
|
|
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,
|
|
#compile_params=[
|
|
# "-ip U8", # entrada em bytes; compila a conversão interna p/ FP16
|
|
#"--mean_values=[123.675,116.28,103.53]",
|
|
#"--scale_values=[58.395,57.12,57.375]",
|
|
#"--reverse_input_channels" # pq você treinou em RGB
|
|
#],
|
|
)
|
|
print(f"Blob salvo em: {blob_path}") |