2034 lines
72 KiB
Python
2034 lines
72 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
_9_test_infer_multihead.py
|
|
|
|
Teste/visualização do SegFormer OAK-FCC-3 Multi-Head.
|
|
|
|
Contrato esperado após normalize + split:
|
|
|
|
dataset/split/val/group/<grupo>/
|
|
tensors/<base>.npy # CHW float32 [R,G,B,RE,NIR]
|
|
masks/<base>.npy # semantic: 0=chao, 1=cana, 2=erva, 255=ignore
|
|
masks_vegetation/<base>.npy # vegetation: 0=background, 1=vegetation, 255=ignore
|
|
masks_cana/<base>.npy # cana: 0=not_cana, 1=cana, 255=ignore
|
|
metas/<base>.json
|
|
previews/<base>.png
|
|
|
|
Mostra:
|
|
- RGB preview do tensor
|
|
- GT/pred/overlay semantic
|
|
- GT/pred/overlay vegetation
|
|
- GT/pred/overlay cana
|
|
- alvo operacional = vegetation == 1 AND cana == 0
|
|
- mapas de confiança das heads binárias
|
|
|
|
Exemplo:
|
|
|
|
python .\_9_test_infer_multihead.py ^
|
|
--config config.json ^
|
|
--split_folder val ^
|
|
--ckpt backup\segformer_b1\test_multi\stacked_raw5_multihead\best_score.pt
|
|
|
|
Controles:
|
|
D / seta direita : próxima amostra
|
|
A / seta esquerda: amostra anterior
|
|
S : salvar painel atual em --out_dir
|
|
SPACE : alterna modo compacto/detalhado
|
|
Q / ESC : sair
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import time
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Sequence, Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from transformers import SegformerForSemanticSegmentation
|
|
|
|
from core.raw_processor_core import RawProcessorCore
|
|
|
|
|
|
# ============================================================
|
|
# Config default
|
|
# ============================================================
|
|
|
|
DEFAULT_HEADS = {
|
|
"semantic": {
|
|
"enabled": True,
|
|
"type": "multiclass",
|
|
"num_classes": 3,
|
|
"mask_dir": "masks",
|
|
"classes": {"chao": 0, "cana": 1, "erva": 2},
|
|
"ignore_index": 255,
|
|
},
|
|
"vegetation": {
|
|
"enabled": True,
|
|
"type": "binary",
|
|
"num_classes": 2,
|
|
"mask_dir": "masks_vegetation",
|
|
"classes": {"background": 0, "vegetation": 1},
|
|
"ignore_index": 255,
|
|
},
|
|
"cana": {
|
|
"enabled": True,
|
|
"type": "binary",
|
|
"num_classes": 2,
|
|
"mask_dir": "masks_cana",
|
|
"classes": {"not_cana": 0, "cana": 1},
|
|
"ignore_index": 255,
|
|
},
|
|
"target": {
|
|
"enabled": True,
|
|
"type": "binary",
|
|
"num_classes": 2,
|
|
"mask_dir": "__derived_target__",
|
|
"classes": {
|
|
"background": 0,
|
|
"target": 1,
|
|
},
|
|
"ignore_index": 255,
|
|
"derived_from": ["vegetation", "cana"],
|
|
},
|
|
}
|
|
|
|
SEMANTIC_COLORS_RGB = {
|
|
0: (85, 85, 85), # chao
|
|
1: (0, 190, 0), # cana
|
|
2: (230, 55, 55), # erva
|
|
}
|
|
|
|
BINARY_COLORS_RGB = {
|
|
0: (30, 30, 30),
|
|
1: (0, 220, 80),
|
|
}
|
|
|
|
CANA_COLORS_RGB = {
|
|
0: (30, 30, 30),
|
|
1: (40, 210, 255),
|
|
}
|
|
|
|
TARGET_COLORS_RGB = {
|
|
0: (30, 30, 30),
|
|
1: (255, 70, 30),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class SampleItem:
|
|
group: str
|
|
base: str
|
|
tensor_path: Optional[Path]
|
|
masks: Dict[str, Optional[Path]]
|
|
meta_path: Optional[Path] = None
|
|
preview_path: Optional[Path] = None
|
|
|
|
# novo
|
|
source_kind: str = "tensor" # "tensor" ou "raw_native_multi"
|
|
raw_group: Optional[dict] = None
|
|
|
|
|
|
# ============================================================
|
|
# Util
|
|
# ============================================================
|
|
|
|
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 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 merge_dict(dst: dict, src: dict) -> dict:
|
|
out = copy.deepcopy(dst)
|
|
|
|
def rec(a, b):
|
|
for k, v in b.items():
|
|
if isinstance(v, dict) and isinstance(a.get(k), dict):
|
|
rec(a[k], v)
|
|
else:
|
|
a[k] = v
|
|
|
|
if isinstance(src, dict):
|
|
rec(out, src)
|
|
return out
|
|
|
|
|
|
def build_heads_config(config: dict, ignore_index: int) -> Dict[str, dict]:
|
|
cfg = merge_dict(DEFAULT_HEADS, config.get("heads", {}) or {})
|
|
active = {}
|
|
for name, hcfg in cfg.items():
|
|
if not bool(hcfg.get("enabled", True)):
|
|
continue
|
|
hcfg.setdefault("ignore_index", ignore_index)
|
|
hcfg["ignore_index"] = int(hcfg.get("ignore_index", ignore_index))
|
|
hcfg["num_classes"] = int(hcfg.get("num_classes", 2))
|
|
hcfg["mask_dir"] = str(hcfg.get("mask_dir", "masks"))
|
|
active[name] = hcfg
|
|
|
|
for required in ("semantic", "vegetation", "cana"):
|
|
if required not in active:
|
|
raise RuntimeError(f"Head obrigatória ausente no config: {required}")
|
|
|
|
return active
|
|
|
|
|
|
def load_labelmap(labelmap_path: Path) -> Tuple[Dict[int, str], Dict[str, int], int, Dict[int, Tuple[int, int, int]]]:
|
|
try:
|
|
from helpers import carregar_labelmap_completo, _infer_ignore_id
|
|
|
|
_cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(str(labelmap_path))
|
|
ignore_id = int(_infer_ignore_id(ignore_rgb, 255))
|
|
|
|
id2label = {
|
|
int(k): str(v)
|
|
for k, v in id_para_nome.items()
|
|
if str(v).lower() not in ("ignore", "void", "background_ignore")
|
|
}
|
|
label2id = {v.lower(): k for k, v in id2label.items()}
|
|
|
|
colormap_rgb = {}
|
|
if isinstance(_colormap_rgb, dict):
|
|
for k, v in _colormap_rgb.items():
|
|
ik = int(k)
|
|
if ik in id2label:
|
|
colormap_rgb[ik] = tuple(map(int, v[:3]))
|
|
elif isinstance(_colormap_rgb, (list, tuple)):
|
|
for ik, v in enumerate(_colormap_rgb):
|
|
if ik in id2label:
|
|
colormap_rgb[ik] = tuple(map(int, v[:3]))
|
|
|
|
return id2label, label2id, ignore_id, colormap_rgb
|
|
|
|
except Exception as e:
|
|
print(f"[WARN] Não consegui usar helpers.carregar_labelmap_completo: {e}")
|
|
print("[WARN] Usando parser simples do labelmap.")
|
|
|
|
id2label: Dict[int, str] = {}
|
|
colormap_rgb: Dict[int, Tuple[int, int, int]] = {}
|
|
ignore_id = 255
|
|
next_id = 0
|
|
|
|
with labelmap_path.open("r", encoding="utf-8") as f:
|
|
for raw_line in f:
|
|
s = raw_line.strip()
|
|
if not s or s.startswith("#"):
|
|
continue
|
|
|
|
name = None
|
|
color = None
|
|
cid = None
|
|
|
|
if ":" in s and not s.split(":", 1)[0].strip().isdigit():
|
|
name_part, rest = s.split(":", 1)
|
|
name = name_part.strip()
|
|
color_txt = rest.split("::", 1)[0].strip().strip(":")
|
|
rgb_parts = [p.strip() for p in color_txt.split(",") if p.strip()]
|
|
if len(rgb_parts) >= 3:
|
|
color = tuple(int(float(p)) for p in rgb_parts[:3])
|
|
else:
|
|
parts = s.replace(",", " ").replace(":", " ").split()
|
|
if len(parts) >= 2 and parts[0].isdigit():
|
|
cid = int(parts[0])
|
|
name = parts[1]
|
|
if len(parts) >= 5:
|
|
color = tuple(int(float(p)) for p in parts[2:5])
|
|
elif len(parts) >= 1:
|
|
name = parts[0]
|
|
|
|
if not name:
|
|
continue
|
|
|
|
if name.lower() in ("ignore", "void", "background_ignore"):
|
|
ignore_id = 255
|
|
continue
|
|
|
|
if cid is None:
|
|
cid = next_id
|
|
next_id = max(next_id, cid + 1)
|
|
|
|
id2label[int(cid)] = str(name)
|
|
if color is not None:
|
|
colormap_rgb[int(cid)] = color
|
|
|
|
if not id2label:
|
|
id2label = {0: "chao", 1: "cana", 2: "erva"}
|
|
|
|
label2id = {v.lower(): k for k, v in id2label.items()}
|
|
for cid in id2label:
|
|
colormap_rgb.setdefault(cid, SEMANTIC_COLORS_RGB.get(cid, (255, 255, 255)))
|
|
|
|
return id2label, label2id, int(ignore_id), colormap_rgb
|
|
|
|
|
|
def find_raw_dataset_layout_root(path: Path) -> Optional[Path]:
|
|
"""
|
|
Detecta layout bruto:
|
|
root/metas
|
|
root/previews
|
|
root/bins
|
|
root/masks opcional
|
|
|
|
Aceita root, root/metas, root/bins, root/previews ou arquivo dentro deles.
|
|
"""
|
|
p = path.resolve()
|
|
candidates = []
|
|
|
|
if p.is_file():
|
|
candidates.append(p.parent)
|
|
candidates.append(p.parent.parent)
|
|
else:
|
|
candidates.append(p)
|
|
candidates.append(p.parent)
|
|
|
|
for c in candidates:
|
|
if not c:
|
|
continue
|
|
|
|
if c.name.lower() in ("metas", "metadata", "jsons", "previews", "bins", "masks"):
|
|
root = c.parent
|
|
else:
|
|
root = c
|
|
|
|
if (root / "metas").is_dir() and (root / "previews").is_dir() and (root / "bins").is_dir():
|
|
return root
|
|
|
|
return None
|
|
|
|
|
|
def resolve_raw_sibling_file(root: Path, subdir: str, stem: str, exts: Tuple[str, ...]) -> Optional[Path]:
|
|
folder = root / subdir
|
|
if not folder.is_dir():
|
|
return None
|
|
|
|
for ext in exts:
|
|
p = folder / f"{stem}{ext}"
|
|
if p.exists():
|
|
return p
|
|
|
|
return None
|
|
|
|
|
|
def resolve_raw_capture_group_from_json(json_path: Path, dataset_root: Path) -> dict:
|
|
"""
|
|
Resolve uma captura bruta:
|
|
metas/<base>.json
|
|
previews/<base>.png
|
|
bins/<arquivos das cameras>
|
|
masks/<base> opcional
|
|
"""
|
|
json_path = json_path.resolve()
|
|
meta = load_json(json_path)
|
|
base_name = json_path.stem
|
|
bins_dir = dataset_root / "bins"
|
|
|
|
preview_path = resolve_raw_sibling_file(
|
|
dataset_root,
|
|
"previews",
|
|
base_name,
|
|
(".png", ".jpg", ".jpeg"),
|
|
)
|
|
|
|
mask_path = resolve_raw_sibling_file(
|
|
dataset_root,
|
|
"masks",
|
|
base_name,
|
|
(".npy", ".png", ".tif", ".tiff"),
|
|
)
|
|
|
|
group = {
|
|
"json": json_path,
|
|
"png": preview_path,
|
|
"mask": mask_path,
|
|
"final_raw": None,
|
|
"cameras": {},
|
|
"dataset_root": dataset_root,
|
|
}
|
|
|
|
if "saved_payload_paths" in meta:
|
|
for cam_id, fname in meta["saved_payload_paths"].items():
|
|
fname_path = Path(fname)
|
|
|
|
candidates = [
|
|
bins_dir / fname_path.name,
|
|
json_path.parent / fname,
|
|
dataset_root / fname,
|
|
bins_dir / f"{base_name}_{cam_id}.bin",
|
|
bins_dir / f"{base_name}_{cam_id}.raw",
|
|
bins_dir / f"{base_name}_{cam_id.lower()}.bin",
|
|
bins_dir / f"{base_name}_{cam_id.lower()}.raw",
|
|
]
|
|
|
|
found = next((c for c in candidates if c.exists()), None)
|
|
if found is not None:
|
|
group["cameras"][cam_id] = found
|
|
else:
|
|
print(f"[WARN] bin não encontrado para {cam_id}: {fname}")
|
|
|
|
return group
|
|
|
|
|
|
def build_multispec_from_raw_native_multi_for_infer(group: dict, meta: dict) -> Tuple[np.ndarray, dict]:
|
|
"""
|
|
Gera tensor MULTISPEC [R,G,B,RE,NIR] em CHW float32 0..1
|
|
a partir dos bins RAW_BRUTO salvos.
|
|
"""
|
|
if meta.get("saved_payload_type") != "raw_native_multi":
|
|
raise RuntimeError(
|
|
f"Captura bruta não suportada aqui: saved_payload_type={meta.get('saved_payload_type')}"
|
|
)
|
|
|
|
stream_meta = meta.get("stream_meta", {}) or {}
|
|
saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
|
|
saved_shapes = meta.get("saved_payload_shapes", {}) or {}
|
|
|
|
frame = {}
|
|
|
|
for cam_id, path in group["cameras"].items():
|
|
saved_dtype = saved_dtypes.get(cam_id)
|
|
saved_shape = saved_shapes.get(cam_id)
|
|
|
|
if saved_dtype is None or saved_shape is None:
|
|
raise RuntimeError(f"Faltam dtype/shape para {cam_id} no JSON: {group['json']}")
|
|
|
|
arr = np.fromfile(str(path), dtype=np.dtype(saved_dtype)).reshape(tuple(saved_shape))
|
|
frame[cam_id] = arr
|
|
|
|
if not frame:
|
|
raise RuntimeError(f"Nenhum bin de câmera encontrado para: {group['json']}")
|
|
|
|
sensor_width = int(meta.get("sensor_width", 1280))
|
|
sensor_height = int(meta.get("sensor_height", 800))
|
|
bayer = meta.get("bayer_pattern", "RGGB")
|
|
|
|
calib_path = meta.get("camera_params_json") or "calibration/module_params.json"
|
|
|
|
if calib_path and not os.path.isfile(calib_path):
|
|
json_dir = Path(group["json"]).parent
|
|
alt = json_dir / calib_path
|
|
|
|
if alt.exists():
|
|
calib_path = str(alt)
|
|
else:
|
|
raise RuntimeError(
|
|
f"module_params.json não encontrado: {calib_path}. "
|
|
f"Sem ele eu não gero MULTISPEC, porque faltaria fusion/radiometria/flatfield."
|
|
)
|
|
|
|
core = RawProcessorCore(
|
|
sensor_width=sensor_width,
|
|
sensor_height=sensor_height,
|
|
bayer_pattern=bayer,
|
|
calibration_json_path=calib_path,
|
|
)
|
|
|
|
processing_meta = dict(stream_meta)
|
|
|
|
if meta.get("actual_camera_controls") is not None:
|
|
processing_meta["actual_camera_controls"] = meta.get("actual_camera_controls")
|
|
|
|
if meta.get("startup_camera_controls") is not None:
|
|
processing_meta["startup_camera_controls"] = meta.get("startup_camera_controls")
|
|
|
|
tensor = core.build_infer_tensor_from_stream(frame, processing_meta, 5)
|
|
|
|
if tensor is None:
|
|
raise RuntimeError(f"RawProcessorCore retornou tensor None para: {group['json']}")
|
|
|
|
tensor = np.asarray(tensor, dtype=np.float32)
|
|
tensor = np.nan_to_num(tensor, nan=0.0, posinf=1.0, neginf=0.0)
|
|
tensor = np.clip(tensor, 0.0, 1.0)
|
|
|
|
processing_info = {
|
|
"fusion_result": getattr(core, "last_fusion_result", None),
|
|
"fusion_config_used": getattr(core, "fusion_config", None),
|
|
"radiometric_normalization_result": getattr(core, "last_radiometric_normalization_result", None),
|
|
"patch_normalization_result": getattr(core, "last_patch_normalization_result", None),
|
|
"frame_quality": getattr(core, "last_frame_quality_result", None),
|
|
}
|
|
|
|
return tensor, processing_info
|
|
|
|
|
|
DEFAULT_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"]
|
|
|
|
def get_input_channel_names(config: dict, channels_override: Optional[int] = None) -> List[str]:
|
|
if "input_channels" in config:
|
|
names = [str(c).upper() for c in config["input_channels"]]
|
|
else:
|
|
n = int(channels_override or config.get("channels", 5))
|
|
names = DEFAULT_CHANNEL_ORDER[:n]
|
|
|
|
invalid = [c for c in names if c not in DEFAULT_CHANNEL_ORDER]
|
|
if invalid:
|
|
raise RuntimeError(f"Canais inválidos em input_channels: {invalid}")
|
|
|
|
return names
|
|
|
|
|
|
def get_input_channel_indices(config: dict, channels_override: Optional[int] = None) -> List[int]:
|
|
names = get_input_channel_names(config, channels_override)
|
|
return [DEFAULT_CHANNEL_ORDER.index(c) for c in names]
|
|
|
|
|
|
# ============================================================
|
|
# Dataset
|
|
# ============================================================
|
|
|
|
def collect_samples(root: Path, heads_config: Dict[str, dict], require_masks: bool = False) -> Tuple[List[SampleItem], bool]:
|
|
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")))
|
|
|
|
# ============================================================
|
|
# MODO 1: dataset já normalizado com tensors/*.npy
|
|
# ============================================================
|
|
if tensor_paths:
|
|
samples: List[SampleItem] = []
|
|
has_any_gt = False
|
|
|
|
for tp in tensor_paths:
|
|
base = tp.stem
|
|
group_name = tp.parent.parent.name if tp.parent.name == "tensors" else "default"
|
|
group_dir = tp.parent.parent if tp.parent.name == "tensors" else tp.parent
|
|
|
|
masks: Dict[str, Optional[Path]] = {}
|
|
missing = []
|
|
|
|
for head_name, hcfg in heads_config.items():
|
|
mask_dir = group_dir / str(hcfg.get("mask_dir", "masks"))
|
|
mask_path = mask_dir / f"{base}.npy"
|
|
|
|
if mask_path.exists():
|
|
masks[head_name] = mask_path
|
|
has_any_gt = True
|
|
else:
|
|
masks[head_name] = None
|
|
missing.append(f"{head_name}:{mask_path}")
|
|
|
|
if require_masks and missing:
|
|
raise RuntimeError(f"Masks ausentes para {tp}: {missing}")
|
|
|
|
meta_path = group_dir / "metas" / f"{base}.json"
|
|
preview_path = group_dir / "previews" / f"{base}.png"
|
|
|
|
samples.append(SampleItem(
|
|
group=group_name,
|
|
base=base,
|
|
tensor_path=tp,
|
|
masks=masks,
|
|
meta_path=meta_path if meta_path.exists() else None,
|
|
preview_path=preview_path if preview_path.exists() else None,
|
|
source_kind="tensor",
|
|
raw_group=None,
|
|
))
|
|
|
|
return samples, has_any_gt
|
|
|
|
# ============================================================
|
|
# MODO 2: dataset bruto com bins/metas/previews/masks
|
|
# ============================================================
|
|
raw_root = find_raw_dataset_layout_root(root)
|
|
|
|
if raw_root is None:
|
|
raise RuntimeError(
|
|
f"Nenhum tensor .npy encontrado e também não detectei layout bruto "
|
|
f"com bins/metas/previews em: {root}"
|
|
)
|
|
|
|
meta_paths = sorted((raw_root / "metas").glob("*.json"))
|
|
if not meta_paths:
|
|
raise RuntimeError(f"Nenhum meta .json encontrado em: {raw_root / 'metas'}")
|
|
|
|
samples: List[SampleItem] = []
|
|
has_any_gt = False
|
|
|
|
for mp in meta_paths:
|
|
base = mp.stem
|
|
meta = load_json(mp)
|
|
|
|
if meta.get("saved_payload_type") != "raw_native_multi":
|
|
print(f"[WARN] pulando {mp.name}: saved_payload_type={meta.get('saved_payload_type')}")
|
|
continue
|
|
|
|
raw_group = resolve_raw_capture_group_from_json(mp, raw_root)
|
|
|
|
masks: Dict[str, Optional[Path]] = {}
|
|
|
|
# Tenta masks específicas multi-head primeiro.
|
|
for head_name, hcfg in heads_config.items():
|
|
mask_dir = raw_root / str(hcfg.get("mask_dir", "masks"))
|
|
mask_path = None
|
|
|
|
mask_dir_name = str(hcfg.get("mask_dir", "masks"))
|
|
|
|
if mask_dir_name == "__derived_target__" or bool(hcfg.get("derived", False)):
|
|
masks[head_name] = None
|
|
continue
|
|
|
|
for ext in (".npy", ".png", ".tif", ".tiff"):
|
|
p = mask_dir / f"{base}{ext}"
|
|
if p.exists():
|
|
mask_path = p
|
|
break
|
|
|
|
masks[head_name] = mask_path
|
|
if mask_path is not None:
|
|
has_any_gt = True
|
|
|
|
# Se só existir masks/<base>.png ou .npy semantic, conecta na head semantic.
|
|
if masks.get("semantic") is None:
|
|
semantic_mask = resolve_raw_sibling_file(
|
|
raw_root,
|
|
"masks",
|
|
base,
|
|
(".npy", ".png", ".tif", ".tiff"),
|
|
)
|
|
if semantic_mask is not None:
|
|
masks["semantic"] = semantic_mask
|
|
has_any_gt = True
|
|
|
|
if require_masks and not any(p is not None for p in masks.values()):
|
|
raise RuntimeError(f"Mask ausente para captura bruta: {mp}")
|
|
|
|
preview_path = resolve_raw_sibling_file(
|
|
raw_root,
|
|
"previews",
|
|
base,
|
|
(".png", ".jpg", ".jpeg"),
|
|
)
|
|
|
|
samples.append(SampleItem(
|
|
group=raw_root.name,
|
|
base=base,
|
|
tensor_path=None,
|
|
masks=masks,
|
|
meta_path=mp,
|
|
preview_path=preview_path,
|
|
source_kind="raw_native_multi",
|
|
raw_group=raw_group,
|
|
))
|
|
|
|
if not samples:
|
|
raise RuntimeError(f"Nenhuma captura raw_native_multi válida encontrada em: {raw_root}")
|
|
|
|
return samples, has_any_gt
|
|
|
|
|
|
def load_tensor(
|
|
path: Path,
|
|
channels: int,
|
|
channel_indices: Optional[List[int]] = None,
|
|
) -> 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")
|
|
|
|
# Normaliza para CHW.
|
|
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}")
|
|
|
|
if channel_indices is not None:
|
|
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}. Shape={chw.shape}"
|
|
)
|
|
chw = chw[channel_indices, :, :]
|
|
else:
|
|
if chw.shape[0] < channels:
|
|
raise RuntimeError(
|
|
f"Tensor {path} tem {chw.shape[0]} canais, "
|
|
f"mas config pediu {channels}."
|
|
)
|
|
chw = chw[:channels, :, :]
|
|
|
|
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 load_sample_tensor(
|
|
sample: SampleItem,
|
|
channels: int,
|
|
channel_indices: Optional[List[int]] = None,
|
|
) -> np.ndarray:
|
|
"""
|
|
Carrega tensor de uma amostra.
|
|
- Se for tensor pronto: lê .npy.
|
|
- Se for RAW_BRUTO: gera MULTISPEC em tempo real a partir dos bins.
|
|
"""
|
|
if sample.source_kind == "tensor":
|
|
if sample.tensor_path is None:
|
|
raise RuntimeError(f"Sample tensor sem tensor_path: {sample.base}")
|
|
return load_tensor(
|
|
sample.tensor_path,
|
|
channels=channels,
|
|
channel_indices=channel_indices,
|
|
)
|
|
|
|
if sample.source_kind == "raw_native_multi":
|
|
if sample.raw_group is None or sample.meta_path is None:
|
|
raise RuntimeError(f"Sample raw sem raw_group/meta_path: {sample.base}")
|
|
|
|
meta = load_json(sample.meta_path)
|
|
tensor, _processing_info = build_multispec_from_raw_native_multi_for_infer(sample.raw_group, meta)
|
|
|
|
if tensor.ndim != 3:
|
|
raise RuntimeError(f"Tensor RAW gerado inválido: {sample.base} shape={tensor.shape}")
|
|
|
|
if channel_indices is not None:
|
|
max_idx = max(channel_indices)
|
|
if tensor.shape[0] <= max_idx:
|
|
raise RuntimeError(
|
|
f"Tensor RAW tem {tensor.shape[0]} canais, "
|
|
f"mas precisa acessar índice {max_idx}."
|
|
)
|
|
tensor = tensor[channel_indices, :, :]
|
|
else:
|
|
if tensor.shape[0] < channels:
|
|
raise RuntimeError(
|
|
f"Tensor RAW com canais incompatíveis: {sample.base} "
|
|
f"shape={tensor.shape}, esperado pelo menos {channels}"
|
|
)
|
|
tensor = tensor[:channels, :, :]
|
|
|
|
return tensor.astype(np.float32)
|
|
|
|
raise RuntimeError(f"source_kind desconhecido: {sample.source_kind}")
|
|
|
|
|
|
def load_mask(path: Optional[Path]) -> Optional[np.ndarray]:
|
|
if path is None:
|
|
return None
|
|
if path.suffix.lower() == ".npy":
|
|
mask = np.load(str(path))
|
|
else:
|
|
mask = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
|
|
if mask is None:
|
|
raise RuntimeError(f"Falha ao ler mask: {path}")
|
|
if mask.ndim == 3:
|
|
mask = mask[:, :, 0]
|
|
return mask.astype(np.int64)
|
|
|
|
|
|
# ============================================================
|
|
# Normalização
|
|
# ============================================================
|
|
|
|
def load_norm_stats(
|
|
path: Optional[Path],
|
|
channels: int,
|
|
channel_indices: Optional[List[int]] = None,
|
|
channel_names: Optional[List[str]] = None,
|
|
) -> Tuple[Optional[List[float]], Optional[List[float]]]:
|
|
if path is None or not path.is_file():
|
|
if path is not None:
|
|
print(f"[NORM] não encontrei norm_stats em {path}. Usando tensor 0..1 sem padronização.")
|
|
return None, None
|
|
|
|
js = load_json(path)
|
|
mean = js.get("mean", None)
|
|
std = js.get("std", None)
|
|
names = js.get("channels", [])
|
|
|
|
if mean is None or std is None:
|
|
raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
|
|
if channel_indices is not None:
|
|
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 = [mean[i] for i in channel_indices]
|
|
std = [std[i] for i in channel_indices]
|
|
|
|
if names:
|
|
names = [names[i] for i in channel_indices]
|
|
elif channel_names:
|
|
names = channel_names
|
|
|
|
else:
|
|
if len(mean) != channels or len(std) != channels:
|
|
raise RuntimeError(
|
|
f"norm_stats incompatível com channels={channels}: "
|
|
f"mean={len(mean)} std={len(std)}"
|
|
)
|
|
|
|
print(f"[NORM] usando {path}")
|
|
print(f"[NORM] channels={names}")
|
|
print(f"[NORM] mean={mean}")
|
|
print(f"[NORM] std ={std}")
|
|
return list(map(float, mean)), list(map(float, std))
|
|
|
|
|
|
# ============================================================
|
|
# Modelo multi-head, igual ao treino
|
|
# ============================================================
|
|
|
|
def patch_segformer_encoder_input_channels(segformer_encoder: nn.Module, in_ch: int):
|
|
if in_ch == 3:
|
|
return segformer_encoder
|
|
|
|
proj = segformer_encoder.encoder.patch_embeddings[0].proj
|
|
|
|
if proj.in_channels == in_ch:
|
|
return segformer_encoder
|
|
|
|
old_weight = proj.weight.data.clone()
|
|
old_bias = proj.bias.data.clone() if proj.bias is not None else None
|
|
|
|
new_proj = nn.Conv2d(
|
|
in_channels=in_ch,
|
|
out_channels=proj.out_channels,
|
|
kernel_size=proj.kernel_size,
|
|
stride=proj.stride,
|
|
padding=proj.padding,
|
|
dilation=proj.dilation,
|
|
groups=proj.groups,
|
|
bias=proj.bias is not None,
|
|
padding_mode=proj.padding_mode,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
if in_ch <= old_weight.shape[1]:
|
|
new_proj.weight.copy_(old_weight[:, :in_ch, :, :])
|
|
else:
|
|
new_proj.weight[:, :old_weight.shape[1], :, :].copy_(old_weight)
|
|
extra = in_ch - old_weight.shape[1]
|
|
mean_w = old_weight.mean(dim=1, keepdim=True)
|
|
new_proj.weight[:, old_weight.shape[1]:, :, :].copy_(mean_w.repeat(1, extra, 1, 1))
|
|
|
|
if old_bias is not None:
|
|
new_proj.bias.copy_(old_bias)
|
|
|
|
segformer_encoder.encoder.patch_embeddings[0].proj = new_proj
|
|
print(f"[MODEL] patch input channels: 3 -> {in_ch}")
|
|
return segformer_encoder
|
|
|
|
|
|
def replace_segformer_decode_classifier(decode_head: nn.Module, num_classes: int):
|
|
old = decode_head.classifier
|
|
if not isinstance(old, nn.Conv2d):
|
|
raise RuntimeError(f"decode_head.classifier não é Conv2d: {type(old)}")
|
|
|
|
new = nn.Conv2d(
|
|
in_channels=old.in_channels,
|
|
out_channels=int(num_classes),
|
|
kernel_size=old.kernel_size,
|
|
stride=old.stride,
|
|
padding=old.padding,
|
|
dilation=old.dilation,
|
|
groups=old.groups,
|
|
bias=old.bias is not None,
|
|
padding_mode=old.padding_mode,
|
|
)
|
|
decode_head.classifier = new
|
|
return decode_head
|
|
|
|
|
|
class MultiHeadSegFormer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
backbone: str,
|
|
channels: int,
|
|
heads_config: Dict[str, dict],
|
|
semantic_id2label: Dict[int, str],
|
|
semantic_label2id: Dict[str, int],
|
|
):
|
|
super().__init__()
|
|
|
|
semantic_classes = int(heads_config["semantic"].get("num_classes", len(semantic_id2label)))
|
|
|
|
base = SegformerForSemanticSegmentation.from_pretrained(
|
|
backbone,
|
|
num_labels=semantic_classes,
|
|
id2label={int(k): str(v) for k, v in semantic_id2label.items()},
|
|
label2id={str(k): int(v) for k, v in semantic_label2id.items()},
|
|
ignore_mismatched_sizes=True,
|
|
)
|
|
|
|
patch_segformer_encoder_input_channels(base.segformer, channels)
|
|
base.config.num_channels = int(channels)
|
|
|
|
self.segformer = base.segformer
|
|
self.decode_heads = nn.ModuleDict()
|
|
self.heads_config = heads_config
|
|
|
|
for head_name, hcfg in heads_config.items():
|
|
h = copy.deepcopy(base.decode_head)
|
|
h = replace_segformer_decode_classifier(h, int(hcfg["num_classes"]))
|
|
self.decode_heads[head_name] = h
|
|
|
|
self.config = base.config
|
|
|
|
def forward(self, pixel_values: torch.Tensor, head_names=None) -> Dict[str, torch.Tensor]:
|
|
outputs = self.segformer(
|
|
pixel_values=pixel_values,
|
|
output_hidden_states=True,
|
|
return_dict=True,
|
|
)
|
|
|
|
hidden_states = outputs.hidden_states
|
|
|
|
if head_names is None:
|
|
selected = list(self.decode_heads.keys())
|
|
else:
|
|
selected = [h for h in head_names if h in self.decode_heads]
|
|
|
|
return {
|
|
head_name: self.decode_heads[head_name](hidden_states)
|
|
for head_name in selected
|
|
}
|
|
|
|
|
|
class MultiHeadTester:
|
|
def __init__(
|
|
self,
|
|
config: dict,
|
|
ckpt_path: Path,
|
|
device: torch.device,
|
|
channels: int,
|
|
heads_config: Dict[str, dict],
|
|
semantic_id2label: Dict[int, str],
|
|
semantic_label2id: Dict[str, int],
|
|
mean: Optional[Sequence[float]],
|
|
std: Optional[Sequence[float]],
|
|
use_amp: bool = True,
|
|
):
|
|
self.config = config
|
|
self.ckpt_path = ckpt_path
|
|
self.device = device
|
|
self.channels = channels
|
|
self.heads_config = heads_config
|
|
self.use_amp = use_amp and device.type == "cuda"
|
|
self.runtime_mode = str(config.get("runtime_mode", "all")).lower()
|
|
|
|
self.mean = None if mean is None else torch.tensor(mean, dtype=torch.float32).view(1, channels, 1, 1).to(device)
|
|
self.std = None if std is None else torch.tensor(std, dtype=torch.float32).view(1, channels, 1, 1).to(device)
|
|
|
|
backbone = config.get("backbone", config.get("pretrained_model", "nvidia/mit-b1"))
|
|
print(f"[MODEL] backbone={backbone}")
|
|
print(f"[MODEL] ckpt={ckpt_path}")
|
|
|
|
self.model = MultiHeadSegFormer(
|
|
backbone=backbone,
|
|
channels=channels,
|
|
heads_config=heads_config,
|
|
semantic_id2label=semantic_id2label,
|
|
semantic_label2id=semantic_label2id,
|
|
)
|
|
|
|
self._load_checkpoint(ckpt_path)
|
|
self.model.to(device)
|
|
self.model.eval()
|
|
|
|
def _load_checkpoint(self, ckpt_path: Path):
|
|
ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
|
|
|
|
if isinstance(ckpt, dict):
|
|
for key in ("model", "model_state", "model_state_dict", "state_dict"):
|
|
if key in ckpt and isinstance(ckpt[key], dict):
|
|
state = ckpt[key]
|
|
break
|
|
else:
|
|
state = ckpt
|
|
else:
|
|
raise RuntimeError(f"Checkpoint em formato inesperado: {type(ckpt)}")
|
|
|
|
clean = {}
|
|
for k, v in state.items():
|
|
nk = k
|
|
for prefix in ("module.", "model."):
|
|
if nk.startswith(prefix):
|
|
nk = nk[len(prefix):]
|
|
clean[nk] = v
|
|
|
|
missing, unexpected = self.model.load_state_dict(clean, strict=False)
|
|
print(f"[MODEL] load_state_dict strict=False | missing={len(missing)} unexpected={len(unexpected)}")
|
|
if missing:
|
|
print("[MODEL] primeiros missing:", missing[:8])
|
|
if unexpected:
|
|
print("[MODEL] primeiros unexpected:", unexpected[:8])
|
|
|
|
def _normalize(self, x: torch.Tensor) -> torch.Tensor:
|
|
if self.mean is not None and self.std is not None:
|
|
return (x - self.mean) / torch.clamp(self.std, min=1e-6)
|
|
return x
|
|
|
|
@torch.inference_mode()
|
|
def infer(self, chw_01: np.ndarray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], float]:
|
|
x = torch.from_numpy(chw_01).unsqueeze(0).to(self.device, non_blocking=True)
|
|
x = self._normalize(x)
|
|
|
|
h, w = int(chw_01.shape[1]), int(chw_01.shape[2])
|
|
|
|
if self.device.type == "cuda":
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
|
|
with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp):
|
|
head_names = None
|
|
if self.runtime_mode in ("target_direct", "target_head"):
|
|
head_names = ["target"]
|
|
elif self.runtime_mode in ("operational", "target_op"):
|
|
head_names = ["vegetation", "cana"]
|
|
|
|
logits_by_head = self.model(pixel_values=x, head_names=head_names)
|
|
|
|
preds = {}
|
|
probs = {}
|
|
|
|
for head_name, logits in logits_by_head.items():
|
|
logits = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
|
|
prob = torch.softmax(logits, dim=1)[0]
|
|
pred = torch.argmax(prob, dim=0)
|
|
preds[head_name] = pred.detach().cpu().numpy().astype(np.uint8)
|
|
probs[head_name] = prob.detach().cpu().numpy().astype(np.float32)
|
|
|
|
if self.device.type == "cuda":
|
|
torch.cuda.synchronize()
|
|
t_ms = (time.perf_counter() - t0) * 1000.0
|
|
|
|
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
|
|
# ============================================================
|
|
|
|
def infer_experiment_tag(config: dict, channels: int) -> str:
|
|
fusion_mode = config.get("fusion_mode", "stacked")
|
|
return f"{fusion_mode}_raw{channels}"
|
|
|
|
|
|
def infer_save_dir(config: dict, config_dir: Path, channels: int) -> Path:
|
|
model_name = config.get("model_name", "test_multi")
|
|
modelo_folder = config.get("modelo", "segformer_b1")
|
|
exp_tag = infer_experiment_tag(config, channels)
|
|
return (config_dir / "backup" / modelo_folder / model_name / exp_tag).resolve()
|
|
|
|
|
|
def find_checkpoint(save_dir: Path, preferred: Optional[str] = None) -> Path:
|
|
if preferred is not None:
|
|
ckpt = Path(preferred)
|
|
if not ckpt.is_absolute():
|
|
ckpt_cwd = (Path.cwd() / ckpt).resolve()
|
|
ckpt_save = (save_dir / ckpt).resolve()
|
|
ckpt = ckpt_cwd if ckpt_cwd.is_file() else ckpt_save
|
|
if not ckpt.is_file():
|
|
raise FileNotFoundError(f"Checkpoint não encontrado: {ckpt}")
|
|
return ckpt
|
|
|
|
candidates = [
|
|
save_dir / "best_score.pt",
|
|
save_dir / "best_target.pt",
|
|
save_dir / "best_cana_head.pt",
|
|
save_dir / "best_semantic_miou.pt",
|
|
save_dir / "last.pt",
|
|
]
|
|
for c in candidates:
|
|
if c.is_file():
|
|
return c
|
|
|
|
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
|
|
# ============================================================
|
|
|
|
def tensor_to_preview_rgb(chw: np.ndarray, gamma: float = 0.85) -> np.ndarray:
|
|
c, h, w = chw.shape
|
|
if c >= 3:
|
|
rgb = np.transpose(chw[:3], (1, 2, 0)).copy()
|
|
else:
|
|
one = chw[0]
|
|
rgb = np.stack([one, one, one], axis=-1)
|
|
|
|
rgb = np.nan_to_num(rgb, nan=0.0, posinf=1.0, neginf=0.0)
|
|
lo = np.percentile(rgb, 1.0)
|
|
hi = np.percentile(rgb, 99.0)
|
|
if hi > lo:
|
|
rgb = (rgb - lo) / (hi - lo)
|
|
rgb = np.clip(rgb, 0.0, 1.0)
|
|
if gamma and gamma > 0:
|
|
rgb = np.power(rgb, gamma)
|
|
return (rgb * 255.0).astype(np.uint8)
|
|
|
|
|
|
def ids_to_rgb(mask: np.ndarray, colormap_rgb: Dict[int, Tuple[int, int, int]], ignore_id: int = 255) -> np.ndarray:
|
|
h, w = mask.shape[:2]
|
|
out = np.zeros((h, w, 3), dtype=np.uint8)
|
|
for cid, color in colormap_rgb.items():
|
|
out[mask == cid] = color
|
|
out[mask == ignore_id] = (0, 0, 0)
|
|
return out
|
|
|
|
|
|
def prob_to_heat_rgb(prob01: np.ndarray) -> np.ndarray:
|
|
p = np.clip(prob01, 0.0, 1.0)
|
|
u8 = (p * 255.0).astype(np.uint8)
|
|
bgr = cv2.applyColorMap(u8, cv2.COLORMAP_TURBO)
|
|
return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
def overlay_rgb(base_rgb: np.ndarray, mask_rgb: np.ndarray, alpha: float) -> np.ndarray:
|
|
return cv2.addWeighted(base_rgb, 1.0 - alpha, mask_rgb, alpha, 0.0)
|
|
|
|
|
|
def put_label(img_rgb: np.ndarray, title: str, subtitle: str = "") -> np.ndarray:
|
|
out = img_rgb.copy()
|
|
cv2.putText(out, title, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.68, (0, 0, 0), 4, cv2.LINE_AA)
|
|
cv2.putText(out, title, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.68, (255, 255, 255), 2, cv2.LINE_AA)
|
|
if subtitle:
|
|
cv2.putText(out, subtitle, (10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 0, 0), 3, cv2.LINE_AA)
|
|
cv2.putText(out, subtitle, (10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 255, 90), 1, cv2.LINE_AA)
|
|
return out
|
|
|
|
|
|
def resize_panel(img: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
|
|
w, h = size
|
|
if img.shape[1] == w and img.shape[0] == h:
|
|
return img
|
|
return cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST)
|
|
|
|
|
|
def compose_grid(panels: List[Tuple[str, np.ndarray, str]], cols: int = 3, max_width: int = 1800) -> np.ndarray:
|
|
if not panels:
|
|
return np.zeros((480, 640, 3), dtype=np.uint8)
|
|
|
|
base_h, base_w = panels[0][1].shape[:2]
|
|
labeled = []
|
|
for title, img, subtitle in panels:
|
|
img = resize_panel(img, (base_w, base_h))
|
|
labeled.append(put_label(img, title, subtitle))
|
|
|
|
rows = []
|
|
blank = np.zeros_like(labeled[0])
|
|
for i in range(0, len(labeled), cols):
|
|
row_imgs = labeled[i:i + cols]
|
|
while len(row_imgs) < cols:
|
|
row_imgs.append(blank.copy())
|
|
rows.append(np.hstack(row_imgs))
|
|
|
|
canvas = np.vstack(rows)
|
|
|
|
if canvas.shape[1] > max_width:
|
|
scale = max_width / canvas.shape[1]
|
|
canvas = cv2.resize(canvas, (int(canvas.shape[1] * scale), int(canvas.shape[0] * scale)), interpolation=cv2.INTER_AREA)
|
|
|
|
return canvas
|
|
|
|
|
|
def class_percent(mask: np.ndarray, class_id: int, ignore_id: int = 255) -> float:
|
|
valid = mask != ignore_id
|
|
den = int(valid.sum())
|
|
if den <= 0:
|
|
return 0.0
|
|
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
|
|
# ============================================================
|
|
|
|
def confusion_matrix_np(pred: np.ndarray, gt: np.ndarray, num_classes: int, ignore_id: int) -> np.ndarray:
|
|
if pred.shape != gt.shape:
|
|
pred = cv2.resize(pred.astype(np.uint8), (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
|
|
valid = gt != ignore_id
|
|
valid &= gt >= 0
|
|
valid &= gt < num_classes
|
|
gt_v = gt[valid].astype(np.int64)
|
|
pred_v = pred[valid].astype(np.int64)
|
|
pred_v = np.clip(pred_v, 0, num_classes - 1)
|
|
|
|
cm = np.bincount(num_classes * gt_v + pred_v, minlength=num_classes * num_classes)
|
|
return cm.reshape(num_classes, num_classes).astype(np.int64)
|
|
|
|
|
|
def metrics_from_cm(cm: np.ndarray) -> Tuple[np.ndarray, float, float]:
|
|
tp = np.diag(cm).astype(np.float64)
|
|
fp = cm.sum(axis=0).astype(np.float64) - tp
|
|
fn = cm.sum(axis=1).astype(np.float64) - tp
|
|
denom = tp + fp + fn
|
|
iou = np.divide(tp, denom, out=np.zeros_like(tp), where=denom > 0)
|
|
miou = float(np.mean(iou)) if len(iou) else 0.0
|
|
acc = float(tp.sum() / max(cm.sum(), 1))
|
|
return iou, miou, acc
|
|
|
|
|
|
def operational_target_mask(veg_mask: np.ndarray, cana_mask: np.ndarray, ignore_id: int = 255) -> np.ndarray:
|
|
out = np.zeros_like(veg_mask, dtype=np.uint8)
|
|
ignore = (veg_mask == ignore_id) | (cana_mask == ignore_id)
|
|
out[(veg_mask == 1) & (cana_mask == 0)] = 1
|
|
out[ignore] = ignore_id
|
|
return out
|
|
|
|
|
|
# ============================================================
|
|
# Main
|
|
# ============================================================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="config.json")
|
|
parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"])
|
|
parser.add_argument("--root_override", default=None)
|
|
parser.add_argument("--test_folder", default=None)
|
|
parser.add_argument("--ckpt", default=None)
|
|
parser.add_argument("--norm_stats", default=None)
|
|
parser.add_argument("--channels", type=int, default=None)
|
|
parser.add_argument("--resize_w", type=int, default=None)
|
|
parser.add_argument("--resize_h", type=int, default=None)
|
|
parser.add_argument("--alpha", type=float, default=0.45)
|
|
parser.add_argument("--ignore_id", type=int, default=None)
|
|
parser.add_argument("--no_amp", action="store_true")
|
|
parser.add_argument("--require_masks", action="store_true")
|
|
parser.add_argument("--out_dir", default="outputs_test_multihead")
|
|
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())
|
|
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)
|
|
config["runtime_mode"] = args.runtime_mode
|
|
|
|
input_channel_names = get_input_channel_names(config, args.channels)
|
|
input_channel_indices = get_input_channel_indices(config, args.channels)
|
|
channels = len(input_channel_names)
|
|
|
|
print(f"Input channels: {input_channel_names} idx={input_channel_indices}")
|
|
res = config.get("resolucao", [1024, 640])
|
|
default_w, default_h = int(res[0]), int(res[1])
|
|
target_w = int(args.resize_w or default_w)
|
|
target_h = int(args.resize_h or default_h)
|
|
|
|
dataset_path = config_dir / "dataset"
|
|
labelmap_path = dataset_path / "labelmap.txt"
|
|
semantic_id2label, semantic_label2id, labelmap_ignore_id, loaded_colormap_rgb = load_labelmap(labelmap_path)
|
|
|
|
ignore_id = int(args.ignore_id if args.ignore_id is not None else labelmap_ignore_id)
|
|
heads_config = build_heads_config(config, ignore_index=ignore_id)
|
|
heads_config["semantic"]["num_classes"] = int(len(semantic_id2label))
|
|
heads_config["semantic"]["ignore_index"] = int(ignore_id)
|
|
|
|
semantic_cmap = dict(SEMANTIC_COLORS_RGB)
|
|
semantic_cmap.update({int(k): tuple(map(int, v)) for k, v in loaded_colormap_rgb.items()})
|
|
|
|
save_dir = infer_save_dir(config, config_dir, channels)
|
|
ckpt_path = find_checkpoint(save_dir, args.ckpt)
|
|
|
|
if args.norm_stats is not None:
|
|
norm_stats_path = resolve_path(args.norm_stats, Path.cwd())
|
|
else:
|
|
# O treino multihead usa stats do dataset normalizado, mas também tentamos alguns fallbacks.
|
|
candidates = [
|
|
dataset_path / f"{default_w}x{default_h}" / "group" / "norm_stats.json",
|
|
save_dir / "norm_stats.json",
|
|
config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test_multi") / config.get("stats_source_tag", "stacked_raw5") / "norm_stats.json",
|
|
]
|
|
norm_stats_path = next((p for p in candidates if p.is_file()), candidates[0])
|
|
|
|
mean, std = load_norm_stats(
|
|
norm_stats_path,
|
|
channels=channels,
|
|
channel_indices=input_channel_indices,
|
|
channel_names=input_channel_names,
|
|
)
|
|
|
|
if args.test_folder is not None:
|
|
root = resolve_path(args.test_folder, Path.cwd())
|
|
elif args.root_override is not None:
|
|
root = resolve_path(args.root_override, Path.cwd())
|
|
else:
|
|
root = (dataset_path / "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, has_gt = collect_samples(root, heads_config=heads_config, require_masks=args.require_masks)
|
|
n = len(samples)
|
|
|
|
print("==========================================")
|
|
print("Teste SegFormer OAK-FCC-3 Multi-Head")
|
|
print(f"Root : {root}")
|
|
print(f"Samples : {n}")
|
|
print(f"GT : {'sim' if has_gt else 'não'}")
|
|
print(f"Resolution : {target_w}x{target_h}")
|
|
print(f"Channels : {channels}")
|
|
print(f"Semantic : {semantic_id2label}")
|
|
print(f"Ignore index: {ignore_id}")
|
|
print("Heads:")
|
|
for name, hcfg in heads_config.items():
|
|
print(f" - {name}: classes={hcfg['num_classes']} mask_dir={hcfg['mask_dir']}")
|
|
print("==========================================")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
tester = MultiHeadTester(
|
|
config=config,
|
|
ckpt_path=ckpt_path,
|
|
device=device,
|
|
channels=channels,
|
|
heads_config=heads_config,
|
|
semantic_id2label=semantic_id2label,
|
|
semantic_label2id=semantic_label2id,
|
|
mean=mean,
|
|
std=std,
|
|
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)
|
|
|
|
idx = max(0, min(args.start_idx, n - 1))
|
|
detailed = False
|
|
|
|
cms_total = {
|
|
name: np.zeros((int(hcfg["num_classes"]), int(hcfg["num_classes"])), dtype=np.int64)
|
|
for name, hcfg in heads_config.items()
|
|
}
|
|
cm_target_total = np.zeros((2, 2), dtype=np.int64)
|
|
visited = set()
|
|
|
|
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]
|
|
chw = load_sample_tensor(
|
|
sample,
|
|
channels=channels,
|
|
channel_indices=input_channel_indices,
|
|
)
|
|
|
|
if (chw.shape[2], chw.shape[1]) != (target_w, target_h):
|
|
hwc = np.transpose(chw, (1, 2, 0))
|
|
hwc = cv2.resize(hwc, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
|
|
chw = np.transpose(hwc, (2, 0, 1)).astype(np.float32)
|
|
|
|
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]
|
|
|
|
pred_sem = preds.get("semantic")
|
|
pred_veg = preds.get("vegetation")
|
|
pred_cana = preds.get("cana")
|
|
|
|
# Target operacional antigo: vegetation AND not cana
|
|
pred_target_op = None
|
|
if pred_veg is not None and pred_cana is not None:
|
|
pred_target_op = operational_target_mask(pred_veg, pred_cana, ignore_id=ignore_id)
|
|
|
|
# Target direta nova: saída da 4ª cabeça
|
|
pred_target_head = preds.get("target")
|
|
|
|
# Fallback visual: se não tiver head target, mostra operacional
|
|
pred_target = pred_target_head if pred_target_head is not None else pred_target_op
|
|
|
|
prob_veg = probs["vegetation"][1] if "vegetation" in probs and probs["vegetation"].shape[0] > 1 else None
|
|
prob_cana = probs["cana"][1] if "cana" in probs and probs["cana"].shape[0] > 1 else None
|
|
|
|
prob_target_op = None
|
|
if prob_veg is not None and prob_cana is not None:
|
|
prob_target_op = np.clip(prob_veg * (1.0 - prob_cana), 0.0, 1.0)
|
|
|
|
prob_target_head = None
|
|
if "target" in probs:
|
|
prob_target_head = probs["target"][1] if probs["target"].shape[0] > 1 else probs["target"][0]
|
|
|
|
prob_target = prob_target_head if prob_target_head is not None else prob_target_op
|
|
|
|
gt_masks = {name: load_mask(path) for name, path in sample.masks.items()}
|
|
for name, gt in list(gt_masks.items()):
|
|
if gt is not None and gt.shape != (pred_h, pred_w):
|
|
gt_masks[name] = cv2.resize(
|
|
gt.astype(np.uint8),
|
|
(pred_w, pred_h),
|
|
interpolation=cv2.INTER_NEAREST,
|
|
)
|
|
|
|
gt_target = None
|
|
if gt_masks.get("target") is not None:
|
|
gt_target = gt_masks["target"]
|
|
elif gt_masks.get("vegetation") is not None and gt_masks.get("cana") is not None:
|
|
gt_target = operational_target_mask(
|
|
gt_masks["vegetation"],
|
|
gt_masks["cana"],
|
|
ignore_id=ignore_id,
|
|
)
|
|
gt_masks["target"] = gt_target
|
|
|
|
# Métricas da amostra.
|
|
metric_lines = []
|
|
sample_metrics = {}
|
|
for head_name, pred in preds.items():
|
|
gt = gt_masks.get(head_name)
|
|
if gt is not None:
|
|
cm = confusion_matrix_np(pred, gt, int(heads_config[head_name]["num_classes"]), ignore_id)
|
|
iou, miou, acc = metrics_from_cm(cm)
|
|
sample_metrics[head_name] = {"iou": iou, "miou": miou, "acc": acc}
|
|
metric_lines.append(f"{head_name}: mIoU={miou:.3f} acc={acc:.3f}")
|
|
|
|
if gt_target is not None:
|
|
cm_t = confusion_matrix_np(pred_target, gt_target, 2, ignore_id)
|
|
iou_t, miou_t, acc_t = metrics_from_cm(cm_t)
|
|
sample_metrics["target_op"] = {"iou": iou_t, "miou": miou_t, "acc": acc_t}
|
|
metric_lines.append(f"target_op: IoU_alvo={iou_t[1]:.3f} acc={acc_t:.3f}")
|
|
if "target" in sample_metrics:
|
|
iou_head = sample_metrics["target"]["iou"]
|
|
metric_lines.append(f"target_head: IoU_alvo={iou_head[1]:.3f}")
|
|
|
|
if idx not in visited:
|
|
for head_name, pred in preds.items():
|
|
gt = gt_masks.get(head_name)
|
|
if gt is not None:
|
|
cms_total[head_name] += confusion_matrix_np(pred, gt, int(heads_config[head_name]["num_classes"]), ignore_id)
|
|
if gt_target is not None and pred_target_op is not None:
|
|
cm_target_total += confusion_matrix_np(pred_target_op, gt_target, 2, ignore_id)
|
|
visited.add(idx)
|
|
|
|
# Visuals.
|
|
panels: List[Tuple[str, np.ndarray, str]] = []
|
|
if pred_sem is not None:
|
|
pred_sem_rgb = ids_to_rgb(pred_sem, semantic_cmap, ignore_id)
|
|
if gt_masks.get("semantic") is not None:
|
|
panels.append(("GT semantic", ids_to_rgb(gt_masks["semantic"], semantic_cmap, ignore_id), "chao/cana/erva"))
|
|
panels.append(("Pred semantic", pred_sem_rgb, f"erva={class_percent(pred_sem, 2, ignore_id):.1f}% cana={class_percent(pred_sem, 1, ignore_id):.1f}%"))
|
|
panels.append(("Overlay semantic", overlay_rgb(preview_rgb, pred_sem_rgb, args.alpha), ""))
|
|
|
|
if pred_veg is not None:
|
|
pred_veg_rgb = ids_to_rgb(pred_veg, BINARY_COLORS_RGB, ignore_id)
|
|
if gt_masks.get("vegetation") is not None:
|
|
panels.append(("GT vegetation", ids_to_rgb(gt_masks["vegetation"], BINARY_COLORS_RGB, ignore_id), "0=fundo 1=veg"))
|
|
panels.append(("Pred vegetation", pred_veg_rgb, f"veg={class_percent(pred_veg, 1, ignore_id):.1f}%"))
|
|
if prob_veg is not None:
|
|
panels.append(("P vegetation", prob_to_heat_rgb(prob_veg), f"mean={float(prob_veg.mean()):.3f}"))
|
|
|
|
if pred_cana is not None:
|
|
pred_cana_rgb = ids_to_rgb(pred_cana, CANA_COLORS_RGB, ignore_id)
|
|
if gt_masks.get("cana") is not None:
|
|
panels.append(("GT cana", ids_to_rgb(gt_masks["cana"], CANA_COLORS_RGB, ignore_id), "0=not_cana 1=cana"))
|
|
panels.append(("Pred cana", pred_cana_rgb, f"cana={class_percent(pred_cana, 1, ignore_id):.1f}%"))
|
|
if prob_cana is not None:
|
|
panels.append(("P cana", prob_to_heat_rgb(prob_cana), f"mean={float(prob_cana.mean()):.3f}"))
|
|
|
|
if detailed:
|
|
if gt_target is not None:
|
|
panels.append(("GT target", ids_to_rgb(gt_target, TARGET_COLORS_RGB, ignore_id), "derivado"))
|
|
|
|
if pred_target_head is not None:
|
|
panels.append(("Pred target HEAD", ids_to_rgb(pred_target_head, TARGET_COLORS_RGB, ignore_id), f"alvo={class_percent(pred_target_head, 1, ignore_id):.1f}%"))
|
|
if prob_target_head is not None:
|
|
panels.append(("P target HEAD", prob_to_heat_rgb(prob_target_head), f"mean={float(prob_target_head.mean()):.3f}"))
|
|
|
|
if pred_target_op is not None:
|
|
panels.append(("Pred target OP", ids_to_rgb(pred_target_op, TARGET_COLORS_RGB, ignore_id), f"alvo={class_percent(pred_target_op, 1, ignore_id):.1f}%"))
|
|
if prob_target_op is not None:
|
|
panels.append(("P target OP", prob_to_heat_rgb(prob_target_op), f"mean={float(prob_target_op.mean()):.3f}"))
|
|
else:
|
|
if pred_target_head is not None:
|
|
rgb = ids_to_rgb(pred_target_head, TARGET_COLORS_RGB, ignore_id)
|
|
panels.append(("Pred target HEAD", rgb, f"alvo={class_percent(pred_target_head, 1, ignore_id):.1f}%"))
|
|
panels.append(("Overlay target HEAD", overlay_rgb(preview_rgb, rgb, args.alpha), "head direta"))
|
|
elif pred_target_op is not None:
|
|
rgb = ids_to_rgb(pred_target_op, TARGET_COLORS_RGB, ignore_id)
|
|
panels.append(("Pred target OP", rgb, f"alvo={class_percent(pred_target_op, 1, ignore_id):.1f}%"))
|
|
panels.append(("Overlay target OP", overlay_rgb(preview_rgb, rgb, args.alpha), "veg & !cana"))
|
|
|
|
canvas = compose_grid(panels, cols=3, max_width=args.max_width)
|
|
|
|
header_h = 78
|
|
header = np.zeros((header_h, canvas.shape[1], 3), dtype=np.uint8)
|
|
header[:] = (25, 25, 25)
|
|
source_name = sample.tensor_path.name if sample.tensor_path is not None else f"{sample.base}.json [{sample.source_kind}]"
|
|
h1 = f"idx {idx + 1}/{n} | {source_name} | inf={t_inf:.1f}ms | {'detalhado' if detailed else 'compacto'}"
|
|
h2 = " | ".join(metric_lines[:3]) if metric_lines else "sem GT"
|
|
cv2.putText(header, h1, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (235, 235, 235), 1, cv2.LINE_AA)
|
|
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):
|
|
break
|
|
elif k in (ord("d"), ord("D"), 83):
|
|
idx = (idx + 1) % n
|
|
elif k in (ord("a"), ord("A"), 81):
|
|
idx = (idx - 1 + n) % n
|
|
elif k == ord(" "):
|
|
detailed = not detailed
|
|
elif k in (ord("s"), ord("S")):
|
|
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:
|
|
print("\n========== RESUMO DOS SAMPLES VISITADOS ==========")
|
|
print(f"visitados={len(visited)}/{n}")
|
|
for head_name, cm in cms_total.items():
|
|
if cm.sum() <= 0:
|
|
continue
|
|
iou, miou, acc = metrics_from_cm(cm)
|
|
print(f"\n[{head_name}] acc={acc:.4f} mIoU={miou:.4f}")
|
|
for i, v in enumerate(iou):
|
|
print(f" IoU {i} = {v:.4f}")
|
|
|
|
if cm_target_total.sum() > 0:
|
|
iou, miou, acc = metrics_from_cm(cm_target_total)
|
|
print(f"\n[target operacional] acc={acc:.4f} mIoU={miou:.4f}")
|
|
print(f" IoU background = {iou[0]:.4f}")
|
|
print(f" IoU alvo = {iou[1]:.4f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|