838 lines
30 KiB
Python
838 lines
30 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Teste/visualização do SegFormer para o OAK-FCC-3 na estrutura nova.
|
|
|
|
Foco deste script:
|
|
1) Testar checkpoint treinado em backup/segformer_b1/test_1/stacked_raw5.
|
|
2) Ler dataset/split/{train,val,test} com tensors .npy e masks.
|
|
3) Ler pasta solta só com tensors .npy, sem GT.
|
|
4) Mostrar preview do tensor, máscara GT, predição e overlay.
|
|
5) Calcular métricas simples por imagem e resumo acumulado quando existe GT.
|
|
|
|
Exemplo, dentro da pasta do dataset oak-fcc-3:
|
|
|
|
python .\_9_test_infer.py ^
|
|
--config config.json ^
|
|
--split_folder val ^
|
|
--ckpt backup\segformer_b1\test_1\stacked_raw5\best_miou.pt
|
|
|
|
Ou usando last/checkpoint:
|
|
|
|
python .\_9_test_infer.py ^
|
|
--config config.json ^
|
|
--split_folder val ^
|
|
--ckpt backup\segformer_b1\test_1\stacked_raw5\last.pt
|
|
|
|
Controles:
|
|
D / seta direita : próxima amostra
|
|
A / seta esquerda: amostra anterior
|
|
S : salvar imagem atual em --out_dir
|
|
Espaço : alterna overlay/predição pura
|
|
Q / ESC : sair
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
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.functional as F
|
|
from transformers import SegformerForSemanticSegmentation
|
|
|
|
|
|
# ============================================================
|
|
# Configs simples
|
|
# ============================================================
|
|
|
|
DEFAULT_CLASS_NAMES = ["chao", "cana", "erva"]
|
|
DEFAULT_COLORMAP_RGB = {
|
|
0: (80, 80, 80), # chao
|
|
1: (0, 180, 0), # cana
|
|
2: (220, 60, 60), # erva
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class SampleItem:
|
|
tensor_path: Path
|
|
mask_path: Optional[Path] = None
|
|
|
|
|
|
# ============================================================
|
|
# Utilitários de path/config
|
|
# ============================================================
|
|
|
|
def _repo_root_from_script() -> Path:
|
|
return Path(__file__).resolve().parent
|
|
|
|
|
|
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 load_json(path: Path) -> dict:
|
|
with path.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def load_labelmap(labelmap_path: Path) -> Tuple[Dict[int, str], Dict[str, int], int, Dict[int, Tuple[int, int, int]]]:
|
|
"""
|
|
Carrega o labelmap do mesmo jeito do treino.
|
|
|
|
Prioridade:
|
|
1) helpers.carregar_labelmap_completo + helpers._infer_ignore_id
|
|
2) parser simples compatível com linhas tipo:
|
|
chao:128,0,0::
|
|
cana:0,0,128::
|
|
erva:0,128,0::
|
|
ignore:255,255,255::
|
|
ou:
|
|
0 chao
|
|
1 cana
|
|
2 erva
|
|
|
|
Importante: ignore/void/background_ignore NÃO entra em id2label.
|
|
Isso evita criar num_classes=4 quando o checkpoint foi treinado com 3 classes.
|
|
"""
|
|
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()}
|
|
id2label = {
|
|
int(k): str(v)
|
|
for k, v in id2label.items()
|
|
if str(v).lower() not in ("ignore", "void", "background_ignore")
|
|
}
|
|
|
|
label2id = {v.lower(): k for k, v in id2label.items()}
|
|
|
|
# O helper pode retornar colormap como:
|
|
# dict -> {0: (128,0,0), 1: (0,0,128), ...}
|
|
# lista -> [(128,0,0), (0,0,128), ...]
|
|
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]))
|
|
|
|
else:
|
|
print(f"[WARN] Tipo inesperado de colormap_rgb vindo do helper: {type(_colormap_rgb)}")
|
|
|
|
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 = dict(enumerate(DEFAULT_CLASS_NAMES))
|
|
|
|
label2id = {v.lower(): k for k, v in id2label.items()}
|
|
for cid in id2label:
|
|
colormap_rgb.setdefault(cid, DEFAULT_COLORMAP_RGB.get(cid, (255, 255, 255)))
|
|
|
|
return id2label, label2id, int(ignore_id), colormap_rgb
|
|
|
|
|
|
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_1")
|
|
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_miou.pt",
|
|
save_dir / "best.pt",
|
|
save_dir / "last.pt",
|
|
save_dir / "checkpoint_last.pt",
|
|
save_dir / "checkpoint_last.pth",
|
|
save_dir / "last.pth",
|
|
]
|
|
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)
|
|
)
|
|
|
|
|
|
# ============================================================
|
|
# Dataset novo: split/tensors + split/masks
|
|
# ============================================================
|
|
|
|
def _possible_mask_paths(tensor_path: Path, split_root: Path) -> List[Path]:
|
|
stem = tensor_path.stem
|
|
candidates: List[Path] = []
|
|
|
|
# Caso clássico:
|
|
# root/tensors/xxx.npy
|
|
# root/masks/xxx.png
|
|
if tensor_path.parent.name == "tensors":
|
|
base = tensor_path.parent.parent
|
|
masks_dir = base / "masks"
|
|
for ext in (".png", ".tif", ".tiff", ".npy"):
|
|
candidates.append(masks_dir / f"{stem}{ext}")
|
|
|
|
# Caso group/gX/tensors/xxx.npy + group/gX/masks/xxx.png
|
|
parts = tensor_path.parts
|
|
if "tensors" in parts:
|
|
tensor_dir = tensor_path.parent
|
|
base = tensor_dir.parent
|
|
masks_dir = base / "masks"
|
|
for ext in (".png", ".tif", ".tiff", ".npy"):
|
|
candidates.append(masks_dir / f"{stem}{ext}")
|
|
|
|
# Fallback amplo dentro do split
|
|
for ext in (".png", ".tif", ".tiff", ".npy"):
|
|
candidates.append(split_root / "masks" / f"{stem}{ext}")
|
|
|
|
return candidates
|
|
|
|
|
|
def collect_samples(root: Path, require_masks: bool = False) -> Tuple[List[SampleItem], bool]:
|
|
"""
|
|
Coleta .npy em:
|
|
root/tensors/*.npy
|
|
root/group/*/tensors/*.npy
|
|
root/**/tensors/*.npy fallback
|
|
|
|
Retorna (samples, has_gt).
|
|
"""
|
|
tensor_paths: List[Path] = []
|
|
|
|
direct = root / "tensors"
|
|
if direct.is_dir():
|
|
tensor_paths.extend(sorted(direct.glob("*.npy")))
|
|
|
|
group = root / "group"
|
|
if group.is_dir():
|
|
for gdir in sorted(group.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")))
|
|
|
|
# Último fallback: qualquer .npy, evitando masks/labels.
|
|
if not tensor_paths:
|
|
for p in sorted(root.glob("**/*.npy")):
|
|
low = str(p).lower()
|
|
if "mask" not in low and "label" not in low:
|
|
tensor_paths.append(p)
|
|
|
|
if not tensor_paths:
|
|
raise RuntimeError(f"Nenhum tensor .npy encontrado em: {root}")
|
|
|
|
samples: List[SampleItem] = []
|
|
gt_count = 0
|
|
|
|
for tp in tensor_paths:
|
|
mask_path = None
|
|
for mp in _possible_mask_paths(tp, root):
|
|
if mp.is_file():
|
|
mask_path = mp
|
|
break
|
|
if mask_path is not None:
|
|
gt_count += 1
|
|
samples.append(SampleItem(tensor_path=tp, mask_path=mask_path))
|
|
|
|
has_gt = gt_count > 0
|
|
if require_masks and gt_count != len(samples):
|
|
raise RuntimeError(f"Máscaras incompletas: {gt_count}/{len(samples)} samples têm GT")
|
|
|
|
return samples, has_gt
|
|
|
|
|
|
def load_tensor(path: Path, channels: int) -> np.ndarray:
|
|
arr = np.load(path).astype(np.float32)
|
|
|
|
# Aceita CHW ou HWC.
|
|
if arr.ndim != 3:
|
|
raise RuntimeError(f"Tensor inválido {path}: shape={arr.shape}, esperado 3D")
|
|
|
|
if arr.shape[0] == channels:
|
|
chw = arr
|
|
elif arr.shape[-1] == channels:
|
|
chw = np.transpose(arr, (2, 0, 1))
|
|
else:
|
|
raise RuntimeError(
|
|
f"Tensor {path.name} incompatível com channels={channels}: shape={arr.shape}"
|
|
)
|
|
|
|
# Segurança: se vier 0..255/0..65535, traz para 0..1.
|
|
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_mask(path: Path) -> np.ndarray:
|
|
if path.suffix.lower() == ".npy":
|
|
mask = np.load(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:
|
|
# Se vier RGB/BGR por acidente, pega primeiro canal.
|
|
mask = mask[:, :, 0]
|
|
return mask.astype(np.int64)
|
|
|
|
|
|
# ============================================================
|
|
# Visualização
|
|
# ============================================================
|
|
|
|
def make_colormap(num_classes: int, class_names: Sequence[str]) -> Dict[int, Tuple[int, int, int]]:
|
|
cmap = dict(DEFAULT_COLORMAP_RGB)
|
|
# Completa classes extras, se existirem.
|
|
rng_colors = [
|
|
(255, 180, 0),
|
|
(0, 160, 220),
|
|
(180, 0, 220),
|
|
(255, 255, 0),
|
|
(0, 255, 180),
|
|
]
|
|
for i in range(num_classes):
|
|
if i not in cmap:
|
|
cmap[i] = rng_colors[i % len(rng_colors)]
|
|
return cmap
|
|
|
|
|
|
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 tensor_to_preview_rgb(chw: np.ndarray, gamma: float = 0.85) -> np.ndarray:
|
|
"""
|
|
Preview visual do tensor [R,G,B,RE,NIR].
|
|
Usa RGB se existir; se channels < 3, replica o primeiro canal.
|
|
Não altera o tensor científico, é só vitrine.
|
|
"""
|
|
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)
|
|
|
|
# Stretch leve por percentil para visualizar melhor sem destruir a inferência.
|
|
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 draw_legend_rgb(colormap_rgb: Dict[int, Tuple[int, int, int]], class_names: Sequence[str], width: int) -> np.ndarray:
|
|
h = 42
|
|
legend = np.zeros((h, width, 3), dtype=np.uint8)
|
|
legend[:] = (25, 25, 25)
|
|
|
|
x = 12
|
|
for cid, name in enumerate(class_names):
|
|
color = colormap_rgb.get(cid, (255, 255, 255))
|
|
cv2.rectangle(legend, (x, 10), (x + 24, 34), color, -1)
|
|
cv2.putText(
|
|
legend,
|
|
f"{cid}:{name}",
|
|
(x + 32, 29),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.55,
|
|
(235, 235, 235),
|
|
1,
|
|
cv2.LINE_AA,
|
|
)
|
|
x += 32 + max(70, len(name) * 11)
|
|
return legend
|
|
|
|
|
|
def put_header(img_rgb: np.ndarray, lines: Sequence[str]) -> np.ndarray:
|
|
out = img_rgb.copy()
|
|
y = 24
|
|
for line in lines:
|
|
cv2.putText(out, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 0), 4, cv2.LINE_AA)
|
|
cv2.putText(out, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 255, 80), 2, cv2.LINE_AA)
|
|
y += 26
|
|
return out
|
|
|
|
|
|
# ============================================================
|
|
# Métricas
|
|
# ============================================================
|
|
|
|
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
|
|
|
|
|
|
# ============================================================
|
|
# Modelo
|
|
# ============================================================
|
|
|
|
class SegformerTester:
|
|
def __init__(
|
|
self,
|
|
config: dict,
|
|
ckpt_path: Path,
|
|
device: torch.device,
|
|
channels: int,
|
|
num_classes: int,
|
|
mean: Optional[Sequence[float]],
|
|
std: Optional[Sequence[float]],
|
|
ignore_id: int = 255,
|
|
use_amp: bool = True,
|
|
) -> None:
|
|
self.config = config
|
|
self.ckpt_path = ckpt_path
|
|
self.device = device
|
|
self.channels = channels
|
|
self.num_classes = num_classes
|
|
self.ignore_id = ignore_id
|
|
self.use_amp = use_amp and device.type == "cuda"
|
|
|
|
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 = SegformerForSemanticSegmentation.from_pretrained(
|
|
backbone,
|
|
num_labels=num_classes,
|
|
num_channels=channels,
|
|
ignore_mismatched_sizes=True,
|
|
)
|
|
|
|
self._load_checkpoint(ckpt_path)
|
|
self.model.to(device)
|
|
self.model.eval()
|
|
|
|
def _load_checkpoint(self, ckpt_path: Path) -> None:
|
|
ckpt = torch.load(str(ckpt_path), map_location="cpu")
|
|
|
|
if isinstance(ckpt, dict):
|
|
for key in ("model_state", "model_state_dict", "state_dict", "model"):
|
|
if key in ckpt and isinstance(ckpt[key], dict):
|
|
state = ckpt[key]
|
|
break
|
|
else:
|
|
# Pode ser o próprio state_dict.
|
|
state = ckpt
|
|
else:
|
|
raise RuntimeError(f"Checkpoint em formato inesperado: {type(ckpt)}")
|
|
|
|
# Remove prefixos comuns.
|
|
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[np.ndarray, float]:
|
|
x = torch.from_numpy(chw_01).unsqueeze(0).to(self.device, non_blocking=True)
|
|
x = self._normalize(x)
|
|
|
|
h, w = chw_01.shape[1], chw_01.shape[2]
|
|
t0 = time.perf_counter()
|
|
with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp):
|
|
out = self.model(pixel_values=x)
|
|
logits = out.logits
|
|
logits = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
|
|
pred = torch.argmax(logits, dim=1)[0]
|
|
if self.device.type == "cuda":
|
|
torch.cuda.synchronize()
|
|
t_ms = (time.perf_counter() - t0) * 1000.0
|
|
return pred.detach().cpu().numpy().astype(np.uint8), t_ms
|
|
|
|
|
|
# ============================================================
|
|
# Norm stats
|
|
# ============================================================
|
|
|
|
def load_norm_stats(path: Optional[Path], channels: int) -> 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 len(mean) != channels or len(std) != channels:
|
|
raise RuntimeError(
|
|
f"norm_stats incompatível com channels={channels}: 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))
|
|
|
|
|
|
# ============================================================
|
|
# Main
|
|
# ============================================================
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="config.json", help="Config usado no treino")
|
|
parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"], help="Split dentro de dataset/split")
|
|
parser.add_argument("--root_override", default=None, help="Pasta direta com tensors/masks")
|
|
parser.add_argument("--test_folder", default=None, help="Pasta avulsa com tensors, com ou sem masks")
|
|
parser.add_argument("--ckpt", default=None, help="Checkpoint .pt/.pth. Se vazio, tenta best_miou.pt/best.pt/last.pt")
|
|
parser.add_argument("--norm_stats", default=None, help="norm_stats.json. Se vazio, tenta no save_dir")
|
|
parser.add_argument("--channels", type=int, default=None, help="Override de canais. Normalmente 5")
|
|
parser.add_argument("--resize_w", type=int, default=None, help="Resize opcional antes da inferência")
|
|
parser.add_argument("--resize_h", type=int, default=None, help="Resize opcional antes da inferência")
|
|
parser.add_argument("--alpha", type=float, default=0.45, help="Alpha do overlay")
|
|
parser.add_argument("--ignore_id", type=int, default=255)
|
|
parser.add_argument("--no_amp", action="store_true", help="Desliga AMP na inferência")
|
|
parser.add_argument("--require_masks", action="store_true", help="Erro se alguma mask estiver faltando")
|
|
parser.add_argument("--out_dir", default="outputs_test", help="Pasta para salvar frames com S")
|
|
parser.add_argument("--start_idx", type=int, default=0)
|
|
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 nova do treino costuma ter:
|
|
# resolucao: [1024, 640]
|
|
# raw_size: [...]
|
|
# channels: 5
|
|
channels = int(args.channels or config.get("channels", 5))
|
|
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"
|
|
id2label, label2id, labelmap_ignore_id, loaded_colormap_rgb = load_labelmap(labelmap_path)
|
|
|
|
# O treino ignora a classe ignore. Logo, num_classes precisa ser só classes reais.
|
|
# No OAK-FCC-3 atual: 0=chao, 1=cana, 2=erva, ignore=255.
|
|
num_classes = int(config.get("num_classes", len(id2label)))
|
|
if num_classes != len(id2label):
|
|
print(f"[WARN] config.num_classes={num_classes}, mas labelmap tem {len(id2label)} classes reais. Usando labelmap.")
|
|
num_classes = len(id2label)
|
|
|
|
class_names = [id2label.get(i, f"class_{i}") for i in range(num_classes)]
|
|
colormap_rgb = make_colormap(num_classes, class_names)
|
|
colormap_rgb.update({int(k): tuple(map(int, v)) for k, v in loaded_colormap_rgb.items() if int(k) < num_classes})
|
|
|
|
# Se o usuário não informou manualmente, usa o ignore_id inferido pelo helper.
|
|
if args.ignore_id == 255:
|
|
args.ignore_id = int(labelmap_ignore_id)
|
|
|
|
save_dir = infer_save_dir(config, config_dir, channels)
|
|
ckpt_path = find_checkpoint(save_dir, args.ckpt)
|
|
|
|
norm_stats_path: Optional[Path]
|
|
if args.norm_stats is not None:
|
|
norm_stats_path = _resolve_path(args.norm_stats, Path.cwd())
|
|
else:
|
|
norm_stats_path = save_dir / "norm_stats.json"
|
|
mean, std = load_norm_stats(norm_stats_path, channels)
|
|
|
|
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, require_masks=args.require_masks)
|
|
n = len(samples)
|
|
|
|
print("==========================================")
|
|
print("Teste SegFormer OAK-FCC-3")
|
|
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"Classes : {num_classes} -> {dict(enumerate(class_names))}")
|
|
print(f"Label2Id : {label2id}")
|
|
print(f"Ignore index: {args.ignore_id}")
|
|
print("==========================================")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
tester = SegformerTester(
|
|
config=config,
|
|
ckpt_path=ckpt_path,
|
|
device=device,
|
|
channels=channels,
|
|
num_classes=num_classes,
|
|
mean=mean,
|
|
std=std,
|
|
ignore_id=args.ignore_id,
|
|
use_amp=not args.no_amp,
|
|
)
|
|
|
|
out_dir = Path(args.out_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
idx = max(0, min(args.start_idx, n - 1))
|
|
show_overlay = True
|
|
cm_total = np.zeros((num_classes, num_classes), dtype=np.int64)
|
|
visited = set()
|
|
|
|
win_name = "OAK-FCC-3 SegFormer Test | D/A navega | S salva | SPACE overlay | Q sai"
|
|
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
|
|
|
|
while True:
|
|
sample = samples[idx]
|
|
chw = load_tensor(sample.tensor_path, channels=channels)
|
|
|
|
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)
|
|
|
|
pred, t_inf = tester.infer(chw)
|
|
preview_rgb = tensor_to_preview_rgb(chw)
|
|
pred_rgb = ids_to_rgb(pred, colormap_rgb, args.ignore_id)
|
|
overlay = cv2.addWeighted(preview_rgb, 1.0 - args.alpha, pred_rgb, args.alpha, 0.0)
|
|
|
|
panels: List[np.ndarray] = [preview_rgb]
|
|
panel_names = ["tensor RGB preview"]
|
|
|
|
metric_line = "sem GT"
|
|
if sample.mask_path is not None:
|
|
gt = load_mask(sample.mask_path)
|
|
if gt.shape != pred.shape:
|
|
gt = cv2.resize(gt.astype(np.uint8), (pred.shape[1], pred.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
gt_rgb = ids_to_rgb(gt, colormap_rgb, args.ignore_id)
|
|
panels.append(gt_rgb)
|
|
panel_names.append("GT")
|
|
|
|
cm = confusion_matrix_np(pred, gt, num_classes, args.ignore_id)
|
|
iou, miou, acc = metrics_from_cm(cm)
|
|
metric_line = " | ".join([f"mIoU={miou:.3f}", f"acc={acc:.3f}"] + [f"{class_names[i]}={iou[i]:.3f}" for i in range(num_classes)])
|
|
|
|
if idx not in visited:
|
|
cm_total += cm
|
|
visited.add(idx)
|
|
|
|
panels.append(overlay if show_overlay else pred_rgb)
|
|
panel_names.append("overlay" if show_overlay else "pred")
|
|
|
|
# Escreve títulos em cada painel.
|
|
titled = []
|
|
for p, name in zip(panels, panel_names):
|
|
pp = p.copy()
|
|
cv2.putText(pp, name, (10, pp.shape[0] - 14), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 0), 4, cv2.LINE_AA)
|
|
cv2.putText(pp, name, (10, pp.shape[0] - 14), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 2, cv2.LINE_AA)
|
|
titled.append(pp)
|
|
|
|
result = np.concatenate(titled, axis=1)
|
|
legend = draw_legend_rgb(colormap_rgb, class_names, result.shape[1])
|
|
result = np.concatenate([result, legend], axis=0)
|
|
|
|
header = [
|
|
f"idx {idx + 1}/{n} | {sample.tensor_path.name} | inf={t_inf:.1f}ms | C={channels}",
|
|
metric_line,
|
|
]
|
|
result = put_header(result, header)
|
|
|
|
# Ajusta ao tamanho da janela sem converter internamente para RGB errado.
|
|
try:
|
|
_, _, win_w, win_h = cv2.getWindowImageRect(win_name)
|
|
except Exception:
|
|
win_w, win_h = 0, 0
|
|
|
|
display_rgb = result
|
|
if win_w > 64 and win_h > 64:
|
|
display_rgb = cv2.resize(result, (win_w, win_h), interpolation=cv2.INTER_NEAREST)
|
|
|
|
cv2.imshow(win_name, cv2.cvtColor(display_rgb, cv2.COLOR_RGB2BGR))
|
|
key = cv2.waitKey(0) & 0xFF
|
|
|
|
if key in (ord("q"), ord("Q"), 27):
|
|
break
|
|
if key in (ord("d"), ord("D"), 83): # 83 seta direita em alguns backends
|
|
idx = (idx + 1) % n
|
|
elif key in (ord("a"), ord("A"), 81):
|
|
idx = (idx - 1 + n) % n
|
|
elif key == ord(" "):
|
|
show_overlay = not show_overlay
|
|
elif key in (ord("s"), ord("S")):
|
|
out_path = out_dir / f"test_{idx:05d}_{sample.tensor_path.stem}.png"
|
|
cv2.imwrite(str(out_path), cv2.cvtColor(result, cv2.COLOR_RGB2BGR))
|
|
print(f"[SAVE] {out_path}")
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
if visited:
|
|
iou, miou, acc = metrics_from_cm(cm_total)
|
|
print("\n========== RESUMO DOS SAMPLES VISITADOS ==========")
|
|
print(f"visitados={len(visited)}/{n}")
|
|
print(f"pixel_acc={acc:.4f}")
|
|
print(f"mIoU={miou:.4f}")
|
|
for i, name in enumerate(class_names):
|
|
print(f"IoU {i}:{name} = {iou[i]:.4f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|