Implementação ONNX/TensorRT

Python/OAK/datasets/_10_export_onnx.py

@@ -0,0 +1,841 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+_10_export_visual_onnx.py
+
+Exporta o checkpoint PyTorch do Visual Worker para ONNX.
+
+Contrato esperado do modelo visual:
+  - Backbone SegFormer Hugging Face, ex: nvidia/mit-b0
+  - Entrada RGB: [N, 3, H, W], float32
+  - Cabeça 1: segmentação semântica do corredor
+  - Cabeça 2: classificação global/status do corredor
+
+Suporta:
+  - ONNX com logits crus de segmentação
+  - ONNX com logits de segmentação redimensionados para HxW
+  - ONNX com máscara argmax de segmentação
+  - ONNX com label_logits ou label_probs
+  - Normalização embutida no grafo: x = (x - mean) / std
+
+Exemplos:
+
+# ONNX recomendado para runtime: recebe RGB 0..1 e já normaliza internamente.
+python _10_export_visual_onnx.py ^
+  --config config.json ^
+  --checkpoint oak-d/backup/segformer_b0/nav_mit_dual_label/best_label.pt ^
+  --out oak-d/backup/segformer_b0/nav_mit_dual_label/best_label.onnx ^
+  --device cuda ^
+  --include-norm ^
+  --semantic-postprocess resize_logits ^
+  --label-postprocess probs
+
+# ONNX cru: runtime precisa enviar tensor já normalizado.
+python _10_export_visual_onnx.py ^
+  --config config.json ^
+  --device cuda ^
+  --semantic-postprocess none ^
+  --label-postprocess logits
+"""
+
+from __future__ import annotations
+
+import json
+import argparse
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import SegformerForSemanticSegmentation
+
+
+# ============================================================
+# Utils
+# ============================================================
+
+
+def load_json(path: str | Path) -> dict:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def save_json(path: str | Path, data: dict):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+
+def ensure_dir(path: str | Path):
+    Path(path).mkdir(parents=True, exist_ok=True)
+
+
+def resolve_path(path_like: Optional[str], base: Optional[Path] = None) -> Optional[Path]:
+    if path_like is None or str(path_like).strip() == "":
+        return None
+
+    p = Path(path_like)
+    if p.is_absolute():
+        return p
+
+    if base is None:
+        base = Path.cwd()
+
+    return (base / p).resolve()
+
+
+def _try_int(text: str) -> Optional[int]:
+    try:
+        return int(str(text).strip())
+    except Exception:
+        return None
+
+
+def clean_label_name(raw_name: str) -> str:
+    """
+    Limpa nomes vindos de labelmap no estilo:
+      naonavegavel:128,0,0::
+      navegavel:0,128,0::
+
+    Mantém apenas o nome lógico da classe.
+    """
+    name = str(raw_name).strip()
+
+    # Remove sufixos visuais comuns do labelmap: classe:R,G,B::
+    if "::" in name:
+        name = name.split("::", 1)[0].strip()
+
+    if ":" in name:
+        left, right = name.split(":", 1)
+        right_clean = right.replace(",", "").replace(" ", "")
+        if right_clean.isdigit():
+            name = left.strip()
+
+    return name.strip()
+
+
+def load_labelmap(labelmap_path: Path) -> Tuple[Dict[int, str], Dict[str, int], int]:
+    """
+    Parser tolerante para labelmap.txt.
+
+    Aceita formatos comuns:
+      navegavel
+      0:navegavel
+      0 navegavel
+      0,navegavel
+      navegavel:0
+      navegavel:0,128,0::
+
+    Retorna:
+      id2label, label2id, ignore_index
+    """
+    if not labelmap_path.is_file():
+        raise FileNotFoundError(f"Labelmap não encontrado: {labelmap_path}")
+
+    id2label: Dict[int, str] = {}
+    ignore_index = 255
+    next_id = 0
+
+    with labelmap_path.open("r", encoding="utf-8") as f:
+        for raw_line in f:
+            line = raw_line.strip()
+            if not line or line.startswith("#"):
+                continue
+
+            lower = line.lower()
+            if lower.startswith("ignore") or lower.startswith("ignore_index"):
+                for sep in ("=", ":", ",", " "):
+                    if sep in line:
+                        maybe = _try_int(line.split(sep)[-1])
+                        if maybe is not None:
+                            ignore_index = maybe
+                            break
+                continue
+
+            cls_id: Optional[int] = None
+            cls_name: Optional[str] = None
+
+            # Caso CVAT/labelmap visual: nome:R,G,B::
+            # Exemplo: navegavel:0,128,0::
+            if "::" in line and ":" in line:
+                before = line.split("::", 1)[0].strip()
+                maybe_name = before.split(":", 1)[0].strip()
+                if maybe_name:
+                    cls_id = next_id
+                    cls_name = maybe_name
+
+            if cls_id is None:
+                for sep in (":", ",", "	", " "):
+                    if sep in line:
+                        parts = [p.strip() for p in line.split(sep) if p.strip()]
+                        if len(parts) >= 2:
+                            left_id = _try_int(parts[0])
+                            right_id = _try_int(parts[-1])
+
+                            if left_id is not None:
+                                cls_id = left_id
+                                cls_name = sep.join(parts[1:]).strip() if sep in (":", ",") else " ".join(parts[1:]).strip()
+                                break
+
+                            if right_id is not None:
+                                cls_id = right_id
+                                cls_name = sep.join(parts[:-1]).strip() if sep in (":", ",") else " ".join(parts[:-1]).strip()
+                                break
+
+            if cls_id is None:
+                cls_id = next_id
+                cls_name = line
+
+            if cls_name is None or cls_name == "":
+                raise RuntimeError(f"Linha inválida no labelmap: {raw_line!r}")
+
+            id2label[int(cls_id)] = clean_label_name(str(cls_name))
+            next_id = max(next_id, int(cls_id) + 1)
+
+    if not id2label:
+        raise RuntimeError(f"Labelmap vazio ou inválido: {labelmap_path}")
+
+    # Garante ids contínuos para o SegFormer.
+    ids_sorted = sorted(id2label.keys())
+    if ids_sorted != list(range(len(ids_sorted))):
+        remap = {old_id: new_id for new_id, old_id in enumerate(ids_sorted)}
+        id2label = {remap[old_id]: name for old_id, name in id2label.items()}
+
+    label2id = {name: idx for idx, name in id2label.items()}
+    return id2label, label2id, ignore_index
+
+
+def get_visual_mode(config: dict) -> str:
+    use_mask2 = bool(config.get("dual_head_mask", config.get("dual_head", False)))
+    use_label = bool(config.get("dual_head_label", False))
+
+    if use_mask2 and use_label:
+        raise RuntimeError("Config inválido: dual_head_mask e dual_head_label ativos juntos.")
+
+    if use_label:
+        return "label"
+    if use_mask2:
+        return "mask2"
+    return "single"
+
+
+def get_save_suffix(mode: str) -> str:
+    if mode == "single":
+        return "_single"
+    if mode == "mask2":
+        return "_dual_mask"
+    if mode == "label":
+        return "_dual_label"
+    raise RuntimeError(f"Modo desconhecido: {mode}")
+
+
+def resolve_default_paths(args, config: dict, config_dir: Path) -> Tuple[Path, Path, str, str, Path]:
+    """
+    Resolve checkpoint e ONNX padrão seguindo o padrão do treino visual:
+      {camera}/backup/{modelo}/{model_name}_{suffix}/{ckpt_name}.pt
+    """
+    camera = str(config.get("camera", "oak-d"))
+    model_key = str(config.get("modelo", "segformer_b0"))
+    model_name = str(config.get("model_name", "visual"))
+    mode = get_visual_mode(config)
+    suffix = get_save_suffix(mode)
+
+    save_dir = config_dir / camera / "backup" / model_key / f"{model_name}{suffix}"
+
+    if mode == "label":
+        default_ckpt_name = "best_label"
+    elif mode == "mask2":
+        default_ckpt_name = "best_mask2"
+    else:
+        default_ckpt_name = "best_main"
+
+    ckpt_name = str(config.get("ckpt_test", default_ckpt_name))
+
+    checkpoint_path = resolve_path(args.checkpoint, Path.cwd())
+    if checkpoint_path is None:
+        checkpoint_path = save_dir / f"{ckpt_name}.pt"
+
+    out_path = resolve_path(args.out, Path.cwd())
+    if out_path is None:
+        out_path = save_dir / f"{ckpt_name}.onnx"
+
+    return checkpoint_path.resolve(), out_path.resolve(), ckpt_name, mode, save_dir.resolve()
+
+
+def resolve_labelmap_path(args, config: dict, config_dir: Path) -> Path:
+    explicit = resolve_path(args.labelmap, Path.cwd())
+    if explicit is not None:
+        return explicit.resolve()
+
+    camera = str(config.get("camera", "oak-d"))
+    return (config_dir / camera / "dataset" / "labelmap.txt").resolve()
+
+
+def resolve_norm_stats_path(args, config: dict, config_dir: Path, save_dir: Path) -> Optional[Path]:
+    explicit = resolve_path(args.norm_stats, Path.cwd())
+    if explicit is not None:
+        return explicit.resolve()
+
+    # Prioridade 1: norm_stats salvo junto ao treino/checkpoint.
+    p = save_dir / "norm_stats.json"
+    if p.is_file():
+        return p.resolve()
+
+    # Prioridade 2: dataset normalizado na resolução do contrato.
+    W, H = config.get("resolucao", [1024, 640])
+    camera = str(config.get("camera", "oak-d"))
+    p = config_dir / camera / "dataset" / f"{int(W)}x{int(H)}" / "group" / "norm_stats.json"
+    if p.is_file():
+        return p.resolve()
+
+    # Prioridade 3: deixa explícito no erro quando --include-norm for usado.
+    return p.resolve()
+
+
+def load_rgb_norm_stats(path: Path) -> Tuple[List[float], List[float], List[str]]:
+    if not path.is_file():
+        raise FileNotFoundError(f"norm_stats não encontrado: {path}")
+
+    js = load_json(path)
+    mean = js.get("mean")
+    std = js.get("std")
+    names = js.get("channels", [])
+
+    if mean is None or std is None:
+        raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
+
+    if names:
+        # Visual worker espera RGB. Aceita stats com canais extras, desde que R,G,B existam.
+        name_to_idx = {str(n).upper(): i for i, n in enumerate(names)}
+        required = ["R", "G", "B"]
+        missing = [ch for ch in required if ch not in name_to_idx]
+        if missing:
+            raise RuntimeError(
+                f"norm_stats incompatível: faltam canais {missing}. "
+                f"channels={names} path={path}"
+            )
+        idx = [name_to_idx[ch] for ch in required]
+        mean_sel = [float(mean[i]) for i in idx]
+        std_sel = [float(std[i]) for i in idx]
+        names_sel = required
+    else:
+        if len(mean) < 3 or len(std) < 3:
+            raise RuntimeError(f"norm_stats precisa de pelo menos 3 valores RGB: {path}")
+        mean_sel = [float(mean[i]) for i in range(3)]
+        std_sel = [float(std[i]) for i in range(3)]
+        names_sel = ["R", "G", "B"]
+
+    return mean_sel, std_sel, names_sel
+
+
+def resolve_label_classes(config: dict, ckpt: Optional[dict] = None) -> Tuple[Dict[int, str], int]:
+    """Resolve classes da cabeça de status/corredor."""
+    label_classes = config.get("label_classes", None)
+
+    if label_classes is not None:
+        label_name_by_id = {i: str(name) for i, name in enumerate(label_classes)}
+        return label_name_by_id, len(label_name_by_id)
+
+    if ckpt is not None:
+        extra = ckpt.get("extra", {}) if isinstance(ckpt, dict) else {}
+        maybe = extra.get("label_name_by_id", None)
+        if isinstance(maybe, dict) and maybe:
+            label_name_by_id = {int(k): str(v) for k, v in maybe.items()}
+            return label_name_by_id, max(label_name_by_id.keys()) + 1
+
+        maybe_config = extra.get("config", {}) if isinstance(extra, dict) else {}
+        maybe_classes = maybe_config.get("label_classes", None) if isinstance(maybe_config, dict) else None
+        if maybe_classes is not None:
+            label_name_by_id = {i: str(name) for i, name in enumerate(maybe_classes)}
+            return label_name_by_id, len(label_name_by_id)
+
+    raise RuntimeError(
+        "Não consegui resolver label_classes. "
+        "Adicione config['label_classes'] com a lista de status do corredor, "
+        "ou use um checkpoint que tenha extra['label_name_by_id']."
+    )
+
+
+# ============================================================
+# Cabeça auxiliar igual ao treino visual
+# ============================================================
+
+
+class LabelHead(nn.Module):
+    """Head de classificação global do frame/status do corredor."""
+
+    def __init__(
+        self,
+        feat_ch: int,
+        num_seg_classes: int,
+        num_label_classes: int,
+        hidden: int = 256,
+        dropout: float = 0.2,
+    ):
+        super().__init__()
+        in_ch = int(feat_ch) + int(num_seg_classes)
+        self.in_ch = in_ch
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.net = nn.Sequential(
+            nn.Linear(in_ch, hidden),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, num_label_classes),
+        )
+
+    def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
+        # Interpola sempre para evitar branch Python durante o trace ONNX.
+        # Se o tamanho já for igual, o resultado é equivalente e o grafo fica estável.
+        feat = F.interpolate(
+            feat,
+            size=logits_seg.shape[-2:],
+            mode="bilinear",
+            align_corners=False,
+        )
+        x = torch.cat([feat, logits_seg], dim=1)
+        x = self.pool(x).flatten(1)
+        return self.net(x)
+
+
+def get_last_feat(out, logits: torch.Tensor) -> torch.Tensor:
+    if hasattr(out, "hidden_states") and out.hidden_states is not None:
+        feat = out.hidden_states[-1]
+    else:
+        feat = logits
+
+    # Interpola sempre para evitar TracerWarning por comparação de shapes.
+    # O contrato do export é resolução fixa, então isso não muda o comportamento útil.
+    feat = F.interpolate(
+        feat,
+        size=logits.shape[-2:],
+        mode="bilinear",
+        align_corners=False,
+    )
+    return feat
+
+
+# ============================================================
+# Modelo composto Visual
+# ============================================================
+
+
+class VisualSegformerDualLabel(nn.Module):
+    """
+    Modelo composto para export:
+      - base_model: SegformerForSemanticSegmentation
+      - label_head: LabelHead treinada junto
+
+    forward retorna dict:
+      semantic: [N, Cseg, h, w]
+      label   : [N, Cstatus]
+    """
+
+    def __init__(self, base_model: nn.Module, label_head: nn.Module):
+        super().__init__()
+        self.base_model = base_model
+        self.label_head = label_head
+
+    def forward(self, pixel_values: torch.Tensor) -> Dict[str, torch.Tensor]:
+        out = self.base_model(pixel_values=pixel_values)
+        logits_seg = out.logits
+        feat = get_last_feat(out, logits_seg)
+        logits_label = self.label_head(feat, logits_seg)
+        return {
+            "semantic": logits_seg,
+            "label": logits_label,
+        }
+
+
+class VisualOnnxWrapper(nn.Module):
+    """
+    Wrapper final exportável para ONNX.
+
+    semantic_postprocess:
+      - none          -> semantic_logits em baixa resolução do decode_head
+      - resize_logits -> semantic_logits em HxW da entrada
+      - argmax_lowres -> semantic_mask em baixa resolução
+      - argmax_fullres-> semantic_mask em HxW da entrada
+
+    label_postprocess:
+      - logits -> label_logits
+      - probs  -> label_probs via softmax
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        include_norm: bool = False,
+        norm_mean: Optional[List[float]] = None,
+        norm_std: Optional[List[float]] = None,
+        semantic_postprocess: str = "none",
+        label_postprocess: str = "logits",
+    ):
+        super().__init__()
+        self.model = model
+        self.include_norm = bool(include_norm)
+        self.semantic_postprocess = str(semantic_postprocess).lower()
+        self.label_postprocess = str(label_postprocess).lower()
+
+        if self.semantic_postprocess not in ("none", "resize_logits", "argmax_lowres", "argmax_fullres"):
+            raise RuntimeError(f"semantic_postprocess inválido: {self.semantic_postprocess}")
+
+        if self.label_postprocess not in ("logits", "probs"):
+            raise RuntimeError(f"label_postprocess inválido: {self.label_postprocess}")
+
+        if self.include_norm:
+            if norm_mean is None or norm_std is None:
+                raise RuntimeError("include_norm=True requer norm_mean/norm_std")
+            mean_t = torch.tensor(norm_mean, dtype=torch.float32).view(1, 3, 1, 1)
+            std_t = torch.tensor(norm_std, dtype=torch.float32).view(1, 3, 1, 1)
+            self.register_buffer("norm_mean", mean_t)
+            self.register_buffer("norm_std", std_t)
+        else:
+            self.register_buffer("norm_mean", torch.empty(0))
+            self.register_buffer("norm_std", torch.empty(0))
+
+    def forward(self, pixel_values: torch.Tensor):
+        input_hw = pixel_values.shape[-2:]
+
+        x = pixel_values
+        if self.include_norm:
+            x = (x - self.norm_mean) / torch.clamp(self.norm_std, min=1e-6)
+
+        outputs = self.model(pixel_values=x)
+        semantic = outputs["semantic"]
+        label = outputs["label"]
+
+        if self.semantic_postprocess == "none":
+            semantic_out = semantic
+        elif self.semantic_postprocess == "resize_logits":
+            semantic_out = F.interpolate(
+                semantic,
+                size=input_hw,
+                mode="bilinear",
+                align_corners=False,
+            )
+        elif self.semantic_postprocess == "argmax_lowres":
+            semantic_out = torch.argmax(semantic, dim=1).to(torch.uint8)
+        elif self.semantic_postprocess == "argmax_fullres":
+            semantic = F.interpolate(
+                semantic,
+                size=input_hw,
+                mode="bilinear",
+                align_corners=False,
+            )
+            semantic_out = torch.argmax(semantic, dim=1).to(torch.uint8)
+        else:
+            raise RuntimeError("semantic_postprocess inválido")
+
+        if self.label_postprocess == "probs":
+            label_out = torch.softmax(label, dim=1)
+        else:
+            label_out = label
+
+        return semantic_out, label_out
+
+
+# ============================================================
+# Build/load
+# ============================================================
+
+
+def build_visual_model(
+    backbone: str,
+    num_seg_classes: int,
+    num_label_classes: int,
+    device: torch.device,
+    input_hw: Tuple[int, int],
+) -> VisualSegformerDualLabel:
+    H, W = input_hw
+
+    base_model = SegformerForSemanticSegmentation.from_pretrained(
+        backbone,
+        num_labels=int(num_seg_classes),
+        ignore_mismatched_sizes=True,
+        use_safetensors=True,
+    )
+    base_model.config.output_hidden_states = True
+    base_model.to(device)
+    base_model.eval()
+
+    with torch.no_grad():
+        dummy = torch.zeros((1, 3, int(H), int(W)), dtype=torch.float32, device=device)
+        out = base_model(pixel_values=dummy)
+        logits = out.logits
+        feat = get_last_feat(out, logits)
+        feat_ch = int(feat.shape[1])
+
+    label_head = LabelHead(
+        feat_ch=feat_ch,
+        num_seg_classes=int(num_seg_classes),
+        num_label_classes=int(num_label_classes),
+        hidden=256,
+        dropout=0.2,
+    ).to(device)
+    label_head.eval()
+
+    return VisualSegformerDualLabel(base_model=base_model, label_head=label_head).to(device)
+
+
+def load_checkpoint_into_model(model: VisualSegformerDualLabel, checkpoint_path: Path):
+    ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
+
+    if "model" not in ckpt:
+        raise RuntimeError(
+            f"Checkpoint não contém chave 'model': {checkpoint_path}. "
+            "Confirme se foi salvo pelo script de treino visual."
+        )
+
+    if "aux_head" not in ckpt:
+        raise RuntimeError(
+            f"Checkpoint não contém chave 'aux_head': {checkpoint_path}. "
+            "Este exportador é para o modelo visual dual_head_label."
+        )
+
+    model.base_model.load_state_dict(ckpt["model"], strict=True)
+    model.label_head.load_state_dict(ckpt["aux_head"], strict=True)
+    return ckpt
+
+
+# ============================================================
+# Main
+# ============================================================
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--config", default="config.json")
+    parser.add_argument("--checkpoint", default="")
+    parser.add_argument("--out", default="")
+    parser.add_argument("--labelmap", default="")
+    parser.add_argument("--norm_stats", default="")
+
+    parser.add_argument("--opset", type=int, default=17)
+    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
+    parser.add_argument("--dynamic-batch", action="store_true")
+
+    parser.add_argument(
+        "--include-norm",
+        action="store_true",
+        help="Inclui normalização RGB x=(x-mean)/std dentro do ONNX. Runtime deve enviar RGB 0..1.",
+    )
+
+    parser.add_argument(
+        "--semantic-postprocess",
+        default="none",
+        choices=["none", "resize_logits", "argmax_lowres", "argmax_fullres"],
+        help="Define a saída semântica exportada.",
+    )
+
+    parser.add_argument(
+        "--label-postprocess",
+        default="logits",
+        choices=["logits", "probs"],
+        help="Define a saída da cabeça de status/corredor.",
+    )
+
+    args = parser.parse_args()
+
+    config_path = resolve_path(args.config, Path.cwd())
+    if config_path is None or not config_path.is_file():
+        raise FileNotFoundError(f"Config não encontrado: {config_path}")
+
+    config_dir = config_path.parent
+    config = load_json(config_path)
+
+    checkpoint_path, out_path, ckpt_name, mode, save_dir = resolve_default_paths(
+        args=args,
+        config=config,
+        config_dir=config_dir,
+    )
+
+    if mode != "label":
+        raise RuntimeError(
+            f"Este exportador foi preparado para dual_head_label. "
+            f"Modo detectado no config: {mode}"
+        )
+
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Checkpoint não encontrado: {checkpoint_path}")
+
+    ensure_dir(out_path.parent)
+
+    labelmap_path = resolve_labelmap_path(args, config, config_dir)
+    semantic_id2label, semantic_label2id, ignore_index = load_labelmap(labelmap_path)
+    num_seg_classes = len(semantic_id2label)
+
+    # Lê ckpt antes para recuperar label_name_by_id, se necessário.
+    ckpt_for_meta = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
+    label_name_by_id, num_label_classes = resolve_label_classes(config, ckpt_for_meta)
+
+    W, H = config.get("resolucao", [1024, 640])
+    W = int(W)
+    H = int(H)
+    backbone = str(config.get("backbone", "nvidia/mit-b0"))
+
+    norm_stats_path = None
+    norm_mean = None
+    norm_std = None
+    norm_channels = ["R", "G", "B"]
+
+    if args.include_norm:
+        norm_stats_path = resolve_norm_stats_path(args, config, config_dir, save_dir)
+        if norm_stats_path is None:
+            raise RuntimeError("include_norm=True, mas norm_stats_path não foi resolvido.")
+        norm_mean, norm_std, norm_channels = load_rgb_norm_stats(norm_stats_path)
+
+    use_cuda = args.device == "cuda" and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    if args.device == "cuda" and not torch.cuda.is_available():
+        print("[WARN] CUDA não disponível. Exportando em CPU.")
+
+    print("==========================================")
+    print("Export Visual Worker SegFormer Dual Label para ONNX")
+    print(f"Config              : {config_path}")
+    print(f"Checkpoint          : {checkpoint_path}")
+    print(f"Output ONNX         : {out_path}")
+    print(f"Save dir            : {save_dir}")
+    print(f"Backbone            : {backbone}")
+    print(f"Resolution          : {W}x{H}")
+    print(f"Input shape         : [1, 3, {H}, {W}]")
+    print(f"Semantic classes    : {num_seg_classes} {semantic_id2label}")
+    print(f"Label classes       : {num_label_classes} {label_name_by_id}")
+    print(f"Include norm        : {args.include_norm}")
+    print(f"Norm stats          : {norm_stats_path}")
+    print(f"Semantic output     : {args.semantic_postprocess}")
+    print(f"Label output        : {args.label_postprocess}")
+    print(f"Dynamic batch       : {args.dynamic_batch}")
+    print(f"Opset               : {args.opset}")
+    print(f"Device              : {device}")
+    print("==========================================")
+
+    print("[MODEL] Montando modelo visual...")
+    model = build_visual_model(
+        backbone=backbone,
+        num_seg_classes=num_seg_classes,
+        num_label_classes=num_label_classes,
+        device=device,
+        input_hw=(H, W),
+    )
+
+    print("[CKPT] Carregando checkpoint...")
+    ckpt = load_checkpoint_into_model(model, checkpoint_path)
+    model.to(device)
+    model.eval()
+
+    wrapper = VisualOnnxWrapper(
+        model=model,
+        include_norm=args.include_norm,
+        norm_mean=norm_mean,
+        norm_std=norm_std,
+        semantic_postprocess=args.semantic_postprocess,
+        label_postprocess=args.label_postprocess,
+    ).to(device)
+    wrapper.eval()
+
+    dummy_input = torch.randn(
+        1,
+        3,
+        H,
+        W,
+        dtype=torch.float32,
+        device=device,
+    )
+
+    input_names = ["pixel_values"]
+
+    if args.semantic_postprocess in ("argmax_lowres", "argmax_fullres"):
+        semantic_output_name = "semantic_mask"
+    else:
+        semantic_output_name = "semantic_logits"
+
+    label_output_name = "label_probs" if args.label_postprocess == "probs" else "label_logits"
+    output_names = [semantic_output_name, label_output_name]
+
+    dynamic_axes = None
+    if args.dynamic_batch:
+        dynamic_axes = {
+            "pixel_values": {0: "batch"},
+            semantic_output_name: {0: "batch"},
+            label_output_name: {0: "batch"},
+        }
+
+    print("[CHECK] Rodando forward PyTorch antes do export...")
+    with torch.no_grad():
+        y = wrapper(dummy_input)
+
+    for name, tensor in zip(output_names, y):
+        print(f"  {name}: shape={tuple(tensor.shape)} dtype={tensor.dtype}")
+
+    print("[EXPORT] Exportando ONNX...")
+    with torch.no_grad():
+        torch.onnx.export(
+            wrapper,
+            dummy_input,
+            str(out_path),
+            export_params=True,
+            opset_version=int(args.opset),
+            do_constant_folding=True,
+            input_names=input_names,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+        )
+
+    print(f"[OK] ONNX salvo em: {out_path}")
+
+    meta_path = out_path.with_suffix(".export_meta.json")
+    save_json(meta_path, {
+        "kind": "visual_worker_segformer_dual_label",
+        "config": str(config_path),
+        "checkpoint": str(checkpoint_path),
+        "onnx": str(out_path),
+        "ckpt_name": ckpt_name,
+        "mode": mode,
+        "backbone": backbone,
+        "input_shape": [1, 3, H, W],
+        "input_channel_names": ["R", "G", "B"],
+        "input_channel_indices": [0, 1, 2],
+        "semantic_output_name": semantic_output_name,
+        "label_output_name": label_output_name,
+        "output_names": output_names,
+        "semantic_postprocess": str(args.semantic_postprocess),
+        "label_postprocess": str(args.label_postprocess),
+        "opset": int(args.opset),
+        "dynamic_batch": bool(args.dynamic_batch),
+        "include_norm": bool(args.include_norm),
+        "norm_stats_path": str(norm_stats_path) if norm_stats_path is not None else None,
+        "norm_channels": norm_channels,
+        "norm_mean": norm_mean,
+        "norm_std": norm_std,
+        "semantic_id2label": semantic_id2label,
+        "semantic_label2id": semantic_label2id,
+        "ignore_index": int(ignore_index),
+        "label_name_by_id": label_name_by_id,
+        "num_label_classes": int(num_label_classes),
+        "checkpoint_epoch": ckpt.get("epoch", None),
+        "checkpoint_bests": ckpt.get("bests", None),
+        "checkpoint_extra": ckpt.get("extra", None),
+    })
+    print(f"[OK] Metadata salvo em: {meta_path}")
+
+    try:
+        import onnx
+        print("[ONNX] Verificando grafo com onnx.checker...")
+        onnx_model = onnx.load(str(out_path))
+        onnx.checker.check_model(onnx_model)
+        print("[OK] onnx.checker passou.")
+    except ImportError:
+        print("[WARN] Pacote onnx não instalado. Pulei onnx.checker.")
+        print("       Instale com: pip install onnx")
+    except Exception as e:
+        print(f"[WARN] onnx.checker encontrou problema: {e}")
+
+    print("\nExport finalizado.")
+
+
+if __name__ == "__main__":
+    main()

