This commit is contained in:
Diego Freitas 2026-05-21 19:57:35 -03:00
parent eeb92bbbc0
commit 26aebaf070
43 changed files with 16545 additions and 5803 deletions

View File

@ -0,0 +1,147 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
from pathlib import Path
from collections import Counter, defaultdict
import argparse
def analisar_dataset(root: Path):
resultado = {}
for grupo_dir in sorted([p for p in root.iterdir() if p.is_dir()]):
labels_dir = grupo_dir / "labels"
if not labels_dir.is_dir():
continue
contagem = Counter()
total = 0
for json_path in sorted(labels_dir.glob("*.json")):
try:
with json_path.open("r", encoding="utf-8") as f:
data = json.load(f)
status = data.get("estado_corredor")
if not status:
status = "SEM_STATUS"
contagem[status] += 1
total += 1
except Exception:
contagem["ERRO_JSON"] += 1
total += 1
resultado[grupo_dir.name] = {
"total": total,
"contagem": contagem
}
return resultado
def imprimir_relatorio(resultado):
todos_status = sorted({
status
for dados in resultado.values()
for status in dados["contagem"].keys()
})
print("\nRELATÓRIO DE STATUS DO CORREDOR")
print("=" * 80)
total_geral = 0
contagem_geral = Counter()
for grupo, dados in resultado.items():
total = dados["total"]
contagem = dados["contagem"]
total_geral += total
contagem_geral.update(contagem)
print(f"\nGrupo: {grupo}")
print(f"Total labels: {total}")
if total == 0:
print(" Nenhum label encontrado.")
continue
for status in todos_status:
qtd = contagem.get(status, 0)
perc = (qtd / total) * 100
print(f" {status:<20} {qtd:>5} {perc:>6.2f}%")
print("\nGERAL")
print("=" * 80)
print(f"Total labels: {total_geral}")
if total_geral > 0:
for status in todos_status:
qtd = contagem_geral.get(status, 0)
perc = (qtd / total_geral) * 100
print(f" {status:<20} {qtd:>5} {perc:>6.2f}%")
def salvar_csv(resultado, out_csv: Path):
import csv
todos_status = sorted({
status
for dados in resultado.values()
for status in dados["contagem"].keys()
})
with out_csv.open("w", newline="", encoding="utf-8") as f:
writer = csv.writer(f)
header = ["grupo", "total"]
for status in todos_status:
header.append(f"{status}_qtd")
header.append(f"{status}_perc")
writer.writerow(header)
for grupo, dados in resultado.items():
total = dados["total"]
contagem = dados["contagem"]
row = [grupo, total]
for status in todos_status:
qtd = contagem.get(status, 0)
perc = (qtd / total) * 100 if total > 0 else 0
row.extend([qtd, round(perc, 2)])
writer.writerow(row)
print(f"\nCSV salvo em: {out_csv}")
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--root",
required=True,
help="Pasta original/group contendo os grupos com labels/"
)
parser.add_argument(
"--csv",
default=None,
help="Opcional: caminho para salvar o relatório CSV"
)
args = parser.parse_args()
root = Path(args.root)
resultado = analisar_dataset(root)
imprimir_relatorio(resultado)
if args.csv:
salvar_csv(resultado, Path(args.csv))
if __name__ == "__main__":
main()

View File

@ -0,0 +1,653 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
corridor_state_labeler.py
-------------------------------------------------
Rotulador rápido de estado global do corredor para dataset de segmentação.
Objetivo
- Percorrer uma raiz de dataset no formato:
dataset_root/
grupo_1/
images/
masks/
labels/ # criado automaticamente
grupo_2/
images/
masks/
labels/
grupo_3/
images/
masks/
labels/
- Mostrar imagem original e máscara lado a lado.
- Permitir classificar cada frame com teclas numéricas.
- Salvar um JSON por imagem dentro de labels/.
- Suportar resume, skip, voltar, avançar e undo.
Exemplo de labels globais:
Direcionando,EntrandoRua,CaminhandoRua,SaindoRua
Dependências:
pip install pillow
Uso com GUI:
python corridor_state_labeler.py
Uso por linha de comando:
python corridor_state_labeler.py --dataset-root "C:/.../oak-d/dataset/original/group" --states Direcionando EntrandoRua CaminhandoRua SaindoRua --resume
"""
import argparse
import json
import os
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import tkinter as tk
from tkinter import filedialog, messagebox
try:
from PIL import Image, ImageTk, ImageOps
except Exception:
print("ERRO: Pillow não encontrado. Instale com: pip install pillow", file=sys.stderr)
raise
IMG_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
MASK_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
IMAGE_DIR_NAMES = ["images", "image", "imgs", "rgb"]
MASK_DIR_NAMES = ["masks", "mask", "segmentacao", "seg", "segmentations"]
LABEL_DIR_NAME = "labels"
@dataclass
class SampleItem:
group: str
image_path: Path
mask_path: Path
label_path: Path
base: str
def normalize_stem(stem: str) -> str:
"""
Remove sufixos comuns para tentar casar image e mask mesmo quando nomes diferem.
Exemplos:
001_rgb -> 001
001_image -> 001
001_mask -> 001
001_segmentacao -> 001
"""
suffixes = [
"_rgb", "_RGB", "_Rgb",
"_image", "_img", "_frame",
"_mask", "_masks",
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
]
out = stem
changed = True
while changed:
changed = False
for sfx in suffixes:
if out.endswith(sfx):
out = out[: -len(sfx)]
changed = True
break
return out
def natural_key(text: str):
"""Ordenação natural simples: 2 antes de 10."""
import re
parts = re.split(r"(\d+)", text)
return [int(p) if p.isdigit() else p.lower() for p in parts]
def find_first_existing_dir(group_dir: Path, candidates: List[str]) -> Optional[Path]:
for name in candidates:
p = group_dir / name
if p.is_dir():
return p
return None
def list_image_files(folder: Path, exts: List[str]) -> List[Path]:
files = []
for p in folder.iterdir():
if p.is_file() and p.suffix.lower() in exts:
files.append(p)
files.sort(key=lambda x: natural_key(x.name))
return files
def collect_samples(dataset_root: Path) -> List[SampleItem]:
"""
Procura grupos dentro de dataset_root.
Cada grupo precisa ter uma pasta images/ e masks/.
Cria labels/ automaticamente.
"""
if not dataset_root.exists() or not dataset_root.is_dir():
raise FileNotFoundError(f"Raiz inválida: {dataset_root}")
samples: List[SampleItem] = []
group_dirs = [p for p in dataset_root.iterdir() if p.is_dir()]
group_dirs.sort(key=lambda x: natural_key(x.name))
for group_dir in group_dirs:
image_dir = find_first_existing_dir(group_dir, IMAGE_DIR_NAMES)
mask_dir = find_first_existing_dir(group_dir, MASK_DIR_NAMES)
if image_dir is None or mask_dir is None:
continue
label_dir = group_dir / LABEL_DIR_NAME
label_dir.mkdir(parents=True, exist_ok=True)
image_files = list_image_files(image_dir, IMG_EXTS)
mask_files = list_image_files(mask_dir, MASK_EXTS)
mask_by_base: Dict[str, Path] = {}
for m in mask_files:
mask_by_base[normalize_stem(m.stem)] = m
for img in image_files:
base = normalize_stem(img.stem)
mask = mask_by_base.get(base)
if mask is None:
# fallback: tenta mesmo stem exato
exact_candidates = [m for m in mask_files if m.stem == img.stem]
mask = exact_candidates[0] if exact_candidates else None
if mask is None:
continue
label_path = label_dir / f"{base}.json"
samples.append(
SampleItem(
group=group_dir.name,
image_path=img,
mask_path=mask,
label_path=label_path,
base=base,
)
)
samples.sort(key=lambda s: (natural_key(s.group), natural_key(s.base)))
return samples
def read_existing_label(label_path: Path) -> Optional[Dict]:
if not label_path.exists():
return None
try:
with label_path.open("r", encoding="utf-8") as f:
return json.load(f)
except Exception:
return None
def safe_relative(path: Path, root: Path) -> str:
try:
return str(path.resolve().relative_to(root.resolve())).replace("\\", "/")
except Exception:
return str(path).replace("\\", "/")
class CorridorStateLabelerApp:
def __init__(
self,
dataset_root: Path,
samples: List[SampleItem],
states: List[str],
resume: bool = True,
display_height: int = 520,
show_existing: bool = False,
):
self.dataset_root = dataset_root
self.states = states
self.display_height = int(display_height)
self.show_existing = show_existing
if resume:
self.samples = [s for s in samples if not s.label_path.exists()]
else:
self.samples = samples
self.idx = 0
self.history: List[Dict] = []
self.root = tk.Tk()
self.root.title("Agrobot - Corridor State Labeler")
self.root.geometry("1600x880")
self.root.bind("<Key>", self.on_key)
self.top_frame = tk.Frame(self.root)
self.top_frame.pack(side=tk.TOP, fill=tk.X)
self.info_label = tk.Label(self.top_frame, text="", font=("Segoe UI", 11), anchor="w")
self.info_label.pack(side=tk.LEFT, padx=10, pady=6, fill=tk.X, expand=True)
self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10), anchor="e")
self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6)
self.image_frame = tk.Frame(self.root)
self.image_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
self.left_frame = tk.LabelFrame(self.image_frame, text="Imagem")
self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)
self.right_frame = tk.LabelFrame(self.image_frame, text="Máscara")
self.right_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)
self.image_label = tk.Label(self.left_frame)
self.image_label.pack(fill=tk.BOTH, expand=True)
self.mask_label = tk.Label(self.right_frame)
self.mask_label.pack(fill=tk.BOTH, expand=True)
self.current_label_var = tk.StringVar(value="")
self.current_label = tk.Label(self.root, textvariable=self.current_label_var, font=("Segoe UI", 11), anchor="w")
self.current_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=2)
self.status_var = tk.StringVar(value="Pronto.")
self.status_label = tk.Label(self.root, textvariable=self.status_var, font=("Segoe UI", 10), anchor="w")
self.status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=4)
self.footer = tk.Label(
self.root,
text="1..9/0=classificar | Espaço=saltar | →/n=próxima | ←/p=anterior | b=desfazer | r=recarregar | q/Esc=sair",
font=("Segoe UI", 10),
)
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
self.render()
def build_legend_text(self) -> str:
parts = []
for i, state in enumerate(self.states, start=1):
key = i if i <= 9 else 0
parts.append(f"[{key}] {state}")
return " | ".join(parts)
def load_and_resize(self, path: Path) -> ImageTk.PhotoImage:
img = Image.open(path).convert("RGB")
img = ImageOps.exif_transpose(img)
w, h = img.size
if h <= 0:
raise ValueError(f"Imagem inválida: {path}")
new_h = self.display_height
new_w = max(1, int(w * (new_h / h)))
img = img.resize((new_w, new_h), Image.BILINEAR)
return ImageTk.PhotoImage(img)
def render(self):
if not self.samples:
messagebox.showinfo("Fim", "Nenhum item para rotular. Talvez tudo já esteja rotulado com resume ligado.")
self.root.destroy()
return
self.idx = max(0, min(self.idx, len(self.samples) - 1))
item = self.samples[self.idx]
try:
self.tk_img = self.load_and_resize(item.image_path)
self.tk_mask = self.load_and_resize(item.mask_path)
self.image_label.configure(image=self.tk_img)
self.mask_label.configure(image=self.tk_mask)
except Exception as e:
self.status_var.set(f"ERRO ao abrir imagem/máscara: {e}")
return
existing = read_existing_label(item.label_path)
existing_txt = ""
if existing:
existing_txt = f" | Já rotulado: {existing.get('estado_corredor')}"
self.info_label.configure(
text=(
f"{self.idx + 1}/{len(self.samples)} | Grupo: {item.group} | Base: {item.base} "
f"| Imagem: {item.image_path.name} | Máscara: {item.mask_path.name}{existing_txt}"
)
)
self.current_label_var.set(
f"Arquivo de label: {safe_relative(item.label_path, self.dataset_root)}"
)
def save_label(self, state_index: int):
if state_index < 0 or state_index >= len(self.states):
return
item = self.samples[self.idx]
state = self.states[state_index]
previous_content = None
previous_existed = item.label_path.exists()
if previous_existed:
try:
previous_content = item.label_path.read_text(encoding="utf-8")
except Exception:
previous_content = None
rec = {
"estado_corredor": state,
"label_id": state_index,
"states": self.states,
"group": item.group,
"base": item.base,
"image": safe_relative(item.image_path, self.dataset_root),
"mask": safe_relative(item.mask_path, self.dataset_root),
"label": safe_relative(item.label_path, self.dataset_root),
"timestamp": datetime.now().isoformat(timespec="seconds"),
"source": "corridor_state_labeler",
}
try:
item.label_path.parent.mkdir(parents=True, exist_ok=True)
with item.label_path.open("w", encoding="utf-8") as f:
json.dump(rec, f, ensure_ascii=False, indent=2)
self.history.append(
{
"action": "label",
"index": self.idx,
"label_path": item.label_path,
"previous_existed": previous_existed,
"previous_content": previous_content,
"new_state": state,
}
)
self.status_var.set(f"Salvo: {item.group}/{item.base}{state}")
self.idx += 1
if self.idx >= len(self.samples):
messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
self.root.destroy()
return
self.render()
except Exception as e:
messagebox.showerror("Erro", f"Falha ao salvar label: {e}")
self.status_var.set(f"ERRO ao salvar label: {e}")
def skip(self):
item = self.samples[self.idx]
self.history.append({"action": "skip", "index": self.idx})
self.status_var.set(f"Pulou: {item.group}/{item.base}")
self.idx += 1
if self.idx >= len(self.samples):
messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
self.root.destroy()
return
self.render()
def undo(self):
if not self.history:
self.status_var.set("Nada para desfazer.")
return
last = self.history.pop()
action = last.get("action")
if action == "label":
label_path: Path = last["label_path"]
try:
if last.get("previous_existed"):
previous_content = last.get("previous_content")
if previous_content is not None:
label_path.write_text(previous_content, encoding="utf-8")
else:
if label_path.exists():
label_path.unlink()
self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
self.status_var.set(f"Desfeito label: {label_path.name}")
self.render()
except Exception as e:
messagebox.showerror("Erro", f"Falha ao desfazer: {e}")
self.status_var.set(f"ERRO no undo: {e}")
elif action == "skip":
self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
self.status_var.set("Desfeito pulo.")
self.render()
def next_item(self):
self.idx = min(len(self.samples) - 1, self.idx + 1)
self.status_var.set("Avançou uma imagem.")
self.render()
def prev_item(self):
self.idx = max(0, self.idx - 1)
self.status_var.set("Voltou uma imagem.")
self.render()
def on_key(self, event):
ch = event.keysym.lower()
if ch in [str(i) for i in range(1, 10)] or ch == "0":
idx = 9 if ch == "0" else int(ch) - 1
self.save_label(idx)
return
if ch in ("space", "s"):
self.skip()
return
if ch in ("right", "n"):
self.next_item()
return
if ch in ("left", "p"):
self.prev_item()
return
if ch == "b":
self.undo()
return
if ch == "r":
self.status_var.set("Recarregado.")
self.render()
return
if ch in ("q", "escape"):
self.root.destroy()
return
def run(self):
self.root.mainloop()
class SetupWindow:
def __init__(self):
self.root = tk.Tk()
self.root.title("Configurar - Corridor State Labeler")
self.root.geometry("760x430")
frm_root = tk.LabelFrame(self.root, text="Raiz do dataset")
frm_root.pack(fill=tk.X, padx=10, pady=8)
self.dataset_root_var = tk.StringVar(value="")
tk.Entry(frm_root, textvariable=self.dataset_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6)
tk.Button(frm_root, text="Escolher...", command=self.choose_dataset_root).pack(side=tk.RIGHT, padx=6, pady=6)
frm_states = tk.LabelFrame(self.root, text="Estados do corredor separados por vírgula")
frm_states.pack(fill=tk.X, padx=10, pady=8)
self.states_var = tk.StringVar(value="Parado,EntrandoRua,CaminhandoRua,SaindoRua,Manobrando,Direcionando,RetornandoBase,Indefinido")
tk.Entry(frm_states, textvariable=self.states_var).pack(fill=tk.X, padx=6, pady=6)
frm_opts = tk.LabelFrame(self.root, text="Opções")
frm_opts.pack(fill=tk.X, padx=10, pady=8)
self.resume_var = tk.BooleanVar(value=True)
self.height_var = tk.IntVar(value=520)
tk.Checkbutton(frm_opts, text="Retomar: pular imagens que já possuem labels/*.json", variable=self.resume_var).pack(anchor="w", padx=6, pady=4)
frm_height = tk.Frame(frm_opts)
frm_height.pack(fill=tk.X, padx=6, pady=4)
tk.Label(frm_height, text="Altura de exibição em px:").pack(side=tk.LEFT)
tk.Entry(frm_height, textvariable=self.height_var, width=8).pack(side=tk.LEFT, padx=6)
self.preview_var = tk.StringVar(value="")
tk.Label(self.root, textvariable=self.preview_var, fg="#555", anchor="w", justify="left").pack(fill=tk.X, padx=12, pady=4)
btn_frame = tk.Frame(self.root)
btn_frame.pack(fill=tk.X, padx=10, pady=10)
tk.Button(btn_frame, text="Verificar dataset", command=self.preview_dataset).pack(side=tk.LEFT, padx=4)
tk.Button(btn_frame, text="Iniciar classificação", command=self.start).pack(side=tk.RIGHT, padx=4)
tk.Label(
self.root,
text="Teclas: 1..9/0=classificar | Espaço=saltar | b=desfazer | ←/→ navegar | q/Esc=sair",
fg="#555",
).pack(pady=4)
self.result = None
def choose_dataset_root(self):
p = filedialog.askdirectory(title="Selecione a raiz do dataset")
if p:
self.dataset_root_var.set(p)
self.preview_dataset()
def parse_states(self) -> List[str]:
return [s.strip() for s in self.states_var.get().split(",") if s.strip()]
def preview_dataset(self):
root_raw = self.dataset_root_var.get().strip()
if not root_raw:
self.preview_var.set("Selecione uma raiz de dataset.")
return
root = Path(root_raw)
try:
samples = collect_samples(root)
groups = sorted(set(s.group for s in samples), key=natural_key)
labeled = sum(1 for s in samples if s.label_path.exists())
self.preview_var.set(
f"Amostras encontradas: {len(samples)} | Já rotuladas: {labeled} | Grupos: {', '.join(groups) if groups else '-'}"
)
except Exception as e:
self.preview_var.set(f"Erro ao verificar dataset: {e}")
def start(self):
root_raw = self.dataset_root_var.get().strip()
states = self.parse_states()
if not root_raw:
messagebox.showwarning("Faltando raiz", "Selecione a raiz do dataset.")
return
if not states:
messagebox.showwarning("Faltando estados", "Informe pelo menos um estado.")
return
if len(states) > 10:
messagebox.showwarning("Muitos estados", "Este rotulador suporta até 10 estados nas teclas 1..9 e 0.")
return
root = Path(root_raw)
try:
samples = collect_samples(root)
except Exception as e:
messagebox.showerror("Erro", f"Falha ao ler dataset: {e}")
return
if not samples:
messagebox.showinfo(
"Sem amostras",
"Nenhum par image/mask encontrado. Verifique se cada grupo possui images/ e masks/.",
)
return
self.result = {
"dataset_root": root,
"states": states,
"resume": self.resume_var.get(),
"display_height": self.height_var.get(),
}
self.root.destroy()
def run(self):
self.root.mainloop()
return self.result
def run_gui_setup():
setup = SetupWindow()
res = setup.run()
if not res:
return
dataset_root = res["dataset_root"]
samples = collect_samples(dataset_root)
app = CorridorStateLabelerApp(
dataset_root=dataset_root,
samples=samples,
states=res["states"],
resume=res["resume"],
display_height=res["display_height"],
)
app.run()
def main():
parser = argparse.ArgumentParser(description="Rotulador de estado global do corredor para dataset de segmentação.")
parser.add_argument("--dataset-root", help="Raiz do dataset contendo grupos com images/ e masks/.")
parser.add_argument("--states", nargs="+", help="Estados do corredor. Ex: Direcionando EntrandoRua CaminhandoRua SaindoRua")
parser.add_argument("--resume", action="store_true", help="Pular itens que já possuem labels/*.json.")
parser.add_argument("--no-resume", action="store_true", help="Não pular itens já rotulados.")
parser.add_argument("--display-height", type=int, default=520, help="Altura de exibição das imagens em px.")
parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir janela de configuração.")
args = parser.parse_args()
need_gui = not args.no_gui_setup and (not args.dataset_root or not args.states)
if need_gui:
run_gui_setup()
return
if not args.dataset_root or not args.states:
print("ERRO: informe --dataset-root e --states, ou rode sem argumentos para abrir a GUI.", file=sys.stderr)
sys.exit(1)
dataset_root = Path(args.dataset_root)
states = args.states
if len(states) > 10:
print("ERRO: máximo de 10 estados: teclas 1..9 e 0.", file=sys.stderr)
sys.exit(1)
samples = collect_samples(dataset_root)
if not samples:
print("ERRO: nenhum par image/mask encontrado.", file=sys.stderr)
sys.exit(1)
resume = True
if args.no_resume:
resume = False
elif args.resume:
resume = True
app = CorridorStateLabelerApp(
dataset_root=dataset_root,
samples=samples,
states=states,
resume=resume,
display_height=args.display_height,
)
app.run()
if __name__ == "__main__":
main()

View File

@ -7,7 +7,7 @@ import argparse
import numpy as np
# ⚙️ Configurações
with open("config_oak.json", "r") as f:
with open("config.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]

View File

@ -1,22 +1,34 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Agrupa imagens e máscaras novas em subpastas por combinação de classes presentes.
Agrupa imagens, máscaras e labels globais em subpastas por combinação de classes presentes.
Estrutura lida (via config.json -> MODELO):
MODELO/dataset/original/new_images/
MODELO/dataset/original/new_masks/
Estrutura lida via config.json -> camera:
MODELO/dataset/original/images/
MODELO/dataset/original/masks/
MODELO/dataset/original/masks2/ opcional, se dual_head_mask=true
MODELO/dataset/original/labels/ opcional/necessário, se dual_head_label=true
Saída:
MODELO/dataset/original/group/<grupo>/images
MODELO/dataset/original/group/<grupo>/masks
Onde <grupo> é os nomes das classes presentes unidos por "_", ex:
chao, erva, cana, chao_erva, erva_cana, chao_erva_cana, etc.
MODELO/dataset/original/group/<grupo>/images/
MODELO/dataset/original/group/<grupo>/masks/
MODELO/dataset/original/group/<grupo>/masks2/ se dual_head_mask=true
MODELO/dataset/original/group/<grupo>/labels/ se dual_head_label=true
Requer: utils.carregar_labelmap_completo(labelmap_path)
O labelmap define mapeamento de cores/ids/nomes das classes.
Onde <grupo> é o nome das classes presentes na máscara unidos por "_", ex:
navegavel
naonavegavel
naonavegavel_navegavel
Observação importante:
- dual_head_mask = segunda cabeça também é máscara pixel-a-pixel.
- dual_head_label = segunda cabeça é label global por frame, vindo de JSON.
O JSON de label é copiado para labels/ com o mesmo base name final da imagem/máscara.
"""
import os
import cv2
import csv
@ -27,42 +39,60 @@ import numpy as np
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ====================== Configurações base ======================
def carregar_config_e_paths():
with open("config_oak.json", "r", encoding="utf-8") as f:
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
EXT_MASKS = (".png", ".jpg", ".jpeg")
EXT_MASKS2 = (".png", ".jpg", ".jpeg")
EXT_LABELS = (".json", ".txt")
MANIFESTO_DEFAULT = "manifest.csv"
def carregar_config_e_paths(modelo_cli=None):
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
MODELO = modelo_cli or config.get("camera")
USE_MASKS2 = bool(config.get("dual_head_mask", False))
USE_LABELS = bool(config.get("dual_head_label", False))
pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
# Pastas origem/destino
PASTA_NEW_IMAGES = os.path.join(pasta_base, "original", "images")
PASTA_NEW_MASKS = os.path.join(pasta_base, "original", "masks")
PASTA_NEW_MASKS2 = os.path.join(pasta_base, "original", "masks2")
PASTA_NEW_LABELS = os.path.join(pasta_base, "original", "labels")
PASTA_FINAL = os.path.join(pasta_base, "original", "group")
return MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_NEW_MASKS2, USE_MASKS2, PASTA_FINAL
return {
"MODELO": MODELO,
"pasta_base": pasta_base,
"labelmap_path": labelmap_path,
"PASTA_NEW_IMAGES": PASTA_NEW_IMAGES,
"PASTA_NEW_MASKS": PASTA_NEW_MASKS,
"PASTA_NEW_MASKS2": PASTA_NEW_MASKS2,
"PASTA_NEW_LABELS": PASTA_NEW_LABELS,
"PASTA_FINAL": PASTA_FINAL,
"USE_MASKS2": USE_MASKS2,
"USE_LABELS": USE_LABELS,
}
# Extensões aceitas
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
EXT_MASKS = (".png", ".jpg", ".jpeg") # prioridade será .png quando houver
EXT_MASKS2 = (".png", ".jpg", ".jpeg") # idem
# Manifesto padrão
MANIFESTO_DEFAULT = "manifest.csv"
# ====================== Utilitários ======================
def garantir_pasta(p):
os.makedirs(p, exist_ok=True)
def nome_disponivel(dest_dir, base_name, ext):
"""Gera nome único em dest_dir com sufixos _001, _002, ... se necessário."""
cand = os.path.join(dest_dir, base_name + ext)
if not os.path.exists(cand):
return cand
i = 1
while True:
cand = os.path.join(dest_dir, f"{base_name}_{i:03d}{ext}")
@ -70,41 +100,69 @@ def nome_disponivel(dest_dir, base_name, ext):
return cand
i += 1
def mapear_masks_por_base(pasta_masks):
"""Retorna {base: caminho_mask}, priorizando .png se houver múltiplas por base."""
def normalizar_base(stem):
"""
Tenta remover sufixos comuns para casar image/mask/label.
Ex:
001_rgb -> 001
001_segmentacao -> 001
001_mask -> 001
"""
sufixos = [
"_rgb", "_RGB", "_Rgb",
"_image", "_img", "_frame",
"_mask", "_masks",
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
"_label", "_labels",
]
out = stem
mudou = True
while mudou:
mudou = False
for sfx in sufixos:
if out.endswith(sfx):
out = out[: -len(sfx)]
mudou = True
break
return out
def mapear_arquivos_por_base(pasta, extensoes, normalizar=True):
"""
Retorna {base: caminho}, priorizando .png para máscaras/imagens e .json para labels.
"""
if not pasta or not os.path.isdir(pasta):
return {}
mapa = {}
for nome in os.listdir(pasta_masks):
prioridade = {
".json": 0,
".png": 1,
".jpg": 2,
".jpeg": 3,
".txt": 4,
}
for nome in os.listdir(pasta):
lower = nome.lower()
if not lower.endswith(EXT_MASKS):
if not lower.endswith(extensoes):
continue
base, ext = os.path.splitext(nome)
cam = os.path.join(pasta_masks, nome)
stem, ext = os.path.splitext(nome)
base = normalizar_base(stem) if normalizar else stem
caminho = os.path.join(pasta, nome)
if base not in mapa:
mapa[base] = cam
mapa[base] = caminho
else:
atual_ext = os.path.splitext(mapa[base])[1].lower()
if atual_ext != ".png" and ext.lower() == ".png":
mapa[base] = cam
if prioridade.get(ext.lower(), 99) < prioridade.get(atual_ext, 99):
mapa[base] = caminho
return mapa
def mapear_masks2_por_base(pasta_masks2):
"""Retorna {base: caminho_mask2}, priorizando .png se houver múltiplas por base."""
if not pasta_masks2 or not os.path.isdir(pasta_masks2):
return {}
mapa = {}
for nome in os.listdir(pasta_masks2):
lower = nome.lower()
if not lower.endswith(EXT_MASKS2):
continue
base, ext = os.path.splitext(nome)
cam = os.path.join(pasta_masks2, nome)
if base not in mapa:
mapa[base] = cam
else:
atual_ext = os.path.splitext(mapa[base])[1].lower()
if atual_ext != ".png" and ext.lower() == ".png":
mapa[base] = cam
return mapa
def localizar_imagem_por_base(pasta_imgs, base):
"""Retorna caminho da imagem correspondente ao base se existir."""
@ -112,119 +170,123 @@ def localizar_imagem_por_base(pasta_imgs, base):
p = os.path.join(pasta_imgs, base + ext)
if os.path.isfile(p):
return p
return None
# fallback com normalização
mapa = mapear_arquivos_por_base(pasta_imgs, EXT_IMAGENS, normalizar=True)
return mapa.get(base)
def inferir_ignore_id(ignore_rgb, cor_para_id):
"""
Tenta inferir o ID da classe ignorada a partir do retorno ignore_rgb e do mapa cor->id.
- Se ignore_rgb for [id] ou (id,), retorna esse id.
- Se ignore_rgb parecer uma cor RGB (len==3), usa cor_para_id[(R,G,B)] se existir.
- Caso não consiga, retorna None.
"""
if ignore_rgb is None:
return None
try:
# caso [id]
if isinstance(ignore_rgb, (list, tuple)) and len(ignore_rgb) == 1:
return int(ignore_rgb[0])
# caso [R,G,B]
if isinstance(ignore_rgb, (list, tuple)) and len(ignore_rgb) == 3:
key = tuple(int(v) for v in ignore_rgb)
return cor_para_id.get(key)
except Exception:
pass
# pode já ser um inteiro simples
if isinstance(ignore_rgb, (int, np.integer)):
return int(ignore_rgb)
return None
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
"""
Versão corrigida e otimizada:
- Se máscara for 1 canal: np.unique direto -> IDs.
- Se for 3 canais:
* se assume_rgb=True: labelmap está em RGB,
mas OpenCV BGR -> convertemos BGR -> RGB.
* se assume_rgb=False: labelmap está em BGR,
mantemos BGR como está.
- Usa converter_mask_rgb_para_ids em amostragem + fallback full-scan.
Retorna: set(ids_presentes)
"""
m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if m is None:
raise RuntimeError(f"Falha ao abrir máscara: {mask_path}")
# ---------------------------
# CASO 1: máscara indexada
# ---------------------------
# Máscara indexada
if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1):
vals = np.unique(m)
return set(int(v) for v in vals)
# ---------------------------
# CASO 2: máscara RGB
# ---------------------------
# Se o labelmap está em RGB (assume_rgb=True),
# convertemos a imagem BGR->RGB para casar com as chaves.
# Máscara RGB/BGR
if assume_rgb:
img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB)
else:
# labelmap já está em BGR; OpenCV entrega BGR; deixa como está
img = m
# Aqui as chaves de cor_para_id estão no MESMO espaço de cor da imagem.
mapa_rgb = cor_para_id
max_classes = len(cor_para_id)
# ---------- AMOSTRAGEM RÁPIDA ----------
step = 8 # pode ajustar para 4 se quiser mais precisão
# Amostragem rápida
step = 8
amostra = img[::step, ::step]
amostra_ids = converter_mask_rgb_para_ids(amostra, mapa_rgb, ignore_id=255)
amostra_ids = converter_mask_rgb_para_ids(amostra, cor_para_id, ignore_id=255)
ids = set(int(x) for x in np.unique(amostra_ids) if x != 255)
if len(ids) >= max_classes:
return ids
# ---------- FULL-SCAN (fallback) ----------
full_ids = converter_mask_rgb_para_ids(img, mapa_rgb, ignore_id=255)
# Full scan
full_ids = converter_mask_rgb_para_ids(img, cor_para_id, ignore_id=255)
ids = set(int(x) for x in np.unique(full_ids) if x != 255)
return ids
def montar_nome_grupo(ids_presentes, id_para_nome):
"""
Constrói o nome do grupo respeitando a ordem natural dos IDs do labelmap.
Ex: {0,1,2} -> chao_cana_obstaculo
"""
#print(f"ids_presentes: {ids_presentes}")
if not ids_presentes:
return "sem_classe"
nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)]
return "_".join(nomes)
def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False, mask2_src=None, dest_mask2_dir=None):
def copiar_arquivo(src, dst, mover=False):
if src is None or dst is None:
return None
garantir_pasta(os.path.dirname(dst))
if mover:
shutil.move(src, dst)
else:
shutil.copy2(src, dst)
return dst
def copiar_ou_mover_pacote(
img_src,
mask_src,
dest_img_dir,
dest_mask_dir,
mover=False,
mask2_src=None,
dest_mask2_dir=None,
label_src=None,
dest_label_dir=None,
):
"""
Copia/move image, mask, mask2 e label usando o MESMO base name final.
Isso é essencial para o loader encontrar os pares depois.
"""
garantir_pasta(dest_img_dir)
garantir_pasta(dest_mask_dir)
base_img = os.path.splitext(os.path.basename(img_src))[0]
img_ext = os.path.splitext(img_src)[1].lower()
img_ext = os.path.splitext(img_src)[1].lower()
mask_ext = os.path.splitext(mask_src)[1].lower()
dst_img = nome_disponivel(dest_img_dir, base_img, img_ext)
new_base = os.path.splitext(os.path.basename(dst_img))[0]
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
dst_mask2 = None
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
# Em caso raro de colisão na mask, recalcula usando ela como referência
if os.path.exists(dst_mask):
# evita colisão invertendo a ordem do "único" para a máscara
dst_mask = nome_disponivel(dest_mask_dir, new_base, mask_ext)
new_base = os.path.splitext(os.path.basename(dst_mask))[0]
dst_img = os.path.join(dest_img_dir, new_base + img_ext)
dst_img = os.path.join(dest_img_dir, new_base + img_ext)
if os.path.exists(dst_img):
dst_img = nome_disponivel(dest_img_dir, new_base, img_ext)
new_base = os.path.splitext(os.path.basename(dst_img))[0]
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
# se tiver mask2, usa o MESMO new_base final
dst_mask2 = None
if mask2_src and dest_mask2_dir:
garantir_pasta(dest_mask2_dir)
mask2_ext = os.path.splitext(mask2_src)[1].lower()
@ -232,50 +294,133 @@ def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False,
if os.path.exists(dst_mask2):
dst_mask2 = nome_disponivel(dest_mask2_dir, new_base, mask2_ext)
if mover:
shutil.move(img_src, dst_img)
shutil.move(mask_src, dst_mask)
if mask2_src and dst_mask2:
shutil.move(mask2_src, dst_mask2)
else:
shutil.copy2(img_src, dst_img)
shutil.copy2(mask_src, dst_mask)
if mask2_src and dst_mask2:
shutil.copy2(mask2_src, dst_mask2)
dst_label = None
if label_src and dest_label_dir:
garantir_pasta(dest_label_dir)
label_ext = os.path.splitext(label_src)[1].lower()
dst_label = os.path.join(dest_label_dir, new_base + label_ext)
if os.path.exists(dst_label):
dst_label = nome_disponivel(dest_label_dir, new_base, label_ext)
copiar_arquivo(img_src, dst_img, mover=mover)
copiar_arquivo(mask_src, dst_mask, mover=mover)
copiar_arquivo(mask2_src, dst_mask2, mover=mover)
copiar_arquivo(label_src, dst_label, mover=mover)
return dst_img, dst_mask, dst_mask2, dst_label
def validar_dimensoes(img_path, mask_path, mask2_path=None, estrito=False):
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
msk = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if img is None or msk is None:
raise RuntimeError("Falha ao abrir img/mask.")
hi, wi = img.shape[:2]
hm, wm = msk.shape[:2]
dim_mismatch = False
if (hi, wi) != (hm, wm):
dim_mismatch = True
msg = f"[AVISO] Dimensões diferem: img {wi}x{hi} vs mask {wm}x{hm} para '{os.path.basename(img_path)}'"
if estrito:
print(msg + " → pulando.")
return False, dim_mismatch
print(msg + " → copiando mesmo assim.")
if mask2_path:
m2 = cv2.imread(mask2_path, cv2.IMREAD_UNCHANGED)
if m2 is None:
print(f"[AVISO] Falha ao abrir mask2: {mask2_path} → ignorando mask2.")
return True, dim_mismatch
h2, w2 = m2.shape[:2]
if (hi, wi) != (h2, w2):
dim_mismatch = True
msg2 = f"[AVISO] Dimensões diferem: img {wi}x{hi} vs mask2 {w2}x{h2} para '{os.path.basename(img_path)}'"
if estrito:
print(msg2 + " → pulando.")
return False, dim_mismatch
print(msg2 + " → copiando mesmo assim.")
return True, dim_mismatch
return dst_img, dst_mask, dst_mask2
# ====================== Pipeline principal ======================
def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT,
validar_dim=True, estrito=False, labelmap_bgr=False):
# carrega config/paths
MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_NEW_MASKS2, USE_MASKS2, PASTA_FINAL = carregar_config_e_paths()
# carrega labelmap completo
def processar(
modelo_cli=None,
mover=False,
manifesto=MANIFESTO_DEFAULT,
validar_dim=True,
estrito=False,
labelmap_bgr=False,
strict_label=False,
):
paths = carregar_config_e_paths(modelo_cli=modelo_cli)
MODELO = paths["MODELO"]
labelmap_path = paths["labelmap_path"]
PASTA_NEW_IMAGES = paths["PASTA_NEW_IMAGES"]
PASTA_NEW_MASKS = paths["PASTA_NEW_MASKS"]
PASTA_NEW_MASKS2 = paths["PASTA_NEW_MASKS2"]
PASTA_NEW_LABELS = paths["PASTA_NEW_LABELS"]
PASTA_FINAL = paths["PASTA_FINAL"]
USE_MASKS2 = paths["USE_MASKS2"]
USE_LABELS = paths["USE_LABELS"]
print(f"[INFO] MODELO: {MODELO}")
print(f"[INFO] dual_head_mask: {USE_MASKS2}")
print(f"[INFO] dual_head_label: {USE_LABELS}")
cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = inferir_ignore_id(ignore_rgb, cor_para_id)
#print(f"cor_para_id: {cor_para_id}, _colormap_rgb: {_colormap_rgb}, id_para_nome: {id_para_nome}")
# garante pastas
garantir_pasta(PASTA_NEW_IMAGES)
garantir_pasta(PASTA_NEW_MASKS)
garantir_pasta(PASTA_FINAL)
usar_masks2 = USE_MASKS2 and os.path.isdir(PASTA_NEW_MASKS2)
if usar_masks2:
garantir_pasta(PASTA_NEW_MASKS2)
print(f"[INFO] masks2 detectada: {PASTA_NEW_MASKS2}")
# indexa máscaras
mapa_masks = mapear_masks_por_base(PASTA_NEW_MASKS)
mapa_masks2 = mapear_masks2_por_base(PASTA_NEW_MASKS2) if usar_masks2 else {}
usar_masks2 = USE_MASKS2 and os.path.isdir(PASTA_NEW_MASKS2)
usar_labels = USE_LABELS and os.path.isdir(PASTA_NEW_LABELS)
if USE_MASKS2 and not usar_masks2:
print(f"[AVISO] dual_head_mask=true, mas pasta masks2 não existe: {PASTA_NEW_MASKS2}")
if USE_LABELS and not usar_labels:
msg = f"[AVISO] dual_head_label=true, mas pasta labels não existe: {PASTA_NEW_LABELS}"
if strict_label:
raise RuntimeError(msg)
print(msg)
if usar_masks2:
print(f"[INFO] Usando masks2: {PASTA_NEW_MASKS2}")
if usar_labels:
print(f"[INFO] Usando labels: {PASTA_NEW_LABELS}")
mapa_masks = mapear_arquivos_por_base(PASTA_NEW_MASKS, EXT_MASKS, normalizar=True)
mapa_masks2 = mapear_arquivos_por_base(PASTA_NEW_MASKS2, EXT_MASKS2, normalizar=True) if usar_masks2 else {}
mapa_labels = mapear_arquivos_por_base(PASTA_NEW_LABELS, EXT_LABELS, normalizar=True) if usar_labels else {}
registros = []
totais = {"total_masks":0, "processados":0, "pulados":0, "sem_imagem":0, "sem_mask2":0, "dim_mismatch":0, "erros":0}
totais = {
"total_masks": 0,
"processados": 0,
"pulados": 0,
"sem_imagem": 0,
"sem_mask2": 0,
"sem_label": 0,
"dim_mismatch": 0,
"erros": 0,
}
por_grupo = {}
for base, mask_path in sorted(mapa_masks.items()):
for base, mask_path in sorted(mapa_masks.items(), key=lambda x: x[0]):
totais["total_masks"] += 1
img_path = localizar_imagem_por_base(PASTA_NEW_IMAGES, base)
if not img_path:
totais["sem_imagem"] += 1
@ -284,114 +429,131 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT,
try:
mask2_path = mapa_masks2.get(base) if usar_masks2 else None
label_path = mapa_labels.get(base) if usar_labels else None
if usar_masks2 and not mask2_path:
totais["sem_mask2"] += 1
print(f"[AVISO] masks2 existe, mas não achei mask2 para base '{base}' (vou agrupar mesmo).")
print(f"[AVISO] masks2 ativada, mas não achei mask2 para base '{base}'. Vou agrupar sem mask2.")
if usar_labels and not label_path:
totais["sem_label"] += 1
msg = f"[AVISO] labels ativada, mas não achei label para base '{base}'."
if strict_label:
print(msg + " → pulando.")
totais["pulados"] += 1
continue
print(msg + " Vou agrupar sem label.")
if validar_dim:
try:
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
msk = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if img is None or msk is None:
raise RuntimeError("Falha ao abrir img/mask.")
hi, wi = img.shape[:2]
hm, wm = msk.shape[:2]
if (hi, wi) != (hm, wm):
totais["dim_mismatch"] += 1
msg = f"[AVISO] Dimensões diferem (img {wi}x{hi} vs mask {wm}x{hm}) para base '{base}'"
if estrito:
print(msg + " → pulando.")
totais["pulados"] += 1
continue
else:
print(msg + " → copiando mesmo assim.")
if mask2_path:
m2 = cv2.imread(mask2_path, cv2.IMREAD_UNCHANGED)
if m2 is None:
print(f"[AVISO] Falha ao abrir mask2: {mask2_path} → ignorando mask2.")
mask2_path = None
else:
h2, w2 = m2.shape[:2]
if (hi, wi) != (h2, w2):
totais["dim_mismatch"] += 1
msg2 = f"[AVISO] Dimensões diferem (img {wi}x{hi} vs mask2 {w2}x{h2}) para base '{base}'"
if estrito:
print(msg2 + " → pulando (por mask2).")
totais["pulados"] += 1
continue
else:
print(msg2 + " → copiando mesmo assim.")
except Exception as e_dim:
print(f"[AVISO] Falha ao validar dimensões: {e_dim} → copiando mesmo assim.")
ok_dim, dim_mismatch = validar_dimensoes(img_path, mask_path, mask2_path=mask2_path, estrito=estrito)
if dim_mismatch:
totais["dim_mismatch"] += 1
if not ok_dim:
totais["pulados"] += 1
continue
ids_presentes = extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=not labelmap_bgr)
# remove classe ignorada, se conhecida
if ignore_id is not None and ignore_id in ids_presentes:
ids_presentes.discard(ignore_id)
# monta nome do grupo
grupo = montar_nome_grupo(ids_presentes, id_para_nome)
# destinos
dest_base = os.path.join(PASTA_FINAL, grupo)
dest_img_dir = os.path.join(dest_base, "images")
dest_img_dir = os.path.join(dest_base, "images")
dest_mask_dir = os.path.join(dest_base, "masks")
dest_mask2_dir = os.path.join(dest_base, "masks2") if usar_masks2 else None
dest_label_dir = os.path.join(dest_base, "labels") if usar_labels else None
dst_img, dst_mask, dst_mask2 = copiar_ou_mover(img_path, mask_path, dest_img_dir, dest_mask_dir, mover=mover, mask2_src=mask2_path, dest_mask2_dir=dest_mask2_dir)
dst_img, dst_mask, dst_mask2, dst_label = copiar_ou_mover_pacote(
img_src=img_path,
mask_src=mask_path,
dest_img_dir=dest_img_dir,
dest_mask_dir=dest_mask_dir,
mover=mover,
mask2_src=mask2_path,
dest_mask2_dir=dest_mask2_dir,
label_src=label_path,
dest_label_dir=dest_label_dir,
)
totais["processados"] += 1
registros.append([img_path, mask_path, mask2_path or "", dst_img, dst_mask, dst_mask2 or "", grupo])
por_grupo[grupo] = por_grupo.get(grupo, 0) + 1
print(f"[OK] {os.path.basename(dst_img)} → grupo: {grupo}")
extra = " +mask2" if dst_mask2 else ""
print(f"[OK] {os.path.basename(dst_img)}{extra} → grupo: {grupo}")
registros.append([
img_path,
mask_path,
mask2_path or "",
label_path or "",
dst_img,
dst_mask,
dst_mask2 or "",
dst_label or "",
grupo,
])
extras = []
if dst_mask2:
extras.append("mask2")
if dst_label:
extras.append("label")
extra_txt = " + " + " + ".join(extras) if extras else ""
print(f"[OK] {os.path.basename(dst_img)}{extra_txt} → grupo: {grupo}")
except Exception as e:
totais["erros"] += 1
print(f"[ERRO] base '{base}': {e}")
# manifesto
if manifesto and registros:
with open(manifesto, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow(["src_image", "src_mask", "src_mask2", "dst_image", "dst_mask", "dst_mask2", "grupo"])
w.writerow([
"src_image",
"src_mask",
"src_mask2",
"src_label",
"dst_image",
"dst_mask",
"dst_mask2",
"dst_label",
"grupo",
])
w.writerows(registros)
print(f"[MANIFESTO] {manifesto} salvo ({len(registros)} entradas).")
# resumo
print("\nResumo: " + " | ".join(f"{k}={v}" for k, v in totais.items()))
if por_grupo:
print("Por grupo:")
for g, c in sorted(por_grupo.items(), key=lambda x: x[0]):
print(f" - {g}: {c}")
# ====================== CLI ======================
def build_cli():
ap = argparse.ArgumentParser(
description="Agrupa pares IMG+MASK por combinação de classes presentes na máscara (a partir de new_*)."
description="Agrupa IMG+MASK e opcionalmente MASK2/LABEL por combinação de classes presentes na máscara."
)
ap.add_argument("--move", action="store_true", help="Move (em vez de copiar) para as pastas de grupo.")
ap.add_argument("--manifest", default=MANIFESTO_DEFAULT, help="CSV de manifesto ('' para não gerar).")
ap.add_argument("--modelo", default=None, help="Sobrescreve MODELO do config.json.")
ap.add_argument("--move", action="store_true", help="Move em vez de copiar para as pastas de grupo.")
ap.add_argument("--manifest", default=MANIFESTO_DEFAULT, help="CSV de manifesto. Use '' para não gerar.")
ap.add_argument("--modelo", default=None, help="Sobrescreve camera do config.json.")
ap.add_argument("--no-validate", action="store_true", help="Não validar dimensões de imagem/máscara.")
ap.add_argument("--strict", action="store_true", help="Se validar e forem diferentes, pular o par.")
ap.add_argument("--labels-bgr", action="store_true",
help="Use se o labelmap estiver em BGR (por padrão assume RGB).")
ap.add_argument("--labels-bgr", action="store_true", help="Use se o labelmap estiver em BGR. Por padrão assume RGB.")
ap.add_argument("--strict-label", action="store_true", help="Se dual_head_label=true e faltar label, pula o item.")
return ap
if __name__ == "__main__":
args = build_cli().parse_args()
manifest = None if (args.manifest.strip() == "") else args.manifest
processar(
modelo_cli=args.modelo,
mover=args.move,
manifesto=manifest,
validar_dim=not args.no_validate,
estrito=args.strict,
labelmap_bgr=args.labels_bgr
labelmap_bgr=args.labels_bgr,
strict_label=args.strict_label,
)

View File

@ -1,113 +1,185 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Augmenta imagens e máscaras *por grupo*.
Augmenta imagens, máscaras e opcionalmente labels globais por grupo.
Entrada (via config.json -> MODELO):
MODELO/dataset/original/group/<grupo>/images
MODELO/dataset/original/group/<grupo>/masks
Entrada via config_oak.json -> camera:
MODELO/dataset/original/group/<grupo>/images/
MODELO/dataset/original/group/<grupo>/masks/
MODELO/dataset/original/group/<grupo>/masks2/ opcional, se dual_head_mask=true
MODELO/dataset/original/group/<grupo>/labels/ opcional, se dual_head_label=true
Saída:
MODELO/dataset/augmented/group/<grupo>/images
MODELO/dataset/augmented/group/<grupo>/masks
Se "original/group" não existir, faz fallback para:
MODELO/dataset/original/{images,masks}
MODELO/dataset/augmented/{images,masks}
MODELO/dataset/augmented/group/<grupo>/images/
MODELO/dataset/augmented/group/<grupo>/masks/
MODELO/dataset/augmented/group/<grupo>/masks2/ se dual_head_mask=true
MODELO/dataset/augmented/group/<grupo>/labels/ se dual_head_label=true
Transf. geométricas (aplicam a img e máscara) e fotométricas (apenas imagem).
Observação:
- A máscara e a mask2 recebem as mesmas transformações geométricas da imagem.
- O label global NÃO é transformado visualmente. Ele é copiado e atualizado para apontar para o novo base augmentado.
Uso:
python _3_augmentation_grouped.py --copies 5
python _3_augmentation_grouped.py --copies 5 --groups chao,chao_erva,cana
python _3_augmentation_grouped_with_labels.py --copies 5
python _3_augmentation_grouped_with_labels.py --copies 5 --groups navegavel,naonavegavel_navegavel
python _3_augmentation_grouped_with_labels.py --copies 5 --strict-label
"""
import os
import json
import argparse
from pathlib import Path
import cv2
from PIL import Image
import albumentations as A
import argparse
# ⚙️ Configurações
with open("config_oak.json", "r", encoding="utf-8") as f:
# ====================== Configurações ======================
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera", ".")
USE_MASKS2 = config.get("dual_head", False)
# Pastas base
MODELO = config.get("camera", ".")
USE_MASKS2 = bool(config.get("dual_head_mask", False))
USE_LABELS = bool(config.get("dual_head_label", False))
DATASET_BASE = os.path.join(MODELO, "dataset")
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
# Fallback (modo antigo, sem grupos)
# Fallback legacy, sem grupos
ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images")
ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks")
ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2")
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
ORIG_OLD_LABELS = os.path.join(DATASET_BASE, "original", "labels")
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
AUG_OLD_LABELS = os.path.join(DATASET_BASE, "augmented", "labels")
# Extensões aceitas
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg") # manter prioridade PNG quando possível
MSK_EXTS = (".png", ".jpg", ".jpeg")
MSK2_EXTS = (".png", ".jpg", ".jpeg")
LABEL_EXTS = (".json", ".txt")
# ====================== Augmentation ======================
train_tf = A.Compose(
[
A.HorizontalFlip(p=0.5),
A.ShiftScaleRotate(
shift_limit=0.01,
scale_limit=0.10,
rotate_limit=5,
border_mode=cv2.BORDER_REFLECT_101,
interpolation=cv2.INTER_LINEAR,
p=0.30,
),
A.OneOf(
[
A.RandomBrightnessContrast(0.2, 0.2, p=1.0),
A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1.0),
A.RandomGamma(gamma_limit=(90, 110), p=1.0),
],
p=0.70,
),
A.OneOf(
[
A.MotionBlur(blur_limit=3, p=1.0),
A.GaussianBlur(blur_limit=3, p=1.0),
],
p=0.20,
),
A.OneOf(
[
A.GaussNoise(var_limit=(5.0, 15.0), p=1.0),
A.ImageCompression(quality_lower=50, quality_upper=85, p=1.0),
],
p=0.20,
),
A.RandomShadow(p=0.10),
A.RandomSunFlare(p=0.10),
A.ChannelShuffle(p=0.05),
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10),
],
additional_targets={"mask2": "mask"},
)
# ====================== Utilitários ======================
def garantir_dir(p):
os.makedirs(p, exist_ok=True)
# Pipeline de augmentations
train_tf = A.Compose([
A.HorizontalFlip(p=0.5),
# Geométricas (aplicam em imagem e máscara)
A.ShiftScaleRotate(
shift_limit=0.01,
scale_limit=0.10,
rotate_limit=5,
border_mode=cv2.BORDER_REFLECT_101,
interpolation=cv2.INTER_LINEAR,
p=0.30
),
def normalizar_base(stem: str) -> str:
sufixos = [
"_rgb", "_RGB", "_Rgb",
"_image", "_img", "_frame",
"_mask", "_masks",
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
"_label", "_labels",
]
# Fotométricas (somente imagem)
A.OneOf([
A.RandomBrightnessContrast(0.2, 0.2, p=1.0),
A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1.0),
A.RandomGamma(gamma_limit=(90, 110), p=1.0),
], p=0.70),
out = stem
mudou = True
while mudou:
mudou = False
for sfx in sufixos:
if out.endswith(sfx):
out = out[: -len(sfx)]
mudou = True
break
return out
A.OneOf([
A.MotionBlur(blur_limit=3, p=1.0),
A.GaussianBlur(blur_limit=3, p=1.0),
], p=0.20),
A.OneOf([
A.GaussNoise(var_limit=(5.0, 15.0), p=1.0),
A.ImageCompression(quality_lower=50, quality_upper=85, p=1.0),
], p=0.20),
def map_files_by_base(folder, exts):
by_base = {}
if not os.path.isdir(folder):
return by_base
A.RandomShadow(p=0.10),
A.RandomSunFlare(p=0.10),
A.ChannelShuffle(p=0.05),
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10),
], additional_targets={
'mask2': 'mask'
})
prioridade = {
".json": 0,
".png": 1,
".jpg": 2,
".jpeg": 3,
".txt": 4,
}
def load_rgb(path):
# cv2 lê BGR → converte para RGB
im = cv2.imread(path, cv2.IMREAD_COLOR)
if im is None:
raise FileNotFoundError(path)
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
for fname in os.listdir(folder):
lower = fname.lower()
if not lower.endswith(exts):
continue
stem, ext = os.path.splitext(fname)
base = normalizar_base(stem)
cand = os.path.join(folder, fname)
if base not in by_base:
by_base[base] = cand
else:
cur_ext = os.path.splitext(by_base[base])[1].lower()
if prioridade.get(ext.lower(), 99) < prioridade.get(cur_ext, 99):
by_base[base] = cand
return by_base
def save_rgb(path, arr_rgb):
Image.fromarray(arr_rgb).save(path)
def list_groups(root):
"""Lista grupos válidos (que contêm subpastas images e masks)."""
if not os.path.isdir(root):
return []
grupos = []
for name in sorted(os.listdir(root)):
gdir = os.path.join(root, name)
@ -117,68 +189,107 @@ def list_groups(root):
grupos.append(name)
return grupos
def map_masks_by_base(msk_dir):
"""Mapeia máscaras por base (prioriza .png)."""
by_base = {}
if not os.path.isdir(msk_dir):
return by_base
for fname in os.listdir(msk_dir):
f_lower = fname.lower()
if not f_lower.endswith(MSK_EXTS):
continue
base, ext = os.path.splitext(fname)
cand = os.path.join(msk_dir, fname)
if base not in by_base:
by_base[base] = cand
else:
# mantém .png se disponível
cur_ext = os.path.splitext(by_base[base])[1].lower()
if cur_ext != ".png" and ext.lower() == ".png":
by_base[base] = cand
return by_base
def map_masks2_by_base(msk2_dir):
"""Mapeia máscaras2 por base (prioriza .png)."""
by_base = {}
if not os.path.isdir(msk2_dir):
return by_base
for fname in os.listdir(msk2_dir):
f_lower = fname.lower()
if not f_lower.endswith(MSK2_EXTS):
continue
base, ext = os.path.splitext(fname)
cand = os.path.join(msk2_dir, fname)
if base not in by_base:
by_base[base] = cand
else:
cur_ext = os.path.splitext(by_base[base])[1].lower()
if cur_ext != ".png" and ext.lower() == ".png":
by_base[base] = cand
return by_base
def load_rgb(path):
im = cv2.imread(path, cv2.IMREAD_COLOR)
if im is None:
raise FileNotFoundError(path)
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
def ensure_aug_dirs(group_name=None, use_masks2=False):
"""Cria diretórios de saída para o grupo ou modo antigo. Se use_masks2, cria masks2."""
def save_rgb(path, arr_rgb):
garantir_dir(os.path.dirname(path))
Image.fromarray(arr_rgb).save(path)
def ensure_aug_dirs(group_name=None, use_masks2=False, use_labels=False):
if group_name:
img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images")
msk_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks")
msk2_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks2") if use_masks2 else None
labels_out = os.path.join(AUG_GROUP_ROOT, group_name, "labels") if use_labels else None
else:
img_out = AUG_OLD_IMG
msk_out = AUG_OLD_MSK
msk2_out = AUG_OLD_MSK2 if use_masks2 else None
labels_out = AUG_OLD_LABELS if use_labels else None
garantir_dir(img_out)
garantir_dir(msk_out)
if use_masks2 and msk2_out:
garantir_dir(msk2_out)
return img_out, msk_out, msk2_out
def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies, msk2_path=None, msk2_out_dir=None):
if use_labels and labels_out:
garantir_dir(labels_out)
return img_out, msk_out, msk2_out, labels_out
def safe_rel(path, root):
try:
return str(Path(path).resolve().relative_to(Path(root).resolve())).replace("\\", "/")
except Exception:
return str(path).replace("\\", "/")
def copiar_label_aug(label_path, out_label_path, new_base, out_img_path, out_msk_path, out_msk2_path=None, group_name=None):
"""
Copia label global para a amostra augmentada.
Se for JSON, atualiza campos úteis.
Se for TXT, copia o conteúdo como está.
"""
if not label_path or not out_label_path:
return
garantir_dir(os.path.dirname(out_label_path))
ext = os.path.splitext(label_path)[1].lower()
if ext == ".json":
try:
with open(label_path, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception:
data = {}
data["base"] = new_base
data["group"] = group_name if group_name is not None else data.get("group")
data["image"] = safe_rel(out_img_path, DATASET_BASE)
data["mask"] = safe_rel(out_msk_path, DATASET_BASE)
data["label"] = safe_rel(out_label_path, DATASET_BASE)
data["is_augmented"] = True
data["source_label"] = safe_rel(label_path, DATASET_BASE)
if out_msk2_path:
data["mask2"] = safe_rel(out_msk2_path, DATASET_BASE)
with open(out_label_path, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2)
else:
with open(label_path, "r", encoding="utf-8") as f:
content = f.read()
with open(out_label_path, "w", encoding="utf-8") as f:
f.write(content)
def augment_pair(
img_path,
msk_path,
img_out_dir,
msk_out_dir,
copies,
msk2_path=None,
msk2_out_dir=None,
label_path=None,
label_out_dir=None,
group_name=None,
):
base_img, img_ext = os.path.splitext(os.path.basename(img_path))
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
_, msk_ext = os.path.splitext(os.path.basename(msk_path))
msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
label_ext = os.path.splitext(os.path.basename(label_path))[1] if label_path else None
# padroniza pelo base da imagem
base = base_img
base = normalizar_base(base_img)
img = load_rgb(img_path)
msk = load_rgb(msk_path)
@ -194,122 +305,223 @@ def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies, msk2_path
img_aug = aug["image"]
msk_aug = aug["mask"]
out_img = os.path.join(img_out_dir, f"{base}_aug_{i:02d}{img_ext}")
out_msk = os.path.join(msk_out_dir, f"{base}_aug_{i:02d}{msk_ext}")
new_base = f"{base}_aug_{i:02d}"
out_img = os.path.join(img_out_dir, f"{new_base}{img_ext}")
out_msk = os.path.join(msk_out_dir, f"{new_base}{msk_ext}")
save_rgb(out_img, img_aug)
save_rgb(out_msk, msk_aug)
out_msk2 = None
if msk2 is not None and msk2_out_dir:
msk2_aug = aug["mask2"]
out_msk2 = os.path.join(msk2_out_dir, f"{base}_aug_{i:02d}{msk2_ext}")
out_msk2 = os.path.join(msk2_out_dir, f"{new_base}{msk2_ext}")
save_rgb(out_msk2, msk2_aug)
if label_path and label_out_dir:
out_label = os.path.join(label_out_dir, f"{new_base}{label_ext}")
copiar_label_aug(
label_path=label_path,
out_label_path=out_label,
new_base=new_base,
out_img_path=out_img,
out_msk_path=out_msk,
out_msk2_path=out_msk2,
group_name=group_name,
)
gen += 1
return gen
def process_group(group_name, copies):
"""Processa um grupo único (images/masks dentro de ORIG_GROUP_ROOT/<group_name>/)."""
# ====================== Processamento ======================
def process_group(group_name, copies, strict_label=False):
img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images")
msk_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks")
msk2_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks2")
labels_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "labels")
if not (os.path.isdir(img_dir) and os.path.isdir(msk_dir)):
print(f"[WARN] Grupo '{group_name}' inválido (sem images/masks). Pulando.")
print(f"[WARN] Grupo '{group_name}' inválido, sem images/masks. Pulando.")
return 0
imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
msk_map = map_masks_by_base(msk_dir)
msk_map = map_files_by_base(msk_dir, MSK_EXTS)
use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
msk2_map = map_masks2_by_base(msk2_dir) if use_masks2 else {}
img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name, use_masks2=use_masks2)
msk2_map = map_files_by_base(msk2_dir, MSK2_EXTS) if use_masks2 else {}
use_labels = USE_LABELS and os.path.isdir(labels_dir)
label_map = map_files_by_base(labels_dir, LABEL_EXTS) if use_labels else {}
if USE_MASKS2 and not use_masks2:
print(f"[WARN] [{group_name}] dual_head_mask=true, mas pasta masks2 não existe.")
if USE_LABELS and not use_labels:
msg = f"[WARN] [{group_name}] dual_head_label=true, mas pasta labels não existe."
if strict_label:
print(msg + " Pulando grupo.")
return 0
print(msg + " Gerando sem labels.")
img_out_dir, msk_out_dir, msk2_out_dir, label_out_dir = ensure_aug_dirs(
group_name,
use_masks2=use_masks2,
use_labels=use_labels,
)
count = 0
sem_mask = 0
sem_mask2 = 0
sem_label = 0
for img_file in sorted(imgs):
base, _ = os.path.splitext(img_file)
msk_file = msk_map.get(base)
base_norm = normalizar_base(base)
msk_file = msk_map.get(base_norm)
if not msk_file:
sem_mask += 1
print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.")
continue
msk2_file = msk2_map.get(base) if use_masks2 else None
msk2_file = msk2_map.get(base_norm) if use_masks2 else None
if use_masks2 and not msk2_file:
sem_mask2 += 1
print(f"[WARN] [{group_name}] mask2 não encontrada para {img_file}, gerando só img+mask.")
label_file = label_map.get(base_norm) if use_labels else None
if use_labels and not label_file:
sem_label += 1
msg = f"[WARN] [{group_name}] label não encontrado para {img_file}."
if strict_label:
print(msg + " Pulando item.")
continue
print(msg + " Gerando augmentation sem label.")
try:
count += augment_pair(
os.path.join(img_dir, img_file),
msk_file,
img_out_dir,
msk_out_dir,
img_path=os.path.join(img_dir, img_file),
msk_path=msk_file,
img_out_dir=img_out_dir,
msk_out_dir=msk_out_dir,
copies=copies,
msk2_path=msk2_file,
msk2_out_dir=msk2_out_dir
msk2_out_dir=msk2_out_dir,
label_path=label_file,
label_out_dir=label_out_dir,
group_name=group_name,
)
except Exception as e:
print(f"[ERRO] [{group_name}] {img_file}: {e}")
print(f"[OK] Grupo '{group_name}'{count} pares gerados.")
print(
f"[OK] Grupo '{group_name}'{count} pares gerados. "
f"sem_mask={sem_mask} | sem_mask2={sem_mask2} | sem_label={sem_label}"
)
return count
def process_legacy(copies):
"""Fallback: modo sem grupos (original/images e original/masks)."""
def process_legacy(copies, strict_label=False):
if not (os.path.isdir(ORIG_OLD_IMG) and os.path.isdir(ORIG_OLD_MSK)):
print("[WARN] Modo legacy não encontrado. Nada a fazer.")
return 0
imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS]
msk_map = map_masks_by_base(ORIG_OLD_MSK)
msk_map = map_files_by_base(ORIG_OLD_MSK, MSK_EXTS)
use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
msk2_map = map_masks2_by_base(ORIG_OLD_MSK2) if use_masks2 else {}
img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name=None, use_masks2=use_masks2)
msk2_map = map_files_by_base(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}
use_labels = USE_LABELS and os.path.isdir(ORIG_OLD_LABELS)
label_map = map_files_by_base(ORIG_OLD_LABELS, LABEL_EXTS) if use_labels else {}
if USE_LABELS and not use_labels and strict_label:
print("[WARN] Legacy com dual_head_label=true, mas sem pasta labels. Pulando.")
return 0
img_out_dir, msk_out_dir, msk2_out_dir, label_out_dir = ensure_aug_dirs(
group_name=None,
use_masks2=use_masks2,
use_labels=use_labels,
)
count = 0
for img_file in sorted(imgs):
base, _ = os.path.splitext(img_file)
msk_file = msk_map.get(base)
base_norm = normalizar_base(base)
msk_file = msk_map.get(base_norm)
if not msk_file:
print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
continue
msk2_file = msk2_map.get(base) if use_masks2 else None
if use_masks2 and not msk2_file:
print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.")
msk2_file = msk2_map.get(base_norm) if use_masks2 else None
label_file = label_map.get(base_norm) if use_labels else None
if use_labels and not label_file and strict_label:
print(f"[WARN] (legacy) Label não encontrado para {img_file}, pulando.")
continue
try:
count += augment_pair(
os.path.join(ORIG_OLD_IMG, img_file),
msk_file,
img_out_dir,
msk_out_dir,
img_path=os.path.join(ORIG_OLD_IMG, img_file),
msk_path=msk_file,
img_out_dir=img_out_dir,
msk_out_dir=msk_out_dir,
copies=copies,
msk2_path=msk2_file,
msk2_out_dir=msk2_out_dir
msk2_out_dir=msk2_out_dir,
label_path=label_file,
label_out_dir=label_out_dir,
group_name=None,
)
except Exception as e:
print(f"[ERRO] (legacy) {img_file}: {e}")
print(f"[OK] Legacy → {count} pares gerados.")
return count
def main(copies=5, groups_csv=None):
def main(copies=5, groups_csv=None, strict_label=False):
total = 0
print(f"[INFO] MODELO={MODELO}")
print(f"[INFO] dual_head_mask={USE_MASKS2}")
print(f"[INFO] dual_head_label={USE_LABELS}")
if os.path.isdir(ORIG_GROUP_ROOT):
grupos = list_groups(ORIG_GROUP_ROOT)
if groups_csv:
# filtra pelos grupos desejados
want = {g.strip() for g in groups_csv.split(",") if g.strip()}
grupos = [g for g in grupos if g in want]
if not grupos:
print("[WARN] Nenhum grupo válido encontrado após filtro.")
if not grupos:
print("[WARN] Nenhum grupo encontrado em original/group. Tentando modo legacy...")
total += process_legacy(copies)
total += process_legacy(copies, strict_label=strict_label)
else:
print(f"Grupos encontrados: {', '.join(grupos)}")
for g in grupos:
total += process_group(g, copies)
total += process_group(g, copies, strict_label=strict_label)
else:
# sem estrutura de grupos
total += process_legacy(copies)
total += process_legacy(copies, strict_label=strict_label)
print(f"\nAugmentation completed! Total: {total} pares gerados.")
if __name__ == "__main__":
ap = argparse.ArgumentParser(description="Augmentação por grupos (images/masks)")
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem (default=5).")
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula (ex: chao,erva_cana).")
ap = argparse.ArgumentParser(description="Augmentação por grupos com suporte a images/masks/masks2/labels.")
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem.")
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula.")
ap.add_argument("--strict-label", action="store_true", help="Se dual_head_label=true e faltar label, pula o item/grupo.")
args = ap.parse_args()
main(copies=args.copies, groups_csv=args.groups)
main(copies=args.copies, groups_csv=args.groups, strict_label=args.strict_label)

View File

@ -1,67 +1,86 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Normaliza/redimensiona imagens e máscaras mantendo a ESTRUTURA POR GRUPO.
Normaliza/redimensiona imagens, máscaras, masks2 e labels mantendo a ESTRUTURA POR GRUPO.
Entradas (via config.json -> MODELO, RESOLUCAO):
- MODELO/dataset/original/group/<grupo>/{images,masks}
- MODELO/dataset/augmented/group/<grupo>/{images,masks}
Entradas via config.json -> camera, resolucao:
- MODELO/dataset/original/group/<grupo>/{images,masks,(masks2),(labels)}
- MODELO/dataset/augmented/group/<grupo>/{images,masks,(masks2),(labels)}
Saídas (por resolução):
- MODELO/dataset/<WxH>/group/<grupo>/{images,masks}
Saídas por resolução:
- MODELO/dataset/<WxH>/group/<grupo>/{images,masks,(masks2),(labels)}
Fallback (modo legado, se não houver "group/"):
- original/{images,masks} e augmented/{images,masks} -> <WxH>/{images,masks}
Fallback legado, se não houver group/:
- original/{images,masks,(masks2),(labels)}
- augmented/{images,masks,(masks2),(labels)}
- saída: <WxH>/{images,masks,(masks2),(labels)}
Conversão de máscara:
- máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids
- Ignora classe "ignore" conforme labelmap (usa índice 255 como padrão quando necessário)
- máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids.
- Ignore conforme labelmap, usando índice 255 como fallback.
Labels:
- Ativado por config['dual_head_label'].
- Copia labels .json/.txt para a saída.
- Se JSON, atualiza image/mask/mask2/label/base/source/normalized.
- Se houver label_id, salva também um .npy com o inteiro para facilitar o Dataset no treino.
"""
import argparse
import os
import json
import cv2
from typing import Dict, List, Tuple
from typing import Dict, List, Tuple, Optional
import numpy as np
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ⚙️ Configurações
with open("config_oak.json", "r", encoding="utf-8") as f:
# ===================== CONFIG =====================
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
USE_MASKS2 = config["dual_head"]
RESOLUCAO = tuple(config["resolucao"]) # [W, H] ou [width, height]
# Compatibilidade:
# - scripts antigos usam dual_head para masks2
# - scripts novos podem usar dual_head_mask e dual_head_label separados
USE_MASKS2 = bool(config.get("dual_head_mask", config.get("dual_head", False)))
USE_LABELS = bool(config.get("dual_head_label", False))
RESOLUCAO = tuple(config["resolucao"]) # [W, H]
pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
mode = "single"
if USE_MASKS2:
mode = "mask2"
elif USE_LABELS:
mode = "label"
suffix = {
"single": "_single",
"mask2": "_dual_mask",
"label": "_dual_label",
}[mode]
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME + suffix)
# Dimensões alvo (pode expandir para múltiplas se quiser)
RESOLUCOES = {
f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1]),
}
# Fontes a processar
FONTES = ["original", "augmented"]
# Extensões aceitas
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg") # preferir .png
MSK2_EXTS = (".png", ".jpg", ".jpeg") # idem
MSK_EXTS = (".png", ".jpg", ".jpeg")
MSK2_EXTS = (".png", ".jpg", ".jpeg")
LABEL_EXTS = (".json", ".txt")
# === Acumuladores globais para mean/std dos canais RAW4 ===
GLOBAL_SUM = None # soma por canal
GLOBAL_SUMSQ = None # soma dos quadrados por canal
GLOBAL_PIXELS = 0 # n de pixels por canal (H*W por imagem)
GLOBAL_SUM = None
GLOBAL_SUMSQ = None
GLOBAL_PIXELS = 0
# ===================== HELPERS =====================
def infer_ignore_id(ignore_rgb, default_id=255):
"""
Tenta inferir o ID de ignore a partir do valor retornado por carregar_labelmap_completo.
- Se for [id] retorna id
- Se for (R,G,B) retorna default_id (tipicamente 255)
- Se for int, retorna direto
"""
if isinstance(ignore_rgb, (list, tuple)):
if len(ignore_rgb) == 1:
try:
@ -74,13 +93,15 @@ def infer_ignore_id(ignore_rgb, default_id=255):
return ignore_rgb
return default_id
def garantir_dir(p):
os.makedirs(p, exist_ok=True)
def list_groups(root) -> List[str]:
"""Lista grupos válidos com subpastas images e masks."""
if not os.path.isdir(root):
return []
grupos = []
for name in sorted(os.listdir(root)):
gdir = os.path.join(root, name)
@ -90,81 +111,130 @@ def list_groups(root) -> List[str]:
grupos.append(name)
return grupos
def map_masks_by_base(msk_dir: str) -> Dict[str, str]:
"""Retorna {base: caminho_mask}, priorizando .png quando houver múltiplas por base."""
def normalizar_base(stem: str) -> str:
sufixos = [
"_rgb", "_RGB", "_Rgb",
"_image", "_img", "_frame",
"_mask", "_masks",
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
"_label", "_labels",
]
out = stem
mudou = True
while mudou:
mudou = False
for sfx in sufixos:
if out.endswith(sfx):
out = out[: -len(sfx)]
mudou = True
break
return out
def map_files_by_base(folder: str, exts: Tuple[str, ...]) -> Dict[str, str]:
by_base = {}
if not os.path.isdir(msk_dir):
if not os.path.isdir(folder):
return by_base
for fname in os.listdir(msk_dir):
f_lower = fname.lower()
if not f_lower.endswith(MSK_EXTS):
prioridade = {
".json": 0,
".png": 1,
".jpg": 2,
".jpeg": 3,
".txt": 4,
}
for fname in os.listdir(folder):
lower = fname.lower()
if not lower.endswith(exts):
continue
base, ext = os.path.splitext(fname)
cand = os.path.join(msk_dir, fname)
stem, ext = os.path.splitext(fname)
base = normalizar_base(stem)
cand = os.path.join(folder, fname)
if base not in by_base:
by_base[base] = cand
else:
cur_ext = os.path.splitext(by_base[base])[1].lower()
if cur_ext != ".png" and ext.lower() == ".png":
if prioridade.get(ext.lower(), 99) < prioridade.get(cur_ext, 99):
by_base[base] = cand
return by_base
def map_masks_by_base(msk_dir: str) -> Dict[str, str]:
return map_files_by_base(msk_dir, MSK_EXTS)
def map_masks2_by_base(msk2_dir: str) -> Dict[str, str]:
"""Retorna {base: caminho_mask2}, priorizando .png quando houver múltiplas por base."""
by_base = {}
if not os.path.isdir(msk2_dir):
return by_base
for fname in os.listdir(msk2_dir):
f_lower = fname.lower()
if not f_lower.endswith(MSK2_EXTS):
continue
base, ext = os.path.splitext(fname)
cand = os.path.join(msk2_dir, fname)
if base not in by_base:
by_base[base] = cand
else:
cur_ext = os.path.splitext(by_base[base])[1].lower()
if cur_ext != ".png" and ext.lower() == ".png":
by_base[base] = cand
return by_base
return map_files_by_base(msk2_dir, MSK2_EXTS)
def normalize_pair(caminho_rgb: str, caminho_mask: str, cor_para_id, ignore_id: int, out_img_dir: str, out_msk_dir: str, dim: Tuple[int,int], prefix: str = ""):
"""Redimensiona e grava a imagem e a máscara (se houver)."""
img_rgb = cv2.imread(caminho_rgb)
if img_rgb is None:
print(f"[!] Erro ao ler imagem: {caminho_rgb}")
return False
# Nome de saída com prefixo para distinguir fonte (ex: original_, augmented_)
nome = os.path.basename(caminho_rgb)
if prefix:
nome_saida_img = f"{prefix}{nome}"
else:
nome_saida_img = nome
nome_saida_msk = nome_saida_img
def map_labels_by_base(label_dir: str) -> Dict[str, str]:
return map_files_by_base(label_dir, LABEL_EXTS)
def trocar_ext_para_png(nome: str) -> str:
for ext in (".jpg", ".jpeg", ".png"):
if nome_saida_msk.lower().endswith(ext):
nome_saida_msk = nome_saida_msk[: -len(ext)] + ".png"
break
if nome.lower().endswith(ext):
return nome[: -len(ext)] + ".png"
return nome + ".png"
# Redimensiona imagem
img_resized = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)
def safe_rel(path: str, root: str) -> str:
try:
return os.path.relpath(path, root).replace("\\", "/")
except Exception:
return str(path).replace("\\", "/")
# ===================== NORMALIZAÇÃO =====================
def normalize_pair(
caminho_rgb: str,
caminho_mask: str,
cor_para_id,
ignore_id: int,
out_img_dir: str,
out_msk_dir: str,
dim: Tuple[int, int],
prefix: str = "",
):
"""Redimensiona e grava imagem + máscara ID."""
img_bgr = cv2.imread(caminho_rgb)
if img_bgr is None:
print(f"[!] Erro ao ler imagem: {caminho_rgb}")
return False, None, None
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
nome = os.path.basename(caminho_rgb)
nome_saida_img = f"{prefix}{nome}" if prefix else nome
nome_saida_msk = trocar_ext_para_png(nome_saida_img)
img_resized_rgb = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
img_float = img_resized.astype(np.float32) / 255.0 # 0-1
img_float = img_resized_rgb.astype(np.float32) / 255.0
h, w, c = img_float.shape
flat = img_float.reshape(-1, c).astype(np.float64)
if (GLOBAL_SUM is None):
if GLOBAL_SUM is None:
GLOBAL_SUM = np.zeros(c, dtype=np.float64)
GLOBAL_SUMSQ = np.zeros(c, dtype=np.float64)
GLOBAL_SUM += flat.sum(axis=0)
GLOBAL_SUMSQ += (flat ** 2).sum(axis=0)
GLOBAL_PIXELS += h * w
garantir_dir(out_img_dir)
cv2.imwrite(os.path.join(out_img_dir, nome_saida_img), img_resized)
out_img_path = os.path.join(out_img_dir, nome_saida_img)
img_resized_bgr = cv2.cvtColor(img_resized_rgb, cv2.COLOR_RGB2BGR)
cv2.imwrite(out_img_path, img_resized_bgr)
# Processa e redimensiona máscara (se existir)
out_msk_path = None
if caminho_mask and os.path.isfile(caminho_mask):
msk_bgr = cv2.imread(caminho_mask, cv2.IMREAD_COLOR)
if msk_bgr is None:
@ -173,184 +243,409 @@ def normalize_pair(caminho_rgb: str, caminho_mask: str, cor_para_id, ignore_id:
msk_rgb = cv2.cvtColor(msk_bgr, cv2.COLOR_BGR2RGB)
mask_ids = converter_mask_rgb_para_ids(msk_rgb, cor_para_id, ignore_id)
mask_resized = cv2.resize(mask_ids, dim, interpolation=cv2.INTER_NEAREST)
garantir_dir(out_msk_dir)
cv2.imwrite(os.path.join(out_msk_dir, nome_saida_msk), mask_resized)
out_msk_path = os.path.join(out_msk_dir, nome_saida_msk)
cv2.imwrite(out_msk_path, mask_resized)
return True, out_img_path, out_msk_path
return True
def normalize_pair_mask2(caminho_rgb: str, caminho_mask2: str,
out_msk2_dir: str, dim: Tuple[int,int], prefix: str = ""):
"""
Redimensiona e grava máscara2 (corredor binário), assumindo que ela é uma máscara "pronta".
- Se for RGB/BGR (3 canais): converte para cinza e faz threshold (0/255) antes de redimensionar.
- Se for 1 canal: mantém, faz threshold (0/255).
- Redimensiona com INTER_NEAREST.
Saída sempre .png com o mesmo nome base/prefixo do arquivo de imagem.
"""
def normalize_pair_mask2(
caminho_rgb: str,
caminho_mask2: str,
out_msk2_dir: str,
dim: Tuple[int, int],
prefix: str = "",
):
if not caminho_mask2 or not os.path.isfile(caminho_mask2):
return False
return False, None
nome = os.path.basename(caminho_rgb)
nome_saida = f"{prefix}{nome}" if prefix else nome
for ext in (".jpg", ".jpeg", ".png"):
if nome_saida.lower().endswith(ext):
nome_saida = nome_saida[: -len(ext)] + ".png"
break
nome_saida = trocar_ext_para_png(nome_saida)
m2 = cv2.imread(caminho_mask2, cv2.IMREAD_UNCHANGED)
if m2 is None:
print(f"[!] Erro ao ler máscara2: {caminho_mask2}")
return False
return False, None
if len(m2.shape) == 3:
# BGR/RGB -> gray
m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY)
else:
m2g = m2
# binariza para 0/255 (evita lixo de compressão)
_, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY)
m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST)
garantir_dir(out_msk2_dir)
cv2.imwrite(os.path.join(out_msk2_dir, nome_saida), m2res)
out_msk2_path = os.path.join(out_msk2_dir, nome_saida)
cv2.imwrite(out_msk2_path, m2res)
return True, out_msk2_path
def normalize_label(
caminho_rgb: str,
caminho_label: Optional[str],
out_label_dir: Optional[str],
out_img_path: Optional[str],
out_msk_path: Optional[str],
out_msk2_path: Optional[str],
source_root: str,
output_root: str,
grupo: Optional[str],
fonte_nome: str,
prefix: str = "",
):
"""Copia/atualiza label global e salva label_id como .npy quando possível."""
if not caminho_label or not out_label_dir or not os.path.isfile(caminho_label):
return False
nome = os.path.basename(caminho_rgb)
nome_saida = f"{prefix}{nome}" if prefix else nome
base_saida, _ = os.path.splitext(nome_saida)
label_ext = os.path.splitext(caminho_label)[1].lower()
garantir_dir(out_label_dir)
out_label_path = os.path.join(out_label_dir, base_saida + label_ext)
out_label_id_path = os.path.join(out_label_dir, base_saida + ".npy")
label_id = None
if label_ext == ".json":
try:
with open(caminho_label, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception:
data = {}
label_id = data.get("label_id")
data["normalized"] = True
data["source"] = fonte_nome
data["group"] = grupo if grupo is not None else data.get("group")
data["base"] = base_saida
data["image"] = safe_rel(out_img_path, output_root) if out_img_path else None
data["mask"] = safe_rel(out_msk_path, output_root) if out_msk_path else None
data["label"] = safe_rel(out_label_path, output_root)
data["source_label"] = safe_rel(caminho_label, source_root)
if out_msk2_path:
data["mask2"] = safe_rel(out_msk2_path, output_root)
with open(out_label_path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2, ensure_ascii=False)
else:
with open(caminho_label, "r", encoding="utf-8") as f:
txt = f.read().strip()
with open(out_label_path, "w", encoding="utf-8") as f:
f.write(txt)
label_id = None
if label_id is not None:
try:
np.save(out_label_id_path, np.array(int(label_id), dtype=np.int64))
except Exception as e:
print(f"[WARN] Não consegui salvar label_id npy para {caminho_label}: {e}")
return True
def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: int, groups_except: str = None):
"""Processa uma raiz do tipo .../<fonte>/group/ agrupando por cada subpasta de grupo."""
# ===================== PROCESSAMENTO =====================
def process_group_root(
fonte_root: str,
fonte_nome: str,
cor_para_id,
ignore_id: int,
groups_except: str = "",
strict_label: bool = False,
):
total = 0
grupos = list_groups(fonte_root)
if not grupos:
return 0
not_want = {g.strip() for g in groups_except.split(",") if g.strip()}
grupos_desconsiderar = [g for g in grupos if g in not_want]
for nome_res, dim in RESOLUCOES.items():
out_root = os.path.join(pasta_base, nome_res, "group")
for grupo in grupos:
if grupo in grupos_desconsiderar:
print(f"[WARN] Grupo desconsiderado nao sera processado: {grupo}")
print(f"[WARN] Grupo desconsiderado não será processado: {grupo}")
continue
in_img_dir = os.path.join(fonte_root, grupo, "images")
in_msk_dir = os.path.join(fonte_root, grupo, "masks")
in_msk2_dir = os.path.join(fonte_root, grupo, "masks2")
in_label_dir = os.path.join(fonte_root, grupo, "labels")
if not (os.path.isdir(in_img_dir) and os.path.isdir(in_msk_dir)):
print(f"[WARN] Grupo inválido (sem images/masks): {grupo}")
print(f"[WARN] Grupo inválido sem images/masks: {grupo}")
continue
out_img_dir = os.path.join(out_root, grupo, "images")
out_msk_dir = os.path.join(out_root, grupo, "masks")
usar_masks2 = USE_MASKS2 and os.path.isdir(in_msk2_dir)
usar_labels = USE_LABELS and os.path.isdir(in_label_dir)
out_msk2_dir = os.path.join(out_root, grupo, "masks2") if usar_masks2 else None
out_label_dir = os.path.join(out_root, grupo, "labels") if usar_labels else None
msk_map = map_masks_by_base(in_msk_dir)
msk2_map = map_masks2_by_base(in_msk2_dir) if usar_masks2 else {}
label_map = map_labels_by_base(in_label_dir) if usar_labels else {}
if USE_LABELS and not usar_labels:
msg = f"[WARN] [{fonte_nome} | {grupo}] dual_head_label=true, mas labels/ não existe."
if strict_label:
print(msg + " Pulando grupo.")
continue
print(msg + " Seguindo sem labels.")
imgs = [f for f in os.listdir(in_img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
n = len(imgs)
for i, fname in enumerate(sorted(imgs), 1):
base, _ = os.path.splitext(fname)
base_norm = normalizar_base(base)
caminho_rgb = os.path.join(in_img_dir, fname)
caminho_mask = msk_map.get(base)
caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
ok = normalize_pair(
caminho_rgb, caminho_mask, cor_para_id, ignore_id,
out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_"
caminho_mask = msk_map.get(base_norm)
caminho_mask2 = msk2_map.get(base_norm) if usar_masks2 else None
caminho_label = label_map.get(base_norm) if usar_labels else None
if usar_labels and not caminho_label:
msg = f"[WARN] [{fonte_nome} | {grupo}] label não encontrado para {fname}."
if strict_label:
print(msg + " Pulando item.")
continue
print(msg + " Seguindo sem label.")
ok, out_img_path, out_msk_path = normalize_pair(
caminho_rgb,
caminho_mask,
cor_para_id,
ignore_id,
out_img_dir,
out_msk_dir,
dim,
prefix=f"{fonte_nome}_",
)
out_msk2_path = None
if usar_masks2 and out_msk2_dir:
if not caminho_mask2:
print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname} (vou seguir).")
print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname}.")
else:
normalize_pair_mask2(
caminho_rgb, caminho_mask2,
out_msk2_dir, dim, prefix=f"{fonte_nome}_"
_, out_msk2_path = normalize_pair_mask2(
caminho_rgb,
caminho_mask2,
out_msk2_dir,
dim,
prefix=f"{fonte_nome}_",
)
if usar_labels and caminho_label and out_label_dir:
normalize_label(
caminho_rgb=caminho_rgb,
caminho_label=caminho_label,
out_label_dir=out_label_dir,
out_img_path=out_img_path,
out_msk_path=out_msk_path,
out_msk2_path=out_msk2_path,
source_root=fonte_root,
output_root=os.path.join(pasta_base, nome_res),
grupo=grupo,
fonte_nome=fonte_nome,
prefix=f"{fonte_nome}_",
)
if ok:
total += 1
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n}{fname}")
return total
def process_legacy_root(legacy_img: str, legacy_msk: str, fonte_nome: str, cor_para_id, ignore_id: int):
"""Processa estrutura legado (sem grupos)."""
def process_legacy_root(
legacy_img: str,
legacy_msk: str,
fonte_nome: str,
cor_para_id,
ignore_id: int,
strict_label: bool = False,
):
if not (os.path.isdir(legacy_img) and os.path.isdir(legacy_msk)):
return 0
total = 0
for nome_res, dim in RESOLUCOES.items():
out_img_dir = os.path.join(pasta_base, nome_res, "images")
out_msk_dir = os.path.join(pasta_base, nome_res, "masks")
legacy_msk2 = os.path.join(os.path.dirname(legacy_msk), "masks2")
legacy_root = os.path.dirname(legacy_msk)
legacy_msk2 = os.path.join(legacy_root, "masks2")
legacy_labels = os.path.join(legacy_root, "labels")
usar_masks2 = USE_MASKS2 and os.path.isdir(legacy_msk2)
usar_labels = USE_LABELS and os.path.isdir(legacy_labels)
out_msk2_dir = os.path.join(pasta_base, nome_res, "masks2") if usar_masks2 else None
out_label_dir = os.path.join(pasta_base, nome_res, "labels") if usar_labels else None
msk_map = map_masks_by_base(legacy_msk)
msk2_map = map_masks2_by_base(legacy_msk2) if usar_masks2 else {}
label_map = map_labels_by_base(legacy_labels) if usar_labels else {}
if USE_LABELS and not usar_labels:
msg = f"[WARN] [{fonte_nome} | legacy] dual_head_label=true, mas labels/ não existe."
if strict_label:
print(msg + " Pulando legacy.")
return total
print(msg + " Seguindo sem labels.")
imgs = [f for f in os.listdir(legacy_img) if os.path.splitext(f.lower())[1] in IMG_EXTS]
n = len(imgs)
for i, fname in enumerate(sorted(imgs), 1):
base, _ = os.path.splitext(fname)
base_norm = normalizar_base(base)
caminho_rgb = os.path.join(legacy_img, fname)
caminho_mask = msk_map.get(base)
caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
ok = normalize_pair(
caminho_rgb, caminho_mask, cor_para_id, ignore_id,
out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_"
caminho_mask = msk_map.get(base_norm)
caminho_mask2 = msk2_map.get(base_norm) if usar_masks2 else None
caminho_label = label_map.get(base_norm) if usar_labels else None
if usar_labels and not caminho_label:
msg = f"[WARN] [{fonte_nome} | legacy] label não encontrado para {fname}."
if strict_label:
print(msg + " Pulando item.")
continue
print(msg + " Seguindo sem label.")
ok, out_img_path, out_msk_path = normalize_pair(
caminho_rgb,
caminho_mask,
cor_para_id,
ignore_id,
out_img_dir,
out_msk_dir,
dim,
prefix=f"{fonte_nome}_",
)
out_msk2_path = None
if usar_masks2 and out_msk2_dir:
if not caminho_mask2:
print(f"[WARN] [{fonte_nome} | legacy] masks2 existe, mas não achei mask2 p/ {fname} (vou seguir).")
print(f"[WARN] [{fonte_nome} | legacy] masks2 existe, mas não achei mask2 p/ {fname}.")
else:
normalize_pair_mask2(
caminho_rgb, caminho_mask2,
out_msk2_dir, dim, prefix=f"{fonte_nome}_"
_, out_msk2_path = normalize_pair_mask2(
caminho_rgb,
caminho_mask2,
out_msk2_dir,
dim,
prefix=f"{fonte_nome}_",
)
if usar_labels and caminho_label and out_label_dir:
normalize_label(
caminho_rgb=caminho_rgb,
caminho_label=caminho_label,
out_label_dir=out_label_dir,
out_img_path=out_img_path,
out_msk_path=out_msk_path,
out_msk2_path=out_msk2_path,
source_root=legacy_root,
output_root=os.path.join(pasta_base, nome_res),
grupo=None,
fonte_nome=fonte_nome,
prefix=f"{fonte_nome}_",
)
if ok:
total += 1
print(f"[{fonte_nome} | legacy | {nome_res}] {i}/{n}{fname}")
return total
# ===================== MAIN =====================
def main(args):
# === Labelmap ===
# Espera tupla na ordem: (cor_para_id, colormap_rgb, id_para_nome, ignore_rgb)
cor_para_id, _colormap_rgb, _id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = infer_ignore_id(ignore_rgb, default_id=255)
print(cor_para_id, _colormap_rgb, _id_para_nome)
print(f"[INFO] cor_para_id: {cor_para_id}")
print(f"[INFO] classes: {_id_para_nome}")
print(f"[INFO] dual_head_mask/masks2: {USE_MASKS2}")
print(f"[INFO] dual_head_label/labels: {USE_LABELS}")
total_geral = 0
# === ORIGINAL ===
orig_group_root = os.path.join(pasta_base, "original", "group")
if os.path.isdir(orig_group_root):
total_geral += process_group_root(orig_group_root, "original", cor_para_id, ignore_id, groups_except=args.groups_except)
total_geral += process_group_root(
orig_group_root,
"original",
cor_para_id,
ignore_id,
groups_except=args.groups_except,
strict_label=args.strict_label,
)
else:
legacy_img = os.path.join(pasta_base, "original", "images")
legacy_msk = os.path.join(pasta_base, "original", "masks")
total_geral += process_legacy_root(legacy_img, legacy_msk, "original", cor_para_id, ignore_id)
total_geral += process_legacy_root(
legacy_img,
legacy_msk,
"original",
cor_para_id,
ignore_id,
strict_label=args.strict_label,
)
# === AUGMENTED ===
aug_group_root = os.path.join(pasta_base, "augmented", "group")
if os.path.isdir(aug_group_root):
total_geral += process_group_root(aug_group_root, "augmented", cor_para_id, ignore_id, groups_except=args.groups_except)
total_geral += process_group_root(
aug_group_root,
"augmented",
cor_para_id,
ignore_id,
groups_except=args.groups_except,
strict_label=args.strict_label,
)
else:
legacy_img = os.path.join(pasta_base, "augmented", "images")
legacy_msk = os.path.join(pasta_base, "augmented", "masks")
total_geral += process_legacy_root(legacy_img, legacy_msk, "augmented", cor_para_id, ignore_id)
total_geral += process_legacy_root(
legacy_img,
legacy_msk,
"augmented",
cor_para_id,
ignore_id,
strict_label=args.strict_label,
)
print(f"\n✅ Concluído! Total normalizados: {total_geral}")
# === calcula mean/std globais e salva em JSON ===
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
if GLOBAL_SUM is not None and GLOBAL_PIXELS > 0:
# média e variância por canal
mean = (GLOBAL_SUM / GLOBAL_PIXELS)
var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean**2
std = np.sqrt(np.maximum(var, 1e-6))
mean = GLOBAL_SUM / GLOBAL_PIXELS
var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean ** 2
std = np.sqrt(np.maximum(var, 1e-6))
# Converte para list pra salvar em JSON
mean_list = mean.tolist()
std_list = std.tolist()
# Se quiser, você pode nomear os canais explicitamente
# dependendo da convenção do raw4:
std_list = std.tolist()
channel_names = ["R", "G", "B"]
stats = {
@ -369,10 +664,21 @@ def main(args):
print(f" mean: {mean_list}")
print(f" std : {std_list}")
else:
print("⚠️ Nenhum RAW processado, não há stats para salvar.")
print("⚠️ Nenhuma imagem processada, não há stats para salvar.")
if __name__ == "__main__":
ap = argparse.ArgumentParser(description="Augmentação por grupos (images/masks)")
ap.add_argument("--groups-except", type=str, default="", help="Lista de grupos para nao usar, separados por vírgula (ex: chao,erva_cana).")
ap = argparse.ArgumentParser(description="Normalização por grupos images/masks/masks2/labels.")
ap.add_argument(
"--groups-except",
type=str,
default="",
help="Lista de grupos para não usar, separados por vírgula. Ex: chao,erva_cana",
)
ap.add_argument(
"--strict-label",
action="store_true",
help="Se dual_head_label=true e faltar label, pula item/grupo.",
)
args = ap.parse_args()
main(args)

View File

@ -1,72 +1,88 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Split estratificado por GRUPO com **val/test do ORIGINAL** e
garantia de NÃO VAZAMENTO entre splits (mesma família não cruza splits).
Split estratificado por GRUPO com val/test do ORIGINAL e garantia de NÃO VAZAMENTO.
de:
MODELO/dataset/<WxH>/group/<grupo>/{images,masks}
MODELO/dataset/<WxH>/group/<grupo>/{images,masks,(masks2),(labels)}
Escreve em:
MODELO/dataset/split/<split>/group/<grupo>/{images,masks}
MODELO/dataset/split/<split>/group/<grupo>/{images,masks,(masks2),(labels)}
Definições:
- "Família" = todas as variações da MESMA base original:
original_<base>.* e augmented_<base>_aug_XX.*
- Val/Test: **original_<base>** (sem augmented)
- Train: original_<base> **e** todos augmented_<base>_aug_XX
- Família = todas as variações da mesma base original:
original_<base>.*
augmented_<base>_aug_XX.*
- Val/Test: somente original_<base>.
- Train: original_<base> + todos augmented_<base>_aug_XX.
Se não houver prefixos (legado), cai para o comportamento antigo (sem família),
mas ainda evita colocar augmented em val/test se detectar sufixo "_aug_XX".
Labels:
- Ativados por config['dual_head_label'].
- Copia labels .json/.txt e .npy quando existirem.
- O pareamento é feito pelo mesmo base name da imagem.
Uso:
python _7_split_grouped_noleak.py
python _7_split_grouped_noleak.py --train 0.7 --val 0.29 --test 0.01 --seed 42
python _7_split_grouped_noleak.py --min-train 1 --min-val 1 --min-test 0
python _7_split_grouped_noleak.py --modelo OAK-1-Lite-W --resolucao 640x384
python _7_split_grouped_noleak_with_labels.py
python _7_split_grouped_noleak_with_labels.py --train 0.7 --val 0.29 --test 0.01 --seed 42
python _7_split_grouped_noleak_with_labels.py --strict-label
python _7_split_grouped_noleak_with_labels.py --cap-train-families "navegavel:300,naonavegavel_navegavel:800"
"""
import os
import re
import json
import shutil
import random
import argparse
from typing import Dict, List, Optional
# ⚙️ Configurações
with open("config_oak.json", "r", encoding="utf-8") as f:
# ===================== CONFIG =====================
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
USE_MASKS2 = bool(config.get("dual_head_mask", config.get("dual_head", False)))
USE_LABELS = bool(config.get("dual_head_label", False))
RESOLUCAO = tuple(config.get("resolucao"))
# Pastas
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
pasta_destino = os.path.join(MODELO, "dataset", "split")
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXT = ".png" # máscaras normalizadas em PNG (recomendado)
MSK_EXT = ".png"
LABEL_EXTS = (".json", ".txt", ".npy")
# Regex para identificar famílias
RE_ORIGINAL_PREFIX = re.compile(r'^original_(.+)$', re.IGNORECASE)
RE_AUGMENTED_FAMILY = re.compile(r'^augmented_(.+?)(?:_aug_\d+)?$', re.IGNORECASE)
RE_AUG_SUFFIX = re.compile(r'_aug_\d+$', re.IGNORECASE)
RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
RE_AUGMENTED_FAMILY = re.compile(r"^augmented_(.+?)(?:_aug_\d+)?$", re.IGNORECASE)
RE_AUG_SUFFIX = re.compile(r"_aug_\d+$", re.IGNORECASE)
# ===================== HELPERS =====================
def garantir(p):
os.makedirs(p, exist_ok=True)
def lista_grupos(root):
if not os.path.isdir(root):
return []
out = []
for g in sorted(os.listdir(root)):
gdir = os.path.join(root, g)
if not os.path.isdir(gdir): continue
if not os.path.isdir(gdir):
continue
if os.path.isdir(os.path.join(gdir, "images")) and os.path.isdir(os.path.join(gdir, "masks")):
out.append(g)
return out
def listar_imagens(img_dir):
if not os.path.isdir(img_dir): return []
if not os.path.isdir(img_dir):
return []
fs = []
for f in os.listdir(img_dir):
ext = os.path.splitext(f.lower())[1]
@ -74,19 +90,27 @@ def listar_imagens(img_dir):
fs.append(f)
return sorted(fs)
def mask_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return base + MSK_EXT
def mask2_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return base + MSK_EXT # masks2 também normalizadas em PNG no normalize
return base + MSK_EXT
def label_candidates_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return [base + ext for ext in LABEL_EXTS]
def classify_source_and_family(filename_no_ext):
"""
Retorna (source, family_key)
source {"original", "augmented", "unknown"}
family_key = base associada ao original (sem prefixo/sufixos), ex: "foo_001"
source {original, augmented, unknown}
family_key = base original sem prefixo/sufixo.
"""
m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
if m:
@ -96,46 +120,66 @@ def classify_source_and_family(filename_no_ext):
if m:
return "augmented", m.group(1)
# legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
if RE_AUG_SUFFIX.search(filename_no_ext):
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
return "augmented", fam
return "unknown", filename_no_ext
def build_family_index(img_dir, msk_dir):
def has_any_label(label_dir: str, img_name: str) -> bool:
if not label_dir or not os.path.isdir(label_dir):
return False
return any(os.path.exists(os.path.join(label_dir, cand)) for cand in label_candidates_from_image_name(img_name))
def build_family_index(img_dir, msk_dir, label_dir=None, require_label=False):
"""
Constroi índice de famílias a partir de img_dir/msk_dir.
Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]}
(strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe)
Constrói índice de famílias a partir de img_dir/msk_dir/labels.
Retorna:
dict family -> {
original: str|None,
augmented: [str],
all: [str]
}
Os nomes são arquivos de imagem.
"""
familias = {}
imgs = listar_imagens(img_dir)
for img_name in imgs:
base_no_ext, ext = os.path.splitext(img_name)
base_no_ext, _ = os.path.splitext(img_name)
mask_name = mask_from_image_name(img_name)
if not os.path.exists(os.path.join(msk_dir, mask_name)):
continue # garante pareamento
continue
if require_label and not has_any_label(label_dir, img_name):
continue
source, fam = classify_source_and_family(base_no_ext)
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
d["all"].append(img_name)
if source == "original":
d["original"] = img_name
elif source == "augmented":
d["augmented"].append(img_name)
else:
# trata como original desconhecido para não perder dado
if d["original"] is None:
d["original"] = img_name
else:
d["augmented"].append(img_name)
return familias
def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
n_train = int(round(n * p_train))
n_val = int(round(n * p_val))
n_test = n - n_train - n_val
n_val = int(round(n * p_val))
n_test = n - n_train - n_val
if n_test < 0:
excesso = -n_test
@ -151,8 +195,8 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
min_sum = min_train + min_val + min_test
if n >= min_sum:
n_train = max(n_train, min_train)
n_val = max(n_val, min_val)
n_test = max(n_test, min_test)
n_val = max(n_val, min_val)
n_test = max(n_test, min_test)
total = n_train + n_val + n_test
while total > n:
@ -165,6 +209,7 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
else:
break
total = n_train + n_val + n_test
while total < n:
if n_train - min_train <= n_val - min_val:
n_train += 1
@ -177,11 +222,9 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
n_val = max(0, min(resto, min_val))
n_test = max(0, resto - n_val)
# ajuste final
diff = n - (n_train + n_val + n_test)
if diff != 0:
if diff > 0:
# adiciona em train, depois val
take = min(diff, n - n_train)
n_train += take
diff -= take
@ -189,7 +232,6 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
n_val += diff
else:
diff = -diff
# tira de test, depois val
take = min(diff, n_test)
n_test -= take
diff -= take
@ -198,38 +240,113 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
return n_train, n_val, n_test
def copiar(nomes, src_img_dir, src_msk_dir, dst_img_dir, dst_msk_dir, src_msk2_dir=None, dst_msk2_dir=None):
garantir(dst_img_dir); garantir(dst_msk_dir)
def copiar_labels_para_item(nome, src_label_dir, dst_label_dir):
if not src_label_dir or not dst_label_dir or not os.path.isdir(src_label_dir):
return 0
garantir(dst_label_dir)
copied = 0
for cand in label_candidates_from_image_name(nome):
src = os.path.join(src_label_dir, cand)
if not os.path.exists(src):
continue
shutil.copy2(src, os.path.join(dst_label_dir, cand))
copied += 1
return copied
def copiar(
nomes,
src_img_dir,
src_msk_dir,
dst_img_dir,
dst_msk_dir,
src_msk2_dir=None,
dst_msk2_dir=None,
src_label_dir=None,
dst_label_dir=None,
strict_label=False,
):
garantir(dst_img_dir)
garantir(dst_msk_dir)
use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir))
use_labels = bool(src_label_dir and dst_label_dir and os.path.isdir(src_label_dir))
if use_msk2:
garantir(dst_msk2_dir)
if use_labels:
garantir(dst_label_dir)
moved = 0
skipped_no_label = 0
for nome in nomes:
mask_name = mask_from_image_name(nome)
src_img = os.path.join(src_img_dir, nome)
src_msk = os.path.join(src_msk_dir, mask_name)
if not (os.path.exists(src_img) and os.path.exists(src_msk)):
continue
if strict_label and use_labels and not has_any_label(src_label_dir, nome):
skipped_no_label += 1
continue
shutil.copy2(src_img, os.path.join(dst_img_dir, nome))
shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name))
if use_msk2:
m2_name = mask2_from_image_name(nome)
src_m2 = os.path.join(src_msk2_dir, m2_name)
if os.path.exists(src_m2):
shutil.copy2(src_m2, os.path.join(dst_msk2_dir, m2_name))
if use_labels:
copied = copiar_labels_para_item(nome, src_label_dir, dst_label_dir)
if strict_label and copied == 0:
skipped_no_label += 1
continue
moved += 1
if skipped_no_label > 0:
print(f"[WARN] {skipped_no_label} itens pulados por falta de label.")
return moved
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
# ===================== SPLIT =====================
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None, strict_label=False):
src_img_dir = os.path.join(pasta_origem, group_name, "images")
src_msk_dir = os.path.join(pasta_origem, group_name, "masks")
src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2")
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
src_label_dir = os.path.join(pasta_origem, group_name, "labels")
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
use_labels = USE_LABELS and os.path.isdir(src_label_dir)
if USE_LABELS and not use_labels:
msg = f"[{group_name}] dual_head_label=true, mas labels/ não existe."
if strict_label:
print(f"[WARN] {msg} Pulando grupo.")
return {"train": 0, "val": 0, "test": 0, "familias": 0}
print(f"[WARN] {msg} Split seguirá sem copiar labels.")
familias = build_family_index(
src_img_dir,
src_msk_dir,
label_dir=src_label_dir if use_labels else None,
require_label=bool(strict_label and use_labels),
)
familias = build_family_index(src_img_dir, src_msk_dir)
# apenas famílias que têm ORIGINAL para participar de val/test
familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
total_familias = len(familias_originais)
if total_familias == 0:
print(f"[{group_name}] 0 famílias com original, pulando.")
return {"train": 0, "val": 0, "test": 0, "familias": 0}
@ -238,81 +355,145 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
rng.shuffle(familias_originais)
n_tr, n_va, n_te = allocate_counts(
total_familias, p_train, p_val, p_test,
mins["train"], mins["val"], mins["test"]
total_familias,
p_train,
p_val,
p_test,
mins["train"],
mins["val"],
mins["test"],
)
fam_train = set(familias_originais[:n_tr])
fam_val = set(familias_originais[n_tr:n_tr+n_va])
fam_test = set(familias_originais[n_tr+n_va: n_tr+n_va+n_te])
fam_val = set(familias_originais[n_tr : n_tr + n_va])
fam_test = set(familias_originais[n_tr + n_va : n_tr + n_va + n_te])
# --- CAP por grupo (apenas no TRAIN) ---
if caps_map and group_name in caps_map:
cap = caps_map[group_name]
if len(fam_train) > cap:
fam_list = list(fam_train)
rng.shuffle(fam_list) # usa o rng já criado com seed
kept = set(fam_list[:cap])
dropped = set(fam_list[cap:])
rng.shuffle(fam_list)
kept = set(fam_list[:cap])
dropped = set(fam_list[cap:])
fam_train = kept
print(f"[{group_name}] cap-train-families={cap} → mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN")
print(
f"[{group_name}] cap-train-families={cap}"
f"mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN"
)
# listas de nomes por split (imagens)
nomes_train, nomes_val, nomes_test = [], [], []
for fam, d in familias.items():
if fam in fam_train:
# train recebe original + todos augmented
if d["original"]:
nomes_train.append(d["original"])
if d["augmented"]:
nomes_train.extend(d["augmented"])
elif fam in fam_val:
# val recebe somente original
if d["original"]:
nomes_val.append(d["original"])
elif fam in fam_test:
# test recebe somente original
if d["original"]:
nomes_test.append(d["original"])
else:
# famílias sem original (não devem cair aqui) ficam fora
pass
# dest dirs
dest_train_img = os.path.join(pasta_destino, "train", "group", group_name, "images")
dest_train_msk = os.path.join(pasta_destino, "train", "group", group_name, "masks")
dest_val_img = os.path.join(pasta_destino, "val", "group", group_name, "images")
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
dest_test_img = os.path.join(pasta_destino, "test", "group", group_name, "images")
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks")
dest_val_img = os.path.join(pasta_destino, "val", "group", group_name, "images")
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
dest_test_img = os.path.join(pasta_destino, "test", "group", group_name, "images")
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks")
dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
m_train = copiar(nomes_train, src_img_dir, src_msk_dir, dest_train_img, dest_train_msk, src_msk2_dir, dest_train_msk2)
m_val = copiar(nomes_val, src_img_dir, src_msk_dir, dest_val_img, dest_val_msk, src_msk2_dir, dest_val_msk2)
m_test = copiar(nomes_test, src_img_dir, src_msk_dir, dest_test_img, dest_test_msk, src_msk2_dir, dest_test_msk2)
dest_train_label = os.path.join(pasta_destino, "train", "group", group_name, "labels") if use_labels else None
dest_val_label = os.path.join(pasta_destino, "val", "group", group_name, "labels") if use_labels else None
dest_test_label = os.path.join(pasta_destino, "test", "group", group_name, "labels") if use_labels else None
print(f"[{group_name}] famílias={total_familias} → train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}")
m_train = copiar(
nomes_train,
src_img_dir,
src_msk_dir,
dest_train_img,
dest_train_msk,
src_msk2_dir,
dest_train_msk2,
src_label_dir,
dest_train_label,
strict_label=strict_label,
)
m_val = copiar(
nomes_val,
src_img_dir,
src_msk_dir,
dest_val_img,
dest_val_msk,
src_msk2_dir,
dest_val_msk2,
src_label_dir,
dest_val_label,
strict_label=strict_label,
)
m_test = copiar(
nomes_test,
src_img_dir,
src_msk_dir,
dest_test_img,
dest_test_msk,
src_msk2_dir,
dest_test_msk2,
src_label_dir,
dest_test_label,
strict_label=strict_label,
)
print(
f"[{group_name}] famílias={total_familias}"
f"train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}"
)
return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}
def main():
ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (val/test só original).")
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).")
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação (default=0.29).")
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste (default=0.01).")
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento (default=42).")
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
# ===================== MAIN =====================
def parse_cap_map(s):
caps = {}
if not s:
return caps
for item in s.split(","):
if not item.strip():
continue
k, v = item.strip().split(":")
caps[k.strip()] = int(v)
return caps
def main():
ap = argparse.ArgumentParser(description="Split estratificado por grupo sem vazamento, com suporte a labels.")
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino.")
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação.")
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste.")
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento.")
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de famílias por grupo em train.")
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de famílias por grupo em val.")
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de famílias por grupo em test.")
ap.add_argument("--modelo", type=str, default=None, help="Sobrescreve MODELO do config.json.")
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 640x480).")
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH.")
ap.add_argument("--cap-train-families", type=str, default="", help="Mapa 'grupo:cap,...' p/ limitar número de FAMÍLIAS no TRAIN. Ex.: 'chao:350'")
ap.add_argument(
"--cap-train-families",
type=str,
default="",
help="Mapa grupo:cap para limitar famílias no TRAIN. Ex: 'navegavel:350'",
)
ap.add_argument(
"--strict-label",
action="store_true",
help="Se dual_head_label=true e faltar label, pula item/grupo.",
)
args = ap.parse_args()
@ -326,14 +507,6 @@ def main():
else:
resolucao = RESOLUCAO
def parse_cap_map(s):
caps = {}
if not s: return caps
for item in s.split(","):
k,v = item.strip().split(":")
caps[k.strip()] = int(v)
return caps
caps_map = parse_cap_map(args.cap_train_families)
global pasta_origem, pasta_destino
@ -341,12 +514,18 @@ def main():
pasta_destino = os.path.join(modelo, "dataset", "split")
soma = args.train + args.val + args.test
if soma <= 0: raise ValueError("Soma de proporções deve ser > 0.")
p_train = args.train / soma
p_val = args.val / soma
p_test = args.test / soma
if soma <= 0:
raise ValueError("Soma de proporções deve ser > 0.")
mins = {"train": max(0, args.min_train), "val": max(0, args.min_val), "test": max(0, args.min_test)}
p_train = args.train / soma
p_val = args.val / soma
p_test = args.test / soma
mins = {
"train": max(0, args.min_train),
"val": max(0, args.min_val),
"test": max(0, args.min_test),
}
garantir(pasta_destino)
@ -357,22 +536,28 @@ def main():
random.seed(args.seed)
total_global = {"train":0, "val":0, "test":0, "familias":0}
total_global = {"train": 0, "val": 0, "test": 0, "familias": 0}
print(f"[INFO] Origem: {pasta_origem}")
print(f"[INFO] Destino: {pasta_destino}")
print(f"[INFO] dual_head_mask/masks2: {USE_MASKS2}")
print(f"[INFO] dual_head_label/labels: {USE_LABELS}")
print(f"Grupos: {', '.join(grupos)}")
print(f"Proporções normalizadas: train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
print(f"Mínimos por grupo (famílias): train={mins['train']} val={mins['val']} test={mins['test']}")
print(f"Mínimos por grupo: train={mins['train']} val={mins['val']} test={mins['test']}")
for g in grupos:
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map)
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map, strict_label=args.strict_label)
for k in total_global.keys():
total_global[k] += res.get(k, 0)
print("\nResumo global (imagens copiadas):")
print("\nResumo global, imagens copiadas:")
print(f" train: {total_global['train']}")
print(f" val: {total_global['val']}")
print(f" test: {total_global['test']}")
print(f" famílias (total): {total_global['familias']}")
print(f" famílias total: {total_global['familias']}")
print("\n✅ Split sem vazamento concluído!")
if __name__ == "__main__":
main()

View File

@ -258,7 +258,7 @@ def run_one_epoch(model: nn.Module,
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="config_oak.json")
parser.add_argument("--config", default="config.json")
parser.add_argument("--epochs", type=int, default=120)
parser.add_argument("--batch", type=int, default=1) # <<< default seguro pra 8GB
parser.add_argument("--lr", type=float, default=6e-5)

View File

@ -1,965 +0,0 @@
# _8_train_segformer_b3_dual_v2.py
# Treino SegFormer-B3 com 2 cabeças:
# - Head 1: segmentação semântica (3 classes + ignore)
# - Head 2: corredor binário (0/1 + ignore)
#
# Features:
# - class weights (opcional, estimado do train)
# - pos_weight (opcional, estimado do train p/ BCE do corredor)
# - warmup + cosine (opcional) + ReduceLROnPlateau
# - early stopping com métrica agregada
# - resume completo (model + corridor_head + optimizer + scaler)
#
# Observação: este script assume que seu ROISegDataset retorna dict com:
# {"image": <tensor/ndarray>, "mask": <tensor/ndarray>, ...}
# e que você já acertou para também expor "image_path"/"mask_path" OU
# que seu DualMaskROISegDataset consiga derivar paths corretamente.
#python _8_train_segformer_b3_dual.py --epochs 80 --batch 2 --num_workers 2 --amp --amp_val --use_weights --auto_pos_weight --warmup_epochs 2 --cosine_epochs 15 --lambda_corr 0.5 --corr_dice_max 0.10 --corr_dice_ramp_epochs 10 --es_metric harmonic --es_patience 12
#python _8_train_segformer_b3_dual.py --epochs 120 --batch 2 --num_workers 4 --amp --amp_val --grad_accum 1 --use_weights --auto_pos_weight --warmup_epochs 2 --cosine_epochs 20 --lambda_corr 0.7 --corr_dice_max 0.15 --corr_dice_ramp_epochs 10 --es_metric harmonic --es_patience 15
import os
import re
import time
import math
import json
import random
import argparse
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.amp import autocast
from torch.cuda.amp import GradScaler
# HuggingFace transformers
from transformers import SegformerForSemanticSegmentation
# Seu dataset (o mesmo que você já usa no projeto)
from roi_seg_dataset import ROISegDataset
# -----------------------------
# Utils
# -----------------------------
def seed_everything(seed: int = 42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
def normalize_img(img: torch.Tensor) -> torch.Tensor:
return (img - IMAGENET_MEAN.to(img.device)) / IMAGENET_STD.to(img.device)
@torch.no_grad()
def update_confusion_matrix(cm: torch.Tensor, preds: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = 255):
preds = preds.view(-1)
labels = labels.view(-1)
valid = labels != ignore_index
preds = preds[valid]
labels = labels[valid]
idx = labels * num_classes + preds
bins = torch.bincount(idx, minlength=num_classes * num_classes)
cm += bins.view(num_classes, num_classes)
@torch.no_grad()
def compute_iou_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Tuple[float, List[float]]:
cm = cm.float()
tp = torch.diag(cm)
fp = cm.sum(0) - tp
fn = cm.sum(1) - tp
denom = tp + fp + fn + eps
iou = (tp / denom).cpu().tolist()
miou = float(np.mean(iou))
return miou, iou
@torch.no_grad()
def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float:
cm = cm.float()
acc = (torch.diag(cm).sum() / (cm.sum() + eps)).item()
return acc
def default_collate_dual(batch):
imgs, m1s, m2s = [], [], []
for item in batch:
# item pode vir como (img, mask1, mask2) ou dict etc.
if isinstance(item, dict):
img = item["image"]
mask1 = item["mask1"]
mask2 = item["mask2"]
else:
img, mask1, mask2 = item
if isinstance(img, np.ndarray):
img = torch.from_numpy(img)
if isinstance(mask1, np.ndarray):
mask1 = torch.from_numpy(mask1)
if isinstance(mask2, np.ndarray):
mask2 = torch.from_numpy(mask2)
if img.ndim == 3 and img.shape[-1] == 3:
img = img.permute(2, 0, 1)
if img.dtype != torch.float32:
img = img.float()
if img.max() > 1.5:
img = img / 255.0
imgs.append(img)
m1s.append(mask1.long())
m2s.append(mask2.long())
return torch.stack(imgs, 0), torch.stack(m1s, 0), torch.stack(m2s, 0)
def estimate_class_weights_from_ds(ds: ROISegDataset, num_classes: int, ignore_index: int = 255, max_samples: int = 800) -> torch.Tensor:
n = min(len(ds), max_samples)
idxs = np.random.choice(len(ds), size=n, replace=False)
counts = np.zeros(num_classes, dtype=np.float64)
for i in idxs:
item = ds[i]
mask = item["mask"] if isinstance(item, dict) else item[1]
m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
m = m.reshape(-1)
m = m[m != ignore_index]
if m.size == 0:
continue
counts += np.bincount(m, minlength=num_classes)[:num_classes]
freq = counts / (counts.sum() + 1e-12)
freq = np.clip(freq, 1e-12, 1.0)
weights = 1.0 / np.log(1.02 + freq)
weights = weights / weights.mean()
return torch.tensor(weights, dtype=torch.float32)
def estimate_pos_weight_from_masks2(ds_dual, ignore_index2: int = 255, max_samples: int = 800) -> Optional[torch.Tensor]:
"""
Estima pos_weight = neg/pos para BCEWithLogitsLoss.
Considera targets binários {0,1} ignorando ignore_index2.
"""
n = min(len(ds_dual), max_samples)
if n <= 0:
return None
idxs = np.random.choice(len(ds_dual), size=n, replace=False)
pos = 0
neg = 0
for i in idxs:
it = ds_dual[i]
m2 = it["mask2"] if isinstance(it, dict) else it[2]
if isinstance(m2, torch.Tensor):
m2 = m2.cpu().numpy()
m2 = np.array(m2).reshape(-1)
m2 = m2[m2 != ignore_index2]
if m2.size == 0:
continue
pos += int((m2 == 1).sum())
neg += int((m2 == 0).sum())
if pos <= 0:
return None
pw = float(neg) / float(pos)
# clamp leve pra não explodir treino
pw = float(np.clip(pw, 1.0, 50.0))
return torch.tensor([pw], dtype=torch.float32)
def dice_loss_with_logits(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = 255, eps: float = 1e-6) -> torch.Tensor:
"""
logits: [B,1,H,W]
targets: [B,H,W] com {0,1} e possivelmente ignore_index
"""
probs = torch.sigmoid(logits)
t = targets.float().unsqueeze(1)
if ignore_index is not None:
valid = (targets != ignore_index).unsqueeze(1)
probs = probs[valid]
t = t[valid]
probs = probs.reshape(-1)
t = t.reshape(-1)
inter = (probs * t).sum()
denom = probs.sum() + t.sum() + eps
dice = (2.0 * inter + eps) / denom
return 1.0 - dice
@torch.no_grad()
def binary_iou_and_acc_from_logits(logits: torch.Tensor, targets: torch.Tensor, thr: float = 0.5, ignore_index: int = 255, eps: float = 1e-6) -> Tuple[float, float]:
probs = torch.sigmoid(logits)
preds = (probs >= thr).long().squeeze(1)
if logits.shape[-2:] != targets.shape[-2:]:
preds = F.interpolate(preds.unsqueeze(1).float(), size=targets.shape[-2:], mode="nearest").long().squeeze(1)
valid = (targets != ignore_index)
if valid.sum().item() == 0:
return 0.0, 0.0
p = preds[valid]
t = targets[valid].long()
tp = ((p == 1) & (t == 1)).sum().float()
fp = ((p == 1) & (t == 0)).sum().float()
fn = ((p == 0) & (t == 1)).sum().float()
iou = (tp / (tp + fp + fn + eps)).item()
acc = ((p == t).sum().float() / (t.numel() + eps)).item()
return iou, acc
def _pretty_iou(iou_list: List[float], class_name_by_id: Dict[int, str]) -> str:
parts = []
for cid, v in enumerate(iou_list):
parts.append(f"{class_name_by_id.get(cid, str(cid))}:{v:.3f}")
return " | ".join(parts)
def save_checkpoint(path: str,
base_model: nn.Module,
corridor_head: nn.Module,
optimizer: torch.optim.Optimizer,
scaler: Optional[GradScaler],
epoch: int,
best_miou: float,
best_main_iou: float,
best_corr_iou: float,
extra: Optional[Dict[str, Any]] = None):
ckpt = {
"epoch": epoch,
"model": base_model.state_dict(),
"corridor_head": corridor_head.state_dict(),
"optimizer": optimizer.state_dict(),
"best_miou": best_miou,
"best_main_iou": best_main_iou,
"best_corr_iou": best_corr_iou,
}
if scaler is not None:
ckpt["scaler"] = scaler.state_dict()
if extra:
ckpt["extra"] = extra
torch.save(ckpt, path)
def load_checkpoint(path: str,
base_model: nn.Module,
corridor_head: nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None,
scaler: Optional[GradScaler] = None,
map_location: str = "cpu") -> Dict[str, Any]:
ckpt = torch.load(path, map_location=map_location, weights_only=False)
base_model.load_state_dict(ckpt["model"], strict=True)
corridor_head.load_state_dict(ckpt["corridor_head"], strict=True)
if optimizer is not None and "optimizer" in ckpt:
optimizer.load_state_dict(ckpt["optimizer"])
if scaler is not None and "scaler" in ckpt:
scaler.load_state_dict(ckpt["scaler"])
return ckpt
# -----------------------------
# Dual Dataset Wrapper
# -----------------------------
class DualMaskROISegDataset(torch.utils.data.Dataset):
def __init__(self, base_ds: ROISegDataset, split_root: str, ignore_index2: int = 255):
self.base_ds = base_ds
self.split_root = split_root
self.ignore_index2 = ignore_index2
# NÃO falha aqui por pasta. Vamos resolver por caminho da mask1.
self.masks2_dir = os.path.join(split_root, "masks2")
self.has_flat_masks2 = os.path.isdir(self.masks2_dir)
def __len__(self):
return len(self.base_ds)
def _get_mask1_path(self, item: Dict[str, Any], idx: int) -> str:
# 1) se vier no dict, usa
for k in ("mask_path", "mask_file", "maskname", "mask_name"):
if k in item and item[k]:
return str(item[k])
# 2) fallback: usa o índice no ROISegDataset (ele tem msk_paths)
if hasattr(self.base_ds, "msk_paths"):
return str(self.base_ds.msk_paths[idx])
raise RuntimeError(
"ROISegDataset não retornou mask_path/mask_file e não possui base_ds.msk_paths. "
"Não dá pra localizar a máscara para buscar a masks2."
)
def _load_mask2_from_mask1_path(self, mask1_path: str, target_shape_hw: Tuple[int, int]) -> torch.Tensor:
# 1) modo robusto: espelha estrutura trocando /masks/ -> /masks2/
norm = os.path.normpath(mask1_path)
parts = norm.split(os.sep)
try:
i = parts.index("masks")
parts[i] = "masks2"
path2 = os.sep.join(parts)
except ValueError:
path2 = ""
# 2) fallback: modo “flat” (split_root/masks2/<filename>)
if (not path2) or (not os.path.exists(path2)):
if not self.has_flat_masks2:
raise FileNotFoundError(
f"Mask2 não encontrada espelhando a mask1.\n"
f"mask1: {mask1_path}\n"
f"tentativa: {path2}\n"
f"E também não existe: {self.masks2_dir}"
)
filename = os.path.basename(mask1_path)
path2 = os.path.join(self.masks2_dir, filename)
if not os.path.exists(path2):
raise FileNotFoundError(f"Mask2 não encontrada: {path2}")
m = np.array(Image.open(path2))
if m.ndim == 3:
m = m[..., 0]
if m.max() > 1:
m = (m >= 128).astype(np.uint8)
H, W = target_shape_hw
if m.shape[0] != H or m.shape[1] != W:
m = np.array(Image.fromarray(m).resize((W, H), resample=Image.NEAREST))
return torch.from_numpy(m).long()
def __getitem__(self, idx):
item = self.base_ds[idx]
# caso venha (dict,) ou [dict]
if isinstance(item, (list, tuple)) and len(item) == 1 and isinstance(item[0], dict):
item = item[0]
if isinstance(item, dict):
pass
elif isinstance(item, (list, tuple)) and len(item) >= 2:
img, mask = item[0], item[1]
item = {"image": img, "mask": mask}
else:
raise RuntimeError(
f"ROISegDataset retornou formato inesperado: type={type(item)} value={item}"
)
img = item["image"]
mask1 = item["mask"]
if isinstance(mask1, torch.Tensor):
h, w = int(mask1.shape[-2]), int(mask1.shape[-1])
else:
m = np.array(mask1)
h, w = m.shape[0], m.shape[1]
mask1_path = self._get_mask1_path(item, idx)
mask2 = self._load_mask2_from_mask1_path(mask1_path, (h, w))
item["mask1"] = mask1
item["mask2"] = mask2
return item
# -----------------------------
# Corridor head
# -----------------------------
class CorridorHead(nn.Module):
"""
Head simples:
feat (B,C,H,W) + logits_seg (B,K,H,W) -> concat -> convs -> 1 canal (logit corredor)
"""
def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1):
super().__init__()
in_ch = feat_ch + num_classes
self.in_ch = in_ch
self.net = nn.Sequential(
nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Dropout2d(dropout),
nn.Conv2d(hidden, hidden, kernel_size=3, padding=1),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Dropout2d(dropout),
nn.Conv2d(hidden, 1, kernel_size=1)
)
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
x = torch.cat([feat, logits_seg], dim=1)
return self.net(x)
# -----------------------------
# One epoch
# -----------------------------
def run_one_epoch_dual(base_model: nn.Module,
corridor_head: nn.Module,
loader: DataLoader,
optimizer: Optional[torch.optim.Optimizer],
device: torch.device,
num_classes: int,
ignore_index1: int,
ignore_index2: int,
criterion_seg: nn.Module,
bce_corr: nn.Module,
lambda_corr: float,
amp: bool,
scaler: Optional[GradScaler],
train: bool,
grad_accum: int = 1,
corridor_thr: float = 0.5,
corr_dice_mix: float = 0.0) -> Dict[str, Any]:
base_model.train(train)
corridor_head.train(train)
total_loss = 0.0
total_loss_seg = 0.0
total_loss_corr = 0.0
cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)
corr_iou_accum = 0.0
corr_acc_accum = 0.0
t0 = time.time()
n_batches = 0
with torch.set_grad_enabled(train):
if train and optimizer is not None:
optimizer.zero_grad(set_to_none=True)
for step, (imgs, masks1, masks2) in enumerate(loader, start=1):
imgs = imgs.to(device, non_blocking=True)
masks1 = masks1.to(device, non_blocking=True)
masks2 = masks2.to(device, non_blocking=True)
#print("masks2 unique:", torch.unique(masks2)[:10])
#print("masks2 max:", masks2.max().item(), "min:", masks2.min().item())
imgs = normalize_img(imgs)
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
out = base_model(pixel_values=imgs)
logits1 = out.logits # [B,K,H,W]
if logits1.shape[-2:] != masks1.shape[-2:]:
logits1 = F.interpolate(logits1, size=masks1.shape[-2:], mode="bilinear", align_corners=False)
loss_seg = criterion_seg(logits1, masks1)
# feat do encoder: pega hidden_states (se disponível), senão reusa logits como proxy
feat = None
if hasattr(out, "hidden_states") and out.hidden_states is not None:
feat = out.hidden_states[-1] # [B, C, h, w]
else:
# fallback: usa logits como feature (não ideal, mas mantém vivo)
feat = logits1
if feat.shape[-2:] != logits1.shape[-2:]:
feat = F.interpolate(feat, size=logits1.shape[-2:], mode="bilinear", align_corners=False)
logits_corr = corridor_head(feat, logits1) # [B,1,H,W]
if logits_corr.shape[-2:] != masks2.shape[-2:]:
logits_corr = F.interpolate(logits_corr, size=masks2.shape[-2:], mode="bilinear", align_corners=False)
#with torch.no_grad():
# p = torch.sigmoid(logits_corr)
# print("corr prob mean:", p.mean().item(), "min:", p.min().item(), "max:", p.max().item())
# BCE (com ignore via masking manual)
valid = (masks2 != ignore_index2)
if valid.sum().item() > 0:
tgt = masks2.float().unsqueeze(1)
bce_val = bce_corr(logits_corr[valid.unsqueeze(1)], tgt[valid.unsqueeze(1)])
if corr_dice_mix > 0:
d_val = dice_loss_with_logits(logits_corr, masks2, ignore_index=ignore_index2)
loss_corr = (1.0 - corr_dice_mix) * bce_val + corr_dice_mix * d_val
else:
loss_corr = bce_val
else:
loss_corr = torch.zeros([], device=device, dtype=loss_seg.dtype)
loss = loss_seg + lambda_corr * loss_corr
if train and grad_accum > 1:
loss = loss / grad_accum
if train and optimizer is not None:
if amp and scaler is not None and device.type == "cuda":
scaler.scale(loss).backward()
if (step % grad_accum) == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
else:
loss.backward()
if (step % grad_accum) == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
total_loss += float(loss.item()) * (grad_accum if (train and grad_accum > 1) else 1.0)
total_loss_seg += float(loss_seg.item())
total_loss_corr += float(loss_corr.item())
n_batches += 1
# metrics seg
preds1 = torch.argmax(logits1.detach(), dim=1)
update_confusion_matrix(cm, preds1, masks1, num_classes, ignore_index=ignore_index1)
# metrics corr
ciou, cacc = binary_iou_and_acc_from_logits(logits_corr.detach(), masks2, thr=corridor_thr, ignore_index=ignore_index2)
corr_iou_accum += ciou
corr_acc_accum += cacc
miou, iou_per_class = compute_iou_from_cm(cm)
acc = compute_pixel_acc_from_cm(cm)
out = {
"loss": total_loss / max(1, n_batches),
"loss_seg": total_loss_seg / max(1, n_batches),
"loss_corr": total_loss_corr / max(1, n_batches),
"acc": acc,
"miou": miou,
"iou_per_class": iou_per_class,
"corr_iou": corr_iou_accum / max(1, n_batches),
"corr_acc": corr_acc_accum / max(1, n_batches),
"time_s": time.time() - t0
}
return out
# -----------------------------
# LR / EarlyStop helpers
# -----------------------------
@dataclass
class EarlyStopState:
best: float = -1e9
bad_epochs: int = 0
stopped: bool = False
def agg_metric(es_metric: str, miou: float, main_iou: float, corr_iou: float) -> float:
if es_metric == "miou":
return float(miou)
if es_metric == "main":
return float(main_iou)
if es_metric == "corr":
return float(corr_iou)
if es_metric == "harmonic":
eps = 1e-6
a = max(eps, float(main_iou))
b = max(eps, float(corr_iou))
return float(2.0 * a * b / (a + b + eps))
# default
return float(miou)
def apply_warmup(optimizer, base_lr: float, epoch: int, warmup_epochs: int):
if warmup_epochs <= 0:
return
if epoch <= warmup_epochs:
wf = epoch / float(warmup_epochs)
for g in optimizer.param_groups:
g["lr"] = base_lr * wf
# -----------------------------
# Main
# -----------------------------
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", default="config_oak.json")
ap.add_argument("--epochs", type=int, default=120)
ap.add_argument("--batch", type=int, default=4)
ap.add_argument("--num_workers", type=int, default=4)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--lr", type=float, default=6e-5)
ap.add_argument("--wd", type=float, default=1e-2)
ap.add_argument("--min_lr", type=float, default=6e-6)
ap.add_argument("--amp", action="store_true")
ap.add_argument("--amp_val", action="store_true")
ap.add_argument("--grad_accum", type=int, default=1)
ap.add_argument("--ignore_index", type=int, default=255)
ap.add_argument("--ignore_index2", type=int, default=255)
ap.add_argument("--lambda_corr", type=float, default=0.5)
ap.add_argument("--corr_thr", type=float, default=0.5)
# class weights
ap.add_argument("--use_weights", action="store_true")
ap.add_argument("--cw_max_samples", type=int, default=800)
# pos_weight corredor
ap.add_argument("--auto_pos_weight", action="store_true")
ap.add_argument("--pw_max_samples", type=int, default=800)
# mistura dice no corredor (rampa)
ap.add_argument("--corr_dice_max", type=float, default=0.15)
ap.add_argument("--corr_dice_ramp_epochs", type=int, default=10)
# schedulers
ap.add_argument("--warmup_epochs", type=int, default=0)
ap.add_argument("--cosine_epochs", type=int, default=0, help="se >0: usa cosine por N épocas e depois troca p/ plateau")
ap.add_argument("--plateau_patience", type=int, default=3)
ap.add_argument("--plateau_factor", type=float, default=0.6)
ap.add_argument("--plateau_cooldown", type=int, default=1)
# early stopping
ap.add_argument("--es_metric", default="harmonic", choices=["miou", "main", "corr", "harmonic"])
ap.add_argument("--es_patience", type=int, default=12)
ap.add_argument("--es_min_delta", type=float, default=2e-4)
# logs/ckpt
ap.add_argument("--save_every", type=int, default=0)
ap.add_argument("--resume", action="store_true")
ap.add_argument("--main_class", type=str, default=None)
args = ap.parse_args()
seed_everything(args.seed)
with open(args.config, "r") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
RESOLUCAO = config["resolucao"]
ROI_INICIO = config["roi_inicio"]
ROI_TAMANHO = config["roi_tamanho"]
BACKBONE = config["backbone"]
MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower()
if args.main_class is not None:
MAIN_CLASS_NAME = args.main_class.lower()
dataset_path = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
# save paths
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME + "_dual")
os.makedirs(save_path, exist_ok=True)
last_ckpt_path = os.path.join(save_path, "last.pt")
best_miou_path = os.path.join(save_path, "best_miou.pt")
best_main_path = os.path.join(save_path, "best_main.pt")
best_corr_path = os.path.join(save_path, "best_corr.pt")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)
# seus splits estão em train/group/... então o ROISegDataset deve saber ler de lá
split_train = os.path.join(dataset_path, "split", "train")
split_val = os.path.join(dataset_path, "split", "val")
ds_train_base = ROISegDataset(
split_train,
save_path, ROI_INICIO, ROI_TAMANHO,
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
)
ds_val_base = ROISegDataset(
split_val,
save_path, ROI_INICIO, ROI_TAMANHO,
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
)
ds_train = DualMaskROISegDataset(ds_train_base, split_train, ignore_index2=args.ignore_index2)
ds_val = DualMaskROISegDataset(ds_val_base, split_val, ignore_index2=args.ignore_index2)
CLASS_NAMES = getattr(ds_train_base, "classes", None)
if CLASS_NAMES is None:
raise RuntimeError("ROISegDataset precisa expor .classes (list ou dict).")
if isinstance(CLASS_NAMES, list):
num_classes = len(CLASS_NAMES)
class_id_by_name = {n.lower(): i for i, n in enumerate(CLASS_NAMES)}
class_name_by_id = {i: n for i, n in enumerate(CLASS_NAMES)}
elif isinstance(CLASS_NAMES, dict):
if all(isinstance(k, int) for k in CLASS_NAMES.keys()):
num_classes = len(CLASS_NAMES)
class_name_by_id = {int(k): str(v) for k, v in CLASS_NAMES.items()}
class_id_by_name = {str(v).lower(): int(k) for k, v in CLASS_NAMES.items()}
else:
class_id_by_name = {str(k).lower(): int(v) for k, v in CLASS_NAMES.items()}
num_classes = len(class_id_by_name)
class_name_by_id = {v: k for k, v in class_id_by_name.items()}
else:
raise RuntimeError("Formato de .classes não reconhecido.")
main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None)
if main_class_id is None:
print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.")
else:
print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}")
dl_train = DataLoader(
ds_train,
batch_size=args.batch,
shuffle=True,
num_workers=args.num_workers,
pin_memory=(device.type == "cuda"),
collate_fn=default_collate_dual
)
dl_val = DataLoader(
ds_val,
batch_size=args.batch,
shuffle=False,
num_workers=max(0, args.num_workers // 2),
pin_memory=(device.type == "cuda"),
collate_fn=default_collate_dual
)
# model base SegFormer
base_model = SegformerForSemanticSegmentation.from_pretrained(
BACKBONE,
num_labels=num_classes,
ignore_mismatched_sizes=True,
use_safetensors=True, # <- evita torch.load do .bin
)
base_model.config.output_hidden_states = True
base_model.to(device)
# descobre feat_ch via dummy forward (mais robusto)
with torch.no_grad():
dummy = torch.zeros((1, 3, 512, 512), device=device)
out = base_model(pixel_values=dummy)
feat_ch = None
if hasattr(out, "hidden_states") and out.hidden_states is not None:
feat_ch = int(out.hidden_states[-1].shape[1])
else:
feat_ch = int(out.logits.shape[1]) # fallback
corridor_head = CorridorHead(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1).to(device)
print(f"[corridor_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {corridor_head.in_ch}")
# losses
class_weights = None
if args.use_weights:
w = estimate_class_weights_from_ds(ds_train_base, num_classes=num_classes, ignore_index=args.ignore_index, max_samples=args.cw_max_samples)
class_weights = w.to(device)
print("[INFO] class_weights:", class_weights.detach().cpu().numpy().round(3).tolist())
criterion_seg = nn.CrossEntropyLoss(ignore_index=args.ignore_index, weight=class_weights)
pos_weight = None
if args.auto_pos_weight:
pw = estimate_pos_weight_from_masks2(ds_train, ignore_index2=args.ignore_index2, max_samples=args.pw_max_samples)
if pw is not None:
pos_weight = pw.to(device)
print("[INFO] pos_weight corredor (neg/pos):", float(pos_weight.item()))
else:
print("[INFO] pos_weight não estimável (sem positivos suficientes), usando None.")
bce_corr = nn.BCEWithLogitsLoss(pos_weight=pos_weight) if pos_weight is not None else nn.BCEWithLogitsLoss()
# optimizer (base + head)
params = list(base_model.parameters()) + list(corridor_head.parameters())
optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wd)
scaler = GradScaler(enabled=(args.amp and device.type == "cuda"))
# schedulers
cosine = None
if args.cosine_epochs and args.cosine_epochs > 0:
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=max(1, args.cosine_epochs),
eta_min=args.min_lr
)
plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode="min",
factor=args.plateau_factor,
patience=args.plateau_patience,
cooldown=args.plateau_cooldown,
min_lr=args.min_lr,
verbose=True
)
active_sched = "cosine" if cosine is not None else "plateau"
# resume
start_epoch = 1
best_miou = -1.0
best_main_iou = -1.0
best_corr_iou = -1.0
if args.resume and os.path.exists(last_ckpt_path):
ckpt = load_checkpoint(last_ckpt_path, base_model, corridor_head, optimizer=optimizer, scaler=scaler, map_location="cpu")
start_epoch = int(ckpt.get("epoch", 0)) + 1
best_miou = float(ckpt.get("best_miou", -1.0))
best_main_iou = float(ckpt.get("best_main_iou", -1.0))
best_corr_iou = float(ckpt.get("best_corr_iou", -1.0))
print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f} best_corr_iou={best_corr_iou:.4f}")
# early stopping
es = EarlyStopState(best=-1e9, bad_epochs=0, stopped=False)
for epoch in range(start_epoch, args.epochs + 1):
lr_now = optimizer.param_groups[0]["lr"]
# warmup
apply_warmup(optimizer, args.lr, epoch, args.warmup_epochs)
lr_now = optimizer.param_groups[0]["lr"]
# corr dice ramp
if args.corr_dice_max > 0 and args.corr_dice_ramp_epochs > 0:
corr_dice_mix = min(args.corr_dice_max, (epoch - 1) / float(args.corr_dice_ramp_epochs) * args.corr_dice_max)
else:
corr_dice_mix = 0.0
print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} | lambda_corr={args.lambda_corr} | corr_dice={corr_dice_mix:.3f} | sched={active_sched} ====")
if device.type == "cuda":
torch.cuda.empty_cache()
tr = run_one_epoch_dual(
base_model=base_model,
corridor_head=corridor_head,
loader=dl_train,
optimizer=optimizer,
device=device,
num_classes=num_classes,
ignore_index1=args.ignore_index,
ignore_index2=args.ignore_index2,
criterion_seg=criterion_seg,
bce_corr=bce_corr,
lambda_corr=args.lambda_corr,
amp=args.amp,
scaler=scaler,
train=True,
grad_accum=max(1, args.grad_accum),
corridor_thr=args.corr_thr,
corr_dice_mix=corr_dice_mix
)
if device.type == "cuda":
torch.cuda.empty_cache()
va = run_one_epoch_dual(
base_model=base_model,
corridor_head=corridor_head,
loader=dl_val,
optimizer=None,
device=device,
num_classes=num_classes,
ignore_index1=args.ignore_index,
ignore_index2=args.ignore_index2,
criterion_seg=criterion_seg,
bce_corr=bce_corr,
lambda_corr=args.lambda_corr,
amp=args.amp_val,
scaler=None,
train=False,
grad_accum=1,
corridor_thr=args.corr_thr,
corr_dice_mix=corr_dice_mix
)
# scheduler step
if cosine is not None and epoch <= args.cosine_epochs:
cosine.step()
else:
active_sched = "plateau"
plateau.step(va["loss"])
main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id]
corr_iou = float(va["corr_iou"])
print(f"TRAIN: loss={tr['loss']:.4f} (seg={tr['loss_seg']:.4f} corr={tr['loss_corr']:.4f}) "
f"acc={tr['acc']:.4f} miou={tr['miou']:.4f} corr_iou={tr['corr_iou']:.4f} (t={tr['time_s']:.1f}s)")
print(f"VAL : loss={va['loss']:.4f} (seg={va['loss_seg']:.4f} corr={va['loss_corr']:.4f}) "
f"acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={float(main_iou):.4f} "
f"corr_iou={va['corr_iou']:.4f} corr_acc={va['corr_acc']:.4f} (t={va['time_s']:.1f}s)")
print("IoU per class:", _pretty_iou(va["iou_per_class"], class_name_by_id))
# save last
save_checkpoint(
last_ckpt_path,
base_model,
corridor_head,
optimizer,
scaler=scaler,
epoch=epoch,
best_miou=best_miou,
best_main_iou=best_main_iou,
best_corr_iou=best_corr_iou,
extra={
"val_loss": va["loss"],
"val_miou": va["miou"],
"val_main_iou": float(main_iou),
"val_corr_iou": float(corr_iou),
"lr": optimizer.param_groups[0]["lr"],
"corr_dice_mix": corr_dice_mix
}
)
# save every
if args.save_every > 0 and (epoch % args.save_every == 0):
save_checkpoint(
os.path.join(save_path, f"epoch_{epoch:04d}.pt"),
base_model, corridor_head, optimizer, scaler=scaler,
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou
)
# bests
if va["miou"] > best_miou:
best_miou = va["miou"]
save_checkpoint(best_miou_path, base_model, corridor_head, optimizer, scaler=scaler,
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou)
print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}")
if float(main_iou) > best_main_iou:
best_main_iou = float(main_iou)
save_checkpoint(best_main_path, base_model, corridor_head, optimizer, scaler=scaler,
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou)
print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}")
if corr_iou > best_corr_iou:
best_corr_iou = corr_iou
save_checkpoint(best_corr_path, base_model, corridor_head, optimizer, scaler=scaler,
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou)
print(f"[BEST CORR] {best_corr_iou:.4f} -> saved: {best_corr_path}")
# early stopping
score = agg_metric(args.es_metric, va["miou"], float(main_iou), corr_iou)
if score > es.best + args.es_min_delta:
es.best = score
es.bad_epochs = 0
print(f"[ES] improved {args.es_metric} -> {score:.6f}")
else:
es.bad_epochs += 1
print(f"[ES] no improve ({es.bad_epochs}/{args.es_patience}) best={es.best:.6f} now={score:.6f}")
if es.bad_epochs >= args.es_patience:
print(f"🛑 Early stopping acionado (metric={args.es_metric}).")
break
if __name__ == "__main__":
main()

File diff suppressed because it is too large Load Diff

View File

@ -215,7 +215,7 @@ def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Lê config do projeto (mesmo padrão do fastscnn)
with open("config_oak.json", "r") as f:
with open("config.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]
@ -261,7 +261,7 @@ def main():
model = load_segformer_from_checkpoint(pt_path, backbone=backbone, num_classes=num_classes, device=device)
# Normalização (ImageNet, padrão de muita coisa; se teu treino usou outro, troca aqui)
from _8_train_segformer_b3 import normalize_img
from _8_train_segformer import normalize_img
if args.camera:
# === Modo câmera (igual estilo do fastscnn) ===

View File

@ -1,570 +0,0 @@
import json
import os
import glob
import argparse
from dataclasses import dataclass
from typing import List, Optional
import cv2
import numpy as np
import torch
import torch.nn as nn
from transformers import SegformerForSemanticSegmentation
# -----------------------------
# Utils
# -----------------------------
def find_device():
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
def load_labelmap_lut_bgr(labelmap_path: str):
"""
labelmap no formato:
classe:r,g,b
Ignora linhas vazias e comentários '#'
Se existir 'ignore', retorna ignore_id=255 e ignore_bgr.
Retorna:
lut_bgr: dict[int, tuple(b,g,r)]
id_to_name: dict[int, str]
ignore_id: int (default 255)
"""
lut_bgr = {}
id_to_name = {}
ignore_id = 255
ignore_bgr = (180, 0, 180) # fallback
idx = 0
with open(labelmap_path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line or line.startswith("#"):
continue
parts = line.split(":")
if len(parts) < 2:
continue
name = parts[0].strip()
rgb_str = parts[1].strip()
r, g, b = map(int, rgb_str.split(","))
bgr = (b, g, r)
if name.lower() == "ignore":
ignore_bgr = bgr
continue
lut_bgr[idx] = bgr
id_to_name[idx] = name
idx += 1
lut_bgr[ignore_id] = ignore_bgr
return lut_bgr, id_to_name, ignore_id
def imread_bgr(path: str) -> np.ndarray:
img = cv2.imread(path, cv2.IMREAD_COLOR)
if img is None:
raise RuntimeError(f"Falha ao ler imagem: {path}")
return img
def imread_gray(path: str) -> np.ndarray:
m = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
if m is None:
raise RuntimeError(f"Falha ao ler máscara: {path}")
return m
def ensure_uint8(x: np.ndarray) -> np.ndarray:
if x.dtype == np.uint8:
return x
x = np.clip(x, 0, 255).astype(np.uint8)
return x
def normalize01(x: np.ndarray, eps: float = 1e-6) -> np.ndarray:
x = x.astype(np.float32)
mn, mx = float(x.min()), float(x.max())
return (x - mn) / (mx - mn + eps)
def colorize_seg(mask_ids: np.ndarray, lut_bgr: dict) -> np.ndarray:
h, w = mask_ids.shape[:2]
out = np.zeros((h, w, 3), dtype=np.uint8)
# pinta cada id presente
for cid, bgr in lut_bgr.items():
out[mask_ids == cid] = bgr
return out
def overlay_mask(img_bgr: np.ndarray, mask_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
return cv2.addWeighted(img_bgr, 1.0 - alpha, mask_bgr, alpha, 0.0)
def put_hud(img: np.ndarray, lines: List[str]) -> np.ndarray:
out = img.copy()
y = 22
for s in lines:
cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (15, 15, 15), 3, cv2.LINE_AA)
cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (245, 245, 245), 1, cv2.LINE_AA)
y += 22
return out
# -----------------------------
# Corridor Head (igual ao treino)
# -----------------------------
class CorridorHead(nn.Module):
"""
Head simples:
feat (B,C,H,W) + logits_seg (B,K,H,W) -> concat -> convs -> 1 canal (logit corredor)
"""
def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1):
super().__init__()
in_ch = feat_ch + num_classes
self.in_ch = in_ch
self.net = nn.Sequential(
nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Dropout2d(dropout),
nn.Conv2d(hidden, hidden, kernel_size=3, padding=1),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Dropout2d(dropout),
nn.Conv2d(hidden, 1, kernel_size=1)
)
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
x = torch.cat([feat, logits_seg], dim=1)
return self.net(x)
# -----------------------------
# Dataset discovery (split/test/group/**)
# -----------------------------
@dataclass
class Sample:
img_path: str
mask_path: Optional[str]
mask2_path: Optional[str]
group_name: str
filename: str
IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
def discover_samples(split_root: str) -> List[Sample]:
"""
Espera estrutura:
split/test/group/<qualquer>/images/*.png
split/test/group/<qualquer>/masks/*.png
split/test/group/<qualquer>/masks2/*.png (opcional, mas recomendado)
"""
group_root = os.path.join(split_root, "group")
if not os.path.isdir(group_root):
raise RuntimeError(f"Não achei pasta: {group_root}")
samples: List[Sample] = []
# pega qualquer images/ em qualquer subpasta de group
img_dirs = glob.glob(os.path.join(group_root, "**", "images"), recursive=True)
img_dirs = [d for d in img_dirs if os.path.isdir(d)]
for idir in img_dirs:
base = os.path.dirname(idir) # .../group/<grupo_ou_subgrupo>
group_name = os.path.relpath(base, group_root).replace("\\", "/")
mdir = os.path.join(base, "masks")
m2dir = os.path.join(base, "masks2")
img_paths = []
for ext in IMG_EXTS:
img_paths.extend(glob.glob(os.path.join(idir, f"*{ext}")))
img_paths = sorted(img_paths)
def find_mask(dir_path: str, filename: str):
stem, _ = os.path.splitext(filename)
for ext in [".png", ".jpg", ".jpeg", ".bmp", ".tif"]:
p = os.path.join(dir_path, stem + ext)
if os.path.exists(p):
return p
return None
for ip in img_paths:
fn = os.path.basename(ip)
mask_path = find_mask(mdir, fn)
mask2_path = find_mask(m2dir, fn)
samples.append(Sample(
img_path=ip,
mask_path=mask_path,
mask2_path=mask2_path,
group_name=group_name,
filename=fn
))
if len(samples) == 0:
raise RuntimeError(f"Nenhuma imagem encontrada em: {group_root}/**/images")
return samples
# -----------------------------
# Model loading
# -----------------------------
def build_models(num_classes: int, device: torch.device, backbone: str):
base_model = SegformerForSemanticSegmentation.from_pretrained(
backbone,
num_labels=num_classes,
ignore_mismatched_sizes=True,
use_safetensors=True, # <- evita torch.load do .bin
)
base_model.to(device)
base_model.config.output_hidden_states = True
# descobre feat_ch real a partir de um forward fake
with torch.no_grad():
dummy = torch.zeros((1, 3, 512, 512), device=device)
out = base_model(pixel_values=dummy, output_hidden_states=True)
feat_ch = int(out.hidden_states[-1].shape[1])
in_ch = feat_ch + num_classes
corridor_head = CorridorHead(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1)
print(f"[corridor_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {in_ch}")
return base_model, corridor_head
def make_corridor_head(in_ch: int, mid: int = 256):
"""
Cria CorridorHead independente de qual assinatura a classe tem no script.
Suporta:
- CorridorHead(in_ch=..., mid=...)
- CorridorHead(in_channels=..., hidden=...)
- CorridorHead(feat_ch=..., num_classes=...) (quando in_ch = feat_ch + num_classes)
"""
# 1) assinatura direta in_ch/mid
try:
return CorridorHead(in_ch=in_ch, mid=mid)
except TypeError:
pass
# 2) variações comuns
try:
return CorridorHead(in_channels=in_ch, mid=mid)
except TypeError:
pass
try:
return CorridorHead(in_channels=in_ch, hidden=mid)
except TypeError:
pass
# 3) assinatura "feat_ch + num_classes"
# aqui só funciona se for exatamente "feat_ch, num_classes"
# e se in_ch for decomponível (ex: 6 = 3+3 ou 515 = 512+3)
for feat_ch_guess, num_classes_guess in [(512, 3), (3, 3)]:
if feat_ch_guess + num_classes_guess == in_ch:
try:
return CorridorHead(feat_ch=feat_ch_guess, num_classes=num_classes_guess, hidden=mid)
except TypeError:
try:
return CorridorHead(feat_ch=feat_ch_guess, num_classes=num_classes_guess)
except TypeError:
pass
raise RuntimeError(
f"Não consegui instanciar CorridorHead para in_ch={in_ch}. "
f"Verifique a assinatura do __init__() da classe CorridorHead no script."
)
def load_dual_checkpoint(ckpt_path, base_model, corridor_head, device):
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
base_model.load_state_dict(ckpt["model"], strict=True)
# >>> pega quantos canais o head do checkpoint espera
sd = ckpt["corridor_head"]
in_ch_ckpt = int(sd["net.0.weight"].shape[1]) # ex: 6 ou 515
# >>> recria head se o in_ch não bater
in_ch_model = int(corridor_head.net[0].in_channels)
if in_ch_model != in_ch_ckpt:
print(f"[corridor_head] rebuild: in_ch_model={in_ch_model} -> in_ch_ckpt={in_ch_ckpt}")
corridor_head = make_corridor_head(in_ch=in_ch_ckpt, mid=256)
corridor_head.load_state_dict(sd, strict=True)
base_model.to(device).eval()
corridor_head.to(device).eval()
meta = ckpt.get("meta", {})
return meta, corridor_head
# -----------------------------
# Inference
# -----------------------------
def run_corridor_head(corridor_head, x_rgb, feat, logits_seg):
import inspect
import torch.nn.functional as F
n_args = len(inspect.signature(corridor_head.forward).parameters)
# descobre in_ch esperado pela primeira conv do head
in_ch = None
if hasattr(corridor_head, "net"):
try:
in_ch = int(corridor_head.net[0].in_channels)
except Exception:
in_ch = None
# prepara entradas conforme o head
if in_ch == 6:
# RGB(3) + probs(3)
probs = torch.softmax(logits_seg, dim=1)
if probs.shape[-2:] != x_rgb.shape[-2:]:
probs = F.interpolate(probs, size=x_rgb.shape[-2:], mode="bilinear", align_corners=False)
feat_in = x_rgb
logits_in = probs
else:
# feat(512) + logits(3)
if logits_seg.shape[-2:] != feat.shape[-2:]:
logits_in = F.interpolate(logits_seg, size=feat.shape[-2:], mode="bilinear", align_corners=False)
else:
logits_in = logits_seg
feat_in = feat
# chama do jeito que o head espera
if n_args == 2:
return corridor_head(feat_in, logits_in)
elif n_args == 1:
x2 = torch.cat([feat_in, logits_in], dim=1)
return corridor_head(x2)
raise RuntimeError(f"Assinatura inesperada: {inspect.signature(corridor_head.forward)}")
@torch.no_grad()
def infer_dual(base_model, corridor_head, img_bgr: np.ndarray, device: torch.device,
input_size: int = 512, corridor_thr: float = 0.5):
"""
Retorna:
pred_ids_full: [H,W] uint8
prob_corr_full: [H,W] float32 (0..1)
bin_corr_full: [H,W] uint8 (0/255)
"""
h0, w0 = img_bgr.shape[:2]
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
img_res = cv2.resize(img_rgb, (input_size, input_size), interpolation=cv2.INTER_AREA)
x = torch.from_numpy(img_res).float() / 255.0
x = x.permute(2, 0, 1).unsqueeze(0).to(device)
out = base_model(pixel_values=x, output_hidden_states=True)
logits_seg = out.logits # [B, 3, H, W]
feat = out.hidden_states[-1] # [B, 512, h, w] (não vamos usar no modo 6)
logits_corr = run_corridor_head(corridor_head, x, feat, logits_seg)
prob_corr = torch.sigmoid(logits_corr)[0, 0].detach().float().cpu().numpy()
bin_corr = (prob_corr >= corridor_thr).astype(np.uint8) * 255
pred_ids = torch.argmax(logits_seg, dim=1)[0].detach().cpu().numpy().astype(np.uint8)
# volta pro tamanho original
pred_ids_full = cv2.resize(pred_ids, (w0, h0), interpolation=cv2.INTER_NEAREST)
prob_corr_full = cv2.resize(prob_corr, (w0, h0), interpolation=cv2.INTER_LINEAR)
bin_corr_full = cv2.resize(bin_corr, (w0, h0), interpolation=cv2.INTER_NEAREST)
return pred_ids_full, prob_corr_full, bin_corr_full
# -----------------------------
# Viewer
# -----------------------------
def make_panel(sample: Sample , lut_bgr: dict, img_bgr: np.ndarray,
gt1: Optional[np.ndarray],
pred1: np.ndarray,
gt2: Optional[np.ndarray],
prob2: np.ndarray,
bin2: np.ndarray,
corridor_thr: float):
h, w = img_bgr.shape[:2]
pred1_col = colorize_seg(pred1, lut_bgr)
pred1_ov = overlay_mask(img_bgr, pred1_col, 0.45)
if gt1 is not None:
gt1_col = colorize_seg(gt1, lut_bgr)
gt1_ov = overlay_mask(img_bgr, gt1_col, 0.45)
else:
gt1_ov = img_bgr.copy()
gt1_col = np.zeros_like(img_bgr)
# --- PRED CORR (prob) sólido: verde=navegável, vermelho=bloqueado
prob = np.clip(prob2, 0.0, 1.0).astype(np.float32)
g = (prob * 255.0).astype(np.uint8)
r = ((1.0 - prob) * 255.0).astype(np.uint8)
b = np.zeros_like(g, dtype=np.uint8)
prob_solid = cv2.merge([b, g, r]) # BGR
# (opcional) desenha contorno do binário por cima do prob
edges = cv2.Canny(bin2, 50, 150)
prob_solid[edges > 0] = (255, 255, 255)
# --- PRED CORR (bin) sólido: branco=navegável, preto=bloqueado
bin2_solid = cv2.cvtColor(bin2, cv2.COLOR_GRAY2BGR)
if gt2 is not None:
gt2_vis = (gt2.copy()).astype(np.uint8)
gt2_vis[gt2_vis == 1] = 255
gt2_bgr = cv2.cvtColor(gt2_vis, cv2.COLOR_GRAY2BGR)
gt2_ov = overlay_mask(img_bgr, gt2_bgr, 0.35)
else:
gt2_ov = img_bgr.copy()
# monta grid
tile_w = 520
tile_h = int(tile_w * h / w)
def fit(im):
return cv2.resize(im, (tile_w, tile_h), interpolation=cv2.INTER_AREA)
row1 = np.concatenate([fit(img_bgr), fit(gt1_ov), fit(pred1_ov)], axis=1)
row2 = np.concatenate([fit(gt2_ov), fit(prob_solid), fit(bin2_solid)], axis=1)
panel = np.concatenate([row1, row2], axis=0)
hud = [
f"{sample.group_name}/{sample.filename}",
f"Corr thr={corridor_thr:.2f} | A/D navega | Q sai",
"Topo: IMG | GT SEG | PRED SEG",
"Baixo: GT CORR | PRED CORR (prob) | PRED CORR (bin)"
]
panel = put_hud(panel, hud)
return panel
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", type=str, default=None, help="Caminho do .pt (se não informar, usa best_miou.pt)")
ap.add_argument("--split", default="test", choices=["test", "val"], help="qual split usar")
ap.add_argument("--input_size", type=int, default=512)
ap.add_argument("--corr_thr", type=float, default=0.5)
args = ap.parse_args()
device = find_device()
print("Device:", device)
with open("config_oak.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"] # ex: "segformer_b3"
BACKBONE = config["backbone"]
modelo_folder = config["modelo"] # ex: "dif"
dataset_path = os.path.join(MODELO, "dataset")
split_root = os.path.join(dataset_path, "split", args.split)
if not os.path.isdir(split_root):
print(f"[WARN] split/{args.split} não existe. Vou usar split/val.")
split_root = os.path.join(dataset_path, "split", "val")
samples = discover_samples(split_root)
print(f"[INFO] samples: {len(samples)} | split_root={split_root}")
labelmap_path = os.path.join(dataset_path, "labelmap.txt") # ajusta se seu arquivo estiver em outro lugar
lut_bgr, id_to_name, ignore_id = load_labelmap_lut_bgr(labelmap_path)
print("[labelmap]", id_to_name)
# ---- Descobre checkpoint ----
if args.ckpt is not None:
ckpt_path = args.ckpt
else:
# Caminho padrão igual ao treino:
# save_path = MODELO/backup/modelo/model_name/raw4
save_path = os.path.join(MODELO, "backup", modelo_folder, f"{MODEL_NAME}_dual")
ckpt_path = os.path.join(save_path, "best_miou.pt")
if not os.path.isfile(ckpt_path):
raise SystemExit(f"Checkpoint não encontrado em: {ckpt_path}")
print(f"[model] ckpt = {ckpt_path}")
num_classes = len(id_to_name)
base_model, corridor_head = build_models(num_classes=num_classes, device=device, backbone=BACKBONE)
meta, corridor_head = load_dual_checkpoint(ckpt_path, base_model, corridor_head, device=device)
print(f"... epoch={meta.get('epoch')} best_miou={meta.get('best_miou')} best_main={meta.get('best_main_iou')} best_corr={meta.get('best_corr_iou')}")
idx = 0
win = "SegFormer Dual Viewer"
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
while True:
s = samples[idx]
img = imread_bgr(s.img_path)
gt1 = imread_gray(s.mask_path) if s.mask_path else None
gt2 = imread_gray(s.mask2_path) if s.mask2_path else None
pred1, prob2, bin2 = infer_dual(
base_model=base_model,
corridor_head=corridor_head,
img_bgr=img,
device=device,
input_size=args.input_size,
corridor_thr=args.corr_thr
)
def dbg_mask(name, m):
if m is None:
print(f"[{name}] None")
else:
u = np.unique(m)
print(f"[{name}] shape={m.shape} dtype={m.dtype} unique={u[:20]}{'...' if len(u)>20 else ''}")
#dbg_mask("gt1(seg)", gt1)
#dbg_mask("gt2(corr)", gt2)
panel = make_panel(
sample=s,
lut_bgr=lut_bgr,
img_bgr=img,
gt1=gt1,
pred1=pred1,
gt2=gt2,
prob2=prob2,
bin2=bin2,
corridor_thr=args.corr_thr
)
cv2.imshow(win, panel)
k = cv2.waitKey(0) & 0xFF
# Q / ESC
if k in (ord('q'), 27):
break
# A (volta)
if k in (ord('a'), ord('A')):
idx = (idx - 1) % len(samples)
continue
# D (avança)
if k in (ord('d'), ord('D')):
idx = (idx + 1) % len(samples)
continue
# W/S muda threshold
if k in (ord('w'), ord('W')):
args.corr_thr = min(0.95, args.corr_thr + 0.05)
continue
if k in (ord('s'), ord('S')):
args.corr_thr = max(0.05, args.corr_thr - 0.05)
continue
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

@ -0,0 +1,998 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Viewer/teste SegFormer com suporte a:
1) single
- segmentação semântica
2) dual_head_mask
- segmentação + mask2 binária
3) dual_head_label
- segmentação + estado global do corredor
Compatível com checkpoint salvo pelo script:
_8_train_segformer_dual_mask_or_label.py
Uso:
python _9_test_segformer_dual_mask_or_label.py --config config.json --split test
python _9_test_segformer_dual_mask_or_label.py --config config.json --split val --ckpt caminho.pt
python _9_test_segformer_dual_mask_or_label.py --config config.json --split test --save_dir outputs_test
Teclas no viewer:
D / seta direita = próxima
A / seta esquerda = anterior
W/S = muda threshold da mask2, apenas modo dual_head_mask
Q / ESC = sair
"""
import os
import re
import cv2
import glob
import json
import argparse
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Any
import time
import depthai as dai
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import SegformerForSemanticSegmentation
# ============================================================
# Utils básicos
# ============================================================
IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
def normalize_img(img: torch.Tensor) -> torch.Tensor:
return (img - IMAGENET_MEAN.to(img.device)) / IMAGENET_STD.to(img.device)
def find_device():
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
def load_labelmap_lut_bgr(labelmap_path: str):
lut_bgr = {}
id_to_name = {}
ignore_id = 255
ignore_bgr = (180, 0, 180)
idx = 0
with open(labelmap_path, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line or line.startswith("#"):
continue
parts = line.split(":")
if len(parts) < 2:
continue
name = parts[0].strip()
rgb_str = parts[1].strip()
r, g, b = map(int, rgb_str.split(","))
bgr = (b, g, r)
if name.lower() == "ignore":
ignore_bgr = bgr
continue
lut_bgr[idx] = bgr
id_to_name[idx] = name
idx += 1
lut_bgr[ignore_id] = ignore_bgr
return lut_bgr, id_to_name, ignore_id
def imread_bgr(path: str) -> np.ndarray:
img = cv2.imread(path, cv2.IMREAD_COLOR)
if img is None:
raise RuntimeError(f"Falha ao ler imagem: {path}")
return img
def imread_gray(path: str) -> np.ndarray:
m = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
if m is None:
raise RuntimeError(f"Falha ao ler máscara: {path}")
return m
def colorize_seg(mask_ids: np.ndarray, lut_bgr: dict) -> np.ndarray:
h, w = mask_ids.shape[:2]
out = np.zeros((h, w, 3), dtype=np.uint8)
for cid, bgr in lut_bgr.items():
out[mask_ids == cid] = bgr
return out
def overlay_mask(img_bgr: np.ndarray, mask_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
return cv2.addWeighted(img_bgr, 1.0 - alpha, mask_bgr, alpha, 0.0)
def put_hud(img: np.ndarray, lines: List[str]) -> np.ndarray:
out = img.copy()
y = 24
for s in lines:
cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.62, (15, 15, 15), 3, cv2.LINE_AA)
cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.62, (245, 245, 245), 1, cv2.LINE_AA)
y += 24
return out
def resize_fit(im: np.ndarray, tile_w: int, tile_h: Optional[int] = None) -> np.ndarray:
h, w = im.shape[:2]
if tile_h is None:
tile_h = int(tile_w * h / max(1, w))
return cv2.resize(im, (tile_w, tile_h), interpolation=cv2.INTER_AREA)
def safe_float(v, default=None):
try:
return float(v)
except Exception:
return default
# ============================================================
# Métricas
# ============================================================
def seg_metrics(pred: np.ndarray, gt: Optional[np.ndarray], num_classes: int, ignore_id: int = 255) -> Optional[Dict[str, Any]]:
if gt is None:
return None
if pred.shape != gt.shape:
pred = cv2.resize(pred, (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST)
valid = gt != ignore_id
if valid.sum() == 0:
return None
ious = []
for c in range(num_classes):
p = (pred == c) & valid
t = (gt == c) & valid
inter = np.logical_and(p, t).sum()
union = np.logical_or(p, t).sum()
ious.append(float(inter / union) if union > 0 else 0.0)
acc = float((pred[valid] == gt[valid]).sum() / max(1, valid.sum()))
return {"miou": float(np.mean(ious)), "acc": acc, "ious": ious}
def mask2_metrics(bin_pred_255: np.ndarray, gt: Optional[np.ndarray]) -> Optional[Dict[str, float]]:
if gt is None:
return None
if bin_pred_255.shape != gt.shape:
bin_pred_255 = cv2.resize(bin_pred_255, (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST)
p = bin_pred_255 >= 128
t = gt >= 128
tp = np.logical_and(p, t).sum()
fp = np.logical_and(p, ~t).sum()
fn = np.logical_and(~p, t).sum()
union = tp + fp + fn
iou = float(tp / union) if union > 0 else 0.0
acc = float((p == t).sum() / p.size)
return {"iou": iou, "acc": acc}
# ============================================================
# Heads iguais ao treino novo
# ============================================================
class Mask2Head(nn.Module):
def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1):
super().__init__()
in_ch = feat_ch + num_classes
self.in_ch = in_ch
self.net = nn.Sequential(
nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Dropout2d(dropout),
nn.Conv2d(hidden, hidden, kernel_size=3, padding=1),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Dropout2d(dropout),
nn.Conv2d(hidden, 1, kernel_size=1),
)
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
x = torch.cat([feat, logits_seg], dim=1)
return self.net(x)
class LabelHead(nn.Module):
def __init__(self, feat_ch: int, num_seg_classes: int, num_label_classes: int, hidden: int = 256, dropout: float = 0.2):
super().__init__()
in_ch = feat_ch + num_seg_classes
self.in_ch = in_ch
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.net = nn.Sequential(
nn.Linear(in_ch, hidden),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(hidden, num_label_classes),
)
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
if feat.shape[-2:] != logits_seg.shape[-2:]:
feat = F.interpolate(feat, size=logits_seg.shape[-2:], mode="bilinear", align_corners=False)
x = torch.cat([feat, logits_seg], dim=1)
x = self.pool(x).flatten(1)
return self.net(x)
# ============================================================
# Discovery
# ============================================================
@dataclass
class Sample:
img_path: str
mask_path: Optional[str]
mask2_path: Optional[str]
label_json_path: Optional[str]
label_npy_path: Optional[str]
group_name: str
filename: str
IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
def find_by_stem(dir_path: str, filename: str, exts: Tuple[str, ...]) -> Optional[str]:
if not dir_path or not os.path.isdir(dir_path):
return None
stem, _ = os.path.splitext(filename)
for ext in exts:
p = os.path.join(dir_path, stem + ext)
if os.path.exists(p):
return p
return None
def discover_samples(split_root: str) -> List[Sample]:
group_root = os.path.join(split_root, "group")
if not os.path.isdir(group_root):
raise RuntimeError(f"Não achei pasta: {group_root}")
samples: List[Sample] = []
img_dirs = glob.glob(os.path.join(group_root, "**", "images"), recursive=True)
img_dirs = [d for d in img_dirs if os.path.isdir(d)]
for idir in img_dirs:
base = os.path.dirname(idir)
group_name = os.path.relpath(base, group_root).replace("\\", "/")
mdir = os.path.join(base, "masks")
m2dir = os.path.join(base, "masks2")
ldir = os.path.join(base, "labels")
img_paths = []
for ext in IMG_EXTS:
img_paths.extend(glob.glob(os.path.join(idir, f"*{ext}")))
img_paths = sorted(img_paths)
for ip in img_paths:
fn = os.path.basename(ip)
mask_path = find_by_stem(mdir, fn, (".png", ".jpg", ".jpeg", ".bmp", ".tif"))
mask2_path = find_by_stem(m2dir, fn, (".png", ".jpg", ".jpeg", ".bmp", ".tif"))
label_npy_path = find_by_stem(ldir, fn, (".npy",))
label_json_path = find_by_stem(ldir, fn, (".json", ".txt"))
samples.append(Sample(
img_path=ip,
mask_path=mask_path,
mask2_path=mask2_path,
label_json_path=label_json_path,
label_npy_path=label_npy_path,
group_name=group_name,
filename=fn,
))
if len(samples) == 0:
raise RuntimeError(f"Nenhuma imagem encontrada em: {group_root}/**/images")
return samples
def discover_image_folder(folder: str) -> List[Sample]:
samples = []
img_paths = []
for ext in IMG_EXTS:
img_paths.extend(glob.glob(os.path.join(folder, f"*{ext}")))
for ip in sorted(img_paths):
samples.append(Sample(
img_path=ip,
mask_path=None,
mask2_path=None,
label_json_path=None,
label_npy_path=None,
group_name="external",
filename=os.path.basename(ip),
))
if not samples:
raise RuntimeError(f"Nenhuma imagem encontrada em: {folder}")
return samples
def read_gt_label(sample: Sample) -> Tuple[Optional[int], Optional[str]]:
if sample.label_npy_path and os.path.exists(sample.label_npy_path):
try:
v = np.load(sample.label_npy_path)
return int(np.array(v).reshape(-1)[0]), None
except Exception:
pass
if sample.label_json_path and os.path.exists(sample.label_json_path):
ext = os.path.splitext(sample.label_json_path)[1].lower()
if ext == ".json":
with open(sample.label_json_path, "r", encoding="utf-8") as f:
data = json.load(f)
lid = data.get("label_id")
lname = data.get("estado_corredor") or data.get("label") or data.get("state")
return (int(lid) if lid is not None else None), lname
if ext == ".txt":
with open(sample.label_json_path, "r", encoding="utf-8") as f:
txt = f.read().strip()
try:
return int(txt), None
except Exception:
return None, txt
return None, None
# ============================================================
# Model loading
# ============================================================
def build_base_model(num_classes: int, device: torch.device, backbone: str):
base_model = SegformerForSemanticSegmentation.from_pretrained(
backbone,
num_labels=num_classes,
ignore_mismatched_sizes=True,
use_safetensors=True,
)
base_model.config.output_hidden_states = True
base_model.to(device).eval()
return base_model
def get_feat_ch(base_model, device, input_h=512, input_w=512) -> int:
with torch.no_grad():
dummy = torch.zeros((1, 3, input_h, input_w), device=device)
out = base_model(pixel_values=dummy, output_hidden_states=True)
return int(out.hidden_states[-1].shape[1])
def infer_mode_from_config(config: dict) -> str:
use_mask = bool(config.get("dual_head_mask", config.get("dual_head", False)))
use_label = bool(config.get("dual_head_label", False))
if use_mask and use_label:
raise RuntimeError("Config com dual_head_mask e dual_head_label ativos ao mesmo tempo. Este viewer espera apenas um.")
if use_mask:
return "mask2"
if use_label:
return "label"
return "single"
def default_save_path(config: dict, mode: str) -> str:
modelo = config["camera"]
model_name = config["model_name"]
modelo_folder = config["modelo"]
suffix = {
"single": "_single",
"mask2": "_dual_mask",
"label": "_dual_label",
}[mode]
return os.path.join(modelo, "backup", modelo_folder, model_name + suffix)
def load_checkpoint(ckpt_path: str, base_model, aux_head, device):
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
base_model.load_state_dict(ckpt["model"], strict=True)
if aux_head is not None:
if "aux_head" not in ckpt:
raise RuntimeError("Checkpoint não possui aux_head. Ele é compatível com single head apenas?")
aux_head.load_state_dict(ckpt["aux_head"], strict=True)
base_model.to(device).eval()
if aux_head is not None:
aux_head.to(device).eval()
return ckpt
# ============================================================
# Inference
# ============================================================
def preprocess_img(img_bgr: np.ndarray, input_size: int, device: torch.device):
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
img_res = cv2.resize(img_rgb, (input_size, input_size), interpolation=cv2.INTER_AREA)
x = torch.from_numpy(img_res).float() / 255.0
x = x.permute(2, 0, 1).unsqueeze(0).to(device)
img_tensor = normalize_img(x)
return img_tensor
@torch.no_grad()
def infer_sample(base_model, aux_head, img_bgr: np.ndarray, device: torch.device, mode: str, input_size: int = 512, mask2_thr: float = 0.5):
h0, w0 = img_bgr.shape[:2]
x = preprocess_img(img_bgr, input_size, device)
out = base_model(pixel_values=x, output_hidden_states=True)
logits_seg = out.logits
pred_ids = torch.argmax(logits_seg, dim=1)[0].detach().cpu().numpy().astype(np.uint8)
pred_ids_full = cv2.resize(pred_ids, (w0, h0), interpolation=cv2.INTER_NEAREST)
result = {
"pred_seg": pred_ids_full,
"mask2_prob": None,
"mask2_bin": None,
"label_id": None,
"label_conf": None,
"label_probs": None,
}
if mode == "single" or aux_head is None:
return result
feat = out.hidden_states[-1]
if feat.shape[-2:] != logits_seg.shape[-2:]:
feat = F.interpolate(feat, size=logits_seg.shape[-2:], mode="bilinear", align_corners=False)
if mode == "mask2":
logits_m2 = aux_head(feat, logits_seg)
prob = torch.sigmoid(logits_m2)[0, 0].detach().float().cpu().numpy()
bin2 = (prob >= mask2_thr).astype(np.uint8) * 255
result["mask2_prob"] = cv2.resize(prob, (w0, h0), interpolation=cv2.INTER_LINEAR)
result["mask2_bin"] = cv2.resize(bin2, (w0, h0), interpolation=cv2.INTER_NEAREST)
return result
if mode == "label":
logits_label = aux_head(feat, logits_seg)
probs = torch.softmax(logits_label, dim=1)[0].detach().cpu().numpy()
lid = int(np.argmax(probs))
result["label_id"] = lid
result["label_conf"] = float(probs[lid])
result["label_probs"] = probs
return result
return result
# ============================================================
# Painéis
# ============================================================
def make_label_prob_bar(probs: np.ndarray, names: Dict[int, str], width: int = 520, row_h: int = 34) -> np.ndarray:
n = len(probs)
h = max(row_h * n, row_h)
img = np.zeros((h, width, 3), dtype=np.uint8) + 35
for i, p in enumerate(probs):
y0 = i * row_h
bar_w = int((width - 180) * float(p))
cv2.rectangle(img, (170, y0 + 8), (170 + bar_w, y0 + row_h - 8), (80, 180, 80), -1)
txt = f"{i} {names.get(i, f'label_{i}')}: {p:.3f}"
cv2.putText(img, txt, (8, y0 + 23), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (245, 245, 245), 1, cv2.LINE_AA)
return img
def make_panel_single(sample: Sample, lut_bgr, img_bgr, gt1, pred1, segm, id_to_name):
pred_col = colorize_seg(pred1, lut_bgr)
pred_ov = overlay_mask(img_bgr, pred_col, 0.45)
if gt1 is not None:
gt_col = colorize_seg(gt1, lut_bgr)
gt_ov = overlay_mask(img_bgr, gt_col, 0.45)
else:
gt_ov = img_bgr.copy()
tile_w = 520
row = np.concatenate([resize_fit(img_bgr, tile_w), resize_fit(gt_ov, tile_w), resize_fit(pred_ov, tile_w)], axis=1)
lines = [
f"{sample.group_name}/{sample.filename}",
"Modo: single | D/A navega | Q sai",
]
if segm:
lines.append(f"SEG acc={segm['acc']:.3f} mIoU={segm['miou']:.3f}")
lines.append("IMG | GT SEG | PRED SEG")
return put_hud(row, lines)
def make_panel_mask2(sample: Sample, lut_bgr, img_bgr, gt1, pred1, gt2, prob2, bin2, mask2_thr, segm, m2m):
pred_col = colorize_seg(pred1, lut_bgr)
pred_ov = overlay_mask(img_bgr, pred_col, 0.45)
if gt1 is not None:
gt_col = colorize_seg(gt1, lut_bgr)
gt_ov = overlay_mask(img_bgr, gt_col, 0.45)
else:
gt_ov = img_bgr.copy()
prob = np.clip(prob2, 0.0, 1.0).astype(np.float32)
g = (prob * 255.0).astype(np.uint8)
r = ((1.0 - prob) * 255.0).astype(np.uint8)
b = np.zeros_like(g, dtype=np.uint8)
prob_solid = cv2.merge([b, g, r])
edges = cv2.Canny(bin2, 50, 150)
prob_solid[edges > 0] = (255, 255, 255)
bin2_solid = cv2.cvtColor(bin2, cv2.COLOR_GRAY2BGR)
if gt2 is not None:
gt2_vis = gt2.copy().astype(np.uint8)
gt2_vis[gt2_vis > 1] = 255
gt2_bgr = cv2.cvtColor(gt2_vis, cv2.COLOR_GRAY2BGR)
gt2_ov = overlay_mask(img_bgr, gt2_bgr, 0.35)
else:
gt2_ov = img_bgr.copy()
tile_w = 520
tile_h = int(tile_w * img_bgr.shape[0] / img_bgr.shape[1])
row1 = np.concatenate([resize_fit(img_bgr, tile_w, tile_h), resize_fit(gt_ov, tile_w, tile_h), resize_fit(pred_ov, tile_w, tile_h)], axis=1)
row2 = np.concatenate([resize_fit(gt2_ov, tile_w, tile_h), resize_fit(prob_solid, tile_w, tile_h), resize_fit(bin2_solid, tile_w, tile_h)], axis=1)
panel = np.concatenate([row1, row2], axis=0)
lines = [
f"{sample.group_name}/{sample.filename}",
f"Modo: dual_head_mask | thr={mask2_thr:.2f} | W/S threshold | D/A navega | Q sai",
]
if segm:
lines.append(f"SEG acc={segm['acc']:.3f} mIoU={segm['miou']:.3f}")
if m2m:
lines.append(f"MASK2 acc={m2m['acc']:.3f} IoU={m2m['iou']:.3f}")
lines.append("Topo: IMG | GT SEG | PRED SEG")
lines.append("Baixo: GT MASK2 | PRED MASK2 prob | PRED MASK2 bin")
return put_hud(panel, lines)
def make_panel_label(sample: Sample, lut_bgr, img_bgr, gt1, pred1, pred_label_id, pred_conf, label_probs, gt_label_id, gt_label_name, label_names, segm):
pred_col = colorize_seg(pred1, lut_bgr)
pred_ov = overlay_mask(img_bgr, pred_col, 0.45)
if gt1 is not None:
gt_col = colorize_seg(gt1, lut_bgr)
gt_ov = overlay_mask(img_bgr, gt_col, 0.45)
else:
gt_ov = img_bgr.copy()
tile_w = 520
tile_h = int(tile_w * img_bgr.shape[0] / img_bgr.shape[1])
row1 = np.concatenate([resize_fit(img_bgr, tile_w, tile_h), resize_fit(gt_ov, tile_w, tile_h), resize_fit(pred_ov, tile_w, tile_h)], axis=1)
probs_panel = make_label_prob_bar(label_probs, label_names, width=tile_w)
blank = np.zeros_like(probs_panel) + 25
row2 = np.concatenate([probs_panel, blank, blank], axis=1)
panel = np.concatenate([row1, row2], axis=0)
pred_name = label_names.get(pred_label_id, f"label_{pred_label_id}") if pred_label_id is not None else "-"
gt_name = gt_label_name or (label_names.get(gt_label_id, f"label_{gt_label_id}") if gt_label_id is not None else "-")
ok_txt = "OK" if (gt_label_id is not None and pred_label_id == gt_label_id) else "MISS" if gt_label_id is not None else "SEM_GT"
lines = [
f"{sample.group_name}/{sample.filename}",
"Modo: dual_head_label | D/A navega | Q sai",
f"GT={gt_name} ({gt_label_id}) | PRED={pred_name} ({pred_label_id}) conf={pred_conf:.3f} | {ok_txt}",
]
if segm:
lines.append(f"SEG acc={segm['acc']:.3f} mIoU={segm['miou']:.3f}")
lines.append("Topo: IMG | GT SEG | PRED SEG | Baixo: probabilidades por estado")
return put_hud(panel, lines)
def compute_roi_indices(h: int, roi_inicio: float, roi_tamanho: float):
y0 = int(h * roi_inicio)
y1 = int(h * min(1.0, roi_inicio + roi_tamanho))
y0 = max(0, min(h - 1, y0))
y1 = max(y0 + 1, min(h, y1))
return y0, y1
def resize_to_config(img_bgr: np.ndarray, resolucao):
w, h = int(resolucao[0]), int(resolucao[1])
return cv2.resize(img_bgr, (w, h), interpolation=cv2.INTER_AREA)
def draw_label_result(img, label_id, label_conf, label_probs, label_names):
if label_id is None:
return img
name = label_names.get(label_id, f"label_{label_id}")
txt = f"STATUS: {name} | conf={label_conf:.3f}"
cv2.putText(img, txt, (12, 64), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (20, 20, 20), 4, cv2.LINE_AA)
cv2.putText(img, txt, (12, 64), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (245, 245, 245), 2, cv2.LINE_AA)
return img
# ============================================================
# Main
# ============================================================
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", type=str, default="config.json")
ap.add_argument("--ckpt", type=str, default=None)
ap.add_argument("--split", default="test", choices=["test", "val", "train"])
ap.add_argument("--input_size", type=int, default=512)
ap.add_argument("--mask2_thr", type=float, default=0.5)
ap.add_argument("--save_dir", type=str, default=None, help="Se informado, salva painéis em disco além de abrir viewer.")
ap.add_argument("--no_view", action="store_true", help="Não abre janela, apenas salva se --save_dir for informado e imprime métricas.")
ap.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal (ex: norm_stats.json).")
ap.add_argument("--test_folder", type=str, default=None, help="pasta para teste")
ap.add_argument("--camera", action="store_true", help="Inferir em tempo real pela OAK-D Lite")
ap.add_argument("--camera_fps", type=int, default=20)
ap.add_argument("--camera_res", type=str, default="720p", choices=["720p", "1080p"])
args = ap.parse_args()
device = find_device()
print("Device:", device)
with open(args.config, "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
BACKBONE = config["backbone"]
RESOLUCAO = config.get("resolucao", [args.input_size, args.input_size])
ROI_INICIO = float(config.get("roi_inicio", 0.0))
ROI_TAMANHO = float(config.get("roi_tamanho", 1.0))
dataset_path = os.path.join(MODELO, "dataset")
mode = infer_mode_from_config(config)
print(f"[INFO] mode={mode}")
if args.camera:
samples = []
print("[INFO] modo camera OAK-D Lite")
elif args.test_folder:
samples = discover_image_folder(args.test_folder)
print(f"[INFO] samples: {len(samples)} | test_folder={args.test_folder}")
else:
split_root = os.path.join(dataset_path, "split", args.split)
if not os.path.isdir(split_root):
print(f"[WARN] split/{args.split} não existe. Vou usar split/val.")
split_root = os.path.join(dataset_path, "split", "val")
samples = discover_samples(split_root)
print(f"[INFO] samples: {len(samples)} | split_root={split_root}")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
lut_bgr, id_to_name, ignore_id = load_labelmap_lut_bgr(labelmap_path)
num_classes = len(id_to_name)
print("[labelmap]", id_to_name)
save_path = default_save_path(config, mode)
if args.ckpt:
ckpt_path = args.ckpt
else:
# Para label, costuma ser interessante ver best_label; para mask2, best_mask2; fallback best_miou
candidates = []
if mode == "label":
candidates += [os.path.join(save_path, "best_label.pt")]
elif mode == "mask2":
candidates += [os.path.join(save_path, "best_mask2.pt")]
candidates += [os.path.join(save_path, "best_miou.pt"), os.path.join(save_path, "last.pt")]
ckpt_path = next((p for p in candidates if os.path.isfile(p)), candidates[0])
if not os.path.isfile(ckpt_path):
raise SystemExit(f"Checkpoint não encontrado: {ckpt_path}")
print(f"[model] ckpt={ckpt_path}")
base_model = build_base_model(num_classes=num_classes, device=device, backbone=BACKBONE)
aux_head = None
label_names = {}
feat_ch = get_feat_ch(base_model, device, input_h=args.input_size, input_w=args.input_size)
# Carrega ckpt antes para descobrir metadados se necessário
ckpt_pre = torch.load(ckpt_path, map_location="cpu", weights_only=False)
extra = ckpt_pre.get("extra", {}) or {}
if mode == "mask2":
aux_head = Mask2Head(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1)
elif mode == "label":
print(extra)
label_names_raw = extra.get("label_name_by_id", {}) or {}
# JSON pode salvar keys como string
label_names = {int(k): str(v) for k, v in label_names_raw.items()} if label_names_raw else {}
print(label_names_raw)
# Se não tiver no ckpt, tenta inferir do dataset de labels
max_label_id = -1
for s in samples:
lid, lname = read_gt_label(s)
if lid is not None:
max_label_id = max(max_label_id, int(lid))
if lname:
label_names[int(lid)] = str(lname)
if max_label_id < 0:
# último fallback: pelo shape da última camada do checkpoint
sd = ckpt_pre.get("aux_head", {})
for k, v in sd.items():
if k.endswith("net.3.weight") or k.endswith("net.3.bias"):
max_label_id = int(v.shape[0]) - 1
break
num_label_classes = max_label_id + 1
if num_label_classes <= 0:
raise RuntimeError("Não consegui inferir num_label_classes para LabelHead.")
for i in range(num_label_classes):
label_names.setdefault(i, f"label_{i}")
aux_head = LabelHead(feat_ch=feat_ch, num_seg_classes=num_classes, num_label_classes=num_label_classes, hidden=256, dropout=0.2)
print(f"[label_head] classes={num_label_classes} names={label_names}")
ckpt = load_checkpoint(ckpt_path, base_model, aux_head, device)
print(f"[ckpt] epoch={ckpt.get('epoch')} bests={ckpt.get('bests')}")
if args.save_dir:
os.makedirs(args.save_dir, exist_ok=True)
# ==========================
# Normalizador (fixo ou dinâmico)
# ==========================
global NORM_MEAN, NORM_STD
norm_stats = None
# Caminho padrão: dentro do dataset, nome do arquivo de stats
# (ajusta aqui pro nome que você realmente usou: norm_stats.json, por ex.)
norm_stats_path = os.path.join(save_path, "norm_stats.json")
if args.norm_stats is not None:
norm_stats_path = args.norm_stats
if norm_stats_path is not None and os.path.exists(norm_stats_path):
with open(norm_stats_path, "r", encoding="utf-8") as f:
norm_stats = json.load(f)
stats_channels = norm_stats.get("channels", [])
stats_mean = norm_stats.get("mean", [])
stats_std = norm_stats.get("std", [])
print(f"[NORM] usando stats fixos de: {norm_stats_path}")
print(f"[NORM] channels={stats_channels}")
print(f"[NORM] mean={stats_mean}")
print(f"[NORM] std ={stats_std}")
# Garante que temos pelo menos R,G,B
idx_by_name = {name: i for i, name in enumerate(stats_channels)}
required = ["R", "G", "B"]
if not all(ch in idx_by_name for ch in required):
print("[NORM] AVISO: norm_stats não contém todos os canais R,G,B. Mantendo normalize imagenet.")
else:
NORM_MEAN = torch.tensor(stats_mean, dtype=torch.float32, device=device).view(3, 1, 1)
NORM_STD = torch.tensor(stats_std, dtype=torch.float32, device=device).view(3, 1, 1).clamp_min(1e-6)
print("[NORM] Normalização fixa por canal ativada para [R,G,B].")
else:
if norm_stats_path:
print(f"[NORM] Caminho de norm_stats não encontrado: {norm_stats_path}. Usando normalize imagenet.")
else:
print("[NORM] norm_stats não informado. Usando normalize imagenet.")
if args.camera:
pipeline = dai.Pipeline()
cam_rgb = pipeline.create(dai.node.Camera).build()
rgb_out = cam_rgb.requestOutput(
size=(RESOLUCAO[0], RESOLUCAO[1]),
type=dai.ImgFrame.Type.BGR888p,
fps=args.camera_fps
)
rgb_queue = rgb_out.createOutputQueue(maxSize=4, blocking=False)
win = "SegFormer Dual - OAK-D Lite"
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
pipeline.start()
prev_time = time.time()
while pipeline.isRunning():
in_rgb = rgb_queue.get()
frame_bgr = in_rgb.getCvFrame()
H, W = frame_bgr.shape[:2]
y0, y1 = compute_roi_indices(H, ROI_INICIO, ROI_TAMANHO)
roi_bgr = frame_bgr[y0:y1, 0:W]
roi_input = resize_to_config(roi_bgr, RESOLUCAO)
res = infer_sample(
base_model=base_model,
aux_head=aux_head,
img_bgr=roi_input,
device=device,
mode=mode,
input_size=args.input_size,
mask2_thr=args.mask2_thr,
)
pred_seg = res["pred_seg"]
pred_col = colorize_seg(pred_seg, lut_bgr)
pred_col_roi = cv2.resize(
pred_col,
(roi_bgr.shape[1], roi_bgr.shape[0]),
interpolation=cv2.INTER_NEAREST
)
overlay = frame_bgr.copy()
overlay[y0:y1, 0:W] = overlay_mask(
overlay[y0:y1, 0:W],
pred_col_roi,
alpha=0.45
)
now = time.time()
fps = 1.0 / max(1e-6, now - prev_time)
prev_time = now
lines = [
f"FPS: {fps:.1f}",
f"Modo: {mode}",
"Q/ESC sair",
]
if mode == "label":
lid = res["label_id"]
conf = res["label_conf"] or 0.0
lname = label_names.get(lid, f"label_{lid}") if lid is not None else "-"
lines.append(f"Label: {lname} ({lid}) conf={conf:.3f}")
elif mode == "mask2" and res["mask2_bin"] is not None:
mask2_bin = cv2.resize(
res["mask2_bin"],
(roi_bgr.shape[1], roi_bgr.shape[0]),
interpolation=cv2.INTER_NEAREST
)
edges = cv2.Canny(mask2_bin, 50, 150)
overlay[y0:y1, 0:W][edges > 0] = (255, 255, 255)
lines.append(f"Mask2 thr={args.mask2_thr:.2f}")
overlay = put_hud(overlay, lines)
cv2.imshow(win, overlay)
k = cv2.waitKey(1) & 0xFF
if k in (ord("q"), ord("Q"), 27):
break
if mode == "mask2" and k in (ord("w"), ord("W")):
args.mask2_thr = min(0.95, args.mask2_thr + 0.05)
if mode == "mask2" and k in (ord("s"), ord("S")):
args.mask2_thr = max(0.05, args.mask2_thr - 0.05)
cv2.destroyAllWindows()
return
idx = 0
win = "SegFormer Test Viewer"
if not args.no_view:
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
global_seg_ious = []
global_seg_accs = []
global_mask2_iou = []
global_mask2_acc = []
label_total = 0
label_correct = 0
while True:
s = samples[idx]
img = imread_bgr(s.img_path)
gt1 = imread_gray(s.mask_path) if s.mask_path else None
gt2 = imread_gray(s.mask2_path) if s.mask2_path else None
gt_label_id, gt_label_name = read_gt_label(s)
res = infer_sample(
base_model=base_model,
aux_head=aux_head,
img_bgr=img,
device=device,
mode=mode,
input_size=args.input_size,
mask2_thr=args.mask2_thr,
)
pred1 = res["pred_seg"]
segm = seg_metrics(pred1, gt1, num_classes=num_classes, ignore_id=ignore_id)
if segm:
global_seg_ious.append(segm["miou"])
global_seg_accs.append(segm["acc"])
if mode == "mask2":
m2m = mask2_metrics(res["mask2_bin"], gt2)
if m2m:
global_mask2_iou.append(m2m["iou"])
global_mask2_acc.append(m2m["acc"])
panel = make_panel_mask2(
sample=s,
lut_bgr=lut_bgr,
img_bgr=img,
gt1=gt1,
pred1=pred1,
gt2=gt2,
prob2=res["mask2_prob"],
bin2=res["mask2_bin"],
mask2_thr=args.mask2_thr,
segm=segm,
m2m=m2m,
)
elif mode == "label":
pred_label_id = res["label_id"]
pred_conf = res["label_conf"] if res["label_conf"] is not None else 0.0
if gt_label_id is not None:
label_total += 1
if pred_label_id == gt_label_id:
label_correct += 1
panel = make_panel_label(
sample=s,
lut_bgr=lut_bgr,
img_bgr=img,
gt1=gt1,
pred1=pred1,
pred_label_id=pred_label_id,
pred_conf=pred_conf,
label_probs=res["label_probs"],
gt_label_id=gt_label_id,
gt_label_name=gt_label_name,
label_names=label_names,
segm=segm,
)
else:
panel = make_panel_single(s, lut_bgr, img, gt1, pred1, segm, id_to_name)
if args.save_dir:
out_name = f"{idx:05d}_{s.group_name.replace('/', '_')}_{os.path.splitext(s.filename)[0]}.jpg"
cv2.imwrite(os.path.join(args.save_dir, out_name), panel)
if args.no_view:
idx += 1
if idx >= len(samples):
break
continue
cv2.imshow(win, panel)
k = cv2.waitKey(0) & 0xFF
if k in (ord("q"), ord("Q"), 27):
break
if k in (ord("a"), ord("A"), 81):
idx = (idx - 1) % len(samples)
continue
if k in (ord("d"), ord("D"), 83):
idx = (idx + 1) % len(samples)
continue
if mode == "mask2" and k in (ord("w"), ord("W")):
args.mask2_thr = min(0.95, args.mask2_thr + 0.05)
continue
if mode == "mask2" and k in (ord("s"), ord("S")):
args.mask2_thr = max(0.05, args.mask2_thr - 0.05)
continue
if not args.no_view:
cv2.destroyAllWindows()
print("\nResumo parcial/geral da sessão:")
if global_seg_ious:
print(f"SEG : mean_acc={np.mean(global_seg_accs):.4f} mean_mIoU={np.mean(global_seg_ious):.4f} n={len(global_seg_ious)}")
if mode == "mask2" and global_mask2_iou:
print(f"MASK2: mean_acc={np.mean(global_mask2_acc):.4f} mean_IoU={np.mean(global_mask2_iou):.4f} n={len(global_mask2_iou)}")
if mode == "label" and label_total > 0:
print(f"LABEL: acc={label_correct / max(1, label_total):.4f} correct={label_correct}/{label_total}")
if __name__ == "__main__":
main()

View File

@ -2,7 +2,8 @@
"camera": "oak-d",
"modelo": "segformer_b0",
"model_name": "nav_mit",
"dual_head": false,
"dual_head_label": true,
"dual_head_mask": false,
"main_class_name": "navegavel",
"es_classes": "",
"model_to_use": "geral",
@ -13,5 +14,15 @@
"shaves": 3,
"channels": 3,
"use_ndvi": false,
"backbone": "nvidia/mit-b0"
"backbone": "nvidia/mit-b0",
"label_classes": [
"Parado",
"EntrandoRua",
"CaminhandoRua",
"SaindoRua",
"Manobrando",
"Direcionando",
"RetornandoBase",
"Indefinido"
]
}

View File

@ -0,0 +1,12 @@
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
modelo = "nvidia/mit-b0"
pasta = r"C:\AgroBaseModels\Backbones\nvidia_mit_b0"
model = SegformerForSemanticSegmentation.from_pretrained(modelo)
processor = SegformerImageProcessor.from_pretrained(modelo)
model.save_pretrained(pasta)
processor.save_pretrained(pasta)
print("Backbone salvo em:", pasta)

File diff suppressed because it is too large Load Diff

View File

@ -216,7 +216,23 @@ class MultiSpectralClient:
def get_next_decoded(self, timeout=2.0, update_radiometry=True):
frame, meta = self.get_next_frame(timeout=timeout)
decoded = self.core.decode_stream_cameras(frame, meta)
frame_type = meta.get("frame_type")
if frame_type == "RAW_BRUTO":
decoded = self.core.decode_stream_cameras(frame, meta)
elif frame_type in ("RGB", "MULTISPEC"):
decoded = {
"cam2": {
"name": "RGB",
"image": frame,
"meta": meta
}
}
else:
raise RuntimeError(f"frame_type não suportado: {frame_type}")
if update_radiometry:
self.update_radiometry(decoded, meta)

View File

@ -230,6 +230,7 @@ def main():
try:
with MultiSpectralClient(
#mx_id="194430108133AC2F00",
width=raw_w,
height=raw_h,
bayer=args.bayer,

View File

@ -82,6 +82,43 @@ DEFAULT_DATASET_BASE = "dataset"
CHANNEL_NAMES = ["R", "G", "B", "RE", "NIR"]
_CORE_CACHE: Dict[tuple, RawProcessorCore] = {}
def get_or_create_core(
raw_size: Tuple[int, int],
bayer_pattern: str,
module_params_path: Optional[str],
) -> RawProcessorCore:
raw_w, raw_h = raw_size
bayer = str(bayer_pattern or "BGGR").upper()
key = (
int(raw_w),
int(raw_h),
bayer,
str(Path(module_params_path).resolve()) if module_params_path else None,
)
core = _CORE_CACHE.get(key)
if core is not None:
return core
print(
"[CORE_CACHE] criando RawProcessorCore "
f"raw={raw_w}x{raw_h} bayer={bayer} module={module_params_path}"
)
core = RawProcessorCore(
sensor_width=int(raw_w),
sensor_height=int(raw_h),
bayer_pattern=bayer,
calibration_json_path=module_params_path,
)
_CORE_CACHE[key] = core
return core
@dataclass
class SampleBundle:
group: str
@ -198,6 +235,13 @@ def collect_samples_from_group(group_dir: Path) -> List[SampleBundle]:
return samples
def copy_json_safe(obj):
try:
return json.loads(json.dumps(obj, default=str))
except Exception:
return None
# ============================================================
# Module params / tensor
# ============================================================
@ -304,14 +348,12 @@ def build_tensor_from_sample(
raw_size: Tuple[int, int],
module_params_path: Optional[str],
) -> Tuple[np.ndarray, dict]:
raw_w, raw_h = raw_size
bayer = str(meta.get("bayer_pattern") or meta.get("bayer") or "RGGB").upper()
bayer = str(meta.get("bayer_pattern") or meta.get("bayer") or "BGGR").upper()
core = RawProcessorCore(
sensor_width=int(raw_w),
sensor_height=int(raw_h),
core = get_or_create_core(
raw_size=raw_size,
bayer_pattern=bayer,
calibration_json_path=module_params_path,
module_params_path=module_params_path,
)
frame = load_frame_from_saved_bins(sample, meta)
@ -329,14 +371,15 @@ def build_tensor_from_sample(
info = {
"module_params": module_params_path,
"bayer_pattern": bayer,
"fusion_result": getattr(core, "last_fusion_result", None),
"patch_normalization_result": getattr(core, "last_patch_normalization_result", None),
"frame_quality": getattr(core, "last_frame_quality_result", None),
"fusion_result": copy_json_safe(getattr(core, "last_fusion_result", None)),
"patch_normalization_result": copy_json_safe(getattr(core, "last_patch_normalization_result", None)),
"frame_quality": copy_json_safe(getattr(core, "last_frame_quality_result", None)),
}
return tensor, info
# ============================================================
# Máscara / preview
# ============================================================
@ -407,6 +450,82 @@ def save_bgr(path: Path, bgr: np.ndarray):
raise RuntimeError(f"Falha ao salvar preview: {path}")
def get_class_ids_from_config() -> dict:
heads = config.get("heads", {}) or {}
semantic = heads.get("semantic", {}) or {}
classes = semantic.get("classes", {}) or {}
return {
"chao": int(classes.get("chao", 0)),
"cana": int(classes.get("cana", 1)),
"erva": int(classes.get("erva", 2)),
}
def build_head_masks_from_semantic(
mask_ids: np.ndarray,
class_ids: dict,
ignore_id: int = 255,
) -> dict:
"""
Gera máscaras auxiliares para multi-head a partir da máscara semântica alinhada.
Entrada:
mask_ids: HW com IDs semânticos.
Saída:
{
"semantic": HW uint8,
"vegetation": HW uint8 com 0/1/ignore,
"cana": HW uint8 com 0/1/ignore
}
"""
chao_id = int(class_ids.get("chao", 0))
cana_id = int(class_ids.get("cana", 1))
erva_id = int(class_ids.get("erva", 2))
semantic = mask_ids.astype(np.uint8, copy=True)
ignore_mask = semantic == ignore_id
vegetation = np.zeros_like(semantic, dtype=np.uint8)
vegetation[(semantic == cana_id) | (semantic == erva_id)] = 1
vegetation[ignore_mask] = ignore_id
cana = np.zeros_like(semantic, dtype=np.uint8)
cana[semantic == cana_id] = 1
cana[ignore_mask] = ignore_id
return {
"semantic": semantic,
"vegetation": vegetation,
"cana": cana,
}
def save_mask_npy_and_debug(mask: np.ndarray, npy_path: Path, png_path: Path):
ensure_dir(npy_path.parent)
ensure_dir(png_path.parent)
mask = mask.astype(np.uint8, copy=False)
np.save(str(npy_path), mask)
# Debug visual:
# 0 -> preto
# 1 -> cinza claro
# 2 -> mais claro, quando existir na semântica
# 255 -> branco
debug = mask.copy()
debug_vis = np.zeros_like(debug, dtype=np.uint8)
debug_vis[debug == 0] = 0
debug_vis[debug == 1] = 120
debug_vis[debug == 2] = 220
debug_vis[debug == 255] = 255
cv2.imwrite(str(png_path), debug_vis)
# ============================================================
# Stats
# ============================================================
@ -467,11 +586,22 @@ def process_group(
out_group = output_root / group_name
out_tensors = out_group / "tensors"
out_masks = out_group / "masks"
out_masks_vegetation = out_group / "masks_vegetation"
out_masks_cana = out_group / "masks_cana"
out_metas = out_group / "metas"
out_previews = out_group / "previews"
for d in (out_tensors, out_masks, out_metas, out_previews):
for d in (
out_tensors,
out_masks,
out_masks_vegetation,
out_masks_cana,
out_metas,
out_previews,
):
ensure_dir(d)
rows = []
@ -515,16 +645,43 @@ def process_group(
# Salva tensor e mask
tensor_path = out_tensors / f"{sample.base}.npy"
mask_path = out_masks / f"{sample.base}.npy"
mask_debug_path = out_masks / f"{sample.base}.png"
preview_path = out_previews / f"{sample.base}.png"
meta_out_path = out_metas / f"{sample.base}.json"
np.save(str(tensor_path), tensor)
np.save(str(mask_path), mask_ids)
# Debug visual da máscara em IDs, só para inspeção rápida.
cv2.imwrite(str(mask_debug_path), mask_ids)
class_ids = get_class_ids_from_config()
head_masks = build_head_masks_from_semantic(
mask_ids=mask_ids,
class_ids=class_ids,
ignore_id=ignore_id,
)
mask_path = out_masks / f"{sample.base}.npy"
mask_debug_path = out_masks / f"{sample.base}.png"
mask_vegetation_path = out_masks_vegetation / f"{sample.base}.npy"
mask_vegetation_debug_path = out_masks_vegetation / f"{sample.base}.png"
mask_cana_path = out_masks_cana / f"{sample.base}.npy"
mask_cana_debug_path = out_masks_cana / f"{sample.base}.png"
save_mask_npy_and_debug(
head_masks["semantic"],
mask_path,
mask_debug_path,
)
save_mask_npy_and_debug(
head_masks["vegetation"],
mask_vegetation_path,
mask_vegetation_debug_path,
)
save_mask_npy_and_debug(
head_masks["cana"],
mask_cana_path,
mask_cana_debug_path,
)
preview_bgr = tensor_to_preview_bgr(tensor)
save_bgr(preview_path, preview_bgr)
@ -542,6 +699,38 @@ def process_group(
"saved_payload_type": "tensor_npy",
"saved_tensor_path": safe_rel(tensor_path, dataset_root),
"saved_mask_path": safe_rel(mask_path, dataset_root),
"head_masks": {
"semantic": {
"path": safe_rel(mask_path, dataset_root),
"shape": list(head_masks["semantic"].shape),
"classes": {
"chao": int(class_ids["chao"]),
"cana": int(class_ids["cana"]),
"erva": int(class_ids["erva"]),
},
"ignore_index": int(ignore_id),
},
"vegetation": {
"path": safe_rel(mask_vegetation_path, dataset_root),
"shape": list(head_masks["vegetation"].shape),
"classes": {
"background": 0,
"vegetation": 1,
},
"positive_from_semantic": ["cana", "erva"],
"ignore_index": int(ignore_id),
},
"cana": {
"path": safe_rel(mask_cana_path, dataset_root),
"shape": list(head_masks["cana"].shape),
"classes": {
"not_cana": 0,
"cana": 1,
},
"positive_from_semantic": ["cana"],
"ignore_index": int(ignore_id),
},
},
"saved_preview_path": safe_rel(preview_path, dataset_root),
"saved_payload_dtype": str(tensor.dtype),
"saved_payload_shape": list(tensor.shape),
@ -575,6 +764,8 @@ def process_group(
"base": sample.base,
"tensor": str(tensor_path),
"mask": str(mask_path),
"mask_vegetation": str(mask_vegetation_path),
"mask_cana": str(mask_cana_path),
"meta": str(meta_out_path),
"preview": str(preview_path),
"quality_status": frame_quality.get("status") if isinstance(frame_quality, dict) else None,
@ -593,7 +784,17 @@ def process_group(
def write_manifest(path: Path, rows: List[dict]):
ensure_dir(path.parent)
fieldnames = ["group", "base", "tensor", "mask", "meta", "preview", "quality_status"]
fieldnames = [
"group",
"base",
"tensor",
"mask",
"mask_vegetation",
"mask_cana",
"meta",
"preview",
"quality_status",
]
with path.open("w", newline="", encoding="utf-8") as f:
w = csv.DictWriter(f, fieldnames=fieldnames)

View File

@ -18,10 +18,15 @@ RESOLUCAO = tuple(config.get("resolucao"))
TENSOR_EXT = ".npy"
MASK_NPY_SUFFIX = ".npy"
AUX_MASK_DIRS = [
"masks_vegetation",
"masks_cana",
]
RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
RE_AUGMENTED_FAMILY = re.compile(r"^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$", re.IGNORECASE)
RE_AUG_SUFFIX = re.compile(r"_aug[a-zA-Z0-9]*_\d+$", re.IGNORECASE)
MULTI_HEAD = bool(config.get("multi_head", False))
def garantir(p):
@ -237,6 +242,46 @@ def copiar_optional(src_dir, dst_dir, base, ext):
return dst
def copiar_mask_dir_optional(src_group_dir, dst_group_dir, mask_dir_name, base):
"""
Copia uma pasta auxiliar de máscara, como:
masks_vegetation/
masks_cana/
Copia:
<base>.npy obrigatório se existir
<base>.png opcional se existir
Retorna caminhos de destino ou None.
"""
src_dir = os.path.join(src_group_dir, mask_dir_name)
dst_dir = os.path.join(dst_group_dir, mask_dir_name)
npy_src = os.path.join(src_dir, base + ".npy")
png_src = os.path.join(src_dir, base + ".png")
result = {
"npy": None,
"png": None,
}
if not os.path.exists(npy_src):
return result
garantir(dst_dir)
npy_dst = os.path.join(dst_dir, base + ".npy")
shutil.copy2(npy_src, npy_dst)
result["npy"] = npy_dst
if os.path.exists(png_src):
png_dst = os.path.join(dst_dir, base + ".png")
shutil.copy2(png_src, png_dst)
result["png"] = png_dst
return result
def copiar(
nomes,
src_group_dir,
@ -304,11 +349,37 @@ def copiar(
shutil.copy2(cand, preview_dst)
break
aux_masks = {}
for aux_dir in AUX_MASK_DIRS:
aux_masks[aux_dir] = copiar_mask_dir_optional(
src_group_dir=src_group_dir,
dst_group_dir=dst_group_dir,
mask_dir_name=aux_dir,
base=base,
)
if MULTI_HEAD:
for aux_dir in AUX_MASK_DIRS:
aux_npy = aux_masks.get(aux_dir, {}).get("npy")
if aux_npy is None:
raise RuntimeError(
f"multi_head=true, mas máscara auxiliar ausente: "
f"{aux_dir}/{base}.npy em {src_group_dir}"
)
rows.append({
"base": base,
"tensor": tensor_dst,
"mask_npy": mask_npy_dst,
"mask_png": mask_png_dst,
"mask_vegetation_npy": aux_masks.get("masks_vegetation", {}).get("npy"),
"mask_vegetation_png": aux_masks.get("masks_vegetation", {}).get("png"),
"mask_cana_npy": aux_masks.get("masks_cana", {}).get("npy"),
"mask_cana_png": aux_masks.get("masks_cana", {}).get("png"),
"meta": meta_dst,
"preview": preview_dst,
})
@ -447,8 +518,16 @@ def write_manifest(path, rows):
"group",
"base",
"tensor",
"mask_npy",
"mask_png",
"mask_vegetation_npy",
"mask_vegetation_png",
"mask_cana_npy",
"mask_cana_png",
"meta",
"preview",
]

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@ -9,7 +9,8 @@
"sensor_height": 800,
"bayer_pattern": "BGGR",
"rgb_processing": {
"mode": "bayer_planes"
"mode": "linear_demosaic",
"demosaic_algorithm": "bilinear"
},
"camera_settings": {
"rgb": {
@ -38,6 +39,9 @@
}
},
"fusion_config": {
"use_remap_cache": true,
"use_remap_for_rgb": false,
"use_remap_for_spec": true,
"alignment_mode": "homography",
"baseline_mm": 75.0,
"reference_camera": "rgb",
@ -49,10 +53,11 @@
},
"nir": {
"dx": 0,
"dy": 1,
"dy": 0,
"theta_deg": 0.0
}
},
"homography_calibration_size": [1280, 800],
"homographies": {
"re_to_rgb": [
[
@ -91,7 +96,7 @@
},
"crop_valid_common": true,
"resize_after_crop": true,
"target_size": null
"target_size": [1024,640]
},
"radiometric_config": {
"enabled": false,
@ -583,35 +588,35 @@
"reference_mode": "fixed",
"reference_controls": {
"rgb": {
"exposure_time_us": 10000,
"sensitivity_iso": 400
"exposure_time_us": 2200,
"sensitivity_iso": 100
},
"re": {
"exposure_time_us": 15000,
"sensitivity_iso": 400
"exposure_time_us": 2500,
"sensitivity_iso": 100
},
"nir": {
"exposure_time_us": 15000,
"sensitivity_iso": 400
"exposure_time_us": 2500,
"sensitivity_iso": 100
}
},
"scale_limits": {
"default": {
"min": 0.15,
"max": 6.0
"max": 3.0
},
"rgb": {
"min": 0.15,
"max": 6.0
"max": 3.0
},
"re": {
"min": 0.15,
"max": 8.0
"max": 2.5
},
"nir": {
"min": 0.15,
"max": 8.0
"max": 2.5
}
},
@ -676,21 +681,21 @@
"rgb_calibration": {
"enabled": true,
"gains": {
"R": 1.2500000000000002,
"R": 1.061,
"G": 1.0,
"B": 1.5500000000000005
"B": 1.452
}
},
"flatfield_config": {
"enabled": true,
"subtract_dark": true,
"subtract_dark": false,
"schema": "multispec_flatfield_v1",
"created_at": "2026-05-08 15:26:18",
"json_file": "calibration/flatfield_maps_v1.json",
"npz_file": "calibration/flatfield_maps_v1.npz",
"apply_before_fusion": true,
"apply_after_decode": true,
"apply_space": "native_camera_space",
"apply_before_fusion": false,
"apply_after_decode": false,
"apply_space": "final_tensor_space",
"map_type": "gain",
"formula": "channel_corrected = max(channel_linear - dark, 0) * gain_map",
"channels": [
@ -772,28 +777,30 @@
"gain_std": 0.23178784549236298
}
},
"exp_gain_correct_during_flat_capture": false,
"exp_gain_correct_during_flat_capture": true,
"smooth_ksize": 31,
"min_gain": 0.25,
"max_gain": 4.0,
"notes": "",
"strength": 0.35,
"strength": 0.25,
"strength_by_channel": {
"R": 0.9,
"G": 0.9,
"B": 0.9,
"R": 0.5,
"G": 0.5,
"B": 0.5,
"RE": 0.25,
"NIR": 0.25
},
"gain_min_runtime": 0.75,
"gain_max_runtime": 1.35,
"runtime_smooth_ksize": 81,
"saturation_guard_enabled": true,
"fast_runtime": true,
"saturation_guard_enabled": false,
"saturation_guard_mode": "fade_strength",
"saturation_guard_threshold": 0.97,
"saturation_guard_soft_start": 0.88,
"saturation_guard_hard": 0.97
"saturation_guard_hard": 0.97,
"clip_output": true
}
}

View File

@ -1,8 +1,7 @@
{
"camera": "oak-fcc-3",
"modelo": "segformer_b1",
"model_name": "test_1",
"dual_head": false,
"model_name": "target_teached",
"main_class_name": "cana",
"es_classes": "",
"model_to_use": "geral",
@ -12,9 +11,69 @@
"roi_tamanho": 1.0,
"shaves": 3,
"channels": 5,
"input_channels": ["R", "G", "B", "RE", "NIR"],
"use_ndvi": false,
"backbone": "nvidia/mit-b1",
"fusion_mode": "stacked",
"stats_source_tag": "stacked_raw5",
"module_params_json": "calibration/module_params.json"
"stats_source_tag": "stacked_raw4",
"module_params_json": "calibration/module_params.json",
"multi_head": true,
"heads": {
"semantic": {
"enabled": true,
"type": "multiclass",
"num_classes": 3,
"mask_dir": "masks",
"classes": {"chao": 0, "cana": 1, "erva": 2},
"ignore_index": 255,
"loss_weight": 0.10
},
"vegetation": {
"enabled": true,
"type": "binary",
"num_classes": 2,
"mask_dir": "masks_vegetation",
"classes": {"background": 0, "vegetation": 1},
"ignore_index": 255,
"loss_weight": 0.20
},
"cana": {
"enabled": true,
"type": "binary",
"num_classes": 2,
"mask_dir": "masks_cana",
"classes": {"not_cana": 0, "cana": 1},
"ignore_index": 255,
"loss_weight": 0.25
},
"target": {
"enabled": true,
"type": "binary",
"num_classes": 2,
"mask_dir": "__derived_target__",
"classes": {"background": 0, "target": 1},
"ignore_index": 255,
"loss_weight": 0.45,
"derived": true
}
},
"target_distillation": {
"enabled": true,
"hard_weight": 0.85,
"distill_weight": 0.15,
"rampup_enabled": true,
"start_epoch": 8,
"rampup_epochs": 12,
"w_sem_erva": 0.45,
"w_veg_not_cana": 0.35,
"w_veg_suppressed": 0.20,
"cana_suppression_power": 1.5,
"teacher_min": 0.0,
"teacher_max": 1.0,
"detach_teacher": true
}
}

View File

@ -0,0 +1,620 @@
import argparse
import json
import sys
import time
from pathlib import Path
import cv2
import numpy as np
try:
import torch
except Exception:
torch = None
# ============================================================
# Ajuste de import local
# ============================================================
THIS_FILE = Path(__file__).resolve()
# Esperado:
# .../Python/Scripts/workers/camera_worker/oak_fcc3_core/benchmark_raw_bruto_scientific.py
WORKERS_DIR = THIS_FILE.parents[2]
if str(WORKERS_DIR) not in sys.path:
sys.path.insert(0, str(WORKERS_DIR))
from camera_worker.oak_fcc3_core.oak_fcc3_client import OakFcc3Client
try:
from camera_worker.oak_fcc3_core.segformer_service import MultiSpecSegformerService
except Exception:
MultiSpecSegformerService = None
# ============================================================
# Utils
# ============================================================
def now_ms():
return time.perf_counter() * 1000.0
def mean(xs):
return float(np.mean(xs)) if xs else 0.0
def p95(xs):
return float(np.percentile(xs, 95)) if xs else 0.0
def maxv(xs):
return float(np.max(xs)) if xs else 0.0
def last_mean(xs, n=30):
return float(np.mean(xs[-n:])) if xs else 0.0
def last_max(xs, n=30):
return float(np.max(xs[-n:])) if xs else 0.0
def load_json_if_exists(path):
if not path:
return None
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def parse_float_list(s, expected=5, default=None):
if s is None:
return default
vals = [float(x.strip()) for x in str(s).split(",") if x.strip() != ""]
if len(vals) != expected:
raise ValueError(f"Esperado {expected} valores, veio {len(vals)}: {s}")
return np.array(vals, dtype=np.float32)
def extract_mean_std_from_model_config(cfg):
if not isinstance(cfg, dict):
return None, None
mean_cfg = cfg.get("mean") or cfg.get("channel_mean") or cfg.get("norm_mean")
std_cfg = cfg.get("std") or cfg.get("channel_std") or cfg.get("norm_std")
norm = cfg.get("normalization") or cfg.get("norm") or {}
if mean_cfg is None and isinstance(norm, dict):
mean_cfg = norm.get("mean") or norm.get("channel_mean")
if std_cfg is None and isinstance(norm, dict):
std_cfg = norm.get("std") or norm.get("channel_std")
if mean_cfg is None or std_cfg is None:
return None, None
mean_arr = np.array(mean_cfg, dtype=np.float32)
std_arr = np.array(std_cfg, dtype=np.float32)
if mean_arr.size != 5 or std_arr.size != 5:
return None, None
return mean_arr, std_arr
def tensor_stats(tensor):
names = ["R", "G", "B", "RE", "NIR"]
out = {}
for i, name in enumerate(names):
ch = tensor[i].astype(np.float32)
out[name] = {
"min": float(np.min(ch)),
"p01": float(np.percentile(ch, 1)),
"p50": float(np.percentile(ch, 50)),
"p99": float(np.percentile(ch, 99)),
"max": float(np.max(ch)),
"mean": float(np.mean(ch)),
"std": float(np.std(ch)),
}
return out
def print_tensor_stats(label, tensor):
print("============================================")
print(f"[{label}] TENSOR")
print(f"shape={tensor.shape} dtype={tensor.dtype}")
stats = tensor_stats(tensor)
for ch, s in stats.items():
print(
f"{ch:>3} | "
f"min={s['min']:.4f} "
f"p01={s['p01']:.4f} "
f"p50={s['p50']:.4f} "
f"p99={s['p99']:.4f} "
f"max={s['max']:.4f} "
f"mean={s['mean']:.4f} "
f"std={s['std']:.4f}"
)
print("============================================")
def to_u8_01(arr):
arr = np.asarray(arr, dtype=np.float32)
arr = np.nan_to_num(arr, nan=0.0, posinf=1.0, neginf=0.0)
arr = np.clip(arr, 0.0, 1.0)
return (arr * 255.0).astype(np.uint8)
def make_panel(tensor):
rgb = np.stack([tensor[0], tensor[1], tensor[2]], axis=2)
rgb_bgr = cv2.cvtColor(to_u8_01(rgb), cv2.COLOR_RGB2BGR)
re_bgr = cv2.cvtColor(to_u8_01(tensor[3]), cv2.COLOR_GRAY2BGR)
nir_bgr = cv2.cvtColor(to_u8_01(tensor[4]), cv2.COLOR_GRAY2BGR)
false_rgb = np.stack(
[
tensor[4], # visual R = NIR
tensor[3], # visual G = RE
tensor[0], # visual B = R real
],
axis=2,
)
false_bgr = cv2.cvtColor(to_u8_01(false_rgb), cv2.COLOR_RGB2BGR)
def title(img, text):
h, w = img.shape[:2]
bar_h = 34
bar = np.zeros((bar_h, w, 3), dtype=np.uint8)
cv2.putText(
bar,
text,
(10, 24),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
(255, 255, 255),
2,
cv2.LINE_AA,
)
return np.vstack([bar, img])
rgb_bgr = title(rgb_bgr, "RGB")
re_bgr = title(re_bgr, "RE")
nir_bgr = title(nir_bgr, "NIR")
false_bgr = title(false_bgr, "Falso color NIR/RE/R")
top = np.hstack([rgb_bgr, re_bgr])
bottom = np.hstack([nir_bgr, false_bgr])
return np.vstack([top, bottom])
# ============================================================
# Benchmark processor
# ============================================================
class ScientificBenchmark:
def __init__(self, args):
self.args = args
self.target_size = (int(args.width), int(args.height))
self.model_cfg = load_json_if_exists(args.model_config_json)
mean_cfg, std_cfg = extract_mean_std_from_model_config(self.model_cfg)
if args.model_mean is not None:
self.mean = parse_float_list(args.model_mean, expected=5)
elif mean_cfg is not None:
self.mean = mean_cfg
else:
self.mean = np.array([0.5, 0.5, 0.5, 0.5, 0.5], dtype=np.float32)
if args.model_std is not None:
self.std = parse_float_list(args.model_std, expected=5)
elif std_cfg is not None:
self.std = std_cfg
else:
self.std = np.array([0.25, 0.25, 0.25, 0.25, 0.25], dtype=np.float32)
self.mean_chw = self.mean[:, None, None].astype(np.float32)
self.std_chw = self.std[:, None, None].astype(np.float32)
self.device = None
if torch is not None and torch.cuda.is_available():
self.device = torch.device("cuda")
elif torch is not None:
self.device = torch.device("cpu")
self.model_svc = None
if args.run_model:
if MultiSpecSegformerService is None:
raise RuntimeError("MultiSpecSegformerService não pôde ser importado.")
if not isinstance(self.model_cfg, dict):
raise RuntimeError("--run_model requer --model_config_json válido.")
self.model_svc = MultiSpecSegformerService(
model_config=self.model_cfg,
mostrar_log=print,
)
dummy = np.zeros((5, args.height, args.width), dtype=np.float32)
for _ in range(max(0, int(args.warmup_model))):
self.model_svc.infer_tensor_fast(dummy, keep_probs=False)
print(f"[BENCH] Warmup modelo concluído: {args.warmup_model}x")
def process_once(self, client):
"""
Mede uma iteração completa do fluxo científico.
"""
times = {
"capture_ms": 0.0,
"decode_ms": 0.0,
"controller_ms": 0.0,
"fuse_total_ms": 0.0,
"fuse_dark_ms": 0.0,
"fuse_radnorm_ms": 0.0,
"fuse_flat_ms": 0.0,
"fuse_prepare_ms": 0.0,
"fuse_warp_ms": 0.0,
"fuse_crop_resize_ms": 0.0,
"fuse_concat_ms": 0.0,
"model_norm_ms": 0.0,
"torch_ms": 0.0,
"infer_ms": 0.0,
"total_ms": 0.0,
"sync_dt_ms": 0.0,
}
t_total0 = now_ms()
# ========================================================
# 1. Captura RAW_BRUTO
# ========================================================
t0 = now_ms()
raw_frame, raw_meta = client.get_next_raw_frame(timeout=self.args.timeout)
times["capture_ms"] = now_ms() - t0
#cp = raw_meta.get("capture_perf", {})
#print(
# "[CAP_ASYNC] "
# f"get_wait={cp.get('async_get_wait_ms',0):.2f}ms "
# f"age={cp.get('async_packet_age_ms',0):.2f}ms "
# f"seq={cp.get('async_packet_seq')} "
# f"thread_wait={cp.get('wait_total_ms',0):.2f}ms "
# f"sleep={cp.get('sleep_ms',0):.2f}ms/{cp.get('sleep_count',0)} "
# f"drain={cp.get('drain_total_ms',0):.2f}ms "
# f"copy={cp.get('drain_frombuffer_copy_ms',0):.2f}ms "
# f"status={cp.get('async_status',{})}"
#)
times["sync_dt_ms"] = float(raw_meta.get("sync_dt_ms", 0.0) or 0.0)
# ========================================================
# 2. Decode RAW10 packed -> float científico por câmera
# ========================================================
t0 = now_ms()
decoded = client.decode_stream_cameras(raw_frame, raw_meta)
times["decode_ms"] = now_ms() - t0
# ========================================================
# 3. RadiometricController update, se ativo
# Isso NÃO é a radiometric_normalization do tensor.
# É o controlador de exposição/ganho.
# ========================================================
t0 = now_ms()
client.update_radiometry(decoded, raw_meta)
times["controller_ms"] = now_ms() - t0
# ========================================================
# 4. Fusão científica no RawProcessorCore
# dark + radnorm + flat + homografia + crop/resize + concat
# ========================================================
t0 = now_ms()
tensor = client.build_infer_tensor_from_decoded(
decoded=decoded,
meta=raw_meta,
channels_expected=5,
target_size=self.target_size,
)
times["fuse_total_ms"] = now_ms() - t0
# Pega detalhamento interno do core
try:
perf = (client.core.last_fusion_result or {}).get("perf", {}) or {}
times["fuse_dark_ms"] = float(perf.get("dark_ms", 0.0) or 0.0)
times["fuse_radnorm_ms"] = float(perf.get("radnorm_ms", 0.0) or 0.0)
times["fuse_flat_ms"] = float(perf.get("flat_ms", 0.0) or 0.0)
times["fuse_prepare_ms"] = float(perf.get("prepare_ms", 0.0) or 0.0)
times["fuse_warp_ms"] = float(perf.get("warp_total_ms", 0.0) or 0.0)
times["fuse_crop_resize_ms"] = float(perf.get("crop_resize_ms", 0.0) or 0.0)
times["fuse_concat_ms"] = float(perf.get("concat_ms", 0.0) or 0.0)
except Exception:
pass
# ========================================================
# 5. Normalização do modelo, opcional
# ========================================================
if self.args.simulate_model_norm:
t0 = now_ms()
tensor = (tensor - self.mean_chw) / np.maximum(self.std_chw, 1e-6)
tensor = np.ascontiguousarray(tensor, dtype=np.float32)
times["model_norm_ms"] = now_ms() - t0
# ========================================================
# 6. Transferência para torch/cuda, opcional
# ========================================================
if self.args.to_torch:
if torch is None:
raise RuntimeError("--to_torch requer torch instalado.")
t0 = now_ms()
x = torch.from_numpy(tensor).unsqueeze(0).to(self.device, non_blocking=True)
if self.device is not None and self.device.type == "cuda":
torch.cuda.synchronize()
times["torch_ms"] = now_ms() - t0
# ========================================================
# 7. Inferência real, opcional
# ========================================================
pred = None
if self.args.run_model:
t0 = now_ms()
pred = self.model_svc.infer_tensor_fast(tensor, keep_probs=False)
times["infer_ms"] = now_ms() - t0
times["total_ms"] = now_ms() - t_total0
return tensor, pred, raw_frame, raw_meta, decoded, times
# ============================================================
# Main
# ============================================================
def main():
ap = argparse.ArgumentParser(
description="Benchmark científico OAK-FCC-3 RAW_BRUTO -> tensor multiespectral final."
)
ap.add_argument(
"--module_params",
default=r"C:\ZendionInc\agrobot_base\Python\OAK\datasets\oak-fcc-3\calibration\module_params.json",
help="Caminho do module_params.json.",
)
ap.add_argument("--width", type=int, default=1024, help="Largura final do tensor.")
ap.add_argument("--height", type=int, default=640, help="Altura final do tensor.")
ap.add_argument("--fps", type=float, default=30.0)
ap.add_argument("--seconds", type=float, default=20.0)
ap.add_argument("--timeout", type=float, default=3.0)
ap.add_argument("--warmup", type=int, default=5)
ap.add_argument("--mx_id", default=None)
ap.add_argument("--sync_tolerance_ms", type=float, default=25.0)
ap.add_argument("--buffer_size", type=int, default=8)
ap.add_argument("--simulate_model_norm", action="store_true")
ap.add_argument("--model_mean", default=None)
ap.add_argument("--model_std", default=None)
ap.add_argument("--to_torch", action="store_true")
ap.add_argument("--run_model", action="store_true")
ap.add_argument("--model_config_json", default=None)
ap.add_argument("--warmup_model", type=int, default=3)
ap.add_argument("--save_debug", action="store_true")
ap.add_argument("--debug_dir", default="raw_bruto_scientific_benchmark")
ap.add_argument("--show", action="store_true")
ap.add_argument("--display_scale", type=float, default=0.65)
args = ap.parse_args()
bench = ScientificBenchmark(args)
client = OakFcc3Client(
width=args.width,
height=args.height,
fps=args.fps,
frame_type="RAW_BRUTO",
output_dtype="uint8",
capture_mode="TRIPLE",
raw_policy="require_triple",
module_calibration_json=args.module_params,
sync_tolerance_ms=args.sync_tolerance_ms,
buffer_size=args.buffer_size,
mx_id=args.mx_id,
)
client.core.warmup_numba_raw10_decode()
samples = {
"capture_ms": [],
"decode_ms": [],
"controller_ms": [],
"fuse_total_ms": [],
"fuse_dark_ms": [],
"fuse_radnorm_ms": [],
"fuse_flat_ms": [],
"fuse_prepare_ms": [],
"fuse_warp_ms": [],
"fuse_crop_resize_ms": [],
"fuse_concat_ms": [],
"model_norm_ms": [],
"torch_ms": [],
"infer_ms": [],
"total_ms": [],
"sync_dt_ms": [],
}
n_frames = 0
t_start = time.perf_counter()
t_last_log = t_start
last_tensor = None
last_meta = None
try:
client.start(print_debug=True)
# Warmup de câmera/controlador/filas
print(f"[BENCH] Warmup frames: {args.warmup}")
for _ in range(max(0, int(args.warmup))):
try:
bench.process_once(client)
except Exception as e:
print(f"[WARN] warmup falhou: {type(e).__name__}: {e}")
time.sleep(0.02)
print("============================================")
print("[BENCH] Iniciando benchmark científico RAW_BRUTO")
print(f"target tensor : (5,{args.height},{args.width})")
print(f"duration : {args.seconds}s")
print("============================================")
while True:
now = time.perf_counter()
elapsed = now - t_start
if elapsed >= args.seconds:
break
try:
tensor, pred, raw_frame, raw_meta, decoded, times = bench.process_once(client)
except TimeoutError as e:
print(f"[TIMEOUT] {e}")
continue
n_frames += 1
last_tensor = tensor
last_meta = raw_meta
for k in samples:
samples[k].append(float(times.get(k, 0.0) or 0.0))
if now - t_last_log >= 1.0:
elapsed = now - t_start
fps = n_frames / max(elapsed, 1e-6)
print(
"[RAW_SCI_PERF] "
f"elapsed={elapsed:.1f}s "
f"frames={n_frames} "
f"fps={fps:.2f} "
f"sync={last_mean(samples['sync_dt_ms']):.2f}ms "
f"capture={last_mean(samples['capture_ms']):.2f}ms "
f"decode={last_mean(samples['decode_ms']):.2f}ms "
f"controller={last_mean(samples['controller_ms']):.2f}ms "
f"fuse={last_mean(samples['fuse_total_ms']):.2f}ms "
f"radnorm={last_mean(samples['fuse_radnorm_ms']):.2f}ms "
f"flat={last_mean(samples['fuse_flat_ms']):.2f}ms "
f"warp={last_mean(samples['fuse_warp_ms']):.2f}ms "
f"crop_resize={last_mean(samples['fuse_crop_resize_ms']):.2f}ms "
f"concat={last_mean(samples['fuse_concat_ms']):.2f}ms "
f"model_norm={last_mean(samples['model_norm_ms']):.2f}ms "
f"torch={last_mean(samples['torch_ms']):.2f}ms "
f"infer={last_mean(samples['infer_ms']):.2f}ms "
f"total={last_mean(samples['total_ms']):.2f}ms "
f"tensor_shape={tuple(tensor.shape)}"
)
t_last_log = now
elapsed_total = time.perf_counter() - t_start
fps_total = n_frames / max(elapsed_total, 1e-6)
print("============================================")
print("RESULTADO FINAL RAW_BRUTO CIENTÍFICO")
print(f"elapsed : {elapsed_total:.2f}s")
print(f"frames : {n_frames}")
print(f"fps : {fps_total:.2f}")
print("--------------------------------------------")
def print_metric(name):
xs = samples[name]
print(
f"{name:18s} "
f"mean={mean(xs):8.2f}ms "
f"p95={p95(xs):8.2f}ms "
f"max={maxv(xs):8.2f}ms"
)
for name in [
"sync_dt_ms",
"capture_ms",
"decode_ms",
"controller_ms",
"fuse_total_ms",
"fuse_dark_ms",
"fuse_radnorm_ms",
"fuse_flat_ms",
"fuse_prepare_ms",
"fuse_warp_ms",
"fuse_crop_resize_ms",
"fuse_concat_ms",
"model_norm_ms",
"torch_ms",
"infer_ms",
"total_ms",
]:
print_metric(name)
print("============================================")
if last_tensor is not None:
print_tensor_stats("LAST RAW_BRUTO SCI", last_tensor)
if args.save_debug:
debug_dir = Path(args.debug_dir)
debug_dir.mkdir(parents=True, exist_ok=True)
np.save(str(debug_dir / "last_tensor.npy"), last_tensor)
stats_path = debug_dir / "last_tensor_stats.json"
with open(stats_path, "w", encoding="utf-8") as f:
json.dump(tensor_stats(last_tensor), f, indent=2, ensure_ascii=False)
panel = make_panel(last_tensor)
cv2.imwrite(str(debug_dir / "last_tensor_panel.png"), panel)
if last_meta is not None:
with open(debug_dir / "last_meta.json", "w", encoding="utf-8") as f:
json.dump(last_meta, f, indent=2, ensure_ascii=False)
print(f"[SAVE] Debug salvo em: {debug_dir}")
if args.show:
panel = make_panel(last_tensor)
if args.display_scale and abs(args.display_scale - 1.0) > 1e-6:
new_w = max(1, int(panel.shape[1] * args.display_scale))
new_h = max(1, int(panel.shape[0] * args.display_scale))
panel = cv2.resize(panel, (new_w, new_h), interpolation=cv2.INTER_AREA)
cv2.imshow("RAW_BRUTO scientific tensor", panel)
print("[INFO] Pressione qualquer tecla para fechar.")
cv2.waitKey(0)
cv2.destroyAllWindows()
finally:
try:
client.stop()
except Exception:
pass
if __name__ == "__main__":
main()

View File

@ -0,0 +1,645 @@
import argparse
import json
import sys
import time
import threading
from pathlib import Path
import cv2
import numpy as np
# ============================================================
# Ajuste de import local
# ============================================================
THIS_FILE = Path(__file__).resolve()
WORKERS_DIR = THIS_FILE.parents[2]
if str(WORKERS_DIR) not in sys.path:
sys.path.insert(0, str(WORKERS_DIR))
from camera_worker.oak_fcc3_core.oak_fcc3_client import OakFcc3Client
try:
from camera_worker.oak_fcc3_core.segformer_service import MultiSpecSegformerService
except Exception:
MultiSpecSegformerService = None
# ============================================================
# FPS / utilidades
# ============================================================
class FpsMeter:
def __init__(self, alpha=0.15):
self.alpha = float(alpha)
self.last_ts = None
self.fps = 0.0
def tick(self):
now = time.perf_counter()
if self.last_ts is not None:
dt = now - self.last_ts
if dt > 1e-9:
inst = 1.0 / dt
self.fps = inst if self.fps <= 0 else (1.0 - self.alpha) * self.fps + self.alpha * inst
self.last_ts = now
return self.fps
def load_json_if_exists(path):
if not path:
return None
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def init_model_service(args):
"""
Mesmo contrato do benchmark científico:
--run_model exige --model_config_json
MultiSpecSegformerService(model_config=cfg).infer_tensor_fast(tensor, keep_probs=False)
"""
if not args.run_model:
return None
if MultiSpecSegformerService is None:
raise RuntimeError("MultiSpecSegformerService não pôde ser importado.")
model_cfg = load_json_if_exists(args.model_config_json)
if not isinstance(model_cfg, dict):
raise RuntimeError("--run_model requer --model_config_json válido.")
svc = MultiSpecSegformerService(
model_config=model_cfg,
mostrar_log=print,
)
dummy = np.zeros((5, int(args.height), int(args.width)), dtype=np.float32)
for _ in range(max(0, int(args.warmup_model))):
svc.infer_tensor_fast(dummy, keep_probs=False)
print(f"[MODEL] Warmup concluído: {args.warmup_model}x")
return svc
def to_u8_01(arr, auto_level=False):
arr = np.asarray(arr, dtype=np.float32)
arr = np.nan_to_num(arr, nan=0.0, posinf=1.0, neginf=0.0)
if auto_level:
p1 = float(np.percentile(arr, 1))
p99 = float(np.percentile(arr, 99))
den = max(p99 - p1, 1e-6)
arr = (arr - p1) / den
arr = np.clip(arr, 0.0, 1.0)
return (arr * 255.0).astype(np.uint8)
def ensure_bgr(img, auto_level=False):
img = np.asarray(img)
if img.ndim == 2:
g = to_u8_01(img, auto_level=auto_level)
return cv2.cvtColor(g, cv2.COLOR_GRAY2BGR)
if img.ndim == 3 and img.shape[2] == 3:
u8 = to_u8_01(img, auto_level=auto_level)
# decoded/tensor RGB vem em RGB; OpenCV mostra BGR
return cv2.cvtColor(u8, cv2.COLOR_RGB2BGR)
raise RuntimeError(f"Imagem inválida para visualização: shape={img.shape}")
def tensor_rgb_to_bgr(tensor, auto_level=False):
rgb = np.stack([tensor[0], tensor[1], tensor[2]], axis=2)
return ensure_bgr(rgb, auto_level=auto_level)
def tensor_channel_to_bgr(tensor, idx, auto_level=False):
return ensure_bgr(tensor[idx], auto_level=auto_level)
def add_title(img, title, color=(255, 255, 255)):
out = img.copy()
h, w = out.shape[:2]
bar_h = 34
bar = np.zeros((bar_h, w, 3), dtype=np.uint8)
cv2.putText(bar, str(title), (10, 23), cv2.FONT_HERSHEY_SIMPLEX, 0.62, color, 2, cv2.LINE_AA)
return np.vstack([bar, out])
def add_hud(panel, lines):
out = panel.copy()
x, y = 12, 45
for line in lines:
cv2.putText(out, line, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.62, (0, 255, 255), 2, cv2.LINE_AA)
y += 24
return out
def resize_tile(img, tile_w, tile_h):
return cv2.resize(img, (int(tile_w), int(tile_h)), interpolation=cv2.INTER_AREA)
def get_cam_by_role(decoded, role):
role = str(role).lower()
for cam_id, item in decoded.items():
r = str(item.get("role") or item.get("meta", {}).get("role") or "").lower()
if r == role:
return cam_id
return None
def decoded_raw_tiles(decoded, tile_w, tile_h, auto_level=False):
"""
Retorna tiles RAW/decoded para RGB, RE, NIR antes da fusão final.
Aqui 'RAW_BRUTO' significa o conteúdo decodificado vindo das câmeras, ainda no espaço nativo.
"""
tiles = {}
rgb_id = get_cam_by_role(decoded, "rgb")
re_id = get_cam_by_role(decoded, "re")
nir_id = get_cam_by_role(decoded, "nir")
if rgb_id is not None:
img = decoded[rgb_id]["image"]
tiles["rgb"] = resize_tile(ensure_bgr(img, auto_level=auto_level), tile_w, tile_h)
else:
tiles["rgb"] = np.zeros((tile_h, tile_w, 3), dtype=np.uint8)
if re_id is not None:
img = decoded[re_id]["image"]
tiles["re"] = resize_tile(ensure_bgr(img, auto_level=auto_level), tile_w, tile_h)
else:
tiles["re"] = np.zeros((tile_h, tile_w, 3), dtype=np.uint8)
if nir_id is not None:
img = decoded[nir_id]["image"]
tiles["nir"] = resize_tile(ensure_bgr(img, auto_level=auto_level), tile_w, tile_h)
else:
tiles["nir"] = np.zeros((tile_h, tile_w, 3), dtype=np.uint8)
return tiles
def tensor_tiles(tensor, tile_w, tile_h, auto_level=False):
return {
"rgb": resize_tile(tensor_rgb_to_bgr(tensor, auto_level=auto_level), tile_w, tile_h),
"re": resize_tile(tensor_channel_to_bgr(tensor, 3, auto_level=auto_level), tile_w, tile_h),
"nir": resize_tile(tensor_channel_to_bgr(tensor, 4, auto_level=auto_level), tile_w, tile_h),
}
def colorize_label_map(label_map, num_classes=None):
label = np.asarray(label_map)
if label.ndim == 3:
label = np.argmax(label, axis=0)
label = label.astype(np.int32)
if num_classes is None:
num_classes = int(max(1, label.max() + 1))
# Paleta simples e estável. BGR.
palette = np.array([
[40, 40, 40],
[60, 180, 60],
[60, 60, 220],
[220, 180, 60],
[180, 60, 180],
[180, 180, 60],
[60, 180, 180],
[220, 220, 220],
], dtype=np.uint8)
out = palette[label % len(palette)]
return out
def try_extract_prediction_tiles(pred, target_w, target_h):
"""
Tentativa genérica. Adapte aqui se o benchmark tiver nomes específicos das cabeças.
Retorna até 3 tiles BGR: semântica/head0/head1.
"""
if pred is None:
blank = np.zeros((target_h, target_w, 3), dtype=np.uint8)
return [blank, blank.copy(), blank.copy()], ["Pred vazio", "Head 1", "Head 2"]
candidates = []
names = []
if isinstance(pred, dict):
# nomes comuns
for key in ("mask", "pred_mask", "class_map", "semantic", "semantic_mask", "segmentation"):
if key in pred:
candidates.append(pred[key])
names.append(key)
heads = pred.get("heads") or pred.get("head_outputs") or pred.get("predictions")
if isinstance(heads, dict):
for k, v in heads.items():
candidates.append(v)
names.append(str(k))
elif isinstance(heads, (list, tuple)):
for i, v in enumerate(heads):
candidates.append(v)
names.append(f"head_{i}")
else:
candidates.append(pred)
names.append("prediction")
tiles = []
out_names = []
for name, arr in zip(names, candidates):
arr = np.asarray(arr)
# remove batch se existir
if arr.ndim == 4 and arr.shape[0] == 1:
arr = arr[0]
if arr.ndim == 3:
# CHW logits/probs ou HWC RGB/probs
if arr.shape[0] <= 32:
vis = colorize_label_map(np.argmax(arr, axis=0), num_classes=arr.shape[0])
elif arr.shape[2] in (1, 3):
vis = ensure_bgr(arr[:, :, 0] if arr.shape[2] == 1 else arr, auto_level=True)
else:
vis = ensure_bgr(np.max(arr, axis=2), auto_level=True)
elif arr.ndim == 2:
# se parecer label map, colore; se parecer float, cinza auto-level
if np.issubdtype(arr.dtype, np.integer) and int(np.max(arr)) <= 64:
vis = colorize_label_map(arr)
else:
vis = ensure_bgr(arr, auto_level=True)
elif arr.ndim == 1:
# vetor de classe/score: desenha texto
vis = np.zeros((target_h, target_w, 3), dtype=np.uint8)
txt = np.array2string(arr[:8], precision=2, separator=", ")
cv2.putText(vis, txt[:80], (10, target_h // 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
else:
continue
vis = resize_tile(vis, target_w, target_h)
tiles.append(vis)
out_names.append(name)
if len(tiles) >= 3:
break
while len(tiles) < 3:
tiles.append(np.zeros((target_h, target_w, 3), dtype=np.uint8))
out_names.append(f"pred_{len(tiles)}")
return tiles[:3], out_names[:3]
def try_run_model(model_svc, tensor):
"""
Inferência real, igual ao benchmark científico.
Mantém esta função isolada para adaptar fácil caso o retorno do modelo mude.
"""
if model_svc is None:
return None, "model_svc_none"
if hasattr(model_svc, "infer_tensor_fast"):
return model_svc.infer_tensor_fast(tensor, keep_probs=False), None
if hasattr(model_svc, "infer"):
return model_svc.infer(tensor), None
return None, "model_svc_sem_infer"
def build_grid(raw_tiles, final_tiles, pred_tiles=None, pred_names=None, tile_w=420, tile_h=260):
rows = []
row_defs = [
("RGB", "rgb"),
("RE", "re"),
("NIR", "nir"),
]
for i, (label, key) in enumerate(row_defs):
left = add_title(raw_tiles[key], f"RAW_BRUTO decoded {label}")
mid = add_title(final_tiles[key], f"Tensor final {label}")
cells = [left, mid]
if pred_tiles is not None:
name = pred_names[i] if pred_names and i < len(pred_names) else f"Pred {i}"
cells.append(add_title(pred_tiles[i], name))
# iguala altura após título
h_min = min(c.shape[0] for c in cells)
norm = [cv2.resize(c, (tile_w, h_min), interpolation=cv2.INTER_AREA) if c.shape[1] != tile_w or c.shape[0] != h_min else c for c in cells]
rows.append(np.hstack(norm))
return np.vstack(rows)
def get_core_perf(client):
for attr in ("raw_processor", "processor", "core", "raw_processor_core"):
obj = getattr(client, attr, None)
if obj is not None and getattr(obj, "last_fusion_result", None) is not None:
return obj.last_fusion_result.get("perf", {}) or {}
return {}
# ============================================================
# Estado compartilhado / workers assíncronos
# ============================================================
class SharedState:
def __init__(self):
self.lock = threading.Lock()
self.running = True
self.latest_decoded = None
self.latest_meta = None
self.latest_tensor = None
self.latest_perf = {}
self.latest_pred = None
self.latest_model_warn = None
self.latest_model_ms = 0.0
self.latest_error = None
self.tensor_fps = FpsMeter()
self.model_fps = FpsMeter()
self.preview_fps = FpsMeter()
self.tensor_seq = 0
self.model_seq = 0
def stop(self):
with self.lock:
self.running = False
def is_running(self):
with self.lock:
return bool(self.running)
def tensor_worker(client, state, args):
"""
Roda no talo: captura RAW_BRUTO, monta tensor final e atualiza cache.
Não depende do FPS da janela.
"""
while state.is_running():
try:
frame, meta, decoded = client.get_next_decoded(timeout=args.timeout)
if not isinstance(frame, dict):
raise RuntimeError(f"RAW_BRUTO esperado como dict. Veio {type(frame)}")
tensor = client.build_infer_tensor_from_decoded(
decoded=decoded,
meta=meta,
channels_expected=5,
target_size=(args.width, args.height),
)
tensor = np.ascontiguousarray(tensor.astype(np.float32, copy=False))
perf = get_core_perf(client)
fps = state.tensor_fps.tick()
with state.lock:
state.latest_decoded = decoded
state.latest_meta = meta
state.latest_tensor = tensor
state.latest_perf = dict(perf or {})
state.tensor_seq += 1
state.latest_error = None
except Exception as e:
with state.lock:
state.latest_error = f"tensor_worker: {type(e).__name__}: {e}"
time.sleep(0.02)
def model_worker(model_svc, state, args):
"""
Opcional: roda inferência no último tensor disponível.
Não bloqueia o worker de tensor nem a janela.
"""
last_seq = -1
while state.is_running():
with state.lock:
tensor = None if state.latest_tensor is None else state.latest_tensor.copy()
seq = state.tensor_seq
if tensor is None or seq == last_seq:
time.sleep(0.005)
continue
last_seq = seq
try:
t0_model = time.perf_counter()
pred, warn = try_run_model(model_svc, tensor)
model_ms = (time.perf_counter() - t0_model) * 1000.0
if warn is None:
state.model_fps.tick()
with state.lock:
state.latest_pred = pred
state.latest_model_warn = warn
state.latest_model_ms = float(model_ms)
state.model_seq += 1
except Exception as e:
with state.lock:
state.latest_model_warn = f"model_worker: {type(e).__name__}: {e}"
time.sleep(0.02)
def snapshot_state(state):
"""
Copia referências do cache para desenhar. A janela roda na cadência do preview.
"""
with state.lock:
return {
"decoded": state.latest_decoded,
"meta": state.latest_meta,
"tensor": state.latest_tensor,
"perf": dict(state.latest_perf or {}),
"pred": state.latest_pred,
"model_warn": state.latest_model_warn,
"model_ms": state.latest_model_ms,
"error": state.latest_error,
"tensor_fps": state.tensor_fps.fps,
"model_fps": state.model_fps.fps,
"tensor_seq": state.tensor_seq,
"model_seq": state.model_seq,
}
# ============================================================
# Main loop assíncrono
# ============================================================
def main():
ap = argparse.ArgumentParser(description="Preview assíncrono RAW_BRUTO decoded vs tensor final multispectral.")
ap.add_argument("--module_params", default=r"C:\ZendionInc\agrobot_base\Python\OAK\datasets\oak-fcc-3\calibration\module_params.json")
ap.add_argument("--width", type=int, default=1024, help="Largura final do tensor.")
ap.add_argument("--height", type=int, default=640, help="Altura final do tensor.")
ap.add_argument("--fps", type=float, default=40.0, help="FPS alvo da câmera.")
ap.add_argument("--preview_fps", type=float, default=5.0, help="FPS da janela OpenCV apenas.")
ap.add_argument("--timeout", type=float, default=2.0)
ap.add_argument("--warmup", type=int, default=5)
ap.add_argument("--mx_id", default=None)
ap.add_argument("--display_scale", type=float, default=0.75)
ap.add_argument("--tile_w", type=int, default=420)
ap.add_argument("--tile_h", type=int, default=260)
ap.add_argument("--auto_level_raw", action="store_true")
ap.add_argument("--auto_level_tensor", action="store_true")
ap.add_argument("--run_model", action="store_true", help="Roda inferência em thread separada usando o último tensor cacheado.")
ap.add_argument("--model_config_json", default=None, help="JSON de configuração do modelo SegFormer, igual ao benchmark.")
ap.add_argument("--warmup_model", type=int, default=3, help="Número de inferências dummy para aquecer o modelo.")
ap.add_argument("--save_last", default=None)
args = ap.parse_args()
client = OakFcc3Client(
width=args.width,
height=args.height,
fps=args.fps,
frame_type="RAW_BRUTO",
capture_mode="TRIPLE",
raw_policy="require_triple",
module_calibration_json=args.module_params,
mx_id=args.mx_id,
)
model_svc = init_model_service(args) if args.run_model else None
state = SharedState()
last_panel = None
last_warn_ts = 0.0
last_seq_drawn = -1
try:
client.start(print_debug=True)
for _ in range(max(0, int(args.warmup))):
try:
client.get_next_decoded(timeout=args.timeout)
except Exception:
pass
time.sleep(0.03)
tw = threading.Thread(target=tensor_worker, args=(client, state, args), daemon=True)
tw.start()
mw = None
if args.run_model:
mw = threading.Thread(target=model_worker, args=(model_svc, state, args), daemon=True)
mw.start()
min_period = 1.0 / max(float(args.preview_fps), 0.1)
next_draw_ts = 0.0
print("[INFO] Preview assíncrono iniciado. Pressione Q ou ESC para sair.")
print("[INFO] Tensor FPS = geração real do tensor. Preview FPS = janela. Model FPS = inferência, se habilitada.")
while True:
now = time.perf_counter()
if now < next_draw_ts:
time.sleep(min(0.005, next_draw_ts - now))
key = cv2.waitKey(1) & 0xFF
if key in (27, ord('q'), ord('Q')):
break
continue
next_draw_ts = now + min_period
snap = snapshot_state(state)
if snap["tensor"] is None or snap["decoded"] is None:
blank = np.zeros((360, 900, 3), dtype=np.uint8)
cv2.putText(blank, "Aguardando primeiro tensor...", (30, 180), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,255,255), 2, cv2.LINE_AA)
cv2.imshow("OAK RAW_BRUTO vs Tensor Final", blank)
key = cv2.waitKey(1) & 0xFF
if key in (27, ord('q'), ord('Q')):
break
continue
# Desenha só usando o cache. Não captura nem monta tensor aqui.
t_draw0 = time.perf_counter()
raw_tiles = decoded_raw_tiles(
snap["decoded"],
tile_w=args.tile_w,
tile_h=args.tile_h,
auto_level=args.auto_level_raw,
)
final_tiles = tensor_tiles(
snap["tensor"],
tile_w=args.tile_w,
tile_h=args.tile_h,
auto_level=args.auto_level_tensor,
)
pred_tiles = None
pred_names = None
if args.run_model:
pred_tiles, pred_names = try_extract_prediction_tiles(
snap["pred"],
target_w=args.tile_w,
target_h=args.tile_h,
)
if snap["model_warn"]:
tnow = time.time()
if tnow - last_warn_ts > 2.0:
print(f"[WARN][MODEL] {snap['model_warn']}")
last_warn_ts = tnow
panel = build_grid(
raw_tiles=raw_tiles,
final_tiles=final_tiles,
pred_tiles=pred_tiles,
pred_names=pred_names,
tile_w=args.tile_w,
tile_h=args.tile_h,
)
preview_fps = state.preview_fps.tick()
draw_ms = (time.perf_counter() - t_draw0) * 1000.0
perf = snap["perf"]
meta = snap["meta"] or {}
tensor_seq = int(snap["tensor_seq"])
dropped_for_preview = max(0, tensor_seq - last_seq_drawn - 1) if last_seq_drawn >= 0 else 0
last_seq_drawn = tensor_seq
hud = [
f"Preview FPS: {preview_fps:.1f} | Tensor FPS: {snap['tensor_fps']:.1f} | Model FPS: {snap['model_fps']:.1f} | infer={snap.get('model_ms', 0.0):.1f}ms",
f"draw={draw_ms:.1f}ms flat={perf.get('flat_ms', 0):.1f} warp={perf.get('warp_total_ms', 0):.1f} crop={perf.get('crop_resize_ms', 0):.1f} fuse={perf.get('total_ms', 0):.1f}",
f"seq={tensor_seq} skipped_preview={dropped_for_preview} frame_type={meta.get('frame_type')} run_model={args.run_model}",
]
if snap["error"]:
hud.append(str(snap["error"])[:120])
panel = add_hud(panel, hud)
last_panel = panel
disp = panel
if args.display_scale and abs(args.display_scale - 1.0) > 1e-6:
disp = cv2.resize(
disp,
(max(1, int(disp.shape[1] * args.display_scale)), max(1, int(disp.shape[0] * args.display_scale))),
interpolation=cv2.INTER_AREA,
)
cv2.imshow("OAK RAW_BRUTO vs Tensor Final", disp)
key = cv2.waitKey(1) & 0xFF
if key in (27, ord('q'), ord('Q')):
break
finally:
state.stop()
time.sleep(0.05)
if args.save_last and last_panel is not None:
out_path = Path(args.save_last)
out_path.parent.mkdir(parents=True, exist_ok=True)
cv2.imwrite(str(out_path), last_panel)
print(f"[SAVE] {out_path}")
try:
client.stop()
except Exception:
pass
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

@ -23,6 +23,7 @@ class OakFcc3Client:
module_calibration_json=None,
sync_mode="best",
sync_tolerance_ms=25.0,
mx_id=None,
**kwargs,
):
self.width = width
@ -37,6 +38,8 @@ class OakFcc3Client:
self.module_params = self._load_module_params(module_calibration_json)
self.fusion_config = self.module_params.get("fusion_config", {}) or {}
self.mx_id = str(mx_id) if mx_id else None
self.svc = OakFcc3Service(
timeout=10,
fps=fps,
@ -48,6 +51,8 @@ class OakFcc3Client:
raw_policy=raw_policy,
sync_mode=sync_mode,
sync_tolerance_ms=sync_tolerance_ms,
mx_id=self.mx_id,
module_calibration_json=module_calibration_json,
**kwargs,
)
@ -152,6 +157,11 @@ class OakFcc3Client:
capture_mode=self.capture_mode,
)
try:
self.mx_id = self.svc.manager.mx_id
except Exception:
pass
applied = self.apply_module_camera_settings()
if print_debug:
@ -181,39 +191,56 @@ class OakFcc3Client:
def get_next_decoded(self, timeout=2.0):
raw_frame, raw_meta = self.get_next_raw_frame(timeout=timeout)
decoded = self.decode_stream_cameras(raw_frame, raw_meta)
self.update_radiometry(decoded, raw_meta)
frame_type = str(raw_meta.get("frame_type", self.frame_type)).upper()
meta = dict(raw_meta)
if frame_type == "RAW_BRUTO":
decoded = self.decode_stream_cameras(raw_frame, raw_meta)
self.update_radiometry(decoded, raw_meta)
frame = raw_frame
return frame, meta, decoded
elif frame_type == "RGB":
decoded = self.decode_stream_cameras(raw_frame, raw_meta)
self.update_radiometry(decoded, raw_meta)
frame = self.build_rgb_tensor(decoded)
meta["output_layout"] = "CHW"
meta["channels"] = ["R", "G", "B"]
meta["shape"] = list(frame.shape)
meta["dtype"] = str(frame.dtype)
return frame, meta, decoded
elif frame_type == "MULTISPEC":
frame = self.build_multispec_tensor(decoded, meta=raw_meta)
decoded = self.decode_oak_aligned_multispec(raw_frame, raw_meta)
self.update_radiometry(decoded, raw_meta)
# Versão inicial segura:
# não chama core.fuse_multispec_cameras(), porque ali teria homografia de novo.
frame = self.build_multispec_tensor_from_oak_aligned(decoded, meta=raw_meta)
meta["output_layout"] = "CHW"
meta["channels"] = ["R", "G", "B", "RE", "NIR"]
meta["shape"] = list(frame.shape)
meta["dtype"] = str(frame.dtype)
meta["aligned_by_oak"] = True
meta["geometry_stage"] = "oak"
return frame, meta, decoded
elif frame_type == "PREVIEW":
decoded = self.decode_stream_cameras(raw_frame, raw_meta)
frame = raw_frame
return frame, meta, decoded
else:
raise RuntimeError(f"frame_type não suportado: {frame_type}")
return frame, meta, decoded
def get_next_tensor_preview(self, timeout=2.0):
frame, meta, decoded = self.get_next_decoded(timeout=timeout)
@ -551,3 +578,95 @@ class OakFcc3Client:
def _gray01_to_bgr(gray01):
g = np.clip(gray01 * 255.0, 0, 255).astype(np.uint8)
return cv2.cvtColor(g, cv2.COLOR_GRAY2BGR)
def decode_oak_aligned_multispec(self, frame, meta):
"""
Decodifica frames alinhados pela OAK.
Entrada esperada:
frame = {
"CAM_A": BGR uint8 HWC,
"CAM_B": GRAY uint8 HW,
"CAM_C": GRAY uint8 HW,
}
Saída:
decoded por role, em float32 0..1.
"""
if not isinstance(frame, dict):
raise RuntimeError("MULTISPEC alinhado esperado como dict de câmeras.")
decoded = {}
camera_info = meta.get("camera_info", {}) or {}
for cam_id, img in frame.items():
info = camera_info.get(cam_id, {}) or {}
role = str(info.get("role", "")).lower()
if not role:
role = str(self.svc.manager.roles.get(cam_id, cam_id)).lower()
img01 = self._frame_to_float01(cam_id, img, role)
decoded[cam_id] = {
"name": role.upper(),
"role": role,
"image": img01,
"meta": {
**info,
"cam_id": cam_id,
"role": role,
"aligned_by_oak": True,
"geometry_stage": "oak",
"homography_applied": role in ("re", "nir"),
"crop_resize_applied": True,
"shape": list(img.shape),
"dtype": str(img.dtype),
},
}
return decoded
def build_multispec_tensor_from_oak_aligned(self, decoded, meta=None):
"""
Monta CHW [R,G,B,RE,NIR] sem reaplicar homografia.
"""
_, rgb_item = self._find_decoded_by_role(decoded, "rgb")
_, re_item = self._find_decoded_by_role(decoded, "re")
_, nir_item = self._find_decoded_by_role(decoded, "nir")
rgb = rgb_item["image"].astype(np.float32, copy=False)
re = re_item["image"].astype(np.float32, copy=False)
nir = nir_item["image"].astype(np.float32, copy=False)
if rgb.ndim != 3 or rgb.shape[2] != 3:
raise RuntimeError(f"RGB alinhado inválido: shape={rgb.shape}")
h, w = rgb.shape[:2]
if re.ndim == 3:
re = re[:, :, 0]
if nir.ndim == 3:
nir = nir[:, :, 0]
if re.shape[:2] != (h, w):
re = cv2.resize(re, (w, h), interpolation=cv2.INTER_LINEAR)
if nir.shape[:2] != (h, w):
nir = cv2.resize(nir, (w, h), interpolation=cv2.INTER_LINEAR)
tensor = np.stack(
[
rgb[:, :, 0], # R
rgb[:, :, 1], # G
rgb[:, :, 2], # B
re,
nir,
],
axis=0,
).astype(np.float32, copy=False)
return np.ascontiguousarray(tensor)

File diff suppressed because it is too large Load Diff

View File

@ -45,6 +45,7 @@ class OakFcc3Service:
def get_config(self):
return {
"mx_id": self.manager.mx_id,
"ok": True,
"fps": self.manager.fps,
"width": self.manager.width,

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,915 @@
# camera_worker/multispec_segformer_service.py
# -*- coding: utf-8 -*-
from __future__ import annotations
import copy
import json
import time
from pathlib import Path
from typing import Dict, 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 SegformerConfig, SegformerForSemanticSegmentation
DEFAULT_HEADS = {
"semantic": {
"enabled": True,
"type": "multiclass",
"num_classes": 3,
"classes": {"chao": 0, "cana": 1, "erva": 2},
"ignore_index": 255,
},
"vegetation": {
"enabled": True,
"type": "binary",
"num_classes": 2,
"classes": {"background": 0, "vegetation": 1},
"ignore_index": 255,
},
"cana": {
"enabled": True,
"type": "binary",
"num_classes": 2,
"classes": {"not_cana": 0, "cana": 1},
"ignore_index": 255,
},
}
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),
}
def load_json(path: str | Path) -> dict:
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
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 = 255) -> 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))
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 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_config = SegformerConfig.from_pretrained(backbone)
base_config.num_labels = semantic_classes
base_config.id2label = {int(k): str(v) for k, v in semantic_id2label.items()}
base_config.label2id = {str(k): int(v) for k, v in semantic_label2id.items()}
base = SegformerForSemanticSegmentation(base_config)
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 = [str(h) for h in head_names if str(h) in self.decode_heads]
return {head_name: self.decode_heads[head_name](hidden_states) for head_name in selected}
class MultiSpecSegformerService:
"""
Service de inferência para tensor multiespectral:
input : CHW float32 [R,G,B,RE,NIR] 0..1
output: dict de predictions multi-head
"""
def __init__(self, model_config: dict, mostrar_log=print):
self.config = model_config or {}
self.mostrar_log = mostrar_log
self.channels = int(self.config.get("channels", 5))
if self.channels != 5:
raise RuntimeError(f"MultiSpecSegformerService espera channels=5, veio {self.channels}")
self.ignore_id = int(self.config.get("ignore_index", 255))
self.heads_config = build_heads_config(self.config, ignore_index=self.ignore_id)
self.semantic_id2label = {
0: "chao",
1: "cana",
2: "erva",
}
self.semantic_label2id = {
"chao": 0,
"cana": 1,
"erva": 2,
}
self.classes = dict(self.semantic_label2id)
# Dict interno, bom para ids_to_rgb
self.colormap_rgb = {
0: (85, 85, 85),
1: (0, 190, 0),
2: (230, 55, 55),
}
# Lista externa, compatível com WeedDetector
self.colormap_list_rgb = [
self.colormap_rgb[0],
self.colormap_rgb[1],
self.colormap_rgb[2],
]
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.use_amp = bool(self.config.get("amp", True)) and self.device.type == "cuda"
self.sync_for_timing = bool(self.config.get("sync_for_timing", False)) and self.device.type == "cuda"
self.runtime_mode = str(self.config.get("runtime_mode", "semantic")).lower()
self.lowres_argmax = bool(self.config.get("lowres_argmax", True))
self.trust_input = bool(self.config.get("trust_input", True))
self.channels_last = bool(self.config.get("channels_last", True)) and self.device.type == "cuda"
self.model_half = bool(self.config.get("model_half", False)) and self.device.type == "cuda"
if self.device.type == "cuda":
torch.backends.cudnn.benchmark = True
try:
torch.set_float32_matmul_precision("high")
except Exception:
pass
self.mean, self.std = self._load_norm_stats_from_config()
backbone = self.config.get("backbone", self.config.get("pretrained_model", "nvidia/mit-b1"))
ckpt_path = self._resolve_checkpoint_path()
self.mostrar_log(f"[MULTIHEAD] device={self.device}")
self.mostrar_log(f"[MULTIHEAD] backbone={backbone}")
self.mostrar_log(f"[MULTIHEAD] ckpt={ckpt_path}")
self.model = MultiHeadSegFormer(
backbone=backbone,
channels=self.channels,
heads_config=self.heads_config,
semantic_id2label=self.semantic_id2label,
semantic_label2id=self.semantic_label2id,
)
self._load_checkpoint(ckpt_path)
self.model.to(self.device)
self.model.eval()
if self.model_half:
self.model.half()
if self.mean is not None:
self.mean = self.mean.half()
if self.std is not None:
self.std = self.std.half()
if self.channels_last:
try:
self.model.to(memory_format=torch.channels_last)
except Exception:
self.channels_last = False
if bool(self.config.get("fold_input_norm", False)):
self._fold_input_normalization_into_first_conv()
if bool(self.config.get("torch_compile", False)):
try:
self.model = torch.compile(
self.model,
mode=str(self.config.get("torch_compile_mode", "reduce-overhead")),
fullgraph=False,
)
self.mostrar_log("[MULTIHEAD][OPT] torch.compile habilitado")
except Exception as e:
self.mostrar_log(f"[MULTIHEAD][OPT] torch.compile falhou: {type(e).__name__}: {e}")
self.mostrar_log(
f"[MULTIHEAD][FAST] runtime_mode={self.runtime_mode} "
f"lowres_argmax={self.lowres_argmax} trust_input={self.trust_input} "
f"channels_last={self.channels_last} model_half={self.model_half} amp={self.use_amp}"
)
self._ultimo_tensor = None
self._ultimo_predictions = None
self._ultimo_probs = None
# ============================================================
# Config/load
# ============================================================
def _resolve_checkpoint_path(self) -> Path:
candidates = []
for key in ("ckpt", "checkpoint", "ia_model_path", "model_path"):
value = self.config.get(key)
if value:
candidates.append(Path(value))
backup_root = self.config.get("backup_root")
modelo = self.config.get("modelo", "segformer_b1")
model_name = self.config.get("model_name", "test_multi")
fusion_mode = self.config.get("fusion_mode", "stacked")
if backup_root:
exp_tag = f"{fusion_mode}_raw{self.channels}_multihead"
candidates.extend([
Path(backup_root) / modelo / model_name / exp_tag / "best_target.pt",
Path(backup_root) / modelo / model_name / exp_tag / "best_score.pt",
Path(backup_root) / modelo / model_name / exp_tag / "last.pt",
])
candidates.extend([
Path("backup") / modelo / model_name / f"{fusion_mode}_raw{self.channels}_multihead" / "best_target.pt",
Path("backup") / modelo / model_name / f"{fusion_mode}_raw{self.channels}_multihead" / "best_score.pt",
Path("backup") / modelo / model_name / f"{fusion_mode}_raw{self.channels}_multihead" / "last.pt",
])
for p in candidates:
p = p.resolve() if not p.is_absolute() else p
if p.is_file():
return p
raise FileNotFoundError(
"Checkpoint multi-head não encontrado. Procurei:\n" +
"\n".join(str(p) for p in candidates)
)
def _load_norm_stats_from_config(self):
path = None
for key in ("norm_stats", "norm_stats_path", "ia_norm_stats_path"):
if self.config.get(key):
path = Path(self.config[key])
break
if path is None:
dataset_root = self.config.get("dataset_root")
ia_resolution = self.config.get("ia_resolution", [1024, 640])
if dataset_root:
w, h = int(ia_resolution[0]), int(ia_resolution[1])
path = Path(dataset_root) / f"{w}x{h}" / "group" / "norm_stats.json"
if path is None or not path.is_file():
self.mostrar_log("[MULTIHEAD][NORM] norm_stats não encontrado. Usando tensor 0..1 sem padronização.")
return None, None
js = load_json(path)
mean = js.get("mean")
std = js.get("std")
names = js.get("channels", [])
if mean is None or std is None:
raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
if len(mean) != self.channels or len(std) != self.channels:
raise RuntimeError(
f"norm_stats incompatível com channels={self.channels}: "
f"mean={len(mean)} std={len(std)} path={path}"
)
self.mostrar_log(f"[MULTIHEAD][NORM] usando {path}")
self.mostrar_log(f"[MULTIHEAD][NORM] channels={names}")
self.mostrar_log(f"[MULTIHEAD][NORM] mean={mean}")
self.mostrar_log(f"[MULTIHEAD][NORM] std ={std}")
mean_t = torch.tensor(mean, dtype=torch.float32).view(1, self.channels, 1, 1).to(self.device)
std_t = torch.tensor(std, dtype=torch.float32).view(1, self.channels, 1, 1).to(self.device)
return mean_t, std_t
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)
self.mostrar_log(
f"[MULTIHEAD] load_state_dict strict=False | "
f"missing={len(missing)} unexpected={len(unexpected)}"
)
if missing:
self.mostrar_log(f"[MULTIHEAD] primeiros missing: {missing[:8]}")
if unexpected:
self.mostrar_log(f"[MULTIHEAD] primeiros unexpected: {unexpected[:8]}")
# ============================================================
# Inferência
# ============================================================
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_tensor(self, tensor5_chw: np.ndarray):
"""
Retorna predictions compatível com o WeedDetector.
Por enquanto:
predictions = semantic mask uint8 HxW
Também mantém:
self._ultimo_predictions_full = dict com semantic/vegetation/cana/target/probs.
"""
if tensor5_chw is None:
return None
chw = np.asarray(tensor5_chw, dtype=np.float32)
if chw.ndim != 3:
raise RuntimeError(f"Tensor inválido: esperado CHW 3D, veio shape={chw.shape}")
if chw.shape[0] != self.channels:
raise RuntimeError(f"Tensor inválido: esperado C={self.channels}, veio shape={chw.shape}")
chw = np.nan_to_num(chw, nan=0.0, posinf=1.0, neginf=0.0)
chw = np.clip(chw, 0.0, 1.0).astype(np.float32, copy=False)
chw = np.ascontiguousarray(chw)
h, w = int(chw.shape[1]), int(chw.shape[2])
x = torch.from_numpy(chw).unsqueeze(0).to(self.device, non_blocking=True)
x = self._normalize(x)
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):
logits_by_head = self.model(pixel_values=x)
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
semantic = preds["semantic"]
vegetation = preds["vegetation"]
cana = preds["cana"]
target = self.operational_target_mask(vegetation, cana, ignore_id=self.ignore_id)
self._ultimo_tensor = chw
self._ultimo_predictions = semantic
self._ultimo_probs = probs
self._ultimo_predictions_full = {
"semantic": semantic,
"vegetation": vegetation,
"cana": cana,
"target": target,
"probs": probs,
"infer_ms": t_ms,
}
return semantic
def infer_tensor_full(self, tensor5_chw: np.ndarray):
semantic = self.infer_tensor(tensor5_chw)
if semantic is None:
return None
return dict(self._ultimo_predictions_full)
@torch.inference_mode()
def infer_tensor_fast(self, tensor5_chw: np.ndarray, keep_probs: bool = False):
if keep_probs:
return self.infer_tensor(tensor5_chw)
return self.infer_tensor_ultrafast(
tensor5_chw,
return_full=bool(self.config.get("return_full_fast", False)),
)
# ============================================================
# Preview/debug
# ============================================================
def preview_infer_cached(self, tensor5_chw=None, predictions=None, alpha=0.5):
"""
Compatível com o uso atual do weed_worker:
rgb_frame, seg_frame, overlay_frame, _, _ = preview_infer_cached(...)
Retorna BGR para OpenCV/TCP.
"""
tensor = tensor5_chw if tensor5_chw is not None else self._ultimo_tensor
pred = predictions if predictions is not None else self._ultimo_predictions
if tensor is None:
return None, None, None, None, None
rgb = self.tensor_to_preview_rgb(tensor)
rgb_bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
if pred is None:
return rgb_bgr, None, rgb_bgr, None, None
seg_rgb = self.ids_to_rgb(pred, self.colormap_rgb, ignore_id=self.ignore_id)
overlay_rgb = cv2.addWeighted(rgb, 1.0 - alpha, seg_rgb, alpha, 0.0)
seg_bgr = cv2.cvtColor(seg_rgb, cv2.COLOR_RGB2BGR)
overlay_bgr = cv2.cvtColor(overlay_rgb, cv2.COLOR_RGB2BGR)
return rgb_bgr, seg_bgr, overlay_bgr, None, None
def get_classes(self):
return self.classes
def get_colormap(self):
return self.colormap_list_rgb
@staticmethod
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
@staticmethod
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)
@staticmethod
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 _runtime_head_names(self):
mode = str(getattr(self, "runtime_mode", "semantic")).lower()
if mode in ("semantic", "sem", "mask"):
return ["semantic"]
if mode in ("target", "spray", "operational"):
return ["vegetation", "cana"]
if mode in ("target_direct", "direct_target", "target_head"):
return ["target"]
if mode in ("all", "full", "debug"):
return ["semantic", "vegetation", "cana", "target"]
return ["semantic"]
def _prepare_input_tensor_fast(self, tensor5_chw: np.ndarray):
if tensor5_chw is None:
return None, None, None
# O pipeline científico já devolve float32 CHW 0..1 contíguo.
# Em runtime, confiar nisso economiza uma passada grande no CPU.
if bool(getattr(self, "trust_input", True)):
chw = tensor5_chw
if not isinstance(chw, np.ndarray):
chw = np.asarray(chw, dtype=np.float32)
if chw.dtype != np.float32 or not chw.flags.c_contiguous:
chw = np.ascontiguousarray(chw, dtype=np.float32)
else:
chw = np.asarray(tensor5_chw, dtype=np.float32)
chw = np.nan_to_num(chw, nan=0.0, posinf=1.0, neginf=0.0)
chw = np.clip(chw, 0.0, 1.0).astype(np.float32, copy=False)
chw = np.ascontiguousarray(chw)
if chw.ndim != 3:
raise RuntimeError(f"Tensor inválido: esperado CHW 3D, veio shape={chw.shape}")
if chw.shape[0] != self.channels:
raise RuntimeError(f"Tensor inválido: esperado C={self.channels}, veio shape={chw.shape}")
h, w = int(chw.shape[1]), int(chw.shape[2])
x = torch.from_numpy(chw).unsqueeze(0).to(self.device, non_blocking=True)
if bool(getattr(self, "channels_last", False)):
try:
x = x.contiguous(memory_format=torch.channels_last)
except Exception:
pass
x = self._normalize(x)
if bool(getattr(self, "model_half", False)):
x = x.half()
return chw, x, (h, w)
def _logits_to_pred_numpy_fast(self, logits: torch.Tensor, out_hw, lowres_argmax: bool):
h, w = int(out_hw[0]), int(out_hw[1])
if lowres_argmax:
# Argmax no mapa pequeno, depois resize da máscara uint8.
# Bem mais barato que interpolar logits CxHxW em float.
pred_small = torch.argmax(logits, dim=1)[0]
pred_np = pred_small.detach().to("cpu", non_blocking=False).numpy().astype(np.uint8)
if pred_np.shape[0] != h or pred_np.shape[1] != w:
pred_np = cv2.resize(pred_np, (w, h), interpolation=cv2.INTER_NEAREST)
return pred_np
# Caminho equivalente ao atual, mas sem softmax.
logits = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
pred = torch.argmax(logits, dim=1)[0]
return pred.detach().to("cpu", non_blocking=False).numpy().astype(np.uint8)
@torch.inference_mode()
def infer_tensor_ultrafast(self, tensor5_chw: np.ndarray, return_full: bool = False):
if tensor5_chw is None:
return None
t_total0 = time.perf_counter()
# ------------------------------------------------------------
# Prepare: numpy -> torch/cuda + layout + normalização
# ------------------------------------------------------------
t_prepare0 = time.perf_counter()
chw, x, out_hw = self._prepare_input_tensor_fast(tensor5_chw)
prepare_ms = (time.perf_counter() - t_prepare0) * 1000.0
if x is None:
return None
head_names = self._runtime_head_names()
lowres_argmax = bool(getattr(self, "lowres_argmax", True))
runtime_mode = str(getattr(self, "runtime_mode", "semantic")).lower()
output_mask_fullres = bool(self.config.get("output_mask_fullres", True))
if bool(getattr(self, "sync_for_timing", False)) and self.device.type == "cuda":
torch.cuda.synchronize()
# ------------------------------------------------------------
# Forward: modelo UMA vez só
# ------------------------------------------------------------
t_forward0 = time.perf_counter()
with torch.autocast(
device_type="cuda",
dtype=torch.float16,
enabled=bool(getattr(self, "use_amp", True)) and self.device.type == "cuda",
):
logits_by_head = self.model(pixel_values=x, head_names=head_names)
if self.device.type == "cuda":
torch.cuda.synchronize()
forward_ms = (time.perf_counter() - t_forward0) * 1000.0
# ------------------------------------------------------------
# Post: argmax / target / cópia CPU
# ------------------------------------------------------------
t_post0 = time.perf_counter()
semantic = None
vegetation = None
cana = None
target = None
if runtime_mode in ("target", "spray", "operational"):
logits_veg = logits_by_head.get("vegetation")
logits_cana = logits_by_head.get("cana")
if logits_veg is None or logits_cana is None:
raise RuntimeError("runtime_mode=target requer heads vegetation e cana")
target = self._target_from_logits_gpu_fast(
logits_veg,
logits_cana,
out_hw=out_hw,
fullres=output_mask_fullres,
)
output = target
else:
preds = {}
if "target" in logits_by_head:
target = self._logits_to_pred_numpy_fast(
logits_by_head["target"],
out_hw=out_hw,
lowres_argmax=lowres_argmax,
)
semantic = None
vegetation = None
cana = None
output = target
else:
for head_name, logits in logits_by_head.items():
preds[head_name] = self._logits_to_pred_numpy_fast(
logits,
out_hw=out_hw,
lowres_argmax=lowres_argmax,
)
semantic = preds.get("semantic")
vegetation = preds.get("vegetation")
cana = preds.get("cana")
if vegetation is not None and cana is not None:
target = self.operational_target_mask(
vegetation,
cana,
ignore_id=self.ignore_id,
)
output = semantic if semantic is not None else target
post_ms = (time.perf_counter() - t_post0) * 1000.0
total_ms = (time.perf_counter() - t_total0) * 1000.0
self._ultimo_tensor = chw
self._ultimo_predictions = output
self._ultimo_probs = None
self._ultimo_predictions_full = {
"semantic": semantic,
"vegetation": vegetation,
"cana": cana,
"target": target,
"probs": None,
"infer_ms": total_ms,
"prepare_ms": prepare_ms,
"forward_ms": forward_ms,
"post_ms": post_ms,
"runtime_mode": runtime_mode,
"lowres_argmax": lowres_argmax,
"heads": list(head_names),
"output_mask_fullres": output_mask_fullres,
}
if return_full:
return dict(self._ultimo_predictions_full)
return output
def _target_from_logits_gpu_fast(self, logits_veg, logits_cana, out_hw, fullres: bool = True):
h, w = int(out_hw[0]), int(out_hw[1])
veg = torch.argmax(logits_veg, dim=1)[0]
cana = torch.argmax(logits_cana, dim=1)[0]
target = ((veg == 1) & (cana == 0)).to(torch.uint8)
target_np = target.detach().to("cpu", non_blocking=False).numpy()
if fullres and (target_np.shape[0] != h or target_np.shape[1] != w):
target_np = cv2.resize(
target_np,
(w, h),
interpolation=cv2.INTER_NEAREST,
)
return target_np.astype(np.uint8, copy=False)
def _fold_input_normalization_into_first_conv(self):
if self.mean is None or self.std is None:
return False
try:
proj = self.model.segformer.encoder.patch_embeddings[0].proj
except Exception:
return False
if not isinstance(proj, nn.Conv2d):
return False
with torch.no_grad():
device = proj.weight.device
dtype = proj.weight.dtype
mean = self.mean.detach().to(device=device, dtype=dtype).view(-1)
std = self.std.detach().to(device=device, dtype=dtype).view(-1)
std = torch.clamp(std, min=1e-6)
w_old = proj.weight.data.clone()
b_old = proj.bias.data.clone() if proj.bias is not None else torch.zeros(
proj.out_channels,
device=device,
dtype=dtype,
)
# W'[:, c] = W[:, c] / std[c]
w_new = w_old / std.view(1, -1, 1, 1)
# bias' = bias - sum(W[:,c,:,:] * mean[c] / std[c])
offset = (w_old * (mean / std).view(1, -1, 1, 1)).sum(dim=(1, 2, 3))
b_new = b_old - offset
proj.weight.data.copy_(w_new)
if proj.bias is None:
proj.bias = nn.Parameter(b_new)
else:
proj.bias.data.copy_(b_new)
self.mean = None
self.std = None
self.mostrar_log("[MULTIHEAD][OPT] normalização foldada na primeira conv")
return True

View File

@ -0,0 +1,99 @@
{
"debug_visual": false,
"frames_consecutivos": 3,
"frames_histerese": 2,
"min_area_px": 400,
"max_area_frac": 0.2,
"ia_roi_begin": 0.0,
"ia_roi_size": 1.0,
"ia_resolution": [1024,640],
"ia_channels": 5,
"ia_use_ndvi": false,
"erva_top_band_frac": 0.30,
"erva_frac_ema": 0.3,
"erva_thresh_vel_gain": 0.4,
"min_frac_erva_global_on": 0.0020,
"min_frac_erva_global_off": 0.0015,
"min_frac_erva_top_on": 0.0015,
"min_frac_erva_top_off": 0.0010,
"min_frac_erva_por_bico": 0.02,
"usar_morfologia": true,
"kernel_morf": 3,
"usar_radar_global_gate": true,
"max_frac_cana_por_bico": 0.009,
"ema_frac_bico": 0.35,
"on_frames_required": 3,
"off_frames_required": 2,
"cana_halo_px": 5,
"min_area_erva_px": 80,
"erva_thresh_vel_gain_local": 0.6,
"k_roi_shift_px_per_vnorm": 24.0,
"fps": 20,
"module_params": "C:\\ZendionInc\\agrobot_base\\Python\\OAK\\datasets\\oak-fcc-3\\calibration\\module_params.json",
"qtd_bicos": 7,
"velocidade_robo": 0,
"ia_model_path": "C:\\AgroBaseModels\\Ervas\\model-3_1.pt",
"ia_labelmap_path": "C:\\AgroBaseModels\\Ervas\\model-3_1.txt",
"ia_norm_stats_path": "C:\\AgroBaseModels\\Ervas\\model-1_1.json",
"ia_backbone": "nvidia/mit-b1",
"faixa_atuacao_bicos": 0.7,
"area_atuacao_bicos": 0.1,
"min_frac_erva_por_bico_on": 0.02,
"min_frac_erva_por_bico_off": 0.01,
"tipo_camera_solo": "multispectral",
"channels": 5,
"backbone": "nvidia/mit-b1",
"ckpt": "C:\\AgroBaseModels\\Ervas\\model-3_1.pt",
"norm_stats_path": "C:\\AgroBaseModels\\Ervas\\model-1_1.json",
"module_calibration_json": "C:\\ZendionInc\\agrobot_base\\Python\\OAK\\datasets\\oak-fcc-3\\calibration\\module_params.json",
"camera_width": 1280,
"camera_height": 800,
"camera_fps": 40,
"amp": true,
"fold_input_norm": true,
"runtime_mode": "target_direct",
"output_mask_fullres": false,
"lowres_argmax": true,
"trust_input": true,
"channels_last": false,
"model_half": true,
"sync_for_timing": false,
"torch_compile": false,
"torch_compile_mode": "reduce-overhead",
"heads": {
"semantic": {
"enabled": true,
"type": "multiclass",
"num_classes": 3,
"classes": {"chao": 0, "cana": 1, "erva": 2},
"ignore_index": 255
},
"vegetation": {
"enabled": true,
"type": "binary",
"num_classes": 2,
"classes": {"background": 0, "vegetation": 1},
"ignore_index": 255
},
"cana": {
"enabled": true,
"type": "binary",
"num_classes": 2,
"classes": {"not_cana": 0, "cana": 1},
"ignore_index": 255
},
"target": {
"enabled": true,
"type": "binary",
"num_classes": 2,
"classes": {"background": 0, "target": 1},
"ignore_index": 255
}
}
}

View File

@ -0,0 +1,5 @@
# label:color_rgb:parts:actions
chao:128,0,0::
cana:0,0,128::
erva:0,128,0::
ignore:255,255,255::

View File

@ -1,37 +1,68 @@
from core.oak_fcc3_service import OakFcc3Service
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
svc = OakFcc3Service(
fps=15,
width=640,
height=400,
frame_type="MULTISPEC",
capture_mode="TRIPLE",
raw_policy="require_triple",
sync_mode="best",
sync_tolerance_ms=25.0,
)
import depthai as dai
svc.connect()
svc.begin(frame_type="RAW_BRUTO", output_dtype="uint8", capture_mode="TRIPLE")
frame, meta = svc.capture_frame(timeout=2.0)
def ler_attr_ou_metodo(obj, nomes, default=None):
for nome in nomes:
try:
valor = getattr(obj, nome, None)
print(meta["camera_info"])
for cam_id, arr in frame.items():
print(cam_id, arr.shape, arr.dtype, arr.size)
if valor is None:
continue
from core.raw_processor_core import RawProcessorCore
from core.raw_processor_preview import RawProcessorPreview
import cv2
if callable(valor):
return valor()
core = RawProcessorCore(sensor_width=1280, sensor_height=800, bayer_pattern="GBRG")
preview = RawProcessorPreview(sensor_width=1280, sensor_height=800, bayer_pattern="GBRG")
return valor
except Exception:
pass
raw16 = core.unpack_raw10_packed(frame["cam0"], sensor_width=1280, sensor_height=800)
img = preview.raw16_to_preview_bgr(raw16, bit_depth=10)
return default
cv2.imwrite("calibration/cam0_raw_preview.png", img)
print(raw16.shape, raw16.dtype, raw16.min(), raw16.max())
svc.stop()
svc.disconnect()
def main():
print("==========================================")
print("Teste simples - listar dispositivos DepthAI")
print("==========================================")
try:
devices = dai.Device.getAllAvailableDevices()
except Exception as e:
print("[ERRO] Falha ao chamar dai.Device.getAllAvailableDevices()")
print("Erro:", e)
return
if not devices:
print("[INFO] Nenhum dispositivo DepthAI/OAK encontrado.")
return
print(f"[OK] Dispositivos encontrados: {len(devices)}")
print()
for i, dev_info in enumerate(devices, start=1):
device_id = ler_attr_ou_metodo(
dev_info,
["getMxId", "mxid", "getDeviceId", "deviceId"],
default=None
)
name = ler_attr_ou_metodo(dev_info, ["name"], default=None)
state = ler_attr_ou_metodo(dev_info, ["state"], default=None)
protocol = ler_attr_ou_metodo(dev_info, ["protocol"], default=None)
platform = ler_attr_ou_metodo(dev_info, ["platform"], default=None)
status = ler_attr_ou_metodo(dev_info, ["status"], default=None)
print(f"Dispositivo #{i}")
print(f" DeviceId : {device_id}")
print(f" Name : {name}")
print(f" State : {state}")
print(f" Protocol : {protocol}")
print(f" Platform : {platform}")
print(f" Status : {status}")
print("-" * 42)
if __name__ == "__main__":
main()

View File

@ -0,0 +1,534 @@
import argparse
import json
import time
from pathlib import Path
import cv2
import depthai as dai
import numpy as np
# ============================================================
# Homografia / crop helpers
# ============================================================
def load_module_params(path: str) -> dict:
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def scale_homography_to_runtime(H, calib_size, runtime_size):
"""
Mesmo conceito do RawProcessorCore:
H calib: ponto_spec_calib -> ponto_rgb_calib
H runtime: ponto_spec_runtime -> ponto_rgb_runtime
"""
H = np.asarray(H, dtype=np.float32)
if calib_size is None:
if abs(H[2, 2]) > 1e-9:
H = H / H[2, 2]
return H.astype(np.float32)
calib_w, calib_h = calib_size
runtime_w, runtime_h = runtime_size
calib_w = float(calib_w)
calib_h = float(calib_h)
runtime_w = float(runtime_w)
runtime_h = float(runtime_h)
sx = runtime_w / calib_w
sy = runtime_h / calib_h
S = np.array(
[
[sx, 0.0, 0.0],
[0.0, sy, 0.0],
[0.0, 0.0, 1.0],
],
dtype=np.float32,
)
S_inv = np.array(
[
[1.0 / sx, 0.0, 0.0],
[0.0, 1.0 / sy, 0.0],
[0.0, 0.0, 1.0],
],
dtype=np.float32,
)
H_runtime = S @ H @ S_inv
if abs(H_runtime[2, 2]) > 1e-9:
H_runtime = H_runtime / H_runtime[2, 2]
return H_runtime.astype(np.float32)
def warp_mask(mask, H, out_w, out_h):
return cv2.warpPerspective(
mask,
H,
(out_w, out_h),
flags=cv2.INTER_NEAREST,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0,
)
def compute_common_crop_box(runtime_w, runtime_h, H_re, H_nir):
"""
Calcula no PC uma vez a área comum válida no espaço RGB.
Depois essa área vira o retângulo final usado na OAK.
"""
base = np.ones((runtime_h, runtime_w), dtype=np.uint8) * 255
rgb_mask = base
re_mask = warp_mask(base, H_re, runtime_w, runtime_h)
nir_mask = warp_mask(base, H_nir, runtime_w, runtime_h)
common = (rgb_mask > 0) & (re_mask > 0) & (nir_mask > 0)
ys, xs = np.where(common)
if xs.size == 0 or ys.size == 0:
raise RuntimeError("Área comum vazia. Verifique as homografias.")
x0 = int(xs.min())
x1 = int(xs.max()) + 1
y0 = int(ys.min())
y1 = int(ys.max()) + 1
return x0, y0, x1, y1
def apply_H_to_point(H, x, y):
p = np.array([float(x), float(y), 1.0], dtype=np.float32)
q = H @ p
if abs(q[2]) < 1e-9:
return float(q[0]), float(q[1])
return float(q[0] / q[2]), float(q[1] / q[2])
def quad_for_output_crop_to_input(H_src_to_rgb, crop_box):
"""
A OAK/ImageManip recebe uma quadrilateral SOURCE que será esticada
para o retângulo de saída.
Como nossa homografia original é source -> RGB, usamos H_inv para descobrir:
canto final no espaço RGB -> ponto correspondente no source.
"""
x0, y0, x1, y1 = crop_box
dst_corners_rgb = [
(x0, y0), # top-left
(x1, y0), # top-right
(x1, y1), # bottom-right
(x0, y1), # bottom-left
]
H_inv = np.linalg.inv(H_src_to_rgb).astype(np.float32)
src_quad = []
for x, y in dst_corners_rgb:
sx, sy = apply_H_to_point(H_inv, x, y)
src_quad.append((sx, sy))
return src_quad
def identity_quad_for_crop(crop_box):
x0, y0, x1, y1 = crop_box
return [
(float(x0), float(y0)),
(float(x1), float(y0)),
(float(x1), float(y1)),
(float(x0), float(y1)),
]
def dai_points(quad):
return [dai.Point2f(float(x), float(y)) for x, y in quad]
# ============================================================
# DepthAI pipeline helpers
# ============================================================
def create_color_camera(pipeline, socket, width, height, fps):
cam = pipeline.create(dai.node.ColorCamera)
cam.setBoardSocket(socket)
# OV9782 costuma ser 800p. Se sua versão não aceitar, ajuste para a resolução suportada.
cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_800_P)
cam.setFps(float(fps))
cam.setInterleaved(False)
cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
cam.setVideoSize(int(width), int(height))
return cam
def create_mono_camera(pipeline, socket, width, height, fps):
cam = pipeline.create(dai.node.MonoCamera)
cam.setBoardSocket(socket)
# OV9282 normalmente suporta 800p.
cam.setResolution(dai.MonoCameraProperties.SensorResolution.THE_800_P)
cam.setFps(float(fps))
return cam
def create_warp_manip(
pipeline,
name,
out_w,
out_h,
src_quad_px,
frame_type=None,
max_output_frame_size=None,
):
manip = pipeline.create(dai.node.ImageManip)
dst_quad_px = [
(0.0, 0.0),
(float(out_w - 1), 0.0),
(float(out_w - 1), float(out_h - 1)),
(0.0, float(out_h - 1)),
]
# Faz crop/warp/resize em uma operação:
# src_quad_px na imagem original -> dst_quad_px na imagem final.
manip.initialConfig.addTransformFourPoints(
dai_points(src_quad_px),
dai_points(dst_quad_px),
False,
)
# NÃO usar setResize nessa versão do DepthAI.
# A saída já deve sair com o tamanho definido pelo dst_quad_px.
manip.setMaxOutputFrameSize(
int(max_output_frame_size)
if max_output_frame_size is not None
else int(out_w * out_h * 3)
)
if frame_type is not None:
manip.initialConfig.setFrameType(frame_type)
xout = pipeline.createXLinkOut()
xout.setStreamName(name)
manip.out.link(xout.input)
return manip, xout
def drain_latest(queue):
"""
Pega o frame mais novo disponível sem acumular fila.
"""
latest = None
while True:
msg = queue.tryGet()
if msg is None:
break
latest = msg
return latest
def clamp_quad(quad, w, h):
out = []
for x, y in quad:
x = max(0.0, min(float(w - 1), float(x)))
y = max(0.0, min(float(h - 1), float(y)))
out.append((x, y))
return out
# ============================================================
# Main benchmark
# ============================================================
def main():
ap = argparse.ArgumentParser(
description="Teste OAK-FFC-3: resize/crop/homografia na OAK e envio de frames alinhados."
)
ap.add_argument("--mxid", default=None)
ap.add_argument("--module_params", required=True)
ap.add_argument("--fps", type=float, default=20.0)
ap.add_argument("--sensor_w", type=int, default=1280)
ap.add_argument("--sensor_h", type=int, default=800)
ap.add_argument("--out_w", type=int, default=1024)
ap.add_argument("--out_h", type=int, default=640)
ap.add_argument("--seconds", type=float, default=20.0)
ap.add_argument("--save_debug", action="store_true")
ap.add_argument("--debug_dir", default="oak_aligned_debug")
args = ap.parse_args()
params = load_module_params(args.module_params)
fusion = params.get("fusion_config", {}) or {}
homographies = fusion.get("homographies", {}) or {}
H_re_raw = homographies.get("re_to_rgb")
H_nir_raw = homographies.get("nir_to_rgb")
if H_re_raw is None or H_nir_raw is None:
raise RuntimeError("module_params precisa conter fusion_config.homographies.re_to_rgb e nir_to_rgb")
calib_size = fusion.get("homography_calibration_size", None)
runtime_size = (args.sensor_w, args.sensor_h)
H_re = scale_homography_to_runtime(H_re_raw, calib_size, runtime_size)
H_nir = scale_homography_to_runtime(H_nir_raw, calib_size, runtime_size)
crop_box = compute_common_crop_box(args.sensor_w, args.sensor_h, H_re, H_nir)
quad_rgb = identity_quad_for_crop(crop_box)
quad_re = quad_for_output_crop_to_input(H_re, crop_box)
quad_nir = quad_for_output_crop_to_input(H_nir, crop_box)
quad_rgb = clamp_quad(quad_rgb, args.sensor_w, args.sensor_h)
quad_re = clamp_quad(quad_re, args.sensor_w, args.sensor_h)
quad_nir = clamp_quad(quad_nir, args.sensor_w, args.sensor_h)
print("============================================")
print("OAK-FCC-3 aligned geometry benchmark")
print(f"module_params : {args.module_params}")
print(f"sensor : {args.sensor_w}x{args.sensor_h}")
print(f"output : {args.out_w}x{args.out_h}")
print(f"fps target : {args.fps}")
print(f"crop_box RGB : {crop_box}")
print(f"quad_rgb : {quad_rgb}")
print(f"quad_re : {quad_re}")
print(f"quad_nir : {quad_nir}")
print("============================================")
pipeline = dai.Pipeline()
# Sockets assumidos:
# CAM_A = RGB, CAM_B = RE, CAM_C = NIR
rgb = create_color_camera(
pipeline,
dai.CameraBoardSocket.CAM_A,
args.sensor_w,
args.sensor_h,
args.fps,
)
re = create_mono_camera(
pipeline,
dai.CameraBoardSocket.CAM_B,
args.sensor_w,
args.sensor_h,
args.fps,
)
nir = create_mono_camera(
pipeline,
dai.CameraBoardSocket.CAM_C,
args.sensor_w,
args.sensor_h,
args.fps,
)
# Tamanhos aproximados de saída.
# RGB BGR888p = 3 canais. Mono GRAY8 = 1 canal.
rgb_manip, _ = create_warp_manip(
pipeline,
"rgb_aligned",
args.out_w,
args.out_h,
quad_rgb,
frame_type=dai.ImgFrame.Type.BGR888p,
max_output_frame_size=args.out_w * args.out_h * 3,
)
re_manip, _ = create_warp_manip(
pipeline,
"re_aligned",
args.out_w,
args.out_h,
quad_re,
frame_type=dai.ImgFrame.Type.GRAY8,
max_output_frame_size=args.out_w * args.out_h,
)
nir_manip, _ = create_warp_manip(
pipeline,
"nir_aligned",
args.out_w,
args.out_h,
quad_nir,
frame_type=dai.ImgFrame.Type.GRAY8,
max_output_frame_size=args.out_w * args.out_h,
)
# Links.
# RGB usa video, RE/NIR usam mono out.
rgb.video.link(rgb_manip.inputImage)
re.out.link(re_manip.inputImage)
nir.out.link(nir_manip.inputImage)
device_info = None
if args.mxid:
device_info = dai.DeviceInfo(args.mxid)
if args.save_debug:
Path(args.debug_dir).mkdir(parents=True, exist_ok=True)
with dai.Device(pipeline, device_info) if device_info is not None else dai.Device(pipeline) as device:
q_rgb = device.getOutputQueue("rgb_aligned", maxSize=2, blocking=False)
q_re = device.getOutputQueue("re_aligned", maxSize=2, blocking=False)
q_nir = device.getOutputQueue("nir_aligned", maxSize=2, blocking=False)
t_start = time.perf_counter()
t_last_log = t_start
n_triplets = 0
n_rgb = 0
n_re = 0
n_nir = 0
last_rgb = None
last_re = None
last_nir = None
lat_samples_ms = []
stack_samples_ms = []
saved = False
while True:
now = time.perf_counter()
if now - t_start >= args.seconds:
break
msg_rgb = drain_latest(q_rgb)
msg_re = drain_latest(q_re)
msg_nir = drain_latest(q_nir)
if msg_rgb is not None:
last_rgb = msg_rgb
n_rgb += 1
if msg_re is not None:
last_re = msg_re
n_re += 1
if msg_nir is not None:
last_nir = msg_nir
n_nir += 1
if last_rgb is not None and last_re is not None and last_nir is not None:
# Mede latência aproximada usando timestamps do dispositivo.
try:
ts_dev = [
last_rgb.getTimestampDevice().total_seconds(),
last_re.getTimestampDevice().total_seconds(),
last_nir.getTimestampDevice().total_seconds(),
]
sync_dt_ms = (max(ts_dev) - min(ts_dev)) * 1000.0
lat_samples_ms.append(sync_dt_ms)
except Exception:
sync_dt_ms = -1.0
# Mede custo de converter e empilhar no PC.
t_stack0 = time.perf_counter()
rgb_bgr = last_rgb.getCvFrame()
re_gray = last_re.getCvFrame()
nir_gray = last_nir.getCvFrame()
# Só para simular o stack final.
# RGB vem BGR, então convertemos para RGB.
rgb_float = cv2.cvtColor(rgb_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
re_float = re_gray.astype(np.float32) / 255.0
nir_float = nir_gray.astype(np.float32) / 255.0
tensor5 = np.empty((5, args.out_h, args.out_w), dtype=np.float32)
tensor5[0] = rgb_float[:, :, 0]
tensor5[1] = rgb_float[:, :, 1]
tensor5[2] = rgb_float[:, :, 2]
tensor5[3] = re_float
tensor5[4] = nir_float
stack_ms = (time.perf_counter() - t_stack0) * 1000.0
stack_samples_ms.append(stack_ms)
n_triplets += 1
if args.save_debug and not saved:
cv2.imwrite(str(Path(args.debug_dir) / "rgb_aligned.png"), rgb_bgr)
cv2.imwrite(str(Path(args.debug_dir) / "re_aligned.png"), re_gray)
cv2.imwrite(str(Path(args.debug_dir) / "nir_aligned.png"), nir_gray)
np.save(str(Path(args.debug_dir) / "tensor5_sample.npy"), tensor5)
saved = True
# Marca como consumidos para não recontar infinitamente o mesmo trio.
last_rgb = None
last_re = None
last_nir = None
if now - t_last_log >= 1.0:
elapsed = now - t_start
fps_triplets = n_triplets / max(elapsed, 1e-6)
fps_rgb = n_rgb / max(elapsed, 1e-6)
fps_re = n_re / max(elapsed, 1e-6)
fps_nir = n_nir / max(elapsed, 1e-6)
sync_mean = float(np.mean(lat_samples_ms[-30:])) if lat_samples_ms else -1.0
sync_max = float(np.max(lat_samples_ms[-30:])) if lat_samples_ms else -1.0
stack_mean = float(np.mean(stack_samples_ms[-30:])) if stack_samples_ms else -1.0
stack_max = float(np.max(stack_samples_ms[-30:])) if stack_samples_ms else -1.0
print(
"[OAK_ALIGN_PERF] "
f"elapsed={elapsed:.1f}s "
f"triplets={n_triplets} "
f"fps_triplets={fps_triplets:.2f} "
f"fps_rgb={fps_rgb:.2f} "
f"fps_re={fps_re:.2f} "
f"fps_nir={fps_nir:.2f} "
f"sync_mean_ms={sync_mean:.2f} "
f"sync_max_ms={sync_max:.2f} "
f"stack_mean_ms={stack_mean:.2f} "
f"stack_max_ms={stack_max:.2f} "
f"tensor_shape={(5, args.out_h, args.out_w)}"
)
t_last_log = now
time.sleep(0.001)
elapsed = time.perf_counter() - t_start
print("============================================")
print("RESULTADO FINAL")
print(f"elapsed : {elapsed:.2f}s")
print(f"triplets : {n_triplets}")
print(f"fps_triplets : {n_triplets / max(elapsed, 1e-6):.2f}")
print(f"rgb frames : {n_rgb} | fps={n_rgb / max(elapsed, 1e-6):.2f}")
print(f"re frames : {n_re} | fps={n_re / max(elapsed, 1e-6):.2f}")
print(f"nir frames : {n_nir} | fps={n_nir / max(elapsed, 1e-6):.2f}")
if lat_samples_ms:
print(f"sync mean/max : {np.mean(lat_samples_ms):.2f} / {np.max(lat_samples_ms):.2f} ms")
if stack_samples_ms:
print(f"stack mean/max: {np.mean(stack_samples_ms):.2f} / {np.max(stack_samples_ms):.2f} ms")
print("============================================")
if __name__ == "__main__":
main()

View File

@ -248,8 +248,11 @@ def build_multispec_from_raw_native_multi(group: dict, meta: dict):
if alt.exists():
calib_path = str(alt)
else:
print(f"[WARN] module_params não encontrado: {calib_path}. Tentando sem calibração.")
calib_path = None
raise RuntimeError(
f"module_params.json não encontrado: {calib_path}. "
f"Não vou gerar MULTISPEC offline sem fusion_config, "
f"porque isso deixaria RE/NIR desalinhados do RGB."
)
core = RawProcessorCore(
sensor_width=sensor_width,
@ -271,6 +274,9 @@ def build_multispec_from_raw_native_multi(group: dict, meta: dict):
tensor = core.build_infer_tensor_from_stream(frame, processing_meta, 5)
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),
}
@ -1135,7 +1141,7 @@ def main():
client = OakFcc3Client(
width=1280,
height=800,
bayer="RGGB",
bayer="BGGR",
frame_type="RAW_BRUTO",
capture_mode="SINGLE",
raw_policy="allow_single",

View File

@ -0,0 +1,708 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
raw_tensor_calibration_tuner.py
Ferramenta visual para calibrar o tensor científico a partir do RAW_BRUTO salvo.
Ela :
<grupo>/bins/<base>_CAM_A.bin
<grupo>/bins/<base>_CAM_B.bin
<grupo>/bins/<base>_CAM_C.bin
<grupo>/metas/<base>.json
E monta o tensor final usando RawProcessorCore + module_params.json,
passando pelo mesmo fluxo real do normalize/runtime:
RAW10 -> linear_demosaic/bayer_planes -> rgb_calibration
-> radiometric_normalization -> homography/crop/resize
-> flatfield final_tensor_space -> tensor [R,G,B,RE,NIR]
Mostra:
esquerda : RGB do tensor base, com module_params original
centro : RGB do tensor ajustado pelos sliders
direita : diff x4 entre base e ajustado
Também mostra p50/mean/p01/p99/std de cada canal do tensor ajustado.
Uso:
python raw_tensor_calibration_tuner.py ^
--src_dir dataset/original/group/chao ^
--module_params calibration/module_params.json ^
--res 1024x640 ^
--raw_size 1280x800
Controles:
A / seta esquerda -> amostra anterior
D / seta direita -> próxima amostra
S -> salva patch JSON
P -> imprime patch JSON
R -> reseta sliders para valores iniciais do module_params
Q / ESC -> sair
Notas:
- Sliders flat_* usam escala 1000 = strength 1.0.
Para flat-field, neutro é 0.
- Sliders rgb_gain_* usam escala 1000 = ganho 1.0.
Para RGB gain, neutro é 1000.
- auto_level/gamma afetam apenas a visualização, não o tensor.
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Tuple, List
import cv2
import numpy as np
# ============================================================
# Import do core
# ============================================================
def add_project_paths():
here = Path.cwd()
candidates = [
here,
here / "Python" / "Scripts" / "workers" / "camera_worker" / "oak_fcc3_core",
here / "Scripts" / "workers" / "camera_worker" / "oak_fcc3_core",
]
for c in candidates:
if c.exists() and str(c) not in sys.path:
sys.path.insert(0, str(c))
add_project_paths()
try:
from core.raw_processor_core import RawProcessorCore
except Exception as e:
raise SystemExit(
"[ERRO] Não consegui importar core.raw_processor_core.RawProcessorCore.\n"
"Rode este script a partir da pasta do projeto onde existe Python/Scripts/workers/camera_worker/oak_fcc3_core,\n"
"ou coloque este script dentro da pasta oak_fcc3_core.\n"
f"Erro original: {type(e).__name__}: {e}"
)
CHANNEL_NAMES = ["R", "G", "B", "RE", "NIR"]
# ============================================================
# IO helpers
# ============================================================
def load_json(path: Path) -> dict:
with path.open("r", encoding="utf-8") as f:
return json.load(f)
def save_json(path: Path, data: dict):
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2)
def parse_size(txt: str) -> Tuple[int, int]:
w, h = [int(x) for x in str(txt).lower().split("x")]
return w, h
@dataclass
class SampleBundle:
base: str
meta_path: Path
bin_paths: Dict[str, Path]
def find_sample_bins(bins_dir: Path, base: str) -> Dict[str, Path]:
out = {}
for cam_id in ("CAM_A", "CAM_B", "CAM_C"):
p = bins_dir / f"{base}_{cam_id}.bin"
if p.exists():
out[cam_id] = p
return out
def collect_samples(src_dir: Path) -> List[SampleBundle]:
metas_dir = src_dir / "metas"
bins_dir = src_dir / "bins"
if not metas_dir.is_dir() or not bins_dir.is_dir():
raise RuntimeError(
f"src_dir precisa conter bins/ e metas/. Veio: {src_dir}"
)
samples = []
for meta_path in sorted(metas_dir.glob("*.json")):
base = meta_path.stem
bin_paths = find_sample_bins(bins_dir, base)
if not all(cam in bin_paths for cam in ("CAM_A", "CAM_B", "CAM_C")):
print(f"[WARN] Pulando {base}: bins incompletos {list(bin_paths.keys())}")
continue
samples.append(SampleBundle(base=base, meta_path=meta_path, bin_paths=bin_paths))
if not samples:
raise RuntimeError(f"Nenhuma amostra válida encontrada em {src_dir}")
return samples
def load_frame_from_saved_bins(sample: SampleBundle, meta: dict) -> Dict[str, np.ndarray]:
saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
saved_shapes = meta.get("saved_payload_shapes", {}) or {}
frame = {}
for cam_id, bin_path in sample.bin_paths.items():
dtype = saved_dtypes.get(cam_id)
shape = saved_shapes.get(cam_id)
if dtype is None or shape is None:
raise RuntimeError(f"Faltam saved_payload_dtypes/shapes para {cam_id} em {sample.base}")
arr = np.fromfile(str(bin_path), dtype=np.dtype(dtype)).reshape(tuple(shape))
frame[cam_id] = arr
return frame
def build_processing_meta(meta: dict) -> dict:
stream_meta = dict(meta.get("stream_meta", {}) or {})
stream_meta["frame_type"] = "RAW_BRUTO"
if "camera_info" not in stream_meta and isinstance(meta.get("camera_info"), dict):
stream_meta["camera_info"] = meta.get("camera_info")
# Mantém compatibilidade com metadados antigos e novos.
for key in (
"actual_camera_controls",
"startup_camera_controls",
"camera_controls",
"frame_controls",
"camera_frames",
):
if meta.get(key) is not None:
stream_meta[key] = meta.get(key)
return stream_meta
def tensor_stats(tensor: np.ndarray) -> Dict[str, dict]:
out = {}
for i, ch in enumerate(CHANNEL_NAMES):
if i >= tensor.shape[0]:
break
vals = tensor[i].astype(np.float32).reshape(-1)
vals = vals[np.isfinite(vals)]
if vals.size == 0:
out[ch] = {}
continue
out[ch] = {
"min": float(np.min(vals)),
"p01": float(np.percentile(vals, 1)),
"p50": float(np.percentile(vals, 50)),
"p99": float(np.percentile(vals, 99)),
"max": float(np.max(vals)),
"mean": float(np.mean(vals)),
"std": float(np.std(vals)),
}
return out
# ============================================================
# Config adjustment
# ============================================================
def deep_merge(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] = copy.deepcopy(v)
rec(out, src)
return out
def load_initial_values(module_params: dict):
ff = module_params.get("flatfield_config", {}) or {}
rgbc = module_params.get("rgb_calibration", {}) or {}
strengths = {ch: float(ff.get("strength", 0.0)) for ch in CHANNEL_NAMES}
for k, v in (ff.get("strength_by_channel", {}) or {}).items():
ku = str(k).upper()
if ku in strengths:
strengths[ku] = float(v)
# Contrato correto do RawProcessorCore: rgb_calibration.gains
gains = {"R": 1.0, "G": 1.0, "B": 1.0}
src = rgbc.get("gains", {}) if isinstance(rgbc, dict) else {}
if isinstance(src, dict):
for ch in ("R", "G", "B"):
if ch in src:
gains[ch] = float(src[ch])
# Compatibilidade com patch antigo do tuner.
src2 = rgbc.get("gain_by_channel", {}) if isinstance(rgbc, dict) else {}
if isinstance(src2, dict):
for ch in ("R", "G", "B"):
if ch in src2:
gains[ch] = float(src2[ch])
return strengths, gains
def make_patch(strengths: Dict[str, float], rgb_gains: Dict[str, float]) -> dict:
return {
"flatfield_config": {
"strength_by_channel": {
"R": float(strengths["R"]),
"G": float(strengths["G"]),
"B": float(strengths["B"]),
"RE": float(strengths["RE"]),
"NIR": float(strengths["NIR"]),
}
},
"rgb_calibration": {
"enabled": True,
"gains": {
"R": float(rgb_gains["R"]),
"G": float(rgb_gains["G"]),
"B": float(rgb_gains["B"]),
}
}
}
def apply_patch_to_core(core: RawProcessorCore, patch: dict):
# Flatfield.
ff_patch = patch.get("flatfield_config", {}) or {}
if isinstance(ff_patch, dict):
core.flatfield_config = deep_merge(core.flatfield_config, ff_patch)
# RGB calibration.
rgb_patch = patch.get("rgb_calibration", {}) or {}
if isinstance(rgb_patch, dict):
core.rgb_calibration = deep_merge(core.rgb_calibration, rgb_patch)
# Limpa caches cujo valor depende do flat gain/strength.
if hasattr(core, "_flatfield_runtime_cache"):
core._flatfield_runtime_cache.clear()
# Geometria/remap não depende do strength nem dos rgb gains,
# então não precisa limpar.
def create_core(raw_size: Tuple[int, int], bayer_pattern: str, module_params: Path) -> RawProcessorCore:
raw_w, raw_h = raw_size
core = RawProcessorCore(
sensor_width=int(raw_w),
sensor_height=int(raw_h),
bayer_pattern=str(bayer_pattern or "RGGB").upper(),
calibration_json_path=str(module_params),
)
return core
# ============================================================
# Tensor generation
# ============================================================
class TensorBuilder:
def __init__(self, module_params: Path, raw_size: Tuple[int, int], res: Tuple[int, int]):
self.module_params = module_params
self.raw_size = raw_size
self.res = res
self.module_params_data = load_json(module_params)
self.base_core = None
self.tune_core = None
self.current_bayer = None
self._sample_cache = {}
self._base_tensor_cache = {}
def ensure_cores(self, bayer_pattern: str):
bayer = str(bayer_pattern or "RGGB").upper()
if self.base_core is not None and self.current_bayer == bayer:
return
print(f"[CORE] criando cores bayer={bayer} raw={self.raw_size} res={self.res}")
self.base_core = create_core(self.raw_size, bayer, self.module_params)
self.tune_core = create_core(self.raw_size, bayer, self.module_params)
self.current_bayer = bayer
def get_frame_meta(self, sample: SampleBundle):
key = sample.base
cached = self._sample_cache.get(key)
if cached is not None:
return cached
meta = load_json(sample.meta_path)
frame = load_frame_from_saved_bins(sample, meta)
processing_meta = build_processing_meta(meta)
bayer = str(meta.get("bayer_pattern") or meta.get("bayer") or "RGGB").upper()
cached = (frame, processing_meta, bayer, meta)
self._sample_cache[key] = cached
return cached
def build_base_tensor(self, sample: SampleBundle):
if sample.base in self._base_tensor_cache:
return self._base_tensor_cache[sample.base]
frame, processing_meta, bayer, _meta = self.get_frame_meta(sample)
self.ensure_cores(bayer)
tensor = self.base_core.build_infer_tensor_from_stream(
frame=frame,
meta=processing_meta,
channels_expected=5,
target_size=self.res,
)
tensor = np.ascontiguousarray(tensor.astype(np.float32, copy=False))
self._base_tensor_cache[sample.base] = tensor
return tensor
def build_tuned_tensor(self, sample: SampleBundle, strengths: Dict[str, float], rgb_gains: Dict[str, float]):
frame, processing_meta, bayer, _meta = self.get_frame_meta(sample)
self.ensure_cores(bayer)
patch = make_patch(strengths, rgb_gains)
apply_patch_to_core(self.tune_core, patch)
tensor = self.tune_core.build_infer_tensor_from_stream(
frame=frame,
meta=processing_meta,
channels_expected=5,
target_size=self.res,
)
return np.ascontiguousarray(tensor.astype(np.float32, copy=False))
# ============================================================
# Display
# ============================================================
def val_to_gain(v: int) -> float:
return float(v) / 1000.0
def gain_to_val(g: float) -> int:
return int(round(max(0.0, min(3.0, float(g))) * 1000.0))
def val_to_strength(v: int) -> float:
return float(v) / 1000.0
def strength_to_val(s: float) -> int:
return int(round(max(0.0, min(2.0, float(s))) * 1000.0))
def read_ui(win: str):
strengths = {
"R": val_to_strength(cv2.getTrackbarPos("flat_R", win)),
"G": val_to_strength(cv2.getTrackbarPos("flat_G", win)),
"B": val_to_strength(cv2.getTrackbarPos("flat_B", win)),
"RE": val_to_strength(cv2.getTrackbarPos("flat_RE", win)),
"NIR": val_to_strength(cv2.getTrackbarPos("flat_NIR", win)),
}
gains = {
"R": val_to_gain(cv2.getTrackbarPos("rgb_gain_R", win)),
"G": val_to_gain(cv2.getTrackbarPos("rgb_gain_G", win)),
"B": val_to_gain(cv2.getTrackbarPos("rgb_gain_B", win)),
}
auto_level = bool(cv2.getTrackbarPos("auto_level", win))
gamma = max(0.05, cv2.getTrackbarPos("gamma_x100", win) / 100.0)
return strengths, gains, auto_level, gamma
def set_ui(win: str, strengths: Dict[str, float], gains: Dict[str, float], gamma: float):
for ch in CHANNEL_NAMES:
cv2.setTrackbarPos(f"flat_{ch}", win, strength_to_val(strengths.get(ch, 0.0)))
for ch in ("R", "G", "B"):
cv2.setTrackbarPos(f"rgb_gain_{ch}", win, gain_to_val(gains.get(ch, 1.0)))
cv2.setTrackbarPos("gamma_x100", win, int(round(gamma * 100)))
def preview_rgb_from_tensor(tensor: np.ndarray, auto_level: bool, gamma: float) -> np.ndarray:
rgb = np.transpose(tensor[:3].astype(np.float32), (1, 2, 0))
rgb = np.nan_to_num(rgb, nan=0.0, posinf=1.0, neginf=0.0)
if auto_level:
# Auto-level visual separado por canal para facilitar inspeção.
out = np.empty_like(rgb)
for c in range(3):
ch = rgb[..., c]
lo = float(np.percentile(ch, 1.0))
hi = float(np.percentile(ch, 99.0))
if hi > lo:
out[..., c] = (ch - lo) / (hi - lo)
else:
out[..., c] = ch
rgb = out
rgb = np.clip(rgb, 0.0, 1.0)
if gamma > 0 and abs(gamma - 1.0) > 1e-6:
rgb = np.power(rgb, gamma)
return (rgb * 255.0).clip(0, 255).astype(np.uint8)
def put_text(img: np.ndarray, lines, x=12, y=25, scale=0.55):
out = img.copy()
yy = y
for line in lines:
cv2.putText(out, str(line), (x, yy), cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), 4, cv2.LINE_AA)
cv2.putText(out, str(line), (x, yy), cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 255, 255), 1, cv2.LINE_AA)
yy += int(22 * scale / 0.55)
return out
def resize_keep_width(img: np.ndarray, target_w: int) -> np.ndarray:
h, w = img.shape[:2]
if w <= target_w:
return img
scale = target_w / float(w)
return cv2.resize(img, (target_w, int(h * scale)), interpolation=cv2.INTER_AREA)
def stats_lines(stats: Dict[str, dict], prefix=""):
lines = []
for ch in CHANNEL_NAMES:
s = stats.get(ch, {}) or {}
if not s:
continue
lines.append(
f"{prefix}{ch}: p50={s['p50']:.4f} mean={s['mean']:.4f} "
f"p01={s['p01']:.4f} p99={s['p99']:.4f} std={s['std']:.4f}"
)
return lines
def build_canvas(
sample: SampleBundle,
idx: int,
total: int,
base_tensor: np.ndarray,
tuned_tensor: np.ndarray,
strengths: Dict[str, float],
gains: Dict[str, float],
auto_level: bool,
gamma: float,
max_width: int,
):
base_rgb = preview_rgb_from_tensor(base_tensor, auto_level=auto_level, gamma=gamma)
tuned_rgb = preview_rgb_from_tensor(tuned_tensor, auto_level=auto_level, gamma=gamma)
diff = np.abs(tuned_rgb.astype(np.int16) - base_rgb.astype(np.int16)).clip(0, 255).astype(np.uint8)
diff = np.clip(diff.astype(np.uint16) * 4, 0, 255).astype(np.uint8)
h = base_rgb.shape[0]
sep = np.zeros((h, 4, 3), dtype=np.uint8)
canvas = np.hstack([base_rgb, sep, tuned_rgb, sep, diff])
st_base = tensor_stats(base_tensor)
st_tuned = tensor_stats(tuned_tensor)
top_lines = [
f"{idx + 1}/{total} | {sample.base}",
"LEFT=base tensor | MID=tuned tensor | RIGHT=diff x4",
f"flat R={strengths['R']:.3f} G={strengths['G']:.3f} B={strengths['B']:.3f} RE={strengths['RE']:.3f} NIR={strengths['NIR']:.3f}",
f"rgb_gain R={gains['R']:.3f} G={gains['G']:.3f} B={gains['B']:.3f} | auto={int(auto_level)} gamma={gamma:.2f}",
"A/D navega | S salva patch | P print | R reset | Q sai",
]
canvas = put_text(canvas, top_lines, x=12, y=24, scale=0.56)
# Painel de stats na parte inferior, com fundo escuro.
panel_h = 190
panel = np.zeros((panel_h, canvas.shape[1], 3), dtype=np.uint8)
stat_lines = ["TUNED tensor stats:"] + stats_lines(st_tuned, "")
# Também mostra relação G/R e G/B, útil para diagnosticar esverdeado.
try:
gr = st_tuned["G"]["mean"] / max(st_tuned["R"]["mean"], 1e-6)
gb = st_tuned["G"]["mean"] / max(st_tuned["B"]["mean"], 1e-6)
stat_lines.append(f"RGB balance mean: G/R={gr:.3f} | G/B={gb:.3f}")
except Exception:
pass
panel = put_text(panel, stat_lines[:8], x=12, y=24, scale=0.58)
full = np.vstack([canvas, panel])
return resize_keep_width(full, max_width)
# ============================================================
# Main
# ============================================================
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src_dir", required=True, help="Grupo original contendo bins/ e metas/. Ex: dataset/original/group/chao")
ap.add_argument("--module_params", required=True, help="calibration/module_params.json")
ap.add_argument("--res", default="1024x640", help="Resolução final WxH do tensor")
ap.add_argument("--raw_size", default="1280x800", help="RAW size WxH")
ap.add_argument("--out_patch", default="raw_tensor_tuned_patch.json")
ap.add_argument("--max_width", type=int, default=1900)
ap.add_argument("--gamma", type=float, default=1.0)
args = ap.parse_args()
src_dir = Path(args.src_dir)
module_params = Path(args.module_params)
if not src_dir.is_dir():
raise SystemExit(f"[ERRO] src_dir não existe: {src_dir}")
if not module_params.is_file():
raise SystemExit(f"[ERRO] module_params não existe: {module_params}")
res = parse_size(args.res)
raw_size = parse_size(args.raw_size)
module_data = load_json(module_params)
initial_strengths, initial_gains = load_initial_values(module_data)
samples = collect_samples(src_dir)
builder = TensorBuilder(
module_params=module_params,
raw_size=raw_size,
res=res,
)
win = "RAW Tensor Calibration Tuner"
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
for ch in CHANNEL_NAMES:
cv2.createTrackbar(f"flat_{ch}", win, 0, 2000, lambda v: None)
for ch in ("R", "G", "B"):
cv2.createTrackbar(f"rgb_gain_{ch}", win, 1000, 3000, lambda v: None)
cv2.createTrackbar("auto_level", win, 0, 1, lambda v: None)
cv2.createTrackbar("gamma_x100", win, 100, 300, lambda v: None)
set_ui(win, initial_strengths, initial_gains, args.gamma)
idx = 0
last_key = None
last_tuned_key = None
base_tensor = None
tuned_tensor = None
print("[INFO] RAW Tensor Calibration Tuner iniciado.")
print("[INFO] O tensor base usa module_params original. O tensor tuned usa sliders.")
print("[INFO] S salva patch JSON. P imprime valores. R reseta.")
print("[INFO] Se ficar lento ao arrastar sliders, solte o mouse: ele recalcula o tensor real a cada mudança.")
while True:
sample = samples[idx]
strengths, gains, auto_level, gamma = read_ui(win)
tuned_key = (
sample.base,
tuple(round(strengths[ch], 4) for ch in CHANNEL_NAMES),
tuple(round(gains[ch], 4) for ch in ("R", "G", "B")),
)
try:
if last_key != sample.base:
t0 = time.perf_counter()
base_tensor = builder.build_base_tensor(sample)
tuned_tensor = None
last_tuned_key = None
last_key = sample.base
print(f"[LOAD] {sample.base} base_tensor={base_tensor.shape} {((time.perf_counter()-t0)*1000):.1f}ms")
if last_tuned_key != tuned_key or tuned_tensor is None:
t0 = time.perf_counter()
tuned_tensor = builder.build_tuned_tensor(sample, strengths, gains)
last_tuned_key = tuned_key
print(f"[TUNE] {sample.base} tuned_tensor={tuned_tensor.shape} {((time.perf_counter()-t0)*1000):.1f}ms")
canvas = build_canvas(
sample=sample,
idx=idx,
total=len(samples),
base_tensor=base_tensor,
tuned_tensor=tuned_tensor,
strengths=strengths,
gains=gains,
auto_level=auto_level,
gamma=gamma,
max_width=args.max_width,
)
cv2.imshow(win, cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))
except Exception as e:
err = np.zeros((600, 1200, 3), dtype=np.uint8)
err = put_text(err, [
f"ERRO em {sample.base}",
f"{type(e).__name__}: {e}",
"A/D navega | Q sai",
], x=20, y=40, scale=0.7)
cv2.imshow(win, cv2.cvtColor(err, cv2.COLOR_RGB2BGR))
k = cv2.waitKey(30) & 0xFF
if k in (ord("q"), ord("Q"), 27):
break
elif k in (ord("d"), ord("D"), 83):
idx = (idx + 1) % len(samples)
last_key = None
last_tuned_key = None
elif k in (ord("a"), ord("A"), 81):
idx = (idx - 1 + len(samples)) % len(samples)
last_key = None
last_tuned_key = None
elif k in (ord("r"), ord("R")):
set_ui(win, initial_strengths, initial_gains, args.gamma)
last_tuned_key = None
elif k in (ord("p"), ord("P")):
patch = make_patch(strengths, gains)
print(json.dumps(patch, ensure_ascii=False, indent=2))
elif k in (ord("s"), ord("S")):
patch = make_patch(strengths, gains)
out = Path(args.out_patch)
save_json(out, patch)
print(f"[SAVE] {out.resolve()}")
print(json.dumps(patch, ensure_ascii=False, indent=2))
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

@ -635,6 +635,7 @@ def main():
offsets_data.setdefault("homographies", {})
offsets_data["schema"] = "manual_multispec_offsets_v2"
offsets_data["reference_camera"] = "rgb"
offsets_data["homography_calibration_size"] = [int(base_w), int(base_h)]
save_offsets_json(args.out_json, offsets_data)
last_msg = f"Offsets salvos em: {args.out_json}"

View File

@ -0,0 +1,92 @@
1. Conectar cartão SD no adaptador e espetar na porta USB
2. Baixar ferramenta de instalação do site
https://www.raspberrypi.com/software/
3. Instalar ferramenta
4. Abrir
5. Configurar Sistema Headless
Choose Device → Raspberry Pi 5
Choose OS → Raspberry Pi OS Other → Raspberry Pi OS Lite 64-bit
Choose Storage → selecione o cartão SD
Clique em Next ou configurações de personalização
Hostname: pi-cam
Usuário: diego
Senha: 1234
Enable SSH → Use password authentication
Country: Brasília - Brasil
Timezone: America/Sao_Paulo
Keyboard: br
Write
6. Remover SD do computador
7. Colocar o SD no Pi
8. Conectar cabo de rede entre Pi e PC
9. Conectar a fonte
10. Espere uns 2 a 5 minutos no primeiro boot
11. Tentar conectar via SSH
ssh diego@pi-cam.local
Se encontrar, digitar yes, e em seguida a senha do usuário para acessar: 1234
12. Atualizar Sistema
sudo apt update
sudo apt full-upgrade -y
sudo reboot
11. Validar configuração do Pi
sudo raspi-config
Interface Options → SSH → Enable
Localisation Options → Timezone → America/Sao_Paulo
System Options → Hostname → pi-cam
12. Conferencias finais
uname -a
cat /etc/os-release
vcgencmd measure_temp
13. Definir IP Fixo (opcional)
sudo nmcli connection modify netplan-eth0 ipv4.addresses 192.168.105.6/24 ipv4.method manual
sudo nmcli connection up netplan-eth0
sudo nmcli device connect eth0
13. Instalar dependencias dev
sudo apt install -y python3 python3-pip python3-venv git
mkdir ~/multispec_module
cd ~/multispec_module
python3 -m venv venv
source venv/bin/activate
pip install numpy opencv-python
sudo apt install -y python3-picamera2 libcamera-apps v4l-utils
14. Conectar no VSCode através do SSH
No VSCode, instalar a extensão Remote - SSH
Apertar Ctrl + Shift + P
Remote-SSH: Connect to Host
+ Add New SSH Host
ssh diego@pi-cam.local
Salvar no config padrão
Ctrl + Shift + P
→ Remote-SSH: Connect to Host
→ pi5-local
Digitar a senha
File → Open Folder
/home/diego/multispec_module
Digitar a senha

View File

@ -216,7 +216,23 @@ class MultiSpectralClient:
def get_next_decoded(self, timeout=2.0, update_radiometry=True):
frame, meta = self.get_next_frame(timeout=timeout)
decoded = self.core.decode_stream_cameras(frame, meta)
frame_type = meta.get("frame_type")
if frame_type == "RAW_BRUTO":
decoded = self.core.decode_stream_cameras(frame, meta)
elif frame_type in ("RGB", "MULTISPEC"):
decoded = {
"cam2": {
"name": "RGB",
"image": frame,
"meta": meta
}
}
else:
raise RuntimeError(f"frame_type não suportado: {frame_type}")
if update_radiometry:
self.update_radiometry(decoded, meta)