Python/OAK/datasets/_9_test_segformer_dual.py

@@ -49,10 +49,14 @@ from transformers import SegformerForSemanticSegmentation
 # Utils básicos
 # ============================================================
 
-IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+# Normalização padrão. Pode ser sobrescrita por norm_stats.json.
-IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+NORM_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+NORM_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+
 def normalize_img(img: torch.Tensor) -> torch.Tensor:
-    return (img - IMAGENET_MEAN.to(img.device)) / IMAGENET_STD.to(img.device)
+    mean = NORM_MEAN.to(img.device)
+    std = NORM_STD.to(img.device).clamp_min(1e-6)
+    return (img - mean) / std
 
 
 def find_device():
@@ -477,6 +481,229 @@ def infer_sample(base_model, aux_head, img_bgr: np.ndarray, device: torch.device
     return result
 
 
+def add_runtime_dll_dirs(trt_home: Optional[str] = None):
+    trt_home = trt_home or os.environ.get("TRT_HOME", r"C:\dev\TensorRT-10.10.0.31")
+
+    dll_dirs = [
+        os.path.join(trt_home, "lib"),
+        os.path.join(trt_home, "bin"),
+    ]
+
+    cuda_home = os.environ.get("CUDA_PATH")
+    if cuda_home:
+        dll_dirs.append(os.path.join(cuda_home, "bin"))
+
+    dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin")
+    dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.3\bin")
+
+    for dll_dir in dll_dirs:
+        if os.path.isdir(dll_dir):
+            try:
+                os.add_dll_directory(dll_dir)
+                print(f"[DLL] add_dll_directory: {dll_dir}")
+            except Exception as e:
+                print(f"[DLL][WARN] falha em {dll_dir}: {e}")
+
+
+def create_onnx_session(onnx_path: str, provider: str = "tensorrt", trt_home: Optional[str] = None):
+    try:
+        import onnxruntime as ort
+    except ImportError:
+        raise ImportError(
+            "onnxruntime não está instalado. Instale com:\n"
+            "  pip install onnxruntime-gpu\n"
+            "ou CPU:\n"
+            "  pip install onnxruntime"
+        )
+
+    provider = provider.lower()
+
+    if provider == "tensorrt":
+        add_runtime_dll_dirs(trt_home)
+
+    available = ort.get_available_providers()
+    print(f"[ONNX] providers disponíveis: {available}")
+
+    sess_options = ort.SessionOptions()
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    if provider == "cuda":
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+    elif provider == "cpu":
+        providers = ["CPUExecutionProvider"]
+
+    elif provider == "tensorrt":
+        cache_dir = os.path.join(os.path.dirname(onnx_path), "trt_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+
+        trt_options = {
+            "device_id": 0,
+            "trt_fp16_enable": True,
+            "trt_engine_cache_enable": True,
+            "trt_engine_cache_path": cache_dir,
+            "trt_timing_cache_enable": True,
+            "trt_timing_cache_path": cache_dir,
+            "trt_max_workspace_size": 4 * 1024 * 1024 * 1024,
+        }
+
+        providers = [
+            ("TensorrtExecutionProvider", trt_options),
+            "CUDAExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+    else:
+        raise RuntimeError(f"Provider desconhecido: {provider}")
+
+    providers_ok = [
+        p for p in providers
+        if (p[0] if isinstance(p, tuple) else p) in available
+    ]
+
+    if not providers_ok:
+        raise RuntimeError(f"Nenhum provider solicitado está disponível. disponível={available}")
+
+    sess = ort.InferenceSession(
+        onnx_path,
+        sess_options=sess_options,
+        providers=providers_ok,
+    )
+
+    active = sess.get_providers()
+    print(f"[ONNX] usando providers: {active}")
+
+    if provider == "tensorrt" and "TensorrtExecutionProvider" not in active:
+        raise RuntimeError(f"TensorRT solicitado, mas não ficou ativo. Providers ativos: {active}")
+
+    if provider == "cuda" and "CUDAExecutionProvider" not in active:
+        raise RuntimeError(f"CUDA solicitado, mas não ficou ativo. Providers ativos: {active}")
+
+    return sess
+
+
+def preprocess_img_onnx(img_bgr: np.ndarray, resolucao, onnx_has_norm: bool):
+    """
+    Retorna input NCHW float32.
+
+    Se ONNX tem normalização embutida:
+      envia RGB 0..1
+
+    Se ONNX NÃO tem normalização embutida:
+      envia RGB já normalizado com NORM_MEAN/NORM_STD
+    """
+    w, h = int(resolucao[0]), int(resolucao[1])
+
+    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    img_res = cv2.resize(img_rgb, (w, h), interpolation=cv2.INTER_AREA)
+
+    x = img_res.astype(np.float32) / 255.0
+    x = np.transpose(x, (2, 0, 1))  # CHW
+
+    if not onnx_has_norm:
+        mean = NORM_MEAN.detach().cpu().numpy().astype(np.float32)
+        std = NORM_STD.detach().cpu().numpy().astype(np.float32)
+        std = np.clip(std, 1e-6, None)
+        x = (x - mean) / std
+
+    return np.expand_dims(x, axis=0).astype(np.float32)
+
+
+def get_onnx_outputs(outputs_dict: Dict[str, np.ndarray]):
+    keys = list(outputs_dict.keys())
+
+    semantic_name = None
+    label_name = None
+
+    for k in ["semantic_logits", "semantic_mask", "semantic", "output_semantic"]:
+        if k in outputs_dict:
+            semantic_name = k
+            break
+
+    for k in ["label_probs", "label_logits", "label", "output_label"]:
+        if k in outputs_dict:
+            label_name = k
+            break
+
+    if semantic_name is None or label_name is None:
+        if len(keys) != 2:
+            raise RuntimeError(f"Outputs ONNX inesperados: {keys}")
+        semantic_name = semantic_name or keys[0]
+        label_name = label_name or keys[1]
+
+    return outputs_dict[semantic_name], outputs_dict[label_name], semantic_name, label_name
+
+
+def infer_sample_onnx(
+    onnx_session,
+    img_bgr: np.ndarray,
+    mode: str,
+    resolucao,
+    onnx_has_norm: bool,
+    mask2_thr: float = 0.5,
+):
+    h0, w0 = img_bgr.shape[:2]
+
+    x = preprocess_img_onnx(
+        img_bgr=img_bgr,
+        resolucao=resolucao,
+        onnx_has_norm=onnx_has_norm,
+    )
+
+    input_name = onnx_session.get_inputs()[0].name
+    output_names = [o.name for o in onnx_session.get_outputs()]
+    raw_outputs = onnx_session.run(None, {input_name: x})
+
+    outputs_dict = {
+        name: arr.astype(np.float32)
+        for name, arr in zip(output_names, raw_outputs)
+    }
+
+    semantic_out, label_out, semantic_name, label_name = get_onnx_outputs(outputs_dict)
+
+    result = {
+        "pred_seg": None,
+        "mask2_prob": None,
+        "mask2_bin": None,
+        "label_id": None,
+        "label_conf": None,
+        "label_probs": None,
+    }
+
+    # semantic_logits: [1,C,H,W]
+    # semantic_mask  : [1,H,W]
+    if semantic_out.ndim == 4:
+        pred_ids = np.argmax(semantic_out, axis=1)[0].astype(np.uint8)
+    elif semantic_out.ndim == 3:
+        pred_ids = semantic_out[0].astype(np.uint8)
+    else:
+        raise RuntimeError(f"Saída semântica ONNX inválida: {semantic_name} shape={semantic_out.shape}")
+
+    pred_ids_full = cv2.resize(pred_ids, (w0, h0), interpolation=cv2.INTER_NEAREST)
+    result["pred_seg"] = pred_ids_full
+
+    # Neste modelo atual não estamos usando mask2 no ONNX.
+    if mode == "mask2":
+        result["mask2_prob"] = None
+        result["mask2_bin"] = None
+        return result
+
+    if mode == "label":
+        if label_name == "label_probs":
+            probs = label_out[0].astype(np.float32)
+        else:
+            logits = label_out.astype(np.float32)
+            logits = logits - np.max(logits, axis=1, keepdims=True)
+            e = np.exp(logits)
+            probs = (e / np.clip(np.sum(e, axis=1, keepdims=True), 1e-12, None))[0]
+
+        lid = int(np.argmax(probs))
+        result["label_id"] = lid
+        result["label_conf"] = float(probs[lid])
+        result["label_probs"] = probs
+
+    return result
+
+
 # ============================================================
 # Painéis
 # ============================================================
@@ -636,6 +863,10 @@ def main():
     ap.add_argument("--camera", action="store_true", help="Inferir em tempo real pela OAK-D Lite")
     ap.add_argument("--camera_fps", type=int, default=20)
     ap.add_argument("--camera_res", type=str, default="720p", choices=["720p", "1080p"])
+    ap.add_argument("--onnx", type=str, default=None, help="Se informado, usa ONNX Runtime em vez de PyTorch.")
+    ap.add_argument("--onnx_provider", type=str, default="tensorrt", choices=["cuda", "cpu", "tensorrt"])
+    ap.add_argument("--onnx_has_norm", action="store_true", help="Use se o ONNX foi exportado com --include-norm.")
+    ap.add_argument("--trt_home", type=str, default=r"C:\dev\TensorRT-10.10.0.31")
     args = ap.parse_args()
 
     device = find_device()
@@ -690,26 +921,36 @@ def main():
         raise SystemExit(f"Checkpoint não encontrado: {ckpt_path}")
     print(f"[model] ckpt={ckpt_path}")
 
-    base_model = build_base_model(num_classes=num_classes, device=device, backbone=BACKBONE)
+    use_onnx = args.onnx is not None and str(args.onnx).strip() != ""
+    onnx_session = None
 
+    base_model = None
     aux_head = None
     label_names = {}
-    feat_ch = get_feat_ch(base_model, device, input_h=args.input_size, input_w=args.input_size)
 
-    # Carrega ckpt antes para descobrir metadados se necessário
+    if use_onnx:
+        if not os.path.isfile(args.onnx):
+            raise SystemExit(f"ONNX não encontrado: {args.onnx}")
+
+        print(f"[runtime] usando ONNX: {args.onnx}")
+        print(f"[runtime] provider: {args.onnx_provider}")
+        print(f"[runtime] onnx_has_norm: {args.onnx_has_norm}")
+
+        onnx_session = create_onnx_session(
+            onnx_path=args.onnx,
+            provider=args.onnx_provider,
+            trt_home=args.trt_home,
+        )
+
+        # Ainda precisamos dos nomes dos labels para HUD.
+        # Primeiro tenta buscar no checkpoint, se existir.
+        if os.path.isfile(ckpt_path):
             ckpt_pre = torch.load(ckpt_path, map_location="cpu", weights_only=False)
             extra = ckpt_pre.get("extra", {}) or {}
-
-    if mode == "mask2":
-        aux_head = Mask2Head(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1)
-    elif mode == "label":
-        print(extra)
             label_names_raw = extra.get("label_name_by_id", {}) or {}
-        # JSON pode salvar keys como string
             label_names = {int(k): str(v) for k, v in label_names_raw.items()} if label_names_raw else {}
-        print(label_names_raw)
 
-        # Se não tiver no ckpt, tenta inferir do dataset de labels
+        # Fallback: tenta inferir dos labels do dataset.
         max_label_id = -1
         for s in samples:
             lid, lname = read_gt_label(s)
@@ -717,19 +958,90 @@ def main():
                 max_label_id = max(max_label_id, int(lid))
                 if lname:
                     label_names[int(lid)] = str(lname)
-        if max_label_id < 0:
+
-            # último fallback: pelo shape da última camada do checkpoint
+        if mode == "label":
+            if max_label_id < 0 and label_names:
+                max_label_id = max(label_names.keys())
+
+            for i in range(max_label_id + 1):
+                label_names.setdefault(i, f"label_{i}")
+
+            # Fallback final para seu caso conhecido.
+            if not label_names:
+                label_names = {
+                    0: "Parado",
+                    1: "EntrandoRua",
+                    2: "CaminhandoRua",
+                    3: "SaindoRua",
+                    4: "Manobrando",
+                    5: "Direcionando",
+                    6: "RetornandoBase",
+                    7: "Indefinido",
+                }
+
+        print(f"[label_names] {label_names}")
+
+    else:
+        print(f"[runtime] usando PyTorch: {ckpt_path}")
+
+        base_model = build_base_model(num_classes=num_classes, device=device, backbone=BACKBONE)
+
+        aux_head = None
+        label_names = {}
+        feat_ch = get_feat_ch(base_model, device, input_h=args.input_size, input_w=args.input_size)
+
+        ckpt_pre = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        extra = ckpt_pre.get("extra", {}) or {}
+
+        if mode == "mask2":
+            aux_head = Mask2Head(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1)
+
+        elif mode == "label":
+            label_names_raw = extra.get("label_name_by_id", {}) or {}
+            label_names = {int(k): str(v) for k, v in label_names_raw.items()} if label_names_raw else {}
+
+            max_label_id = -1
+            for s in samples:
+                lid, lname = read_gt_label(s)
+                if lid is not None:
+                    max_label_id = max(max_label_id, int(lid))
+                    if lname:
+                        label_names[int(lid)] = str(lname)
+
+            # Fonte mais confiável: shape da última camada salva no checkpoint.
+            ckpt_num_label_classes = None
             sd = ckpt_pre.get("aux_head", {})
+
             for k, v in sd.items():
                 if k.endswith("net.3.weight") or k.endswith("net.3.bias"):
-                    max_label_id = int(v.shape[0]) - 1
+                    ckpt_num_label_classes = int(v.shape[0])
                     break
+
+            if ckpt_num_label_classes is not None and ckpt_num_label_classes > 0:
+                num_label_classes = ckpt_num_label_classes
+            else:
+                # Fallback: labels vistos no dataset atual.
                 num_label_classes = max_label_id + 1
+
             if num_label_classes <= 0:
                 raise RuntimeError("Não consegui inferir num_label_classes para LabelHead.")
+
             for i in range(num_label_classes):
                 label_names.setdefault(i, f"label_{i}")
-        aux_head = LabelHead(feat_ch=feat_ch, num_seg_classes=num_classes, num_label_classes=num_label_classes, hidden=256, dropout=0.2)
+            if num_label_classes <= 0:
+                raise RuntimeError("Não consegui inferir num_label_classes para LabelHead.")
+
+            for i in range(num_label_classes):
+                label_names.setdefault(i, f"label_{i}")
+
+            aux_head = LabelHead(
+                feat_ch=feat_ch,
+                num_seg_classes=num_classes,
+                num_label_classes=num_label_classes,
+                hidden=256,
+                dropout=0.2,
+            )
+
             print(f"[label_head] classes={num_label_classes} names={label_names}")
 
         ckpt = load_checkpoint(ckpt_path, base_model, aux_head, device)
@@ -809,6 +1121,16 @@ def main():
             roi_bgr = frame_bgr[y0:y1, 0:W]
             roi_input = resize_to_config(roi_bgr, RESOLUCAO)
 
+            if use_onnx:
+                res = infer_sample_onnx(
+                    onnx_session=onnx_session,
+                    img_bgr=roi_input,
+                    mode=mode,
+                    resolucao=RESOLUCAO,
+                    onnx_has_norm=args.onnx_has_norm,
+                    mask2_thr=args.mask2_thr,
+                )
+            else:
                 res = infer_sample(
                     base_model=base_model,
                     aux_head=aux_head,
@@ -896,6 +1218,16 @@ def main():
         gt2 = imread_gray(s.mask2_path) if s.mask2_path else None
         gt_label_id, gt_label_name = read_gt_label(s)
 
+        if use_onnx:
+            res = infer_sample_onnx(
+                onnx_session=onnx_session,
+                img_bgr=img,
+                mode=mode,
+                resolucao=RESOLUCAO,
+                onnx_has_norm=args.onnx_has_norm,
+                mask2_thr=args.mask2_thr,
+            )
+        else:
             res = infer_sample(
                 base_model=base_model,
                 aux_head=aux_head,

Python/OAK/datasets/config.json

@@ -7,6 +7,7 @@
     "main_class_name": "navegavel", 
     "es_classes": "", 
     "model_to_use": "geral",
+    "ckpt_test": "best_main",
     "raw_size": [1296, 1028], 
     "resolucao": [1024, 576], 
     "roi_inicio": 0.0,

Python/OAK/datasets/oak-fcc-3/_10_export_onnx.py

@@ -0,0 +1,508 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+_10_export_onnx.py
+
+Exporta o checkpoint PyTorch do SegFormer Multi-Head OAK-FCC-3 para ONNX.
+
+Exemplo:
+
+python _10_export_onnx.py --config config.json --checkpoint backup/segformer_b1/target_teached/stacked_raw5/best_score.pt --out backup/segformer_b1/target_teached/stacked_raw5/best_score.onnx --train-script _8_train_multihead.py --opset 17 --device cuda
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import argparse
+import importlib.util
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+
+
+# ============================================================
+# Utils
+# ============================================================
+
+def load_json(path: str | Path) -> dict:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def ensure_dir(path: Path):
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def import_train_module(train_script_path: str | Path):
+    """
+    Importa o script de treinamento como módulo, para reaproveitar:
+      - load_labelmap
+      - build_heads_config
+      - get_input_channel_names
+      - get_input_channel_indices
+      - build_model
+      - load_checkpoint
+    """
+    train_script_path = Path(train_script_path)
+
+    if not train_script_path.exists():
+        raise FileNotFoundError(f"Script de treino não encontrado: {train_script_path}")
+
+    spec = importlib.util.spec_from_file_location(
+        "train_multihead_module",
+        str(train_script_path.resolve())
+    )
+
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Não consegui importar o script: {train_script_path}")
+
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def save_export_metadata(path: Path, data: dict):
+    ensure_dir(path.parent)
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+
+def resolve_norm_stats_path(config: dict, config_path: Path) -> Path | None:
+    fusion_mode = config.get("fusion_mode", "stacked")
+    model = config.get("modelo")
+    model_name = config.get("model_name")
+    ch = config.get("channels")
+
+    normstats_path = Path(f"backup/{model}/{model_name}/{fusion_mode}_raw{ch}/norm_stats.json")
+    if normstats_path.exists():
+        return normstats_path
+
+    ia_resolution = config.get("resolucao", [1024, 640])
+    w, h = int(ia_resolution[0]), int(ia_resolution[1])
+    p = Path("dataset") / f"{w}x{h}" / "group" / "norm_stats.json"
+    if not p.is_absolute():
+        p = (config_path.parent / p).resolve()
+    return p
+        
+
+    return None
+
+
+def load_selected_norm_stats(config: dict, config_path: Path, input_channel_indices: List[int]):
+    path = resolve_norm_stats_path(config, config_path)
+
+    if path is None or not path.is_file():
+        raise FileNotFoundError(f"norm_stats não encontrado: {path}")
+
+    js = load_json(path)
+
+    mean = js.get("mean")
+    std = js.get("std")
+    names = js.get("channels", [])
+
+    if mean is None or std is None:
+        raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
+
+    max_idx = max(input_channel_indices)
+
+    if len(mean) <= max_idx or len(std) <= max_idx:
+        raise RuntimeError(
+            f"norm_stats incompatível: precisa índices={input_channel_indices}, "
+            f"mean={len(mean)} std={len(std)} path={path}"
+        )
+
+    mean_sel = [float(mean[i]) for i in input_channel_indices]
+    std_sel = [float(std[i]) for i in input_channel_indices]
+
+    if names:
+        names_sel = [names[i] for i in input_channel_indices]
+    else:
+        names_sel = []
+
+    return path, mean_sel, std_sel, names_sel
+
+
+# ============================================================
+# Wrapper ONNX
+# ============================================================
+
+class MultiHeadOnnxWrapper(nn.Module):
+    """
+    Wrapper ONNX para exportar o SegFormer Multi-Head.
+
+    Pode exportar:
+      - logits crus
+      - logits redimensionados
+      - argmax em baixa resolução
+      - argmax em resolução da entrada
+
+    Também pode embutir a normalização:
+      x = (x - mean) / std
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        output_heads: List[str],
+        resize_to_input: bool = False,
+        include_norm: bool = False,
+        norm_mean: List[float] | None = None,
+        norm_std: List[float] | None = None,
+        postprocess: str = "none",
+    ):
+        super().__init__()
+        self.model = model
+        self.output_heads = list(output_heads)
+        self.resize_to_input = bool(resize_to_input)
+        self.include_norm = bool(include_norm)
+        self.postprocess = str(postprocess).lower()
+
+        if self.postprocess not in ("none", "resize_logits", "argmax_lowres", "argmax_fullres"):
+            raise RuntimeError(f"postprocess inválido: {self.postprocess}")
+
+        if self.postprocess == "resize_logits":
+            self.resize_to_input = True
+
+        if self.include_norm:
+            if norm_mean is None or norm_std is None:
+                raise RuntimeError("include_norm=True requer norm_mean e norm_std")
+
+            mean_t = torch.tensor(norm_mean, dtype=torch.float32).view(1, len(norm_mean), 1, 1)
+            std_t = torch.tensor(norm_std, dtype=torch.float32).view(1, len(norm_std), 1, 1)
+
+            self.register_buffer("norm_mean", mean_t)
+            self.register_buffer("norm_std", std_t)
+        else:
+            self.register_buffer("norm_mean", torch.empty(0))
+            self.register_buffer("norm_std", torch.empty(0))
+
+    def forward(self, pixel_values: torch.Tensor):
+        input_hw = pixel_values.shape[-2:]
+
+        x = pixel_values
+
+        if self.include_norm:
+            x = (x - self.norm_mean) / torch.clamp(self.norm_std, min=1e-6)
+
+        outputs: Dict[str, torch.Tensor] = self.model(pixel_values=x)
+
+        result = []
+
+        for head_name in self.output_heads:
+            if head_name not in outputs:
+                raise RuntimeError(f"Head ausente no modelo: {head_name}")
+
+            logits = outputs[head_name]
+
+            if self.postprocess in ("none", "resize_logits"):
+                if self.resize_to_input:
+                    logits = torch.nn.functional.interpolate(
+                        logits,
+                        size=input_hw,
+                        mode="bilinear",
+                        align_corners=False,
+                    )
+
+                result.append(logits)
+
+            elif self.postprocess == "argmax_lowres":
+                mask = torch.argmax(logits, dim=1).to(torch.uint8)
+                result.append(mask)
+
+            elif self.postprocess == "argmax_fullres":
+                logits = torch.nn.functional.interpolate(
+                    logits,
+                    size=input_hw,
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                mask = torch.argmax(logits, dim=1).to(torch.uint8)
+                result.append(mask)
+
+        return tuple(result)
+
+
+# ============================================================
+# Main
+# ============================================================
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--config", default="config.json")
+    parser.add_argument("--checkpoint", default="")
+    parser.add_argument("--out", default="")
+
+    parser.add_argument(
+        "--train-script",
+        default="_8_train_multihead.py",
+        help="Script de treino usado para montar exatamente o mesmo modelo.",
+    )
+
+    parser.add_argument("--labelmap", default="dataset/labelmap.txt")
+    parser.add_argument("--opset", type=int, default=17)
+    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
+
+    parser.add_argument(
+        "--resize-to-input",
+        action="store_true",
+        help=(
+            "Se ativo, exporta cada head já redimensionada para HxW da entrada. "
+            "Se desativo, exporta logits brutos do decode_head."
+        ),
+    )
+
+    parser.add_argument(
+        "--include-norm",
+        action="store_true",
+        help="Inclui normalização x=(x-mean)/std dentro do grafo ONNX.",
+    )
+
+    parser.add_argument(
+        "--postprocess",
+        default="none",
+        choices=["none", "resize_logits", "argmax_lowres", "argmax_fullres"],
+        help=(
+            "Define o pós-processamento exportado no ONNX. "
+            "none=logits crus; resize_logits=logits em HxW; "
+            "argmax_lowres=máscara na resolução do decode_head; "
+            "argmax_fullres=máscara HxW pronta."
+        ),
+    )
+
+    parser.add_argument(
+        "--dynamic-batch",
+        action="store_true",
+        help="Permite batch dinâmico no ONNX. H e W continuam fixos.",
+    )
+
+    args = parser.parse_args()
+
+    config_path = Path(args.config)
+    checkpoint_path = Path(args.checkpoint)
+    out_path = Path(args.out)
+    labelmap_path = Path(args.labelmap)
+
+    if not config_path.exists():
+        raise FileNotFoundError(f"Config não encontrado: {config_path}")
+
+    if args.checkpoint != "" and not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint não encontrado: {checkpoint_path}")
+
+    if not labelmap_path.exists():
+        raise FileNotFoundError(f"Labelmap não encontrado: {labelmap_path}")
+
+    if args.out != "":
+        ensure_dir(out_path.parent)
+
+    train_mod = import_train_module(args.train_script)
+
+    config = load_json(config_path)
+
+    W, H = config["resolucao"]
+    backbone = config.get("backbone", "nvidia/mit-b1")
+    fusion_mode = config.get("fusion_mode", "stacked")
+    model = config.get("modelo")
+    model_name = config.get("model_name")
+    ch = config.get("channels")
+    ckpt_name = config.get("ckpt_test", "best_score")
+
+    if args.checkpoint == "":
+        checkpoint_path = Path(f"backup/{model}/{model_name}/{fusion_mode}_raw{ch}/{ckpt_name}.pt")
+    if args.out == "":
+        out_path = Path(f"backup/{model}/{model_name}/{fusion_mode}_raw{ch}/{ckpt_name}.onnx")
+
+    if fusion_mode != "stacked":
+        raise RuntimeError("Este exportador foi pensado para fusion_mode='stacked'.")
+
+    input_channel_names = train_mod.get_input_channel_names(config)
+    input_channel_indices = train_mod.get_input_channel_indices(config)
+    channels = len(input_channel_names)
+
+    semantic_id2label, semantic_label2id, ignore_from_labelmap = train_mod.load_labelmap(
+        str(labelmap_path)
+    )
+
+    heads_config = train_mod.build_heads_config(
+        config,
+        ignore_index=int(ignore_from_labelmap)
+    )
+
+    # Mesmo ajuste feito no treino:
+    # semantic usa o número real do labelmap.
+    heads_config["semantic"]["num_classes"] = int(len(semantic_id2label))
+    heads_config["semantic"]["ignore_index"] = int(ignore_from_labelmap)
+
+    output_heads = list(heads_config.keys())
+
+    print("==========================================")
+    print("Export SegFormer Multi-Head para ONNX")
+    print(f"Config       : {config_path}")
+    print(f"Checkpoint   : {checkpoint_path}")
+    print(f"Output ONNX  : {out_path}")
+    print(f"Backbone     : {backbone}")
+    print(f"Resolution   : {W}x{H}")
+    print(f"Input shape  : [1, {channels}, {H}, {W}]")
+    print(f"Channels     : {input_channel_names} idx={input_channel_indices}")
+    print(f"Heads        : {output_heads}")
+    print(f"Resize output: {args.resize_to_input}")
+    print(f"Include norm : {args.include_norm}")
+    print(f"Postprocess  : {args.postprocess}")
+    print("==========================================")
+
+    use_cuda = args.device == "cuda" and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    if args.device == "cuda" and not torch.cuda.is_available():
+        print("[WARN] CUDA não disponível. Exportando em CPU.")
+
+    model = train_mod.build_model(
+        backbone=backbone,
+        channels=channels,
+        heads_config=heads_config,
+        semantic_id2label=semantic_id2label,
+        semantic_label2id=semantic_label2id,
+    )
+
+    print("[CKPT] Carregando checkpoint...")
+    ckpt = torch.load(
+        str(checkpoint_path),
+        map_location="cpu",
+        weights_only=False,
+    )
+
+    if "model" not in ckpt:
+        raise RuntimeError(
+            "Checkpoint não contém a chave 'model'. "
+            "Confirme se é um checkpoint salvo pelo script de treino."
+        )
+
+    model.load_state_dict(ckpt["model"], strict=True)
+    model.to(device)
+    model.eval()
+
+    norm_stats_path = None
+    norm_mean = None
+    norm_std = None
+    norm_channels = []
+
+    if args.include_norm:
+        norm_stats_path, norm_mean, norm_std, norm_channels = load_selected_norm_stats(
+            config=config,
+            config_path=config_path,
+            input_channel_indices=input_channel_indices,
+        )
+
+        print(f"[NORM] Embutindo normalização no ONNX: {norm_stats_path}")
+        print(f"[NORM] channels={norm_channels if norm_channels else input_channel_names}")
+        print(f"[NORM] mean={norm_mean}")
+        print(f"[NORM] std ={norm_std}")
+
+    wrapper = MultiHeadOnnxWrapper(
+        model=model,
+        output_heads=output_heads,
+        resize_to_input=args.resize_to_input,
+        include_norm=args.include_norm,
+        norm_mean=norm_mean,
+        norm_std=norm_std,
+        postprocess=args.postprocess,
+    )
+    wrapper.to(device)
+    wrapper.eval()
+
+    dummy_input = torch.randn(
+        1,
+        channels,
+        int(H),
+        int(W),
+        dtype=torch.float32,
+        device=device,
+    )
+
+    input_names = ["pixel_values"]
+    if args.postprocess in ("argmax_lowres", "argmax_fullres"):
+        output_names = [f"{name}_mask" for name in output_heads]
+    else:
+        output_names = [f"{name}_logits" for name in output_heads]
+
+    dynamic_axes = None
+    if args.dynamic_batch:
+        dynamic_axes = {
+            "pixel_values": {0: "batch"},
+        }
+        for out_name in output_names:
+            dynamic_axes[out_name] = {0: "batch"}
+
+    print("[CHECK] Rodando forward PyTorch antes do export...")
+    with torch.no_grad():
+        y = wrapper(dummy_input)
+
+    for name, tensor in zip(output_names, y):
+        print(f"  {name}: shape={tuple(tensor.shape)} dtype={tensor.dtype}")
+
+    print("[EXPORT] Exportando ONNX...")
+    with torch.no_grad():
+        torch.onnx.export(
+            wrapper,
+            dummy_input,
+            str(out_path),
+            export_params=True,
+            opset_version=int(args.opset),
+            do_constant_folding=True,
+            input_names=input_names,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+        )
+
+    print(f"[OK] ONNX salvo em: {out_path}")
+
+    meta_path = out_path.with_suffix(".export_meta.json")
+    save_export_metadata(meta_path, {
+        "config": str(config_path),
+        "checkpoint": str(checkpoint_path),
+        "onnx": str(out_path),
+        "backbone": backbone,
+        "input_shape": [1, channels, int(H), int(W)],
+        "input_channel_names": input_channel_names,
+        "input_channel_indices": input_channel_indices,
+        "heads": output_heads,
+        "output_names": output_names,
+        "resize_to_input": bool(args.resize_to_input),
+        "opset": int(args.opset),
+        "dynamic_batch": bool(args.dynamic_batch),
+        "semantic_id2label": semantic_id2label,
+        "heads_config": heads_config,
+        "checkpoint_epoch": ckpt.get("epoch", None),
+        "checkpoint_best": ckpt.get("best", None),
+        "include_norm": bool(args.include_norm),
+        "norm_stats_path": str(norm_stats_path) if norm_stats_path is not None else None,
+        "norm_channels": norm_channels if norm_channels else input_channel_names,
+        "norm_mean": norm_mean,
+        "norm_std": norm_std,
+        "postprocess": str(args.postprocess),
+    })
+
+    print(f"[OK] Metadata salvo em: {meta_path}")
+
+    # Verificação opcional do pacote onnx, se instalado.
+    try:
+        import onnx
+        print("[ONNX] Verificando grafo com onnx.checker...")
+        onnx_model = onnx.load(str(out_path))
+        onnx.checker.check_model(onnx_model)
+        print("[OK] onnx.checker passou.")
+    except ImportError:
+        print("[WARN] Pacote onnx não instalado. Pulei onnx.checker.")
+        print("       Instale com: pip install onnx")
+    except Exception as e:
+        print(f"[WARN] onnx.checker encontrou problema: {e}")
+
+    print("\nExport finalizado.")
+
+
+if __name__ == "__main__":
+    main()

Python/OAK/datasets/oak-fcc-3/_12_benchmark_onnx.py

@@ -0,0 +1,946 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+_12_benchmark_onnx.py
+
+Benchmark PyTorch vs ONNX Runtime para SegFormer OAK-FCC-3 Multi-Head.
+
+Mede:
+  - PyTorch FP32
+  - PyTorch AMP/FP16
+  - ONNX Runtime CUDA ou CPU
+
+Exemplos:
+
+Benchmark ONNX cru 160x256:
+
+python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider cuda
+
+Benchmark ONNX resized 640x1024:
+
+python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider cuda
+
+
+TensorRT
+python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider tensorrt --skip_torch_fp32 --skip_torch_amp
+
+"""
+
+from __future__ import annotations
+
+import gc
+import csv
+import json
+import time
+import argparse
+import importlib.util
+from pathlib import Path
+from typing import Optional, List, Dict, Tuple
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# ============================================================
+# Utils
+# ============================================================
+
+def load_json(path: str | Path) -> dict:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def save_json(path: str | Path, data: dict):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+
+def save_csv(path: str | Path, rows: List[dict]):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if not rows:
+        return
+
+    keys = list(rows[0].keys())
+
+    with path.open("w", newline="", encoding="utf-8") as f:
+        w = csv.DictWriter(f, fieldnames=keys)
+        w.writeheader()
+        w.writerows(rows)
+
+
+def resolve_path(path_like: Optional[str], base: Optional[Path] = None) -> Optional[Path]:
+    if path_like is None:
+        return None
+
+    p = Path(path_like)
+    if p.is_absolute():
+        return p
+
+    if base is None:
+        base = Path.cwd()
+
+    return (base / p).resolve()
+
+
+def import_train_module(train_script_path: str | Path):
+    train_script_path = Path(train_script_path)
+
+    if not train_script_path.exists():
+        raise FileNotFoundError(f"Script de treino não encontrado: {train_script_path}")
+
+    spec = importlib.util.spec_from_file_location(
+        "train_multihead_module",
+        str(train_script_path.resolve())
+    )
+
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Não consegui importar o script: {train_script_path}")
+
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def synchronize_if_cuda(device: torch.device):
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+
+
+def clear_cuda():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def percentile(values: List[float], p: float) -> float:
+    if not values:
+        return 0.0
+    return float(np.percentile(np.asarray(values, dtype=np.float64), p))
+
+
+def summarize_times(times_ms: List[float]) -> dict:
+    arr = np.asarray(times_ms, dtype=np.float64)
+
+    if arr.size == 0:
+        return {
+            "n": 0,
+            "mean_ms": 0.0,
+            "median_ms": 0.0,
+            "min_ms": 0.0,
+            "max_ms": 0.0,
+            "p95_ms": 0.0,
+            "p99_ms": 0.0,
+            "fps_mean": 0.0,
+            "fps_p95_latency": 0.0,
+        }
+
+    mean_ms = float(arr.mean())
+    p95_ms = float(np.percentile(arr, 95))
+    p99_ms = float(np.percentile(arr, 99))
+
+    return {
+        "n": int(arr.size),
+        "mean_ms": mean_ms,
+        "median_ms": float(np.median(arr)),
+        "min_ms": float(arr.min()),
+        "max_ms": float(arr.max()),
+        "p95_ms": p95_ms,
+        "p99_ms": p99_ms,
+        "fps_mean": float(1000.0 / mean_ms) if mean_ms > 0 else 0.0,
+        "fps_p95_latency": float(1000.0 / p95_ms) if p95_ms > 0 else 0.0,
+    }
+
+
+def resolve_model_artifact_paths(
+    args,
+    config: dict,
+    config_dir: Path,
+    channels: int,
+) -> Tuple[Path, Path, str]:
+    """
+    Resolve checkpoint e ONNX.
+
+    Se --checkpoint ou --onnx forem informados, usa os caminhos informados.
+    Se ficarem vazios, monta a partir do config:
+
+      backup/{modelo}/{model_name}/{fusion_mode}_raw{channels}/{ckpt_name}.pt
+      backup/{modelo}/{model_name}/{fusion_mode}_raw{channels}/{ckpt_name}.onnx
+
+    ckpt_name vem de:
+      config["ckpt_test"] ou "best_score"
+    """
+    model = config.get("modelo", "segformer_b1")
+    model_name = config.get("model_name", "target_teached")
+    fusion_mode = config.get("fusion_mode", "stacked")
+    ckpt_name = config.get("ckpt_test", "best_score")
+
+    # Usa o número real de canais selecionados,
+    # não necessariamente config["channels"].
+    ch = int(channels)
+
+    base_dir = config_dir / "backup" / model / model_name / f"{fusion_mode}_raw{ch}"
+
+    if args.checkpoint:
+        checkpoint_path = resolve_path(args.checkpoint, Path.cwd())
+    else:
+        checkpoint_path = base_dir / f"{ckpt_name}.pt"
+
+    if args.onnx:
+        onnx_path = resolve_path(args.onnx, Path.cwd())
+    else:
+        onnx_path = base_dir / f"{ckpt_name}.onnx"
+
+    if checkpoint_path is None or not checkpoint_path.is_file():
+        raise FileNotFoundError(
+            f"Checkpoint não encontrado: {checkpoint_path}\n"
+            f"Dica: informe --checkpoint ou ajuste config['ckpt_test']."
+        )
+
+    if onnx_path is None or not onnx_path.is_file():
+        raise FileNotFoundError(
+            f"ONNX não encontrado: {onnx_path}\n"
+            f"Dica: informe --onnx ou ajuste config['ckpt_test']."
+        )
+
+    return checkpoint_path.resolve(), onnx_path.resolve(), str(ckpt_name)
+
+
+# ============================================================
+# Dataset / normalização
+# ============================================================
+
+def collect_tensor_samples(root: Path, max_samples: int = 50, start_idx: int = 0) -> List[Path]:
+    tensor_paths: List[Path] = []
+
+    direct = root / "tensors"
+    if direct.is_dir():
+        tensor_paths.extend(sorted(direct.glob("*.npy")))
+
+    group_root = root / "group"
+    if group_root.is_dir():
+        for gdir in sorted(group_root.iterdir()):
+            tdir = gdir / "tensors"
+            if tdir.is_dir():
+                tensor_paths.extend(sorted(tdir.glob("*.npy")))
+
+    if not tensor_paths:
+        tensor_paths.extend(sorted(root.glob("**/tensors/*.npy")))
+
+    if not tensor_paths:
+        raise RuntimeError(f"Nenhum tensor .npy encontrado em: {root}")
+
+    start_idx = max(0, int(start_idx))
+    selected = tensor_paths[start_idx:]
+
+    if max_samples > 0:
+        selected = selected[:int(max_samples)]
+
+    return selected
+
+
+def load_tensor(path: Path, channels: int, channel_indices: List[int]) -> np.ndarray:
+    arr = np.load(str(path)).astype(np.float32)
+
+    if arr.ndim != 3:
+        raise RuntimeError(f"Tensor inválido {path}: shape={arr.shape}, esperado 3D")
+
+    if arr.shape[0] in (3, 4, 5):
+        chw = arr
+    elif arr.shape[-1] in (3, 4, 5):
+        chw = np.transpose(arr, (2, 0, 1))
+    else:
+        raise RuntimeError(f"Tensor com layout inesperado: {path} shape={arr.shape}")
+
+    max_idx = max(channel_indices)
+    if chw.shape[0] <= max_idx:
+        raise RuntimeError(
+            f"Tensor {path} tem {chw.shape[0]} canais, "
+            f"mas precisa acessar índice {max_idx}."
+        )
+
+    chw = chw[channel_indices, :, :]
+
+    finite = np.isfinite(chw)
+    if finite.any():
+        mx = float(np.nanmax(chw[finite]))
+        if mx > 2.0 and mx <= 255.0:
+            chw = chw / 255.0
+        elif mx > 255.0:
+            chw = chw / 65535.0
+
+    chw = np.nan_to_num(chw, nan=0.0, posinf=1.0, neginf=0.0)
+    return np.clip(chw, 0.0, 1.0).astype(np.float32)
+
+
+def find_norm_stats(config: dict, config_dir: Path, save_dir: Path, explicit: Optional[str]) -> Optional[Path]:
+    if explicit:
+        return resolve_path(explicit, Path.cwd())
+
+    W, H = config.get("resolucao", [1024, 640])
+    dataset_path = config_dir / "dataset"
+
+    candidates = [
+        dataset_path / f"{int(W)}x{int(H)}" / "group" / "norm_stats.json",
+        save_dir / "norm_stats.json",
+        config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test") / config.get("stats_source_tag", "stacked_raw5") / "norm_stats.json",
+    ]
+
+    for p in candidates:
+        if p.is_file():
+            return p
+
+    return candidates[0]
+
+
+def load_norm_stats(
+    path: Optional[Path],
+    channel_indices: List[int],
+    channel_names: List[str],
+) -> Tuple[Optional[List[float]], Optional[List[float]], Optional[str]]:
+    if path is None or not path.is_file():
+        print("[NORM] Sem norm_stats. Usando tensor 0..1 sem padronização.")
+        return None, None, None
+
+    js = load_json(path)
+    mean = js.get("mean")
+    std = js.get("std")
+    names = js.get("channels", [])
+
+    if mean is None or std is None:
+        raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
+
+    max_idx = max(channel_indices)
+    if len(mean) <= max_idx or len(std) <= max_idx:
+        raise RuntimeError(
+            f"norm_stats incompatível: precisa índices={channel_indices}, "
+            f"mean={len(mean)} std={len(std)}"
+        )
+
+    mean_sel = [float(mean[i]) for i in channel_indices]
+    std_sel = [float(std[i]) for i in channel_indices]
+
+    if names:
+        names_sel = [names[i] for i in channel_indices]
+    else:
+        names_sel = channel_names
+
+    print(f"[NORM] usando {path}")
+    print(f"[NORM] channels={names_sel}")
+    print(f"[NORM] mean={mean_sel}")
+    print(f"[NORM] std ={std_sel}")
+
+    return mean_sel, std_sel, str(path)
+
+
+def normalize_numpy_chw(chw: np.ndarray, mean: Optional[List[float]], std: Optional[List[float]]) -> np.ndarray:
+    if mean is None or std is None:
+        return chw.astype(np.float32)
+
+    mean_np = np.asarray(mean, dtype=np.float32).reshape(-1, 1, 1)
+    std_np = np.asarray(std, dtype=np.float32).reshape(-1, 1, 1)
+    std_np = np.clip(std_np, 1e-6, None)
+
+    return ((chw.astype(np.float32) - mean_np) / std_np).astype(np.float32)
+
+
+def load_inputs_as_numpy(
+    samples: List[Path],
+    channels: int,
+    channel_indices: List[int],
+    mean: Optional[List[float]],
+    std: Optional[List[float]],
+    target_hw: Tuple[int, int],
+    normalize_input: bool = True,
+) -> List[np.ndarray]:
+    H, W = target_hw
+    xs = []
+
+    for p in samples:
+        chw01 = load_tensor(
+            p,
+            channels=channels,
+            channel_indices=channel_indices,
+        )
+
+        if chw01.shape[-2:] != (H, W):
+            hwc = np.transpose(chw01, (1, 2, 0))
+            hwc = cv2.resize(hwc, (W, H), interpolation=cv2.INTER_LINEAR)
+            chw01 = np.transpose(hwc, (2, 0, 1)).astype(np.float32)
+
+        if normalize_input:
+            chw = normalize_numpy_chw(chw01, mean=mean, std=std)
+        else:
+            chw = chw01.astype(np.float32, copy=False)
+
+        x = np.expand_dims(chw, axis=0).astype(np.float32)
+        xs.append(x)
+
+    return xs
+
+
+# ============================================================
+# PyTorch
+# ============================================================
+
+class TorchTupleWrapper(nn.Module):
+    def __init__(self, model: nn.Module, output_heads: List[str]):
+        super().__init__()
+        self.model = model
+        self.output_heads = list(output_heads)
+
+    def forward(self, pixel_values: torch.Tensor):
+        outputs = self.model(pixel_values=pixel_values)
+        return tuple(outputs[h] for h in self.output_heads)
+
+
+@torch.inference_mode()
+def benchmark_torch(
+    model: nn.Module,
+    inputs_np: List[np.ndarray],
+    device: torch.device,
+    warmup: int,
+    repeat: int,
+    amp: bool,
+    label: str,
+) -> Tuple[dict, List[dict]]:
+    model.eval()
+
+    times = []
+    rows = []
+
+    # Precarrega tensors na GPU para medir só inferência do modelo.
+    inputs_t = [
+        torch.from_numpy(x).to(device, non_blocking=True)
+        for x in inputs_np
+    ]
+
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+
+    print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}")
+
+    # Warmup
+    for i in range(max(0, warmup)):
+        x = inputs_t[i % len(inputs_t)]
+        with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=amp and device.type == "cuda"):
+            _ = model(x)
+
+    synchronize_if_cuda(device)
+
+    # Medição
+    total_iter = len(inputs_t) * max(1, repeat)
+    idx = 0
+
+    for r in range(max(1, repeat)):
+        for sample_idx, x in enumerate(inputs_t):
+            synchronize_if_cuda(device)
+            t0 = time.perf_counter()
+
+            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=amp and device.type == "cuda"):
+                _ = model(x)
+
+            synchronize_if_cuda(device)
+            dt_ms = (time.perf_counter() - t0) * 1000.0
+
+            times.append(dt_ms)
+            rows.append({
+                "engine": label,
+                "repeat": r,
+                "sample_idx": sample_idx,
+                "iter_idx": idx,
+                "latency_ms": dt_ms,
+            })
+
+            idx += 1
+
+            if idx % 25 == 0 or idx == total_iter:
+                print(f"  {idx:04d}/{total_iter:04d} | last={dt_ms:.2f}ms")
+
+    summary = summarize_times(times)
+    summary["engine"] = label
+
+    return summary, rows
+
+
+# ============================================================
+# ONNX Runtime
+# ============================================================
+
+def create_onnx_session(onnx_path: Path, provider: str):
+    import os
+
+    trt_home = os.environ.get("TRT_HOME", r"C:\dev\TensorRT-10.10.0.31")
+
+    for dll_dir in [
+        os.path.join(trt_home, "lib"),
+        os.path.join(trt_home, "bin"),
+        r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin",
+    ]:
+        if os.path.isdir(dll_dir):
+            try:
+                os.add_dll_directory(dll_dir)
+                print(f"[DLL] add_dll_directory: {dll_dir}")
+            except Exception as e:
+                print(f"[DLL][WARN] falha em {dll_dir}: {e}")
+                
+    try:
+        import onnxruntime as ort
+    except ImportError:
+        raise ImportError(
+            "onnxruntime não está instalado. Instale com:\n"
+            "  pip install onnxruntime-gpu\n"
+            "ou CPU:\n"
+            "  pip install onnxruntime"
+        )
+
+    available = ort.get_available_providers()
+    print(f"[ONNX] providers disponíveis: {available}")
+
+    provider = provider.lower()
+
+    sess_options = ort.SessionOptions()
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    if provider == "cuda":
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+    elif provider == "cpu":
+        providers = ["CPUExecutionProvider"]
+
+    elif provider == "tensorrt":
+        cache_dir = onnx_path.parent / "trt_cache"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+        trt_options = {
+            "device_id": 0,
+
+            # FP16: o ponto principal do nosso teste.
+            "trt_fp16_enable": True,
+
+            # Cache: evita rebuild do engine a cada execução.
+            "trt_engine_cache_enable": True,
+            "trt_engine_cache_path": str(cache_dir),
+
+            # Timing cache ajuda a acelerar builds futuros.
+            "trt_timing_cache_enable": True,
+            "trt_timing_cache_path": str(cache_dir),
+
+            # Workspace. 4GB é razoável para RTX 3070, ajuste se faltar VRAM.
+            "trt_max_workspace_size": 4 * 1024 * 1024 * 1024,
+        }
+
+        providers = [
+            ("TensorrtExecutionProvider", trt_options),
+            "CUDAExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+
+    else:
+        providers = [provider]
+
+    # Checagem de disponibilidade, lidando com provider tuple.
+    requested_names = [
+        p[0] if isinstance(p, tuple) else p
+        for p in providers
+    ]
+
+    providers_ok = [
+        p for p in providers
+        if (p[0] if isinstance(p, tuple) else p) in available
+    ]
+
+    if not providers_ok:
+        raise RuntimeError(
+            f"Nenhum provider solicitado está disponível. "
+            f"Solicitado={requested_names}, disponível={available}"
+        )
+
+    session = ort.InferenceSession(
+        str(onnx_path),
+        sess_options=sess_options,
+        providers=providers_ok,
+    )
+
+    print(f"[ONNX] usando providers: {session.get_providers()}")
+
+    active_providers = session.get_providers()
+
+    if provider == "tensorrt" and "TensorrtExecutionProvider" not in active_providers:
+        raise RuntimeError(
+            "TensorRTExecutionProvider foi solicitado, mas não ficou ativo. "
+            f"Providers ativos: {active_providers}. "
+            "Provável causa: TensorRT não instalado, DLLs fora do PATH, "
+            "ou versão incompatível com onnxruntime-gpu."
+        )
+
+    if provider == "cuda" and "CUDAExecutionProvider" not in active_providers:
+        raise RuntimeError(
+            "CUDAExecutionProvider foi solicitado, mas não ficou ativo. "
+            f"Providers ativos: {active_providers}."
+        )
+
+    return session
+
+
+def benchmark_onnx(
+    session,
+    inputs_np: List[np.ndarray],
+    warmup: int,
+    repeat: int,
+    label: str,
+) -> Tuple[dict, List[dict]]:
+    input_name = session.get_inputs()[0].name
+
+    times = []
+    rows = []
+
+    print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}")
+
+    # Warmup
+    for i in range(max(0, warmup)):
+        x = inputs_np[i % len(inputs_np)]
+        _ = session.run(None, {input_name: x})
+
+    # Medição
+    total_iter = len(inputs_np) * max(1, repeat)
+    idx = 0
+
+    for r in range(max(1, repeat)):
+        for sample_idx, x in enumerate(inputs_np):
+            t0 = time.perf_counter()
+            _ = session.run(None, {input_name: x})
+            dt_ms = (time.perf_counter() - t0) * 1000.0
+
+            times.append(dt_ms)
+            rows.append({
+                "engine": label,
+                "repeat": r,
+                "sample_idx": sample_idx,
+                "iter_idx": idx,
+                "latency_ms": dt_ms,
+            })
+
+            idx += 1
+
+            if idx % 25 == 0 or idx == total_iter:
+                print(f"  {idx:04d}/{total_iter:04d} | last={dt_ms:.2f}ms")
+
+    summary = summarize_times(times)
+    summary["engine"] = label
+
+    return summary, rows
+
+
+# ============================================================
+# Main
+# ============================================================
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--config", default="config.json")
+    parser.add_argument("--checkpoint", default="")
+    parser.add_argument("--onnx", default="")
+    parser.add_argument("--train-script", default="_8_train_multihead.py")
+    parser.add_argument("--labelmap", default="dataset/labelmap.txt")
+
+    parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"])
+    parser.add_argument("--root_override", default=None)
+    parser.add_argument("--norm_stats", default=None)
+
+    parser.add_argument("--max_samples", type=int, default=50)
+    parser.add_argument("--start_idx", type=int, default=0)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--repeat", type=int, default=5)
+
+    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
+    parser.add_argument("--onnx_provider", default="cuda", choices=["cuda", "cpu", "tensorrt"])
+
+    parser.add_argument("--skip_torch_fp32", action="store_true")
+    parser.add_argument("--skip_torch_amp", action="store_true")
+    parser.add_argument("--skip_onnx", action="store_true")
+
+    parser.add_argument(
+        "--onnx_has_norm",
+        action="store_true",
+        help="Use quando o ONNX já inclui normalização interna. Nesse caso o ONNX recebe tensor 0..1 cru.",
+    )
+
+    parser.add_argument("--out_dir", default=None)
+
+    args = parser.parse_args()
+
+    config_path = resolve_path(args.config, Path.cwd())
+    train_script_path = resolve_path(args.train_script, Path.cwd())
+    labelmap_path = resolve_path(args.labelmap, Path.cwd())
+
+    if config_path is None or not config_path.is_file():
+        raise FileNotFoundError(f"Config não encontrado: {config_path}")
+    if train_script_path is None or not train_script_path.is_file():
+        raise FileNotFoundError(f"Train script não encontrado: {train_script_path}")
+    if labelmap_path is None or not labelmap_path.is_file():
+        raise FileNotFoundError(f"Labelmap não encontrado: {labelmap_path}")
+
+    config_dir = config_path.parent
+    config = load_json(config_path)
+
+    train_mod = import_train_module(train_script_path)
+
+    W, H = config.get("resolucao", [1024, 640])
+    W = int(W)
+    H = int(H)
+
+    backbone = config.get("backbone", "nvidia/mit-b1")
+    input_channel_names = train_mod.get_input_channel_names(config)
+    input_channel_indices = train_mod.get_input_channel_indices(config)
+    channels = len(input_channel_names)
+
+    checkpoint_path, onnx_path, ckpt_name = resolve_model_artifact_paths(
+        args=args,
+        config=config,
+        config_dir=config_dir,
+        channels=channels,
+    )
+
+    semantic_id2label, semantic_label2id, ignore_from_labelmap = train_mod.load_labelmap(
+        str(labelmap_path)
+    )
+
+    heads_config = train_mod.build_heads_config(
+        config,
+        ignore_index=int(ignore_from_labelmap)
+    )
+
+    heads_config["semantic"]["num_classes"] = int(len(semantic_id2label))
+    heads_config["semantic"]["ignore_index"] = int(ignore_from_labelmap)
+
+    output_heads = list(heads_config.keys())
+
+    save_dir = (
+        config_dir
+        / "backup"
+        / config.get("modelo", "segformer_b1")
+        / config.get("model_name", "test")
+        / f"{config.get('fusion_mode', 'stacked')}_raw{channels}"
+    )
+
+    norm_stats_path = find_norm_stats(
+        config=config,
+        config_dir=config_dir,
+        save_dir=save_dir,
+        explicit=args.norm_stats,
+    )
+
+    mean, std, norm_stats_used = load_norm_stats(
+        norm_stats_path,
+        channel_indices=input_channel_indices,
+        channel_names=input_channel_names,
+    )
+
+    if args.root_override:
+        root = resolve_path(args.root_override, Path.cwd())
+    else:
+        root = (config_dir / "dataset" / "split" / args.split_folder).resolve()
+
+    if root is None or not root.is_dir():
+        raise FileNotFoundError(f"Root de dados não encontrado: {root}")
+
+    samples = collect_tensor_samples(
+        root=root,
+        max_samples=args.max_samples,
+        start_idx=args.start_idx,
+    )
+
+    if args.out_dir:
+        out_dir = resolve_path(args.out_dir, Path.cwd())
+    else:
+        out_dir = onnx_path.parent / "benchmarks"
+
+    assert out_dir is not None
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    use_cuda = args.device == "cuda" and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    if args.device == "cuda" and not torch.cuda.is_available():
+        print("[WARN] CUDA indisponível. Usando CPU no PyTorch.")
+
+    print("==========================================")
+    print("Benchmark PyTorch vs ONNX")
+    print(f"Config       : {config_path}")
+    print(f"Checkpoint   : {checkpoint_path}")
+    print(f"ONNX         : {onnx_path}")
+    print(f"Root         : {root}")
+    print(f"Samples      : {len(samples)}")
+    print(f"Warmup       : {args.warmup}")
+    print(f"Repeat       : {args.repeat}")
+    print(f"Backbone     : {backbone}")
+    print(f"Input shape  : [1, {channels}, {H}, {W}]")
+    print(f"Channels     : {input_channel_names} idx={input_channel_indices}")
+    print(f"Heads        : {output_heads}")
+    print(f"Device       : {device}")
+    print(f"ONNX provider: {args.onnx_provider}")
+    print(f"ONNX has norm: {args.onnx_has_norm}")
+    print(f"Out dir      : {out_dir}")
+    print("==========================================")
+
+    if args.onnx_has_norm:
+        print("\n[DATA] Carregando inputs 0..1 crus na RAM...")
+    else:
+        print("\n[DATA] Carregando inputs normalizados na RAM...")
+    inputs_np = load_inputs_as_numpy(
+        samples=samples,
+        channels=channels,
+        channel_indices=input_channel_indices,
+        mean=mean,
+        std=std,
+        target_hw=(H, W),
+        normalize_input=not args.onnx_has_norm,
+    )
+    print(f"[DATA] Inputs carregados: {len(inputs_np)}")
+
+    summaries = []
+    all_rows = []
+
+    # ========================================================
+    # PyTorch
+    # ========================================================
+    need_torch = not args.skip_torch_fp32 or not args.skip_torch_amp
+
+    if need_torch:
+        print("\n[MODEL] Montando PyTorch...")
+        model = train_mod.build_model(
+            backbone=backbone,
+            channels=channels,
+            heads_config=heads_config,
+            semantic_id2label=semantic_id2label,
+            semantic_label2id=semantic_label2id,
+        )
+
+        ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
+        if "model" not in ckpt:
+            raise RuntimeError("Checkpoint não contém chave 'model'.")
+
+        model.load_state_dict(ckpt["model"], strict=True)
+        model.to(device)
+        model.eval()
+
+        torch_model = TorchTupleWrapper(
+            model=model,
+            output_heads=output_heads,
+        ).to(device)
+        torch_model.eval()
+
+        clear_cuda()
+
+        if not args.skip_torch_fp32:
+            summary, rows = benchmark_torch(
+                model=torch_model,
+                inputs_np=inputs_np,
+                device=device,
+                warmup=args.warmup,
+                repeat=args.repeat,
+                amp=False,
+                label="torch_fp32",
+            )
+            summaries.append(summary)
+            all_rows.extend(rows)
+
+        clear_cuda()
+
+        if not args.skip_torch_amp:
+            summary, rows = benchmark_torch(
+                model=torch_model,
+                inputs_np=inputs_np,
+                device=device,
+                warmup=args.warmup,
+                repeat=args.repeat,
+                amp=True,
+                label="torch_amp_fp16",
+            )
+            summaries.append(summary)
+            all_rows.extend(rows)
+
+        del torch_model
+        del model
+        clear_cuda()
+
+    # ========================================================
+    # ONNX
+    # ========================================================
+    if not args.skip_onnx:
+        print("\n[ONNX] Carregando sessão...")
+        session = create_onnx_session(
+            onnx_path=onnx_path,
+            provider=args.onnx_provider,
+        )
+
+        summary, rows = benchmark_onnx(
+            session=session,
+            inputs_np=inputs_np,
+            warmup=args.warmup,
+            repeat=args.repeat,
+            label=f"onnx_{args.onnx_provider}",
+        )
+        summaries.append(summary)
+        all_rows.extend(rows)
+
+    # ========================================================
+    # Relatório
+    # ========================================================
+    print("\n========== RESUMO ==========")
+
+    for s in summaries:
+        print(
+            f"{s['engine']:<16} "
+            f"n={s['n']:<4} "
+            f"mean={s['mean_ms']:.3f}ms "
+            f"median={s['median_ms']:.3f}ms "
+            f"p95={s['p95_ms']:.3f}ms "
+            f"p99={s['p99_ms']:.3f}ms "
+            f"fps_mean={s['fps_mean']:.2f} "
+            f"fps_p95={s['fps_p95_latency']:.2f}"
+        )
+
+    base_name = f"{onnx_path.stem}_{args.onnx_provider}"
+    report_json = out_dir / f"{base_name}_benchmark_report.json"
+    report_csv = out_dir / f"{base_name}_benchmark_rows.csv"
+
+    report = {
+        "config": str(config_path),
+        "checkpoint": str(checkpoint_path),
+        "ckpt_name": ckpt_name,
+        "onnx": str(onnx_path),
+        "onnx_has_norm": bool(args.onnx_has_norm),
+        "root": str(root),
+        "samples": len(samples),
+        "warmup": int(args.warmup),
+        "repeat": int(args.repeat),
+        "input_shape": [1, channels, H, W],
+        "input_channel_names": input_channel_names,
+        "input_channel_indices": input_channel_indices,
+        "heads": output_heads,
+        "norm_stats_used": norm_stats_used,
+        "onnx_provider": args.onnx_provider,
+        "device": str(device),
+        "summaries": summaries,
+    }
+
+    save_json(report_json, report)
+    save_csv(report_csv, all_rows)
+
+    print(f"\n[OK] JSON salvo em: {report_json}")
+    print(f"[OK] CSV salvo em : {report_csv}")
+    print("\nBenchmark finalizado.")
+
+
+if __name__ == "__main__":
+    main()

Python/OAK/datasets/oak-fcc-3/_6_normalize.py

@@ -814,7 +814,11 @@ def main():
         description="Normaliza RAW_BRUTO OAK-FCC-3 para tensor MULTISPEC final de treino."
     )
 
-    ap.add_argument("--src-root", default="dataset/original/group")
+    ap.add_argument(
+        "--src-roots",
+        default="dataset/original/group;dataset/augmented/group",
+        help="Raízes de entrada separadas por ';'. Ex: dataset/original/group;dataset/augmented/group"
+    )
     ap.add_argument("--out-root", default="dataset")
     ap.add_argument("--module-params", default=DEFAULT_MODULE_PARAMS)
     ap.add_argument("--res", default=f"{DEFAULT_RES[0]}x{DEFAULT_RES[1]}", help="Resolução final WxH.")
@@ -827,11 +831,23 @@ def main():
 
     args = ap.parse_args()
 
-    src_root = Path(args.src_root)
+    src_roots = [
+        Path(x.strip())
+        for x in str(args.src_roots).split(";")
+        if x.strip()
+    ]
+
     out_dataset_root = Path(args.out_root)
 
-    if not src_root.is_dir():
+    valid_src_roots = []
-        raise SystemExit(f"[ERRO] src-root não encontrado: {src_root}")
+    for src_root in src_roots:
+        if src_root.is_dir():
+            valid_src_roots.append(src_root)
+        else:
+            print(f"[WARN] src-root não encontrado, ignorando: {src_root}")
+
+    if not valid_src_roots:
+        raise SystemExit(f"[ERRO] Nenhum src-root válido encontrado: {src_roots}")
 
     try:
         res_w, res_h = [int(x) for x in args.res.lower().split("x")]
@@ -856,29 +872,38 @@ def main():
     cor_para_id, _, _, ignore_rgb = carregar_labelmap_completo(str(labelmap_path))
     ignore_id = _infer_ignore_id(ignore_rgb, 255)
 
-    all_groups = list_groups(src_root)
+    groups_by_root = []
+
+    for src_root in valid_src_roots:
+        groups = list_groups(src_root)
 
         if args.groups:
             want = {g.strip() for g in args.groups.split(",") if g.strip()}
-        all_groups = [g for g in all_groups if g in want]
+            groups = [g for g in groups if g in want]
 
-    if not all_groups:
+        if groups:
+            groups_by_root.append((src_root, groups))
+
+    if not groups_by_root:
         raise SystemExit("[ERRO] Nenhum grupo encontrado.")
 
     print("============================================")
     print("Normalize OAK-FCC-3")
-    print(f"SRC          : {src_root}")
+    print(f"SRC          : {[str(x) for x in valid_src_roots]}")
     print(f"OUT          : {output_root}")
     print(f"MODULE PARAM : {args.module_params}")
     print(f"RES          : {res}")
     print(f"RAW SIZE     : {raw_size}")
-    print(f"GROUPS       : {', '.join(all_groups)}")
+    print("GROUPS       :")
+    for root, groups in groups_by_root:
+        print(f"  - {root}: {', '.join(groups)}")
     print(f"SKIP BAD     : {args.skip_bad_quality}")
     print("============================================")
 
     running_stats = RunningStats()
     all_rows = []
 
+    for src_root, all_groups in groups_by_root:
         for group_name in all_groups:
             rows = process_group(
                 group_name=group_name,

Python/OAK/datasets/oak-fcc-3/_9_test_multihead.py

@@ -1043,6 +1043,241 @@ class MultiHeadTester:
         return preds, probs, t_ms
 
 
+class OnnxMultiHeadTester:
+    def __init__(
+        self,
+        config: dict,
+        onnx_path: Path,
+        provider: str,
+        channels: int,
+        heads_config: Dict[str, dict],
+        mean: Optional[Sequence[float]],
+        std: Optional[Sequence[float]],
+        trt_home: Optional[str] = None,
+        trt_fp16: bool = True,
+    ):
+        self.config = config
+        self.onnx_path = onnx_path
+        self.provider = provider.lower()
+        self.channels = int(channels)
+        self.heads_config = heads_config
+        self.runtime_mode = str(config.get("runtime_mode", "all")).lower()
+
+        self.mean = None if mean is None else np.asarray(mean, dtype=np.float32).reshape(1, channels, 1, 1)
+        self.std = None if std is None else np.asarray(std, dtype=np.float32).reshape(1, channels, 1, 1)
+
+        print(f"[ONNX_MODEL] onnx={onnx_path}")
+        print(f"[ONNX_MODEL] provider={provider}")
+
+        self.session = self._create_session(
+            onnx_path=onnx_path,
+            provider=provider,
+            trt_home=trt_home,
+            trt_fp16=trt_fp16,
+        )
+
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_names = [o.name for o in self.session.get_outputs()]
+        print(f"[ONNX_MODEL] input={self.input_name}")
+        print(f"[ONNX_MODEL] outputs={self.output_names}")
+
+    def _create_session(
+        self,
+        onnx_path: Path,
+        provider: str,
+        trt_home: Optional[str],
+        trt_fp16: bool,
+    ):
+        try:
+            import onnxruntime as ort
+        except ImportError:
+            raise ImportError(
+                "onnxruntime não está instalado. Use:\n"
+                "  pip install onnxruntime-gpu"
+            )
+
+        provider = provider.lower()
+
+        if provider == "tensorrt":
+            trt_home = trt_home or os.environ.get("TRT_HOME", r"C:\dev\TensorRT-10.10.0.31")
+
+            dll_dirs = [
+                os.path.join(trt_home, "lib"),
+                os.path.join(trt_home, "bin"),
+            ]
+
+            cuda_home = os.environ.get("CUDA_PATH")
+            if cuda_home:
+                dll_dirs.append(os.path.join(cuda_home, "bin"))
+
+            # fallback comum que vocês estão usando
+            dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.3\bin")
+            dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin")
+
+            for dll_dir in dll_dirs:
+                if os.path.isdir(dll_dir):
+                    try:
+                        os.add_dll_directory(dll_dir)
+                        print(f"[DLL] add_dll_directory: {dll_dir}")
+                    except Exception as e:
+                        print(f"[DLL][WARN] falha em {dll_dir}: {e}")
+
+        available = ort.get_available_providers()
+        print(f"[ONNX] providers disponíveis: {available}")
+
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if provider == "cpu":
+            providers = ["CPUExecutionProvider"]
+
+        elif provider == "cuda":
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        elif provider == "tensorrt":
+            cache_dir = onnx_path.parent / "trt_cache"
+            cache_dir.mkdir(parents=True, exist_ok=True)
+
+            trt_options = {
+                "device_id": 0,
+                "trt_fp16_enable": bool(trt_fp16),
+                "trt_engine_cache_enable": True,
+                "trt_engine_cache_path": str(cache_dir),
+                "trt_timing_cache_enable": True,
+                "trt_timing_cache_path": str(cache_dir),
+                "trt_max_workspace_size": 4 * 1024 * 1024 * 1024,
+            }
+
+            providers = [
+                ("TensorrtExecutionProvider", trt_options),
+                "CUDAExecutionProvider",
+                "CPUExecutionProvider",
+            ]
+
+        else:
+            raise RuntimeError(f"Provider ONNX inválido: {provider}")
+
+        providers_ok = [
+            p for p in providers
+            if (p[0] if isinstance(p, tuple) else p) in available
+        ]
+
+        if not providers_ok:
+            raise RuntimeError(f"Nenhum provider ONNX disponível. Pedido={providers}, disponíveis={available}")
+
+        session = ort.InferenceSession(
+            str(onnx_path),
+            sess_options=sess_options,
+            providers=providers_ok,
+        )
+
+        active = session.get_providers()
+        print(f"[ONNX] usando providers: {active}")
+
+        if provider == "tensorrt" and "TensorrtExecutionProvider" not in active:
+            raise RuntimeError(f"TensorRT solicitado, mas não ficou ativo. Providers ativos: {active}")
+
+        if provider == "cuda" and "CUDAExecutionProvider" not in active:
+            raise RuntimeError(f"CUDA solicitado, mas não ficou ativo. Providers ativos: {active}")
+
+        return session
+
+    def _normalize(self, x: np.ndarray) -> np.ndarray:
+        if self.mean is not None and self.std is not None:
+            return ((x - self.mean[0]) / np.clip(self.std[0], 1e-6, None)).astype(np.float32)
+        return x.astype(np.float32)
+
+    def _selected_heads(self) -> Optional[List[str]]:
+        if self.runtime_mode in ("target_direct", "target_head"):
+            return ["target"]
+        if self.runtime_mode in ("operational", "target_op"):
+            return ["vegetation", "cana"]
+        return None
+
+    @staticmethod
+    def _softmax_np(logits: np.ndarray, axis: int = 1) -> np.ndarray:
+        x = logits.astype(np.float32)
+        x = x - np.max(x, axis=axis, keepdims=True)
+        e = np.exp(x)
+        return e / np.clip(np.sum(e, axis=axis, keepdims=True), 1e-12, None)
+
+    @staticmethod
+    def _resize_logits_nchw(logits: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
+        n, c, h, w = logits.shape
+        th, tw = target_hw
+
+        if (h, w) == (th, tw):
+            return logits
+
+        out = np.empty((n, c, th, tw), dtype=np.float32)
+        for bi in range(n):
+            for ci in range(c):
+                out[bi, ci] = cv2.resize(
+                    logits[bi, ci].astype(np.float32),
+                    (tw, th),
+                    interpolation=cv2.INTER_LINEAR,
+                )
+        return out
+
+    def _map_outputs(self, outputs: List[np.ndarray]) -> Dict[str, np.ndarray]:
+        raw = {
+            name: arr.astype(np.float32)
+            for name, arr in zip(self.output_names, outputs)
+        }
+
+        mapped = {}
+        for head in self.heads_config.keys():
+            candidates = [
+                head,
+                f"{head}_logits",
+                f"output_{head}",
+            ]
+
+            found = None
+            for c in candidates:
+                if c in raw:
+                    found = c
+                    break
+
+            if found is not None:
+                mapped[head] = raw[found]
+
+        return mapped
+
+    def infer(self, chw_01: np.ndarray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], float]:
+        h, w = int(chw_01.shape[1]), int(chw_01.shape[2])
+
+        x = self._normalize(chw_01)
+        x = np.expand_dims(x, axis=0).astype(np.float32)
+
+        t0 = time.perf_counter()
+        outputs = self.session.run(None, {self.input_name: x})
+        t_ms = (time.perf_counter() - t0) * 1000.0
+
+        logits_by_head = self._map_outputs(outputs)
+
+        selected = self._selected_heads()
+        if selected is not None:
+            logits_by_head = {
+                hname: logits
+                for hname, logits in logits_by_head.items()
+                if hname in selected
+            }
+
+        preds = {}
+        probs = {}
+
+        for head_name, logits in logits_by_head.items():
+            logits = self._resize_logits_nchw(logits, (h, w))
+            prob = self._softmax_np(logits, axis=1)[0]
+            pred = np.argmax(prob, axis=0).astype(np.uint8)
+
+            preds[head_name] = pred
+            probs[head_name] = prob.astype(np.float32)
+
+        return preds, probs, t_ms
+
+
 # ============================================================
 # Checkpoints / paths
 # ============================================================
@@ -1084,6 +1319,43 @@ def find_checkpoint(save_dir: Path, preferred: Optional[str] = None) -> Path:
     raise FileNotFoundError("Nenhum checkpoint encontrado. Procurei:\n" + "\n".join(str(c) for c in candidates))
 
 
+def find_onnx_model(save_dir: Path, ckpt_path: Path, preferred: Optional[str] = None) -> Path:
+    """
+    Resolve o .onnx.
+
+    Se preferred for informado, usa ele.
+    Caso contrário, usa o mesmo stem do checkpoint:
+      best_score.pt -> best_score.onnx
+    """
+    if preferred:
+        p = Path(preferred)
+        if not p.is_absolute():
+            p_cwd = (Path.cwd() / p).resolve()
+            p_save = (save_dir / p).resolve()
+            p = p_cwd if p_cwd.is_file() else p_save
+
+        if not p.is_file():
+            raise FileNotFoundError(f"ONNX não encontrado: {p}")
+
+        return p.resolve()
+
+    p = ckpt_path.with_suffix(".onnx")
+
+    if not p.is_file():
+        alt = save_dir / f"{ckpt_path.stem}.onnx"
+        p = alt
+
+    if not p.is_file():
+        raise FileNotFoundError(
+            f"ONNX não encontrado para checkpoint {ckpt_path.name}. Procurei:\n"
+            f"  {ckpt_path.with_suffix('.onnx')}\n"
+            f"  {save_dir / (ckpt_path.stem + '.onnx')}\n"
+            f"Informe manualmente com --onnx."
+        )
+
+    return p.resolve()
+
+
 # ============================================================
 # Visualização
 # ============================================================
@@ -1179,6 +1451,37 @@ def class_percent(mask: np.ndarray, class_id: int, ignore_id: int = 255) -> floa
     return float(((mask == class_id) & valid).sum() * 100.0 / den)
 
 
+def compare_pred_equal_percent(a: Optional[np.ndarray], b: Optional[np.ndarray]) -> Optional[float]:
+    if a is None or b is None:
+        return None
+
+    if a.shape != b.shape:
+        b = cv2.resize(
+            b.astype(np.uint8),
+            (a.shape[1], a.shape[0]),
+            interpolation=cv2.INTER_NEAREST,
+        )
+
+    return float(np.mean(a == b) * 100.0)
+
+
+def diff_mask_rgb(a: Optional[np.ndarray], b: Optional[np.ndarray]) -> Optional[np.ndarray]:
+    if a is None or b is None:
+        return None
+
+    if a.shape != b.shape:
+        b = cv2.resize(
+            b.astype(np.uint8),
+            (a.shape[1], a.shape[0]),
+            interpolation=cv2.INTER_NEAREST,
+        )
+
+    diff = (a != b).astype(np.uint8) * 255
+    rgb = np.zeros((diff.shape[0], diff.shape[1], 3), dtype=np.uint8)
+    rgb[:, :, 0] = diff  # vermelho em RGB
+    return rgb
+
+
 # ============================================================
 # Métricas numpy
 # ============================================================
@@ -1240,6 +1543,10 @@ def main():
     parser.add_argument("--start_idx", type=int, default=0)
     parser.add_argument("--max_width", type=int, default=1800)
     parser.add_argument("--runtime_mode", default="all", choices=["all", "target_direct", "operational"])
+    parser.add_argument("--onnx", default="")
+    parser.add_argument("--onnx_provider", default="", choices=["", "cpu", "cuda", "tensorrt"])
+    parser.add_argument("--trt_home", default=None)
+    parser.add_argument("--trt_no_fp16", action="store_true")
     args = parser.parse_args()
 
     config_path = resolve_path(args.config, Path.cwd())
@@ -1336,6 +1643,32 @@ def main():
         use_amp=not args.no_amp,
     )
 
+    onnx_tester = None
+    onnx_path = None
+
+    if args.onnx_provider:
+        onnx_path = find_onnx_model(
+            save_dir=save_dir,
+            ckpt_path=ckpt_path,
+            preferred=args.onnx,
+        )
+
+        onnx_tester = OnnxMultiHeadTester(
+            config=config,
+            onnx_path=onnx_path,
+            provider=args.onnx_provider,
+            channels=channels,
+            heads_config=heads_config,
+            mean=mean,
+            std=std,
+            trt_home=args.trt_home,
+            trt_fp16=not args.trt_no_fp16,
+        )
+
+    if onnx_tester is not None:
+        print(f"ONNX        : {onnx_path}")
+        print(f"ONNX provider: {args.onnx_provider}")
+
     out_dir = Path(args.out_dir)
     ensure_dir(out_dir)
 
@@ -1351,6 +1684,10 @@ def main():
 
     win_name = "OAK-FCC-3 MultiHead Test | D/A navega | S salva | SPACE detalhado | Q sai"
     cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
+    win_name_onnx = None
+    if onnx_tester is not None:
+        win_name_onnx = "ONNX/TensorRT MultiHead Test | comparação visual"
+        cv2.namedWindow(win_name_onnx, cv2.WINDOW_NORMAL)
 
     while True:
         sample = samples[idx]
@@ -1368,6 +1705,13 @@ def main():
         preds, probs, t_inf = tester.infer(chw)
         preview_rgb = tensor_to_preview_rgb(chw)
 
+        onnx_preds = None
+        onnx_probs = None
+        t_onnx = None
+
+        if onnx_tester is not None:
+            onnx_preds, onnx_probs, t_onnx = onnx_tester.infer(chw)
+
         first_pred = next(iter(preds.values()))
         pred_h, pred_w = first_pred.shape[:2]
 
@@ -1508,7 +1852,143 @@ def main():
         cv2.putText(header, h2, (12, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0, 255, 120), 1, cv2.LINE_AA)
         canvas = np.vstack([header, canvas])
 
+        onnx_canvas = None
+        if onnx_tester is not None and onnx_preds is not None and onnx_probs is not None:
+            onnx_pred_sem = onnx_preds.get("semantic")
+            onnx_pred_veg = onnx_preds.get("vegetation")
+            onnx_pred_cana = onnx_preds.get("cana")
+
+            onnx_pred_target_op = None
+            if onnx_pred_veg is not None and onnx_pred_cana is not None:
+                onnx_pred_target_op = operational_target_mask(
+                    onnx_pred_veg,
+                    onnx_pred_cana,
+                    ignore_id=ignore_id,
+                )
+
+            onnx_pred_target_head = onnx_preds.get("target")
+            onnx_pred_target = onnx_pred_target_head if onnx_pred_target_head is not None else onnx_pred_target_op
+
+            onnx_prob_veg = onnx_probs["vegetation"][1] if "vegetation" in onnx_probs and onnx_probs["vegetation"].shape[0] > 1 else None
+            onnx_prob_cana = onnx_probs["cana"][1] if "cana" in onnx_probs and onnx_probs["cana"].shape[0] > 1 else None
+
+            onnx_prob_target_op = None
+            if onnx_prob_veg is not None and onnx_prob_cana is not None:
+                onnx_prob_target_op = np.clip(onnx_prob_veg * (1.0 - onnx_prob_cana), 0.0, 1.0)
+
+            onnx_prob_target_head = None
+            if "target" in onnx_probs:
+                onnx_prob_target_head = onnx_probs["target"][1] if onnx_probs["target"].shape[0] > 1 else onnx_probs["target"][0]
+
+            onnx_prob_target = onnx_prob_target_head if onnx_prob_target_head is not None else onnx_prob_target_op
+
+            eq_sem = compare_pred_equal_percent(pred_sem, onnx_pred_sem)
+            eq_veg = compare_pred_equal_percent(pred_veg, onnx_pred_veg)
+            eq_cana = compare_pred_equal_percent(pred_cana, onnx_pred_cana)
+            eq_target = compare_pred_equal_percent(pred_target, onnx_pred_target)
+
+            onnx_panels: List[Tuple[str, np.ndarray, str]] = []
+
+            if onnx_pred_sem is not None:
+                onnx_sem_rgb = ids_to_rgb(onnx_pred_sem, semantic_cmap, ignore_id)
+                onnx_panels.append((
+                    "ONNX semantic",
+                    onnx_sem_rgb,
+                    "" if eq_sem is None else f"igual PT={eq_sem:.3f}%"
+                ))
+                onnx_panels.append((
+                    "ONNX overlay semantic",
+                    overlay_rgb(preview_rgb, onnx_sem_rgb, args.alpha),
+                    ""
+                ))
+
+                d = diff_mask_rgb(pred_sem, onnx_pred_sem)
+                if d is not None:
+                    onnx_panels.append(("Diff semantic", d, "vermelho=diferente"))
+
+            if onnx_pred_veg is not None:
+                onnx_veg_rgb = ids_to_rgb(onnx_pred_veg, BINARY_COLORS_RGB, ignore_id)
+                onnx_panels.append((
+                    "ONNX vegetation",
+                    onnx_veg_rgb,
+                    "" if eq_veg is None else f"igual PT={eq_veg:.3f}%"
+                ))
+
+                if onnx_prob_veg is not None:
+                    onnx_panels.append((
+                        "ONNX P vegetation",
+                        prob_to_heat_rgb(onnx_prob_veg),
+                        f"mean={float(onnx_prob_veg.mean()):.3f}"
+                    ))
+
+            if onnx_pred_cana is not None:
+                onnx_cana_rgb = ids_to_rgb(onnx_pred_cana, CANA_COLORS_RGB, ignore_id)
+                onnx_panels.append((
+                    "ONNX cana",
+                    onnx_cana_rgb,
+                    "" if eq_cana is None else f"igual PT={eq_cana:.3f}%"
+                ))
+
+                if onnx_prob_cana is not None:
+                    onnx_panels.append((
+                        "ONNX P cana",
+                        prob_to_heat_rgb(onnx_prob_cana),
+                        f"mean={float(onnx_prob_cana.mean()):.3f}"
+                    ))
+
+            if onnx_pred_target is not None:
+                onnx_target_rgb = ids_to_rgb(onnx_pred_target, TARGET_COLORS_RGB, ignore_id)
+                title = "ONNX target HEAD" if onnx_pred_target_head is not None else "ONNX target OP"
+
+                onnx_panels.append((
+                    title,
+                    onnx_target_rgb,
+                    "" if eq_target is None else f"igual PT={eq_target:.3f}%"
+                ))
+                onnx_panels.append((
+                    "ONNX overlay target",
+                    overlay_rgb(preview_rgb, onnx_target_rgb, args.alpha),
+                    f"inf={t_onnx:.1f}ms"
+                ))
+
+                if onnx_prob_target is not None:
+                    onnx_panels.append((
+                        "ONNX P target",
+                        prob_to_heat_rgb(onnx_prob_target),
+                        f"mean={float(onnx_prob_target.mean()):.3f}"
+                    ))
+
+                d = diff_mask_rgb(pred_target, onnx_pred_target)
+                if d is not None:
+                    onnx_panels.append(("Diff target", d, "vermelho=diferente"))
+
+            onnx_canvas = compose_grid(onnx_panels, cols=3, max_width=args.max_width)
+
+            header_h_onnx = 78
+            header_onnx = np.zeros((header_h_onnx, onnx_canvas.shape[1], 3), dtype=np.uint8)
+            header_onnx[:] = (18, 18, 35)
+
+            h1_onnx = f"ONNX {args.onnx_provider} | {source_name} | inf={t_onnx:.1f}ms"
+            h2_parts = []
+            if eq_sem is not None:
+                h2_parts.append(f"sem={eq_sem:.3f}%")
+            if eq_veg is not None:
+                h2_parts.append(f"veg={eq_veg:.3f}%")
+            if eq_cana is not None:
+                h2_parts.append(f"cana={eq_cana:.3f}%")
+            if eq_target is not None:
+                h2_parts.append(f"target={eq_target:.3f}%")
+            h2_onnx = "igual PyTorch: " + " | ".join(h2_parts) if h2_parts else "comparação indisponível"
+
+            cv2.putText(header_onnx, h1_onnx, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (235, 235, 235), 1, cv2.LINE_AA)
+            cv2.putText(header_onnx, h2_onnx, (12, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0, 255, 120), 1, cv2.LINE_AA)
+
+            onnx_canvas = np.vstack([header_onnx, onnx_canvas])
+
         cv2.imshow(win_name, cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))
+        if onnx_canvas is not None and win_name_onnx is not None:
+            cv2.imshow(win_name_onnx, cv2.cvtColor(onnx_canvas, cv2.COLOR_RGB2BGR))
+        
         k = cv2.waitKey(0) & 0xFF
 
         if k in (ord("q"), ord("Q"), 27):
@@ -1520,10 +2000,15 @@ def main():
         elif k == ord(" "):
             detailed = not detailed
         elif k in (ord("s"), ord("S")):
-            out_path = out_dir / f"multihead_{idx:05d}_{sample.base}.png"
+            out_path = out_dir / f"multihead_pytorch_{idx:05d}_{sample.base}.png"
             cv2.imwrite(str(out_path), cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))
             print(f"[SAVE] {out_path}")
 
+            if onnx_canvas is not None:
+                out_path_onnx = out_dir / f"multihead_onnx_{args.onnx_provider}_{idx:05d}_{sample.base}.png"
+                cv2.imwrite(str(out_path_onnx), cv2.cvtColor(onnx_canvas, cv2.COLOR_RGB2BGR))
+                print(f"[SAVE] {out_path_onnx}")
+
     cv2.destroyAllWindows()
 
     if visited:

Python/OAK/datasets/oak-fcc-3/config.json

@@ -1,7 +1,7 @@
 {
     "camera": "oak-fcc-3",
     "modelo": "segformer_b1",
-    "model_name": "target_teached",
+    "model_name": "target_aug",
     "main_class_name": "cana", 
     "es_classes": "", 
     "model_to_use": "geral",
@@ -15,8 +15,9 @@
     "use_ndvi": false,
     "backbone": "nvidia/mit-b1",
     "fusion_mode": "stacked",
-    "stats_source_tag": "stacked_raw4",
+    "stats_source_tag": "stacked_raw5",
     "module_params_json": "calibration/module_params.json",
+    "ckpt_test": "best_target",
     "multi_head": true,
     "heads": {
         "semantic": {