diff --git a/Python/OAK/datasets/_12_check_percent_status.py b/Python/OAK/datasets/_12_check_percent_status.py new file mode 100644 index 000000000..11b2f50cb --- /dev/null +++ b/Python/OAK/datasets/_12_check_percent_status.py @@ -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() \ No newline at end of file diff --git a/Python/OAK/datasets/_1_dataset_labeler.py b/Python/OAK/datasets/_1_dataset_labeler.py new file mode 100644 index 000000000..88fef5105 --- /dev/null +++ b/Python/OAK/datasets/_1_dataset_labeler.py @@ -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("", 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() diff --git a/Python/OAK/datasets/_2_create_full_mask.py b/Python/OAK/datasets/_2_create_full_mask.py index 43108692a..e98693e02 100644 --- a/Python/OAK/datasets/_2_create_full_mask.py +++ b/Python/OAK/datasets/_2_create_full_mask.py @@ -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"] diff --git a/Python/OAK/datasets/_4_group_images_by_class.py b/Python/OAK/datasets/_4_group_images_by_class.py index 392a46626..ed388afcc 100644 --- a/Python/OAK/datasets/_4_group_images_by_class.py +++ b/Python/OAK/datasets/_4_group_images_by_class.py @@ -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//images - MODELO/dataset/original/group//masks -Onde é os nomes das classes presentes unidos por "_", ex: - chao, erva, cana, chao_erva, erva_cana, chao_erva_cana, etc. + MODELO/dataset/original/group//images/ + MODELO/dataset/original/group//masks/ + MODELO/dataset/original/group//masks2/ se dual_head_mask=true + MODELO/dataset/original/group//labels/ se dual_head_label=true -Requer: utils.carregar_labelmap_completo(labelmap_path) -O labelmap define mapeamento de cores/ids/nomes das classes. +Onde é 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 lê 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, ) diff --git a/Python/OAK/datasets/_5_augmentation.py b/Python/OAK/datasets/_5_augmentation.py index 99613f26d..1c6d0c95c 100644 --- a/Python/OAK/datasets/_5_augmentation.py +++ b/Python/OAK/datasets/_5_augmentation.py @@ -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//images - MODELO/dataset/original/group//masks +Entrada via config_oak.json -> camera: + + MODELO/dataset/original/group//images/ + MODELO/dataset/original/group//masks/ + MODELO/dataset/original/group//masks2/ opcional, se dual_head_mask=true + MODELO/dataset/original/group//labels/ opcional, se dual_head_label=true Saída: - MODELO/dataset/augmented/group//images - MODELO/dataset/augmented/group//masks -Se "original/group" não existir, faz fallback para: - MODELO/dataset/original/{images,masks} - MODELO/dataset/augmented/{images,masks} + MODELO/dataset/augmented/group//images/ + MODELO/dataset/augmented/group//masks/ + MODELO/dataset/augmented/group//masks2/ se dual_head_mask=true + MODELO/dataset/augmented/group//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//).""" + +# ====================== 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) diff --git a/Python/OAK/datasets/_6_normalize.py b/Python/OAK/datasets/_6_normalize.py index 0590b873e..f9d678cfb 100644 --- a/Python/OAK/datasets/_6_normalize.py +++ b/Python/OAK/datasets/_6_normalize.py @@ -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//{images,masks} - - MODELO/dataset/augmented/group//{images,masks} +Entradas via config.json -> camera, resolucao: + - MODELO/dataset/original/group//{images,masks,(masks2),(labels)} + - MODELO/dataset/augmented/group//{images,masks,(masks2),(labels)} -Saídas (por resolução): - - MODELO/dataset//group//{images,masks} +Saídas por resolução: + - MODELO/dataset//group//{images,masks,(masks2),(labels)} -Fallback (modo legado, se não houver "group/"): - - original/{images,masks} e augmented/{images,masks} -> /{images,masks} +Fallback legado, se não houver group/: + - original/{images,masks,(masks2),(labels)} + - augmented/{images,masks,(masks2),(labels)} + - saída: /{images,masks,(masks2),(labels)} Conversão de máscara: - - Lê 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) + - Lê 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 já é 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 ...//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) diff --git a/Python/OAK/datasets/_7_split.py b/Python/OAK/datasets/_7_split.py index b3d2df071..4a5e47824 100644 --- a/Python/OAK/datasets/_7_split.py +++ b/Python/OAK/datasets/_7_split.py @@ -1,72 +1,88 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ -Split estratificado por GRUPO com **val/test só do ORIGINAL** e -garantia de NÃO VAZAMENTO entre splits (mesma família não cruza splits). +Split estratificado por GRUPO com val/test só do ORIGINAL e garantia de NÃO VAZAMENTO. Lê de: - MODELO/dataset//group//{images,masks} + MODELO/dataset//group//{images,masks,(masks2),(labels)} Escreve em: - MODELO/dataset/split//group//{images,masks} + MODELO/dataset/split//group//{images,masks,(masks2),(labels)} Definições: -- "Família" = todas as variações da MESMA base original: - original_.* e augmented__aug_XX.* -- Val/Test: só **original_** (sem augmented) -- Train: original_ **e** todos augmented__aug_XX +- Família = todas as variações da mesma base original: + original_.* + augmented__aug_XX.* +- Val/Test: somente original_. +- Train: original_ + todos augmented__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() diff --git a/Python/OAK/datasets/_8_train_segformer_b3.py b/Python/OAK/datasets/_8_train_segformer.py similarity index 99% rename from Python/OAK/datasets/_8_train_segformer_b3.py rename to Python/OAK/datasets/_8_train_segformer.py index 020053d5f..533b47e66 100644 --- a/Python/OAK/datasets/_8_train_segformer_b3.py +++ b/Python/OAK/datasets/_8_train_segformer.py @@ -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) diff --git a/Python/OAK/datasets/_8_train_segformer_b3_dual.py b/Python/OAK/datasets/_8_train_segformer_b3_dual.py deleted file mode 100644 index 0a912fa01..000000000 --- a/Python/OAK/datasets/_8_train_segformer_b3_dual.py +++ /dev/null @@ -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": , "mask": , ...} -# 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/) - 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() diff --git a/Python/OAK/datasets/_8_train_segformer_dual.py b/Python/OAK/datasets/_8_train_segformer_dual.py new file mode 100644 index 000000000..58fcd31f9 --- /dev/null +++ b/Python/OAK/datasets/_8_train_segformer_dual.py @@ -0,0 +1,1319 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Treino SegFormer com suporte a: + +1) Single head + - Head principal: segmentação semântica. + +2) Dual head mask + - Head principal: segmentação semântica. + - Head secundária: máscara binária/pixel-a-pixel em masks2/. + - Compatível com a lógica antiga de dual_head. + +3) Dual head label + - Head principal: segmentação semântica. + - Head secundária: classificação global do estado do corredor em labels/*.json ou labels/*.npy. + +Config esperado: + dual_head_mask: bool # nova chave para segunda cabeça de máscara + dual_head_label: bool # nova chave para segunda cabeça de classificação global + +Compatibilidade: + dual_head: bool # se existir e dual_head_mask não existir, será tratado como dual_head_mask + +Exemplo label JSON: + { + "estado_corredor": "CaminhandoRua", + "label_id": 2, + "states": ["Direcionando", "EntrandoRua", "CaminhandoRua", "SaindoRua"] + } + +Uso exemplo: + python _8_train_segformer_dual_mask_or_label.py --config config.json --epochs 80 --batch 4 --amp --amp_val --use_weights + +Label head: + python _8_train_segformer_dual_mask_or_label.py --config config.json --epochs 100 --batch 4 --amp --amp_val --use_weights --auto_label_weights --lambda_label 0.4 --es_metric harmonic_label + +Mask head: + python _8_train_segformer_dual_mask_or_label.py --config config.json --epochs 100 --batch 4 --amp --amp_val --use_weights --auto_pos_weight --lambda_mask2 0.5 +""" + +import os +import re +import time +import json +import math +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 + +from transformers import SegformerForSemanticSegmentation + +from roi_seg_dataset import ROISegDataset + + +# ============================================================ +# Utils gerais +# ============================================================ + +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) +NORM_MEAN = IMAGENET_MEAN.clone() +NORM_STD = IMAGENET_STD.clone() + +def normalize_img(img: torch.Tensor) -> torch.Tensor: + return (img - NORM_MEAN.to(img.device)) / NORM_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] + + if labels.numel() == 0: + return + + 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() + return float((torch.diag(cm).sum() / (cm.sum() + eps)).item()) + + +def _pretty_iou(iou_list: List[float], class_name_by_id: Dict[int, str]) -> str: + return " | ".join(f"{class_name_by_id.get(cid, str(cid))}:{v:.3f}" for cid, v in enumerate(iou_list)) + + +# ============================================================ +# Métricas classificação global +# ============================================================ + +@torch.no_grad() +def update_cls_confusion_matrix(cm: torch.Tensor, logits: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = -100): + preds = torch.argmax(logits, dim=1).view(-1) + labels = labels.view(-1) + valid = labels != ignore_index + preds = preds[valid] + labels = labels[valid] + + if labels.numel() == 0: + return + + 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_cls_metrics_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Dict[str, Any]: + cm = cm.float() + tp = torch.diag(cm) + fp = cm.sum(0) - tp + fn = cm.sum(1) - tp + support = cm.sum(1) + + precision = tp / (tp + fp + eps) + recall = tp / (tp + fn + eps) + f1 = 2 * precision * recall / (precision + recall + eps) + + total = cm.sum() + eps + acc = tp.sum() / total + + valid_classes = support > 0 + if valid_classes.any(): + macro_f1 = f1[valid_classes].mean() + macro_recall = recall[valid_classes].mean() + else: + macro_f1 = torch.tensor(0.0, device=cm.device) + macro_recall = torch.tensor(0.0, device=cm.device) + + return { + "acc": float(acc.item()), + "macro_f1": float(macro_f1.item()), + "macro_recall": float(macro_recall.item()), + "precision_per_class": precision.cpu().tolist(), + "recall_per_class": recall.cpu().tolist(), + "f1_per_class": f1.cpu().tolist(), + "support_per_class": support.cpu().tolist(), + } + + +def _pretty_cls(values: List[float], names: Dict[int, str], title: str) -> str: + parts = [] + for i, v in enumerate(values): + parts.append(f"{names.get(i, str(i))}:{v:.3f}") + return f"{title}: " + " | ".join(parts) + + +# ============================================================ +# Dataset/collate +# ============================================================ + +def ensure_img_tensor(img): + if isinstance(img, np.ndarray): + img = torch.from_numpy(img) + 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 + return img + + +def ensure_mask_tensor(mask): + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(mask) + return mask.long() + + +def collate_seg(batch): + imgs, masks = [], [] + for item in batch: + if isinstance(item, dict): + img = item["image"] + mask = item["mask"] + else: + img, mask = item[:2] + imgs.append(ensure_img_tensor(img)) + masks.append(ensure_mask_tensor(mask)) + return torch.stack(imgs, 0), torch.stack(masks, 0) + + +def collate_mask2(batch): + imgs, m1s, m2s = [], [], [] + for item in batch: + if isinstance(item, dict): + img = item["image"] + mask1 = item["mask1"] + mask2 = item["mask2"] + else: + img, mask1, mask2 = item + imgs.append(ensure_img_tensor(img)) + m1s.append(ensure_mask_tensor(mask1)) + m2s.append(ensure_mask_tensor(mask2)) + return torch.stack(imgs, 0), torch.stack(m1s, 0), torch.stack(m2s, 0) + + +def collate_label(batch): + imgs, masks, labels = [], [], [] + for item in batch: + if isinstance(item, dict): + img = item["image"] + mask = item["mask"] + label = item["label"] + else: + img, mask, label = item + imgs.append(ensure_img_tensor(img)) + masks.append(ensure_mask_tensor(mask)) + labels.append(int(label)) + return torch.stack(imgs, 0), torch.stack(masks, 0), torch.tensor(labels, dtype=torch.long) + + +class DualMaskROISegDataset(torch.utils.data.Dataset): + """Wrapper para segunda cabeça pixel-a-pixel masks2/.""" + 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 + 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: + for k in ("mask_path", "mask_file", "maskname", "mask_name"): + if k in item and item[k]: + return str(item[k]) + if hasattr(self.base_ds, "msk_paths"): + return str(self.base_ds.msk_paths[idx]) + raise RuntimeError("ROISegDataset não expõe mask_path nem base_ds.msk_paths.") + + def _load_mask2_from_mask1_path(self, mask1_path: str, target_shape_hw: Tuple[int, int]) -> torch.Tensor: + 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 = "" + + if (not path2) or (not os.path.exists(path2)): + if not self.has_flat_masks2: + raise FileNotFoundError(f"Mask2 não encontrada para mask1={mask1_path}") + 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] + if isinstance(item, (list, tuple)) and len(item) == 1 and isinstance(item[0], dict): + item = item[0] + if not isinstance(item, dict): + if isinstance(item, (list, tuple)) and len(item) >= 2: + item = {"image": item[0], "mask": item[1]} + else: + raise RuntimeError(f"Formato inesperado do ROISegDataset: {type(item)}") + + 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 + + +class DualLabelROISegDataset(torch.utils.data.Dataset): + """Wrapper para segunda cabeça de classificação global labels/.""" + def __init__(self, base_ds: ROISegDataset, split_root: str, strict_label: bool = True): + self.base_ds = base_ds + self.split_root = split_root + self.strict_label = strict_label + self.label_names_by_id: Dict[int, str] = {} + self.label_ids_found: List[int] = [] + + self.valid_indices = [] + for idx in range(len(base_ds)): + try: + mask_path = self._get_mask_path_from_base(idx) + label_path = self._label_path_from_mask_path(mask_path) + if label_path and os.path.exists(label_path): + label_id, label_name = self._read_label(label_path) + if label_id is not None: + self.valid_indices.append(idx) + self.label_ids_found.append(int(label_id)) + if label_name is not None: + self.label_names_by_id[int(label_id)] = str(label_name) + elif not strict_label: + self.valid_indices.append(idx) + elif not strict_label: + self.valid_indices.append(idx) + except Exception: + if not strict_label: + self.valid_indices.append(idx) + + if strict_label and len(self.valid_indices) == 0: + raise RuntimeError("Nenhuma amostra com label encontrada. Verifique split/.../labels.") + + def __len__(self): + return len(self.valid_indices) if self.strict_label else len(self.base_ds) + + def _real_idx(self, idx): + return self.valid_indices[idx] if self.strict_label else idx + + def _get_mask_path_from_item(self, item: Dict[str, Any], idx: int) -> str: + for k in ("mask_path", "mask_file", "maskname", "mask_name"): + if k in item and item[k]: + return str(item[k]) + if hasattr(self.base_ds, "msk_paths"): + return str(self.base_ds.msk_paths[idx]) + raise RuntimeError("ROISegDataset não expõe mask_path nem base_ds.msk_paths.") + + def _get_mask_path_from_base(self, idx: int) -> str: + if hasattr(self.base_ds, "msk_paths"): + return str(self.base_ds.msk_paths[idx]) + item = self.base_ds[idx] + if isinstance(item, dict): + return self._get_mask_path_from_item(item, idx) + raise RuntimeError("Não foi possível obter mask_path para localizar label.") + + def _label_path_from_mask_path(self, mask_path: str) -> Optional[str]: + norm = os.path.normpath(mask_path) + parts = norm.split(os.sep) + try: + i = parts.index("masks") + parts[i] = "labels" + base = os.path.splitext(os.path.basename(mask_path))[0] + root = os.sep.join(parts[:-1]) + for ext in (".npy", ".json", ".txt"): + cand = os.path.join(root, base + ext) + if os.path.exists(cand): + return cand + return os.path.join(root, base + ".json") + except ValueError: + base = os.path.splitext(os.path.basename(mask_path))[0] + for ext in (".npy", ".json", ".txt"): + cand = os.path.join(self.split_root, "labels", base + ext) + if os.path.exists(cand): + return cand + return None + + def _read_label(self, label_path: str) -> Tuple[Optional[int], Optional[str]]: + ext = os.path.splitext(label_path)[1].lower() + if ext == ".npy": + v = np.load(label_path) + return int(np.array(v).reshape(-1)[0]), None + if ext == ".json": + with open(label_path, "r", encoding="utf-8") as f: + data = json.load(f) + label_id = data.get("label_id") + label_name = data.get("estado_corredor") or data.get("label") or data.get("state") + if label_id is None: + return None, label_name + return int(label_id), label_name + if ext == ".txt": + with open(label_path, "r", encoding="utf-8") as f: + txt = f.read().strip() + try: + return int(txt), None + except Exception: + return None, txt + return None, None + + def __getitem__(self, idx): + real_idx = self._real_idx(idx) + item = self.base_ds[real_idx] + if isinstance(item, (list, tuple)) and len(item) == 1 and isinstance(item[0], dict): + item = item[0] + if not isinstance(item, dict): + if isinstance(item, (list, tuple)) and len(item) >= 2: + item = {"image": item[0], "mask": item[1]} + else: + raise RuntimeError(f"Formato inesperado do ROISegDataset: {type(item)}") + + mask_path = self._get_mask_path_from_item(item, real_idx) + label_path = self._label_path_from_mask_path(mask_path) + label_id = -100 + label_name = None + if label_path and os.path.exists(label_path): + lid, lname = self._read_label(label_path) + if lid is not None: + label_id = int(lid) + label_name = lname + + item["label"] = label_id + item["label_path"] = label_path + if label_name is not None: + item["label_name"] = label_name + return item + + +# ============================================================ +# Heads +# ============================================================ + +class Mask2Head(nn.Module): + """Head para segunda máscara pixel-a-pixel.""" + 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): + """Head para classificação global do frame/estado do corredor.""" + def __init__(self, feat_ch: int, num_seg_classes: int, num_label_classes: int, hidden: int = 256, dropout: float = 0.2): + super().__init__() + in_ch = 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) + + +# ============================================================ +# Loss helpers +# ============================================================ + + +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]: + 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) + pw = float(np.clip(pw, 1.0, 50.0)) + return torch.tensor([pw], dtype=torch.float32) + + +def estimate_label_weights(ds_label: DualLabelROISegDataset, num_label_classes: int, max_samples: int = 5000) -> torch.Tensor: + n = min(len(ds_label), max_samples) + counts = np.zeros(num_label_classes, dtype=np.float64) + for i in range(n): + it = ds_label[i] + lid = int(it.get("label", -100)) + if 0 <= lid < num_label_classes: + counts[lid] += 1 + 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 dice_loss_with_logits(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = 255, eps: float = 1e-6) -> torch.Tensor: + 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 + + +# ============================================================ +# Checkpoint +# ============================================================ + +def save_checkpoint(path: str, base_model: nn.Module, aux_head: Optional[nn.Module], optimizer, scaler, epoch: int, bests: Dict[str, float], extra: Optional[Dict[str, Any]] = None): + ckpt = { + "epoch": epoch, + "model": base_model.state_dict(), + "optimizer": optimizer.state_dict(), + "bests": bests, + } + if aux_head is not None: + ckpt["aux_head"] = aux_head.state_dict() + if scaler is not None: + ckpt["scaler"] = scaler.state_dict() + ckpt["extra"] = extra if extra is not None else {} + torch.save(ckpt, path) + + +def load_checkpoint(path: str, base_model: nn.Module, aux_head: Optional[nn.Module], optimizer=None, scaler=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) + if aux_head is not None and "aux_head" in ckpt: + aux_head.load_state_dict(ckpt["aux_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 + + +# ============================================================ +# Epochs +# ============================================================ + +def get_last_feat(out, logits): + if hasattr(out, "hidden_states") and out.hidden_states is not None: + feat = out.hidden_states[-1] + else: + feat = logits + if feat.shape[-2:] != logits.shape[-2:]: + feat = F.interpolate(feat, size=logits.shape[-2:], mode="bilinear", align_corners=False) + return feat + + +def run_epoch_single(base_model, loader, optimizer, device, num_classes, ignore_index, criterion_seg, amp, scaler, train, grad_accum=1): + base_model.train(train) + total_loss = 0.0 + cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device) + n_batches = 0 + t0 = time.time() + + with torch.set_grad_enabled(train): + if train and optimizer is not None: + optimizer.zero_grad(set_to_none=True) + for step, (imgs, masks) in enumerate(loader, start=1): + imgs = normalize_img(imgs.to(device, non_blocking=True)) + masks = masks.to(device, non_blocking=True) + with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): + out = base_model(pixel_values=imgs) + logits = out.logits + if logits.shape[-2:] != masks.shape[-2:]: + logits = F.interpolate(logits, size=masks.shape[-2:], mode="bilinear", align_corners=False) + loss = criterion_seg(logits, masks) + 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) + n_batches += 1 + preds = torch.argmax(logits.detach(), dim=1) + update_confusion_matrix(cm, preds, masks, num_classes, ignore_index) + + if train and optimizer is not None and (n_batches % grad_accum != 0): + if amp and scaler is not None and device.type == "cuda": + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + miou, iou_per_class = compute_iou_from_cm(cm) + return { + "loss": total_loss / max(1, n_batches), + "loss_seg": total_loss / max(1, n_batches), + "acc": compute_pixel_acc_from_cm(cm), + "miou": miou, + "iou_per_class": iou_per_class, + "time_s": time.time() - t0, + } + + +def run_epoch_mask2(base_model, mask2_head, loader, optimizer, device, num_classes, ignore_index1, ignore_index2, criterion_seg, bce_mask2, lambda_mask2, amp, scaler, train, grad_accum=1, mask2_thr=0.5, mask2_dice_mix=0.0): + base_model.train(train) + mask2_head.train(train) + total_loss = 0.0 + total_loss_seg = 0.0 + total_loss_mask2 = 0.0 + cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device) + mask2_iou_accum = 0.0 + mask2_acc_accum = 0.0 + n_batches = 0 + t0 = time.time() + + 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 = normalize_img(imgs.to(device, non_blocking=True)) + masks1 = masks1.to(device, non_blocking=True) + masks2 = masks2.to(device, non_blocking=True) + + with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): + out = base_model(pixel_values=imgs) + logits1 = out.logits + 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 = get_last_feat(out, logits1) + logits_m2 = mask2_head(feat, logits1) + if logits_m2.shape[-2:] != masks2.shape[-2:]: + logits_m2 = F.interpolate(logits_m2, size=masks2.shape[-2:], mode="bilinear", align_corners=False) + valid = masks2 != ignore_index2 + if valid.sum().item() > 0: + tgt = masks2.float().unsqueeze(1) + bce_val = bce_mask2(logits_m2[valid.unsqueeze(1)], tgt[valid.unsqueeze(1)]) + if mask2_dice_mix > 0: + d_val = dice_loss_with_logits(logits_m2, masks2, ignore_index=ignore_index2) + loss_mask2 = (1.0 - mask2_dice_mix) * bce_val + mask2_dice_mix * d_val + else: + loss_mask2 = bce_val + else: + loss_mask2 = torch.zeros([], device=device, dtype=loss_seg.dtype) + loss = loss_seg + lambda_mask2 * loss_mask2 + 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_mask2 += float(loss_mask2.item()) + n_batches += 1 + preds1 = torch.argmax(logits1.detach(), dim=1) + update_confusion_matrix(cm, preds1, masks1, num_classes, ignore_index1) + mi2, ma2 = binary_iou_and_acc_from_logits(logits_m2.detach(), masks2, thr=mask2_thr, ignore_index=ignore_index2) + mask2_iou_accum += mi2 + mask2_acc_accum += ma2 + + if train and optimizer is not None and (n_batches % grad_accum != 0): + if amp and scaler is not None and device.type == "cuda": + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + miou, iou_per_class = compute_iou_from_cm(cm) + return { + "loss": total_loss / max(1, n_batches), + "loss_seg": total_loss_seg / max(1, n_batches), + "loss_aux": total_loss_mask2 / max(1, n_batches), + "acc": compute_pixel_acc_from_cm(cm), + "miou": miou, + "iou_per_class": iou_per_class, + "aux_iou": mask2_iou_accum / max(1, n_batches), + "aux_acc": mask2_acc_accum / max(1, n_batches), + "time_s": time.time() - t0, + } + + +def run_epoch_label(base_model, label_head, loader, optimizer, device, num_classes, num_label_classes, ignore_index, criterion_seg, criterion_label, lambda_label, amp, scaler, train, grad_accum=1): + base_model.train(train) + label_head.train(train) + total_loss = 0.0 + total_loss_seg = 0.0 + total_loss_label = 0.0 + cm_seg = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device) + cm_cls = torch.zeros((num_label_classes, num_label_classes), dtype=torch.int64, device=device) + n_batches = 0 + t0 = time.time() + + with torch.set_grad_enabled(train): + if train and optimizer is not None: + optimizer.zero_grad(set_to_none=True) + for step, (imgs, masks, labels) in enumerate(loader, start=1): + imgs = normalize_img(imgs.to(device, non_blocking=True)) + masks = masks.to(device, non_blocking=True) + labels = labels.to(device, non_blocking=True) + + with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): + out = base_model(pixel_values=imgs) + logits_seg = out.logits + if logits_seg.shape[-2:] != masks.shape[-2:]: + logits_seg = F.interpolate(logits_seg, size=masks.shape[-2:], mode="bilinear", align_corners=False) + loss_seg = criterion_seg(logits_seg, masks) + feat = get_last_feat(out, logits_seg) + logits_label = label_head(feat, logits_seg) + loss_label = criterion_label(logits_label, labels) + loss = loss_seg + lambda_label * loss_label + 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_label += float(loss_label.item()) + n_batches += 1 + + preds_seg = torch.argmax(logits_seg.detach(), dim=1) + update_confusion_matrix(cm_seg, preds_seg, masks, num_classes, ignore_index) + update_cls_confusion_matrix(cm_cls, logits_label.detach(), labels, num_label_classes, ignore_index=-100) + + if train and optimizer is not None and (n_batches % grad_accum != 0): + if amp and scaler is not None and device.type == "cuda": + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + miou, iou_per_class = compute_iou_from_cm(cm_seg) + cls_metrics = compute_cls_metrics_from_cm(cm_cls) + return { + "loss": total_loss / max(1, n_batches), + "loss_seg": total_loss_seg / max(1, n_batches), + "loss_aux": total_loss_label / max(1, n_batches), + "acc": compute_pixel_acc_from_cm(cm_seg), + "miou": miou, + "iou_per_class": iou_per_class, + "aux_acc": cls_metrics["acc"], + "aux_macro_f1": cls_metrics["macro_f1"], + "aux_macro_recall": cls_metrics["macro_recall"], + "aux_f1_per_class": cls_metrics["f1_per_class"], + "aux_recall_per_class": cls_metrics["recall_per_class"], + "aux_support_per_class": cls_metrics["support_per_class"], + "time_s": time.time() - t0, + } + + +# ============================================================ +# LR / Early stop +# ============================================================ + +@dataclass +class EarlyStopState: + best: float = -1e9 + bad_epochs: int = 0 + + +def agg_metric(es_metric: str, miou: float, main_iou: float, aux_value: float) -> float: + if es_metric == "miou": + return float(miou) + if es_metric == "main": + return float(main_iou) + if es_metric in ("aux", "mask2", "label"): + return float(aux_value) + if es_metric in ("harmonic", "harmonic_label", "harmonic_mask2"): + eps = 1e-6 + a = max(eps, float(main_iou)) + b = max(eps, float(aux_value)) + return float(2.0 * a * b / (a + b + eps)) + 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.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_mask2", type=float, default=0.5) + ap.add_argument("--lambda_label", type=float, default=0.4) + ap.add_argument("--mask2_thr", type=float, default=0.5) + + ap.add_argument("--use_weights", action="store_true") + ap.add_argument("--cw_max_samples", type=int, default=800) + ap.add_argument("--auto_pos_weight", action="store_true") + ap.add_argument("--pw_max_samples", type=int, default=800) + ap.add_argument("--auto_label_weights", action="store_true") + ap.add_argument("--label_weight_max_samples", type=int, default=5000) + + ap.add_argument("--mask2_dice_max", type=float, default=0.15) + ap.add_argument("--mask2_dice_ramp_epochs", type=int, default=10) + + ap.add_argument("--warmup_epochs", type=int, default=0) + ap.add_argument("--cosine_epochs", type=int, default=0) + 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) + + ap.add_argument("--es_metric", default="harmonic_label", choices=["miou", "main", "aux", "mask2", "label", "harmonic", "harmonic_label", "harmonic_mask2"]) + ap.add_argument("--es_patience", type=int, default=12) + ap.add_argument("--es_min_delta", type=float, default=2e-4) + + 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) + ap.add_argument("--strict_label", action="store_true", help="Para dual_head_label, exige label em todas as amostras.") + ap.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal (ex: norm_stats.json).") + + args = ap.parse_args() + seed_everything(args.seed) + + with open(args.config, "r", encoding="utf-8") 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"] + + USE_MASK2 = bool(config.get("dual_head_mask", config.get("dual_head", False))) + USE_LABEL = bool(config.get("dual_head_label", False)) + if USE_MASK2 and USE_LABEL: + raise RuntimeError("Ative apenas um: dual_head_mask OU dual_head_label. Os dois juntos ainda não estão implementados neste script.") + + mode = "single" + if USE_MASK2: + mode = "mask2" + elif USE_LABEL: + mode = "label" + + 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") + + suffix = { + "single": "_single", + "mask2": "_dual_mask", + "label": "_dual_label", + }[mode] + save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME + suffix) + 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_aux_path = os.path.join(save_path, f"best_{mode}.pt") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("Device:", device) + print(f"[INFO] train mode: {mode}") + print(f"[INFO] save_path: {save_path}") + + 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) + + CLASS_NAMES = getattr(ds_train_base, "classes", None) + if CLASS_NAMES is None: + raise RuntimeError("ROISegDataset precisa expor .classes.") + + 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) + if main_class_id is None: + print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main usa mIoU.") + else: + print(f"Main class: {MAIN_CLASS_NAME} -> id={main_class_id}") + + aux_head = None + num_label_classes = 0 + label_name_by_id: Dict[int, str] = {} + + if mode == "mask2": + 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) + collate_fn = collate_mask2 + elif mode == "label": + ds_train = DualLabelROISegDataset(ds_train_base, split_train, strict_label=args.strict_label) + ds_val = DualLabelROISegDataset(ds_val_base, split_val, strict_label=args.strict_label) + + # 🔥 PRIORIDADE 1: config.json + config_labels = config.get("label_classes", None) + + if config_labels is not None: + label_name_by_id = {i: name for i, name in enumerate(config_labels)} + num_label_classes = len(config_labels) + + else: + # 🔁 fallback automático (dataset) + max_id = max(ds_train.label_ids_found + ds_val.label_ids_found) + num_label_classes = int(max_id) + 1 + + label_name_by_id = {i: f"label_{i}" for i in range(num_label_classes)} + label_name_by_id.update(ds_train.label_names_by_id) + label_name_by_id.update(ds_val.label_names_by_id) + + collate_fn = collate_label + + print(f"[INFO] label classes: {label_name_by_id}") + print(f"[INFO] train samples with labels: {len(ds_train)} | val: {len(ds_val)}") + else: + ds_train = ds_train_base + ds_val = ds_val_base + collate_fn = collate_seg + + dl_train = DataLoader(ds_train, batch_size=args.batch, shuffle=True, num_workers=args.num_workers, pin_memory=(device.type == "cuda"), collate_fn=collate_fn) + 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=collate_fn) + + 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) + + feat_ch = None + if mode in ("mask2", "label"): + with torch.no_grad(): + dummy_h = int(RESOLUCAO[1]) if RESOLUCAO else 512 + dummy_w = int(RESOLUCAO[0]) if RESOLUCAO else 512 + dummy = torch.zeros((1, 3, dummy_h, dummy_w), device=device) + out = base_model(pixel_values=dummy) + logits = out.logits + feat = get_last_feat(out, logits) + feat_ch = int(feat.shape[1]) + + if mode == "mask2": + aux_head = Mask2Head(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1).to(device) + print(f"[mask2_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {aux_head.in_ch}") + elif mode == "label": + aux_head = LabelHead(feat_ch=feat_ch, num_seg_classes=num_classes, num_label_classes=num_label_classes, hidden=256, dropout=0.2).to(device) + print(f"[label_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {aux_head.in_ch} | classes={num_label_classes}") + + 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) + + bce_mask2 = None + criterion_label = None + if mode == "mask2": + 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 mask2 neg/pos:", float(pos_weight.item())) + bce_mask2 = nn.BCEWithLogitsLoss(pos_weight=pos_weight) if pos_weight is not None else nn.BCEWithLogitsLoss() + elif mode == "label": + label_weights = None + if args.auto_label_weights: + label_weights = estimate_label_weights(ds_train, num_label_classes=num_label_classes, max_samples=args.label_weight_max_samples).to(device) + print("[INFO] label_weights:", label_weights.detach().cpu().numpy().round(3).tolist()) + criterion_label = nn.CrossEntropyLoss(weight=label_weights, ignore_index=-100) + + params = list(base_model.parameters()) + if aux_head is not None: + params += list(aux_head.parameters()) + optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wd) + scaler = GradScaler(enabled=(args.amp and device.type == "cuda")) + + # ========================== + # Normalizador (fixo ou ImageNet) + # ========================== + global NORM_MEAN, NORM_STD + + 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: + if not os.path.exists(norm_stats_path): + raise FileNotFoundError(f"[NORM] norm_stats informado, mas não encontrado: {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}") + + required = ["R", "G", "B"] + + if stats_channels != required: + raise ValueError( + f"[NORM] norm_stats inválido. Esperado channels={required}, " + f"mas veio channels={stats_channels}. " + "Este treino espera imagem em RGB." + ) + + if len(stats_mean) != 3 or len(stats_std) != 3: + raise ValueError( + f"[NORM] mean/std inválidos. Esperado 3 valores, " + f"recebido mean={len(stats_mean)}, std={len(stats_std)}." + ) + + 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 ativada para ordem [R,G,B].") + + else: + print("[NORM] norm_stats não informado. Usando normalização ImageNet.") + + 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" + + start_epoch = 1 + bests = {"miou": -1.0, "main": -1.0, "aux": -1.0} + if args.resume and os.path.exists(last_ckpt_path): + ckpt = load_checkpoint(last_ckpt_path, base_model, aux_head, optimizer=optimizer, scaler=scaler, map_location="cpu") + start_epoch = int(ckpt.get("epoch", 0)) + 1 + bests.update(ckpt.get("bests", {})) + print(f"[RESUME] epoch={start_epoch} bests={bests}") + + es = EarlyStopState() + + for epoch in range(start_epoch, args.epochs + 1): + apply_warmup(optimizer, args.lr, epoch, args.warmup_epochs) + lr_now = optimizer.param_groups[0]["lr"] + + if args.mask2_dice_max > 0 and args.mask2_dice_ramp_epochs > 0: + mask2_dice_mix = min(args.mask2_dice_max, (epoch - 1) / float(args.mask2_dice_ramp_epochs) * args.mask2_dice_max) + else: + mask2_dice_mix = 0.0 + + print(f"\n{time.time()} - ==== Epoch {epoch}/{args.epochs} | mode={mode} | lr={lr_now:.2e} | sched={active_sched} ====") + + if device.type == "cuda": + torch.cuda.empty_cache() + + if mode == "single": + tr = run_epoch_single(base_model, dl_train, optimizer, device, num_classes, args.ignore_index, criterion_seg, args.amp, scaler, True, max(1, args.grad_accum)) + va = run_epoch_single(base_model, dl_val, None, device, num_classes, args.ignore_index, criterion_seg, args.amp_val, None, False, 1) + aux_value = va["miou"] + elif mode == "mask2": + tr = run_epoch_mask2(base_model, aux_head, dl_train, optimizer, device, num_classes, args.ignore_index, args.ignore_index2, criterion_seg, bce_mask2, args.lambda_mask2, args.amp, scaler, True, max(1, args.grad_accum), args.mask2_thr, mask2_dice_mix) + va = run_epoch_mask2(base_model, aux_head, dl_val, None, device, num_classes, args.ignore_index, args.ignore_index2, criterion_seg, bce_mask2, args.lambda_mask2, args.amp_val, None, False, 1, args.mask2_thr, mask2_dice_mix) + aux_value = float(va["aux_iou"]) + else: + tr = run_epoch_label(base_model, aux_head, dl_train, optimizer, device, num_classes, num_label_classes, args.ignore_index, criterion_seg, criterion_label, args.lambda_label, args.amp, scaler, True, max(1, args.grad_accum)) + va = run_epoch_label(base_model, aux_head, dl_val, None, device, num_classes, num_label_classes, args.ignore_index, criterion_seg, criterion_label, args.lambda_label, args.amp_val, None, False, 1) + aux_value = float(va["aux_macro_f1"]) + + 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] + + print(f"TRAIN: loss={tr['loss']:.4f} seg={tr.get('loss_seg', 0):.4f} aux={tr.get('loss_aux', 0):.4f} acc={tr['acc']:.4f} miou={tr['miou']:.4f} t={tr['time_s']:.1f}s") + print(f"VAL : loss={va['loss']:.4f} seg={va.get('loss_seg', 0):.4f} aux={va.get('loss_aux', 0):.4f} acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={float(main_iou):.4f} t={va['time_s']:.1f}s") + print("IoU per class:", _pretty_iou(va["iou_per_class"], class_name_by_id)) + + if mode == "mask2": + print(f"MASK2: val_iou={va['aux_iou']:.4f} val_acc={va['aux_acc']:.4f}") + elif mode == "label": + print(f"LABEL: acc={va['aux_acc']:.4f} macro_f1={va['aux_macro_f1']:.4f} macro_recall={va['aux_macro_recall']:.4f}") + print(_pretty_cls(va["aux_f1_per_class"], label_name_by_id, "F1")) + print(_pretty_cls(va["aux_recall_per_class"], label_name_by_id, "Recall")) + + _extras = { + "mode": mode, + "epoch": epoch, + "config": config, + "args": vars(args), + "norm_stats_path": norm_stats_path, + "norm_mean": NORM_MEAN.detach().cpu().view(-1).tolist(), + "norm_std": NORM_STD.detach().cpu().view(-1).tolist(), + "val_loss": va["loss"], + "val_miou": va["miou"], + "val_main_iou": float(main_iou), + "val_aux": float(aux_value), + "lr": optimizer.param_groups[0]["lr"], + "label_name_by_id": label_name_by_id, + } + save_checkpoint( + last_ckpt_path, + base_model, + aux_head, + optimizer, + scaler, + epoch, + bests, + extra=_extras, + ) + + 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, aux_head, optimizer, scaler, epoch, bests, _extras) + + if va["miou"] > bests["miou"]: + bests["miou"] = float(va["miou"]) + save_checkpoint(best_miou_path, base_model, aux_head, optimizer, scaler, epoch, bests, _extras) + print(f"[BEST mIoU] {bests['miou']:.4f} -> {best_miou_path}") + + if float(main_iou) > bests["main"]: + bests["main"] = float(main_iou) + save_checkpoint(best_main_path, base_model, aux_head, optimizer, scaler, epoch, bests, _extras) + print(f"[BEST MAIN] {bests['main']:.4f} -> {best_main_path}") + + if float(aux_value) > bests["aux"]: + bests["aux"] = float(aux_value) + save_checkpoint(best_aux_path, base_model, aux_head, optimizer, scaler, epoch, bests, _extras) + print(f"[BEST AUX/{mode}] {bests['aux']:.4f} -> {best_aux_path}") + + score = agg_metric(args.es_metric, va["miou"], float(main_iou), float(aux_value)) + 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() diff --git a/Python/OAK/datasets/_9_test_segformer_b3.py b/Python/OAK/datasets/_9_test_segformer.py similarity index 99% rename from Python/OAK/datasets/_9_test_segformer_b3.py rename to Python/OAK/datasets/_9_test_segformer.py index 5c9c459ef..e01a30344 100644 --- a/Python/OAK/datasets/_9_test_segformer_b3.py +++ b/Python/OAK/datasets/_9_test_segformer.py @@ -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) === diff --git a/Python/OAK/datasets/_9_test_segformer_b3_dual.py b/Python/OAK/datasets/_9_test_segformer_b3_dual.py deleted file mode 100644 index 924947ccb..000000000 --- a/Python/OAK/datasets/_9_test_segformer_b3_dual.py +++ /dev/null @@ -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): - """ - Lê 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//images/*.png - split/test/group//masks/*.png - split/test/group//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/ - 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() diff --git a/Python/OAK/datasets/_9_test_segformer_dual.py b/Python/OAK/datasets/_9_test_segformer_dual.py new file mode 100644 index 000000000..9cfa42e0a --- /dev/null +++ b/Python/OAK/datasets/_9_test_segformer_dual.py @@ -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() diff --git a/Python/OAK/datasets/config_oak.json b/Python/OAK/datasets/config.json similarity index 53% rename from Python/OAK/datasets/config_oak.json rename to Python/OAK/datasets/config.json index 340983cbb..c87c21af3 100644 --- a/Python/OAK/datasets/config_oak.json +++ b/Python/OAK/datasets/config.json @@ -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" + ] } \ No newline at end of file diff --git a/Python/OAK/datasets/download_backbone.py b/Python/OAK/datasets/download_backbone.py new file mode 100644 index 000000000..3c6e1912b --- /dev/null +++ b/Python/OAK/datasets/download_backbone.py @@ -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) \ No newline at end of file diff --git a/Python/OAK/datasets/multiespec_module/_5_augmentation.py b/Python/OAK/datasets/multiespec_module/_5_augmentation.py index 675b7a3ed..64316fecc 100644 --- a/Python/OAK/datasets/multiespec_module/_5_augmentation.py +++ b/Python/OAK/datasets/multiespec_module/_5_augmentation.py @@ -1,613 +1,550 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ -Augmentação por grupos para o módulo multiespectral (opção 1: bins como canais sincronizados). +Augmenta imagens, máscaras e opcionalmente labels globais por grupo. -Entrada: - dataset/original/group//previews - dataset/original/group//metas - dataset/original/group//bins - dataset/original/group//masks - (opcional) dataset/original/group//masks2 +Entrada via config_oak.json -> camera: + + MODELO/dataset/original/group//images/ + MODELO/dataset/original/group//masks/ + MODELO/dataset/original/group//masks2/ opcional, se dual_head_mask=true + MODELO/dataset/original/group//labels/ opcional, se dual_head_label=true Saída: - dataset/augmented/group//previews - dataset/augmented/group//metas - dataset/augmented/group//bins - dataset/augmented/group//masks - (opcional) dataset/augmented/group//masks2 -Amostra esperada: - .png # preview - .json # meta - _cam0.bin # bin câmera 0 - _cam1.bin # bin câmera 1 - _cam2.bin # bin câmera 2 (opcional) - .png # mask + MODELO/dataset/augmented/group//images/ + MODELO/dataset/augmented/group//masks/ + MODELO/dataset/augmented/group//masks2/ se dual_head_mask=true + MODELO/dataset/augmented/group//labels/ se dual_head_label=true -Estratégia: -- Geometria sincronizada em preview + masks + todos os bins. -- Blur / ruído / ganho apenas nos bins. -- Preview de saída recebe a mesma geometria; não recebe blur pesado para continuar útil como inspeção visual. -- Meta é copiado e marcado como augmentado. +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. -Importante: -- Este script assume que os .bin são RAW10 packed, um arquivo por câmera. -- A largura/altura do bin é lida do meta.json quando possível; se não existir, cai para config['raw_size']. +Uso: + 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 re -import cv2 import json -import math -import shutil import argparse -import csv import random -from copy import deepcopy -from typing import Dict, List, Optional, Tuple +from pathlib import Path -import numpy as np +import cv2 from PIL import Image +import albumentations as A -from pi.raw_processor_core import RawProcessorCore +# ====================== 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) -RAW_SIZE = config.get("raw_size", [1296, 1028]) # [W, H] +USE_MASKS2 = bool(config.get("dual_head_mask", False)) +USE_LABELS = bool(config.get("dual_head_label", False)) -DATASET_BASE = os.path.join("dataset") +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") -PREVIEW_EXTS = (".jpg", ".jpeg", ".png") -MASK_EXTS = (".png", ".jpg", ".jpeg") -MASK2_EXTS = (".png", ".jpg", ".jpeg") -META_EXTS = (".json",) -BIN_RE = re.compile(r"^(?P.+)_cam(?P\d+)\.bin$", re.IGNORECASE) -MANIFESTO_DEFAULT = "manifest_aug.csv" +# 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") +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") + +IMG_EXTS = (".jpg", ".jpeg", ".png") +MSK_EXTS = (".png", ".jpg", ".jpeg") +MSK2_EXTS = (".png", ".jpg", ".jpeg") +LABEL_EXTS = (".json", ".txt") -# ============================================================ -# Helpers básicos -# ============================================================ +# ====================== Augmentation ====================== -def garantir_dir(p: str): +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) -def save_rgb(path: str, arr_rgb: np.ndarray): - Image.fromarray(arr_rgb).save(path) +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 load_rgb(path: str) -> np.ndarray: +def map_files_by_base(folder, exts): + by_base = {} + if not os.path.isdir(folder): + return by_base + + 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 + + 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 list_groups(root): + if not os.path.isdir(root): + return [] + + grupos = [] + for name in sorted(os.listdir(root)): + gdir = os.path.join(root, name) + 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")): + grupos.append(name) + return grupos + + +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 load_mask_any(path: str) -> np.ndarray: - m = cv2.imread(path, cv2.IMREAD_UNCHANGED) - if m is None: - raise FileNotFoundError(path) - if m.ndim == 2: - return m - if m.shape[2] == 1: - return m[:, :, 0] - return cv2.cvtColor(m, cv2.COLOR_BGR2RGB) +def save_rgb(path, arr_rgb): + garantir_dir(os.path.dirname(path)) + Image.fromarray(arr_rgb).save(path) -def save_mask_any(path: str, mask: np.ndarray): - if mask.ndim == 2: - cv2.imwrite(path, mask) +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: - bgr = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR) - cv2.imwrite(path, bgr) + 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) -def list_groups(root: str) -> List[str]: - if not os.path.isdir(root): - return [] - grupos = [] - for name in sorted(os.listdir(root)): - gdir = os.path.join(root, name) - if not os.path.isdir(gdir): - continue - has_prev = os.path.isdir(os.path.join(gdir, "previews")) - has_meta = os.path.isdir(os.path.join(gdir, "metas")) - has_bins = os.path.isdir(os.path.join(gdir, "bins")) - has_masks = os.path.isdir(os.path.join(gdir, "masks")) - if has_prev and has_meta and has_bins and has_masks: - grupos.append(name) - return grupos + if use_masks2 and msk2_out: + garantir_dir(msk2_out) + if use_labels and labels_out: + garantir_dir(labels_out) -def map_by_base_priorizando_png(folder: str, exts: Tuple[str, ...]) -> Dict[str, str]: - if not os.path.isdir(folder): - return {} - by_base = {} - for fname in os.listdir(folder): - if not fname.lower().endswith(exts): - continue - base, ext = os.path.splitext(fname) - full = os.path.join(folder, fname) - if base not in by_base: - by_base[base] = full - else: - cur_ext = os.path.splitext(by_base[base])[1].lower() - if cur_ext != ".png" and ext.lower() == ".png": - by_base[base] = full - return by_base + return img_out, msk_out, msk2_out, labels_out -def map_bins_by_base(folder: str) -> Dict[str, List[str]]: - by_base = {} - if not os.path.isdir(folder): - return by_base - for fname in os.listdir(folder): - m = BIN_RE.match(fname) - if not m: - continue - base = m.group("base") - cam = int(m.group("cam")) - by_base.setdefault(base, []).append((cam, os.path.join(folder, fname))) - for base in list(by_base.keys()): - by_base[base] = [p for _, p in sorted(by_base[base], key=lambda x: x[0])] - return by_base +def safe_rel(path, root): + try: + return str(Path(path).resolve().relative_to(Path(root).resolve())).replace("\\", "/") + except Exception: + return str(path).replace("\\", "/") -def ensure_aug_dirs(group_name: str, use_masks2: bool): - base = os.path.join(AUG_GROUP_ROOT, group_name) - prev_out = os.path.join(base, "previews") - meta_out = os.path.join(base, "metas") - bins_out = os.path.join(base, "bins") - mask_out = os.path.join(base, "masks") - mask2_out = os.path.join(base, "masks2") if use_masks2 else None - - garantir_dir(prev_out) - garantir_dir(meta_out) - garantir_dir(bins_out) - garantir_dir(mask_out) - if use_masks2 and mask2_out: - garantir_dir(mask2_out) - - return prev_out, meta_out, bins_out, mask_out, mask2_out - - -def make_raw_core() -> RawProcessorCore: - return RawProcessorCore(sensor_width=RAW_SIZE[0], sensor_height=RAW_SIZE[1]) - - -def load_all_bins(meta_path, bins_paths): - import json - - with open(meta_path, "r", encoding="utf-8") as f: - meta = json.load(f) - - core = make_raw_core() - - bins_data = [] - bins_meta = [] - - for path in bins_paths: - filename = os.path.basename(path) - cam_id = filename.split("_")[-1].replace(".bin", "") - - cam_meta = core.extract_camera_meta(meta, cam_id) - data = core.load_native_bin(path, cam_meta) - - bins_data.append(data) - bins_meta.append(cam_meta) - - return bins_data, bins_meta - - -def save_all_bins(core, bins_data, bins_meta, base_name, out_dir): - paths = [] - - for data, meta in zip(bins_data, bins_meta): - cam_id = meta["camera_id"] - out_path = os.path.join(out_dir, f"{base_name}_{cam_id}.bin") - - core.save_native_bin(out_path, data, meta) - paths.append(out_path) - - return paths - - -# ============================================================ -# Meta / resolução -# ============================================================ - -def infer_bin_hw_from_meta(meta_path: str) -> Tuple[int, int]: - with open(meta_path, "r", encoding="utf-8") as f: - meta = json.load(f) - - # tenta campos mais prováveis - width = None - height = None - - for k in ("sensor_width", "width", "raw_width"): - if k in meta: - width = int(meta[k]) - break - for k in ("sensor_height", "height", "raw_height"): - if k in meta: - height = int(meta[k]) - break - - if (width is None or height is None) and "raw_size" in meta and isinstance(meta["raw_size"], (list, tuple)) and len(meta["raw_size"]) == 2: - width, height = int(meta["raw_size"][0]), int(meta["raw_size"][1]) - - if width is None or height is None: - width, height = int(RAW_SIZE[0]), int(RAW_SIZE[1]) - - return width, height - - -def build_augmented_meta(meta_path: str, source_group: str, source_base: str, - aug_base: str, aug_index: int, params: Dict) -> Dict: - with open(meta_path, "r", encoding="utf-8") as f: - meta = json.load(f) - - meta_aug = deepcopy(meta) - meta_aug["augmented"] = True - meta_aug["augmentation"] = { - "source_group": source_group, - "source_base": source_base, - "aug_base": aug_base, - "aug_index": aug_index, - "params": params, - } - return meta_aug - - -# ============================================================ -# Geometria sincronizada -# ============================================================ - -def sample_geom_params() -> Dict: - do_hflip = np.random.rand() < 0.5 - shift_x_frac = float(np.random.uniform(-0.01, 0.01)) - shift_y_frac = float(np.random.uniform(-0.01, 0.01)) - scale = float(np.random.uniform(0.92, 1.08)) - angle = float(np.random.uniform(-5.0, 5.0)) - return { - "hflip": do_hflip, - "shift_x_frac": shift_x_frac, - "shift_y_frac": shift_y_frac, - "scale": scale, - "angle": angle, - } - - -def build_affine_matrix(width: int, height: int, params: Dict) -> np.ndarray: - cx = (width - 1) / 2.0 - cy = (height - 1) / 2.0 - M = cv2.getRotationMatrix2D((cx, cy), params["angle"], params["scale"]) - M[0, 2] += params["shift_x_frac"] * width - M[1, 2] += params["shift_y_frac"] * height - return M - - -def apply_geom_to_image(img: np.ndarray, params: Dict, is_mask: bool = False) -> np.ndarray: - out = img - if params["hflip"]: - out = cv2.flip(out, 1) - - h, w = out.shape[:2] - M = build_affine_matrix(w, h, params) - - interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_LINEAR - if out.ndim == 2: - warped = cv2.warpAffine(out, M, (w, h), flags=interp, borderMode=cv2.BORDER_REFLECT_101) - else: - warped = cv2.warpAffine(out, M, (w, h), flags=interp, borderMode=cv2.BORDER_REFLECT_101) - return warped - - -def apply_effects_rgb(img): - out = img.astype(np.float32) - - if np.random.rand() < 0.2: - out = cv2.GaussianBlur(out, (3,3), 0) - - gain = np.random.uniform(0.97, 1.03) - out *= gain - - noise = np.random.normal(0, 2, out.shape) - out += noise - - return np.clip(out, 0, 255).astype(np.uint8) - - -# ============================================================ -# Blur / ruído / ganho coerentes nos bins -# ============================================================ - -def motion_blur_kernel(ksize=5, angle=0.0): - ksize = int(ksize) - if ksize < 3: - ksize = 3 - if ksize % 2 == 0: - ksize += 1 - - kernel = np.zeros((ksize, ksize), dtype=np.float32) - kernel[ksize // 2, :] = 1.0 - center = (ksize / 2.0 - 0.5, ksize / 2.0 - 0.5) - M = cv2.getRotationMatrix2D(center, angle, 1.0) - kernel = cv2.warpAffine(kernel, M, (ksize, ksize)) - s = kernel.sum() - if s > 0: - kernel /= s - return kernel - - -def sample_bin_effects() -> Dict: - effect = {"kind": "none"} - r = np.random.rand() - if r < 0.15: - effect["kind"] = "motion" - effect["ksize"] = int(np.random.choice([3, 5, 7])) - effect["angle"] = float(np.random.uniform(-20.0, 20.0)) - elif r < 0.30: - effect["kind"] = "gaussian" - effect["ksize"] = int(np.random.choice([3, 5, 7])) - - # leves variações por bin (mantendo coerência física e sem enlouquecer) - effect["gain_min"] = float(np.random.uniform(0.97, 0.995)) - effect["gain_max"] = float(np.random.uniform(1.005, 1.03)) - effect["noise_sigma"] = float(np.random.uniform(0.0, 2.0)) # escala raw10 - return effect - - -def apply_effects_to_bin(bin_img: np.ndarray, effect: Dict, gain: float) -> np.ndarray: - out = bin_img.astype(np.float32) - - if effect["kind"] == "motion": - kernel = motion_blur_kernel(effect["ksize"], effect["angle"]) - out = cv2.filter2D(out, ddepth=-1, kernel=kernel, borderType=cv2.BORDER_REFLECT_101) - elif effect["kind"] == "gaussian": - k = effect["ksize"] - if k % 2 == 0: - k += 1 - out = cv2.GaussianBlur(out, (k, k), 0, borderType=cv2.BORDER_REFLECT_101) - - out *= gain - - sigma = effect.get("noise_sigma", 0.0) - if sigma > 0: - noise = np.random.normal(0.0, sigma, out.shape).astype(np.float32) - out += noise - - out = np.clip(out, 0.0, 1023.0) - return np.round(out).astype(np.uint16) - - -# ============================================================ -# Núcleo da augmentação -# ============================================================ - -def make_preview_from_augmented_preview(preview_geom: np.ndarray) -> np.ndarray: +def copiar_label_aug(label_path, out_label_path, new_base, out_img_path, out_msk_path, out_msk2_path=None, group_name=None): """ - Por enquanto, o preview final é o preview original com a mesma geometria. - Mantemos isso simples nesta etapa para inspeção humana. + Copia label global para a amostra augmentada. + Se for JSON, atualiza campos úteis. + Se for TXT, copia o conteúdo como está. """ - return preview_geom - - -def augment_sample(group_name: str, base: str, - preview_path: str, meta_path: str, bins_paths: List[str], mask_path: str, - preview_out_dir: str, meta_out_dir: str, bins_out_dir: str, mask_out_dir: str, - copies: int, mask2_path: Optional[str] = None, mask2_out_dir: Optional[str] = None, - aug_suffix: str = "aug") -> int: - - preview_ext = os.path.splitext(preview_path)[1].lower() - meta_ext = os.path.splitext(meta_path)[1].lower() - mask_ext = os.path.splitext(mask_path)[1].lower() - mask2_ext = os.path.splitext(mask2_path)[1].lower() if mask2_path else None - - preview = load_rgb(preview_path) - mask = load_mask_any(mask_path) - mask2 = load_mask_any(mask2_path) if mask2_path else None - - raw_w, raw_h = infer_bin_hw_from_meta(meta_path) - bins_imgs, bins_meta = load_all_bins(meta_path, bins_paths) - core = make_raw_core() - - generated = 0 - for i in range(copies): - params = sample_geom_params() - effects = sample_bin_effects() - - preview_g = apply_geom_to_image(preview, params, is_mask=False) - mask_g = apply_geom_to_image(mask, params, is_mask=True) - mask2_g = apply_geom_to_image(mask2, params, is_mask=True) if mask2 is not None else None - - bins_g = [apply_geom_to_image(b, params, is_mask=False) for b in bins_imgs] - - gains = [float(np.random.uniform(effects["gain_min"], effects["gain_max"])) for _ in bins_g] - bins_aug = [] - - for b, meta in zip(bins_g, bins_meta): - if meta["channels"] == 3: - # RGB (cam2) - out = apply_effects_rgb(b) - else: - # RAW mono (cam0, cam1) - gain = float(np.random.uniform(effects["gain_min"], effects["gain_max"])) - out = apply_effects_to_bin(b, effects, gain) - - bins_aug.append(out) - - preview_aug = make_preview_from_augmented_preview(preview_g) - - aug_base = f"{base}_{aug_suffix}_{i:02d}" - out_preview = os.path.join(preview_out_dir, aug_base + preview_ext) - out_meta = os.path.join(meta_out_dir, aug_base + meta_ext) - out_mask = os.path.join(mask_out_dir, aug_base + mask_ext) - - save_rgb(out_preview, preview_aug) - save_mask_any(out_mask, mask_g) - - meta_aug = build_augmented_meta(meta_path, group_name, base, aug_base, i, { - "geometry": params, - "bin_effects": effects, - "bin_gains": gains, - }) - with open(out_meta, "w", encoding="utf-8") as f: - json.dump(meta_aug, f, ensure_ascii=False, indent=2) - - out_bins = save_all_bins(core, bins_aug, bins_meta, aug_base, bins_out_dir) - - if mask2_g is not None and mask2_out_dir: - out_mask2 = os.path.join(mask2_out_dir, aug_base + mask2_ext) - save_mask_any(out_mask2, mask2_g) - - generated += 1 - - return generated - - -# ============================================================ -# Processamento por grupo -# ============================================================ - -def process_group(group_name: str, copies: int, limit: Optional[int] = None, - seed: int = 42, aug_suffix: str = "aug") -> Tuple[int, List[List[str]]]: - gdir = os.path.join(ORIG_GROUP_ROOT, group_name) - previews_dir = os.path.join(gdir, "previews") - metas_dir = os.path.join(gdir, "metas") - bins_dir = os.path.join(gdir, "bins") - masks_dir = os.path.join(gdir, "masks") - masks2_dir = os.path.join(gdir, "masks2") - - if not (os.path.isdir(previews_dir) and os.path.isdir(metas_dir) and os.path.isdir(bins_dir) and os.path.isdir(masks_dir)): - print(f"[WARN] Grupo '{group_name}' inválido. Precisa de previews/metas/bins/masks.") - return 0, [] - - use_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir) - - previews_map = map_by_base_priorizando_png(previews_dir, PREVIEW_EXTS) - metas_map = map_by_base_priorizando_png(metas_dir, META_EXTS) - masks_map = map_by_base_priorizando_png(masks_dir, MASK_EXTS) - bins_map = map_bins_by_base(bins_dir) - masks2_map = map_by_base_priorizando_png(masks2_dir, MASK2_EXTS) if use_masks2 else {} - - bases = sorted(set(previews_map.keys()) & set(metas_map.keys()) & set(masks_map.keys()) & set(bins_map.keys())) - if limit is not None and limit > 0 and limit < len(bases): - rng = np.random.default_rng(seed) - idx = sorted(rng.choice(len(bases), size=limit, replace=False).tolist()) - bases = [bases[i] for i in idx] - - preview_out, meta_out, bins_out, mask_out, mask2_out = ensure_aug_dirs(group_name, use_masks2) - - registros = [] - count = 0 - for base in bases: - try: - gen = augment_sample( - group_name=group_name, - base=base, - preview_path=previews_map[base], - meta_path=metas_map[base], - bins_paths=bins_map[base], - mask_path=masks_map[base], - preview_out_dir=preview_out, - meta_out_dir=meta_out, - bins_out_dir=bins_out, - mask_out_dir=mask_out, - copies=copies, - mask2_path=masks2_map.get(base), - mask2_out_dir=mask2_out, - aug_suffix=aug_suffix, - ) - count += gen - registros.append([ - group_name, - base, - previews_map[base], - metas_map[base], - json.dumps(bins_map[base], ensure_ascii=False), - masks_map[base], - gen, - ]) - except Exception as e: - print(f"[ERRO] [{group_name}] {base}: {e}") - - print(f"[OK] Grupo '{group_name}' -> {count} amostras geradas.") - return count, registros - - -# ============================================================ -# Main -# ============================================================ - -def main(copies: int = 5, groups_csv: Optional[str] = None, - limit: Optional[int] = None, seed: int = 42, - suffix: str = "aug", manifesto: str = MANIFESTO_DEFAULT): - total = 0 - all_records = [] - - grupos = list_groups(ORIG_GROUP_ROOT) - if groups_csv: - 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.") - return - - if not grupos: - print("[WARN] Nenhum grupo encontrado em dataset/original/group.") + if not label_path or not out_label_path: return - print(f"Grupos encontrados: {', '.join(grupos)}") - for g in grupos: - count, records = process_group(g, copies, limit=limit, seed=seed, aug_suffix=suffix) - total += count - all_records.extend(records) + garantir_dir(os.path.dirname(out_label_path)) + ext = os.path.splitext(label_path)[1].lower() - if manifesto: - with open(manifesto, "w", newline="", encoding="utf-8") as f: - w = csv.writer(f) - w.writerow([ - "grupo", - "base", - "src_preview", - "src_meta", - "src_bins_json", - "src_mask", - "generated_copies", - ]) - w.writerows(all_records) + if ext == ".json": + try: + with open(label_path, "r", encoding="utf-8") as f: + data = json.load(f) + except Exception: + data = {} - print(f"\nAugmentation completed! Total: {total} amostras geradas.") + 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)) + _, 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 + + base = normalizar_base(base_img) + + img = load_rgb(img_path) + msk = load_rgb(msk_path) + msk2 = load_rgb(msk2_path) if msk2_path else None + + gen = 0 + for i in range(copies): + if msk2 is not None and msk2_out_dir: + aug = train_tf(image=img, mask=msk, mask2=msk2) + else: + aug = train_tf(image=img, mask=msk) + + img_aug = aug["image"] + msk_aug = aug["mask"] + + 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"{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 + + +# ====================== Processamento ====================== + + +def process_group(group_name, copies, strict_label=False, limit=None, seed=42): + 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.") + return 0 + + imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS] + imgs = sorted(imgs) + + if limit is not None and limit > 0 and limit < len(imgs): + rng = random.Random(seed) + imgs = sorted(rng.sample(imgs, limit)) + print(f"[INFO] [{group_name}] Limit aplicado: {limit} amostras originais selecionadas.") + + msk_map = map_files_by_base(msk_dir, MSK_EXTS) + + use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir) + 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 imgs: + base, _ = os.path.splitext(img_file) + 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_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( + 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, + 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. " + f"sem_mask={sem_mask} | sem_mask2={sem_mask2} | sem_label={sem_label}" + ) + return count + + +def process_legacy(copies, strict_label=False, limit=None, seed=42): + 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] + imgs = sorted(imgs) + + if limit is not None and limit > 0 and limit < len(imgs): + rng = random.Random(seed) + imgs = sorted(rng.sample(imgs, limit)) + print(f"[INFO] Legacy limit aplicado: {limit} amostras originais selecionadas.") + + 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_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 imgs: + base, _ = os.path.splitext(img_file) + 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_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( + 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, + 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, strict_label=False, limit=None, seed=42): + 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: + 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, strict_label=strict_label, limit=limit, seed=seed) + else: + print(f"Grupos encontrados: {', '.join(grupos)}") + for g in grupos: + total += process_group(g, copies, strict_label=strict_label, limit=limit, seed=seed) + else: + total += process_legacy(copies, strict_label=strict_label, limit=limit, seed=seed) + + print(f"\nAugmentation completed! Total: {total} pares gerados.") if __name__ == "__main__": - ap = argparse.ArgumentParser(description="Augmentação por grupos para o módulo multiespectral usando bins sincronizados.") - ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por amostra (default=5).") + 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("--limit", type=int, default=None, help="Quantidade máxima de amostras originais do grupo a augmentar.") + ap.add_argument("--strict-label", action="store_true", help="Se dual_head_label=true e faltar label, pula o item/grupo.") + ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de amostras originais por grupo para augmentar.") ap.add_argument("--seed", type=int, default=42, help="Seed para seleção reproduzível quando usar --limit.") - ap.add_argument("--suffix", type=str, default="aug", help="Sufixo usado no nome dos arquivos gerados.") - ap.add_argument("--manifest", type=str, default=MANIFESTO_DEFAULT, help="CSV de manifesto.") args = ap.parse_args() random.seed(args.seed) - np.random.seed(args.seed) main( copies=args.copies, groups_csv=args.groups, + strict_label=args.strict_label, limit=args.limit, seed=args.seed, - suffix=args.suffix, - manifesto=args.manifest, ) diff --git a/Python/OAK/datasets/multiespec_module/multispectral_client.py b/Python/OAK/datasets/multiespec_module/multispectral_client.py index 6544b9fc2..f773c045c 100644 --- a/Python/OAK/datasets/multiespec_module/multispectral_client.py +++ b/Python/OAK/datasets/multiespec_module/multispectral_client.py @@ -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) diff --git a/Python/OAK/datasets/oak-fcc-3/_0_capture.py b/Python/OAK/datasets/oak-fcc-3/_0_capture.py index 534d24fa9..4edc51d1f 100644 --- a/Python/OAK/datasets/oak-fcc-3/_0_capture.py +++ b/Python/OAK/datasets/oak-fcc-3/_0_capture.py @@ -230,6 +230,7 @@ def main(): try: with MultiSpectralClient( + #mx_id="194430108133AC2F00", width=raw_w, height=raw_h, bayer=args.bayer, diff --git a/Python/OAK/datasets/oak-fcc-3/_6_normalize.py b/Python/OAK/datasets/oak-fcc-3/_6_normalize.py index 41c64b75c..840aae7e4 100644 --- a/Python/OAK/datasets/oak-fcc-3/_6_normalize.py +++ b/Python/OAK/datasets/oak-fcc-3/_6_normalize.py @@ -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) diff --git a/Python/OAK/datasets/oak-fcc-3/_7_split.py b/Python/OAK/datasets/oak-fcc-3/_7_split.py index b00d7c7e4..0398a2d97 100644 --- a/Python/OAK/datasets/oak-fcc-3/_7_split.py +++ b/Python/OAK/datasets/oak-fcc-3/_7_split.py @@ -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: + .npy obrigatório se existir + .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", ] diff --git a/Python/OAK/datasets/oak-fcc-3/_8_train_multihead.py b/Python/OAK/datasets/oak-fcc-3/_8_train_multihead.py new file mode 100644 index 000000000..0ea831d3b --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/_8_train_multihead.py @@ -0,0 +1,1884 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +_8_train_segformer_oak_fcc3_multihead.py + +Treina SegFormer OAK-FCC-3 com tensor multispectral [R,G,B,RE,NIR] +e contrato multi-head: + + 1) semantic_head: + máscara: masks/.npy + classes: 0=chao, 1=cana, 2=erva, 255=ignore + + 2) vegetation_head: + máscara: masks_vegetation/.npy + classes: 0=nao_vegetacao, 1=vegetacao, 255=ignore + + 3) cana_head: + máscara: masks_cana/.npy + classes: 0=nao_cana, 1=cana, 255=ignore + +Decisão operacional futura: + alvo/pulverizavel = vegetation == 1 AND cana == 0 + +Entrada esperada após normalize + split: + + dataset/split/train/group// + tensors/.npy + masks/.npy + masks_vegetation/.npy + masks_cana/.npy + metas/.json + previews/.png + + dataset/split/val/group//... + +Exemplo: + +python _8_train_segformer_oak_fcc3_multihead.py ^ + --epochs 50 --batch 2 --lr 3e-5 --wd 0.01 ^ + --num_workers 2 --amp --amp_val --grad_accum 4 ^ + --class_weights auto +""" + +from __future__ import annotations + +import os +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128") + +import csv +import copy +import json +import time +import argparse +import random +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader +from torch.amp import autocast, GradScaler + +from transformers import SegformerForSemanticSegmentation + + +# ============================================================ +# Seed / util +# ============================================================ + +def set_seed(seed: int = 42): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def ensure_dir(path: Path): + path.mkdir(parents=True, exist_ok=True) + + +def load_json(path: str | Path) -> dict: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +def save_json(path: Path, data: dict): + ensure_dir(path.parent) + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + + +def safe_float(x, default=0.0) -> float: + try: + return float(x) + except Exception: + return float(default) + + +# ============================================================ +# Labelmap +# ============================================================ + +def load_labelmap(labelmap_path: str): + """ + Tenta usar helpers.carregar_labelmap_completo. + Fallback simples para labelmap com uma classe por linha. + """ + try: + from helpers import carregar_labelmap_completo, _infer_ignore_id + + _cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path) + ignore_id = _infer_ignore_id(ignore_rgb, 255) + + id2label = {int(k): str(v) for k, v in id_para_nome.items()} + label2id = {v.lower(): k for k, v in id2label.items()} + + return id2label, label2id, int(ignore_id) + + except Exception as e: + print(f"[WARN] Não consegui usar helpers.carregar_labelmap_completo: {e}") + print("[WARN] Usando parser simples: uma classe por linha.") + + id2label = {} + with open(labelmap_path, "r", encoding="utf-8") as f: + for line in f: + s = line.strip() + if not s or s.startswith("#"): + continue + + parts = s.replace(",", " ").split() + + if len(parts) >= 2 and parts[0].isdigit(): + cid = int(parts[0]) + name = parts[1] + else: + cid = len(id2label) + name = parts[0] + + if name.lower() in ("ignore", "void", "background_ignore"): + continue + + id2label[cid] = name + + label2id = {v.lower(): k for k, v in id2label.items()} + return id2label, label2id, 255 + + +# ============================================================ +# Head config +# ============================================================ + +DEFAULT_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_from": ["vegetation", "cana"], + }, +} + + +def merge_dict(dst: dict, src: dict) -> dict: + out = copy.deepcopy(dst) + + def rec(a, b): + for k, v in b.items(): + if isinstance(v, dict) and isinstance(a.get(k), dict): + rec(a[k], v) + else: + a[k] = v + + if isinstance(src, dict): + rec(out, src) + + return out + + +def build_heads_config(config: dict, ignore_index: int) -> Dict[str, dict]: + cfg = merge_dict(DEFAULT_HEADS, config.get("heads", {}) or {}) + + # Compatibilidade com configs antigas. + if not bool(config.get("multi_head", True)): + print("[WARN] config.multi_head não está true. Este script vai treinar multi-head mesmo assim.") + + for name, hcfg in cfg.items(): + hcfg.setdefault("enabled", True) + hcfg.setdefault("ignore_index", ignore_index) + hcfg["num_classes"] = int(hcfg.get("num_classes", 2)) + hcfg["loss_weight"] = float(hcfg.get("loss_weight", 1.0)) + hcfg["mask_dir"] = str(hcfg.get("mask_dir", "masks")) + + active = {k: v for k, v in cfg.items() if bool(v.get("enabled", True))} + + if "semantic" not in active: + raise RuntimeError("A head 'semantic' deve estar habilitada neste contrato.") + if "vegetation" not in active: + raise RuntimeError("A head 'vegetation' deve estar habilitada neste contrato.") + if "cana" not in active: + raise RuntimeError("A head 'cana' deve estar habilitada neste contrato.") + + total_w = sum(float(v.get("loss_weight", 0.0)) for v in active.values()) + if total_w <= 0: + raise RuntimeError("Soma de loss_weight das heads deve ser > 0.") + + # Normaliza os pesos para não bagunçar magnitude da loss. + for v in active.values(): + v["loss_weight_norm"] = float(v.get("loss_weight", 0.0)) / total_w + + return active + + +# ============================================================ +# Dataset OAK-FCC-3 multi-head +# ============================================================ + +class OakFcc3TensorMultiHeadDataset(Dataset): + """ + Lê o contrato pós-normalização/split: + + root/ + group//tensors/.npy + group//masks/.npy + group//masks_vegetation/.npy + group//masks_cana/.npy + """ + + def __init__( + self, + root: str | Path, + heads_config: Dict[str, dict], + channels: int = 5, + channel_indices: Optional[List[int]] = None, + strict_channels: bool = False, + resize_hw: Optional[Tuple[int, int]] = None, + ): + self.root = Path(root) + self.heads_config = heads_config + self.channels = int(channels) + self.strict_channels = bool(strict_channels) + self.resize_hw = resize_hw + self.channel_indices = channel_indices + + self.samples = self._collect_samples() + + if not self.samples: + raise RuntimeError(f"Nenhuma amostra encontrada em: {self.root}") + + def _collect_samples(self): + samples = [] + + group_root = self.root / "group" + if not group_root.is_dir(): + raise RuntimeError(f"Pasta group não encontrada em: {self.root}") + + for group_dir in sorted(group_root.iterdir()): + if not group_dir.is_dir(): + continue + + tensors_dir = group_dir / "tensors" + metas_dir = group_dir / "metas" + previews_dir = group_dir / "previews" + + if not tensors_dir.is_dir(): + continue + + for tensor_path in sorted(tensors_dir.glob("*.npy")): + base = tensor_path.stem + + masks = {} + missing = [] + + for head_name, hcfg in self.heads_config.items(): + mask_dir_name = str(hcfg.get("mask_dir")) + + # Head target derivada: não precisa arquivo .npy físico + if mask_dir_name == "__derived_target__" or bool(hcfg.get("derived", False)): + masks[head_name] = None + continue + + mask_dir = group_dir / mask_dir_name + mask_path = mask_dir / f"{base}.npy" + + if not mask_path.exists(): + missing.append(f"{head_name}:{mask_path}") + else: + masks[head_name] = mask_path + + if missing: + print(f"[WARN] Pulando {tensor_path}, masks ausentes: {missing}") + continue + + meta_path = metas_dir / f"{base}.json" + preview_path = previews_dir / f"{base}.png" + + samples.append({ + "group": group_dir.name, + "base": base, + "tensor": tensor_path, + "masks": masks, + "meta": meta_path if meta_path.exists() else None, + "preview": preview_path if preview_path.exists() else None, + }) + + return samples + + def __len__(self): + return len(self.samples) + + def _resize_tensor(self, x: torch.Tensor): + if self.resize_hw is None: + return x + + h, w = self.resize_hw + + if x.shape[-2:] != (h, w): + x = F.interpolate( + x.unsqueeze(0), + size=(h, w), + mode="bilinear", + align_corners=False, + ).squeeze(0) + + return x + + def _resize_mask(self, y: torch.Tensor): + if self.resize_hw is None: + return y + + h, w = self.resize_hw + + if y.shape[-2:] != (h, w): + y = F.interpolate( + y.unsqueeze(0).unsqueeze(0).float(), + size=(h, w), + mode="nearest", + ).squeeze(0).squeeze(0).long() + + return y + + def __getitem__(self, idx): + s = self.samples[idx] + + x = np.load(str(s["tensor"])).astype(np.float32) + + if x.ndim != 3: + raise RuntimeError(f"Tensor inválido {s['tensor']}: shape={x.shape}") + + if self.channel_indices is not None: + max_idx = max(self.channel_indices) + if x.shape[0] <= max_idx: + raise RuntimeError( + f"Tensor {s['tensor']} tem {x.shape[0]} canais, " + f"mas precisa acessar índice {max_idx}. Shape={x.shape}" + ) + x = x[self.channel_indices, :, :] + else: + if x.shape[0] < self.channels: + raise RuntimeError( + f"Tensor {s['tensor']} tem {x.shape[0]} canais, " + f"mas config pediu {self.channels}." + ) + x = x[:self.channels] + + xt = torch.from_numpy(np.ascontiguousarray(x)).float() + xt = self._resize_tensor(xt) + + masks = {} + + # Primeiro carrega as heads com arquivo físico + for head_name, path in s["masks"].items(): + hcfg = self.heads_config[head_name] + + if path is None: + continue + + y = np.load(str(path)).astype(np.int64) + yt = torch.from_numpy(np.ascontiguousarray(y)).long() + yt = self._resize_mask(yt) + masks[head_name] = yt + + # Depois deriva a target, se existir no contrato + if "target" in self.heads_config: + if "vegetation" not in masks or "cana" not in masks: + raise RuntimeError("Head target requer masks vegetation e cana carregadas.") + + ignore_index = int(self.heads_config["target"].get("ignore_index", 255)) + + veg = masks["vegetation"] + cana = masks["cana"] + + valid = (veg != ignore_index) & (cana != ignore_index) + + target = torch.zeros_like(veg, dtype=torch.long) + target[(veg == 1) & (cana == 0)] = 1 + target[~valid] = ignore_index + + masks["target"] = target + + return { + "image": xt, + "masks": masks, + "group": s["group"], + "base": s["base"], + } + + +def collate_fn(batch): + imgs = torch.stack([b["image"] for b in batch], dim=0) + + head_names = list(batch[0]["masks"].keys()) + masks = { + h: torch.stack([b["masks"][h] for b in batch], dim=0) + for h in head_names + } + + meta = { + "group": [b["group"] for b in batch], + "base": [b["base"] for b in batch], + } + + return imgs, masks, meta + + +# ============================================================ +# Normalização +# ============================================================ + +class FixedNormalizer(nn.Module): + def __init__(self, mean: List[float], std: List[float]): + super().__init__() + mean_t = torch.tensor(mean, dtype=torch.float32).view(1, -1, 1, 1) + std_t = torch.tensor(std, dtype=torch.float32).view(1, -1, 1, 1) + self.register_buffer("mean", mean_t) + self.register_buffer("std", torch.clamp(std_t, min=1e-6)) + + def forward(self, x): + return (x - self.mean) / self.std + + +def normalize_per_batch(x: torch.Tensor, eps: float = 1e-6): + mean = x.mean(dim=(0, 2, 3), keepdim=True) + std = x.std(dim=(0, 2, 3), keepdim=True).clamp_min(eps) + return (x - mean) / std + + +def build_normalizer(config: dict, args, device: torch.device): + input_channel_names = get_input_channel_names(config) + input_channel_indices = get_input_channel_indices(config) + channels = len(input_channel_names) + model_family = config.get("modelo", "segformer") + model_name = config.get("model_name", "model") + stats_source_tag = config.get("stats_source_tag", f"stacked_raw{channels}") + + candidates = [] + + if args.norm_stats: + candidates.append(Path(args.norm_stats)) + + candidates.append(Path("dataset") / f"{config['resolucao'][0]}x{config['resolucao'][1]}" / "group" / "norm_stats.json") + candidates.append(Path("backup") / model_family / model_name / stats_source_tag / "norm_stats.json") + + for p in candidates: + if p.exists(): + stats = load_json(p) + mean = stats.get("mean", []) + std = stats.get("std", []) + stat_channels = stats.get("channels", []) + + if len(mean) < max(input_channel_indices) + 1 or len(std) < max(input_channel_indices) + 1: + raise RuntimeError( + f"norm_stats incompatível: {p} | " + f"precisa índices={input_channel_indices}, " + f"mean={len(mean)}, std={len(std)}" + ) + + mean = [mean[i] for i in input_channel_indices] + std = [std[i] for i in input_channel_indices] + + if stat_channels: + stat_channels = [stat_channels[i] for i in input_channel_indices] + + print(f"[NORM] usando stats fixos: {p}") + print(f"[NORM] selected_channels={stat_channels or input_channel_names}") + return FixedNormalizer(mean, std).to(device), str(p) + + print("[NORM] norm_stats não encontrado. Usando normalize_per_batch.") + return None, None + + +DEFAULT_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"] + +def get_input_channel_names(config: dict) -> List[str]: + """ + Define quais canais entram no modelo. + + Padrão: + channels=3 -> R,G,B + channels=4 -> R,G,B,RE + channels=5 -> R,G,B,RE,NIR + + Futuramente permite: + "input_channels": ["R", "G", "B", "NIR"] + """ + if "input_channels" in config: + names = [str(c).upper() for c in config["input_channels"]] + else: + n = int(config.get("channels", 5)) + names = DEFAULT_CHANNEL_ORDER[:n] + + invalid = [c for c in names if c not in DEFAULT_CHANNEL_ORDER] + if invalid: + raise RuntimeError(f"Canais inválidos em input_channels: {invalid}") + + return names + + +def get_input_channel_indices(config: dict) -> List[int]: + names = get_input_channel_names(config) + return [DEFAULT_CHANNEL_ORDER.index(c) for c in names] + + +# ============================================================ +# Modelo multi-head +# ============================================================ + +def patch_segformer_encoder_input_channels(segformer_encoder: nn.Module, in_ch: int): + """ + Altera o primeiro patch embedding do SegFormer para aceitar C canais. + Inicializa canais extras pela média dos pesos RGB. + """ + 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): + """ + Troca o classifier final da decode_head do HuggingFace SegFormer. + """ + if not hasattr(decode_head, "classifier"): + raise RuntimeError("decode_head sem atributo classifier. Estrutura SegFormer inesperada.") + + 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, + ) + + nn.init.xavier_uniform_(new.weight) + if new.bias is not None: + nn.init.zeros_(new.bias) + + decode_head.classifier = new + return decode_head + + +class MultiHeadSegFormer(nn.Module): + def __init__( + self, + backbone: str, + channels: int, + heads_config: Dict[str, dict], + semantic_id2label: Dict[int, str], + semantic_label2id: Dict[str, int], + ): + super().__init__() + + semantic_classes = int(heads_config["semantic"].get("num_classes", len(semantic_id2label))) + + base = SegformerForSemanticSegmentation.from_pretrained( + backbone, + num_labels=semantic_classes, + id2label={int(k): str(v) for k, v in semantic_id2label.items()}, + label2id={str(k): int(v) for k, v in semantic_label2id.items()}, + ignore_mismatched_sizes=True, + ) + + patch_segformer_encoder_input_channels(base.segformer, channels) + base.config.num_channels = int(channels) + + self.segformer = base.segformer + self.heads_config = heads_config + self.decode_heads = nn.ModuleDict() + + 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) -> Dict[str, torch.Tensor]: + outputs = self.segformer( + pixel_values=pixel_values, + output_hidden_states=True, + return_dict=True, + ) + + hidden_states = outputs.hidden_states + logits = {} + + for head_name, head in self.decode_heads.items(): + logits[head_name] = head(hidden_states) + + return logits + + +def build_model( + backbone: str, + channels: int, + heads_config: Dict[str, dict], + semantic_id2label: Dict[int, str], + semantic_label2id: Dict[str, int], +): + model = MultiHeadSegFormer( + backbone=backbone, + channels=channels, + heads_config=heads_config, + semantic_id2label=semantic_id2label, + semantic_label2id=semantic_label2id, + ) + return model + + +# ============================================================ +# Métricas / loss +# ============================================================ + +@torch.no_grad() +def update_confusion_matrix(cm, preds, labels, num_classes, ignore_index=255): + preds = preds.reshape(-1) + labels = labels.reshape(-1) + + valid = labels != ignore_index + preds = preds[valid] + labels = labels[valid] + + valid2 = (labels >= 0) & (labels < num_classes) & (preds >= 0) & (preds < num_classes) + preds = preds[valid2] + labels = labels[valid2] + + if labels.numel() == 0: + return + + 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, eps=1e-6): + 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, eps=1e-6): + cm = cm.float() + return float(torch.diag(cm).sum() / (cm.sum() + eps)) + + +def dice_loss(logits, target, num_classes, ignore_index=255, smooth=1.0): + probs = torch.softmax(logits, dim=1) + valid = target != ignore_index + + if valid.sum() == 0: + return logits.new_tensor(0.0) + + target_clamped = target.clone() + target_clamped[~valid] = 0 + target_clamped = target_clamped.long() + + target_1h = F.one_hot(target_clamped, num_classes=num_classes) + target_1h = target_1h.permute(0, 3, 1, 2).float() + + valid_f = valid.unsqueeze(1).float() + probs = probs * valid_f + target_1h = target_1h * valid_f + + dims = (0, 2, 3) + inter = (probs * target_1h).sum(dims) + den = probs.sum(dims) + target_1h.sum(dims) + + dice = (2.0 * inter + smooth) / (den + smooth) + return 1.0 - dice.mean() + + +def build_target_teacher_from_heads( + logits_by_head: Dict[str, torch.Tensor], + target_shape: Tuple[int, int], + config: dict, + erva_id: int = 2, +): + cfg = config.get("target_distillation", {}) or {} + + if not bool(cfg.get("enabled", False)): + return None + + required = ("semantic", "vegetation", "cana") + if any(k not in logits_by_head for k in required): + return None + + sem = logits_by_head["semantic"] + veg = logits_by_head["vegetation"] + cana = logits_by_head["cana"] + + if sem.shape[-2:] != target_shape: + sem = F.interpolate(sem, size=target_shape, mode="bilinear", align_corners=False) + if veg.shape[-2:] != target_shape: + veg = F.interpolate(veg, size=target_shape, mode="bilinear", align_corners=False) + if cana.shape[-2:] != target_shape: + cana = F.interpolate(cana, size=target_shape, mode="bilinear", align_corners=False) + + p_sem = torch.softmax(sem, dim=1) + p_veg = torch.softmax(veg, dim=1) + p_cana = torch.softmax(cana, dim=1) + + p_sem_erva = p_sem[:, int(erva_id), :, :] + p_veg_pos = p_veg[:, 1, :, :] + p_cana_pos = p_cana[:, 1, :, :] + + w_sem_erva = float(cfg.get("w_sem_erva", 0.45)) + w_veg_not_cana = float(cfg.get("w_veg_not_cana", 0.35)) + w_veg_suppressed = float(cfg.get("w_veg_suppressed", 0.20)) + power = float(cfg.get("cana_suppression_power", 1.5)) + + not_cana = torch.clamp(1.0 - p_cana_pos, 0.0, 1.0) + + teacher = ( + w_sem_erva * p_sem_erva + + w_veg_not_cana * p_veg_pos * not_cana + + w_veg_suppressed * p_veg_pos * torch.pow(not_cana, power) + ) + + teacher_min = float(cfg.get("teacher_min", 0.0)) + teacher_max = float(cfg.get("teacher_max", 1.0)) + teacher = torch.clamp(teacher, teacher_min, teacher_max) + + if bool(cfg.get("detach_teacher", True)): + teacher = teacher.detach() + + return teacher + + +def estimate_head_class_weights( + ds: Dataset, + head_name: str, + num_classes: int, + ignore_index: int = 255, + max_samples: int = 800, + seed: int = 42, +): + rng = np.random.default_rng(seed) + n = min(len(ds), max_samples) + idxs = rng.choice(len(ds), size=n, replace=False) + + counts = np.zeros(num_classes, dtype=np.float64) + + for i in idxs: + item = ds[i] + m = item["masks"][head_name].numpy().reshape(-1) + m = m[m != ignore_index] + m = m[(m >= 0) & (m < num_classes)] + + if m.size > 0: + 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), counts + + +def build_criterions( + ds_train: Dataset, + heads_config: Dict[str, dict], + class_weights_mode: str, + device: torch.device, + seed: int, +): + criterions = {} + weight_debug = {} + + for head_name, hcfg in heads_config.items(): + num_classes = int(hcfg["num_classes"]) + ignore_index = int(hcfg.get("ignore_index", 255)) + + mode = str(class_weights_mode).lower() + weights = None + counts = None + + if mode == "none": + weights = None + + elif mode == "auto": + w, counts = estimate_head_class_weights( + ds_train, + head_name=head_name, + num_classes=num_classes, + ignore_index=ignore_index, + seed=seed, + ) + weights = w.to(device) + print(f"[LOSS:{head_name}] class counts:", counts.astype(int).tolist()) + print(f"[LOSS:{head_name}] class weights:", w.cpu().numpy().round(3).tolist()) + + else: + # Formato opcional: + # semantic=1,2,3;vegetation=1,2;cana=1,4 + # ou, para compatibilidade, uma lista aplicada só na semantic. + parsed = parse_class_weights_string(class_weights_mode) + if head_name in parsed: + parts = parsed[head_name] + elif head_name == "semantic" and "__single__" in parsed: + parts = parsed["__single__"] + else: + parts = None + + if parts is not None: + if len(parts) != num_classes: + raise RuntimeError( + f"Pesos da head {head_name} precisam ter {num_classes} valores. Veio: {parts}" + ) + weights = torch.tensor(parts, dtype=torch.float32, device=device) + + criterions[head_name] = nn.CrossEntropyLoss( + weight=weights, + ignore_index=ignore_index, + ) + + weight_debug[head_name] = { + "weights": weights.detach().cpu().tolist() if weights is not None else None, + "counts": counts.astype(int).tolist() if counts is not None else None, + } + + return criterions, weight_debug + + +def parse_class_weights_string(s: str) -> Dict[str, List[float]]: + out = {} + txt = str(s).strip() + + if not txt: + return out + + if "=" not in txt: + out["__single__"] = [float(x) for x in txt.split(",") if x.strip()] + return out + + for block in txt.split(";"): + block = block.strip() + if not block: + continue + k, v = block.split("=", 1) + out[k.strip()] = [float(x) for x in v.split(",") if x.strip()] + + return out + + +@torch.no_grad() +def update_operational_target_cm( + cm_target, + pred_veg: torch.Tensor, + pred_cana: torch.Tensor, + gt_veg: torch.Tensor, + gt_cana: torch.Tensor, + ignore_index: int = 255, +): + """ + Métrica operacional: + target/alvo = vegetação viva e não cana + + Usa heads vegetation e cana. + Classe 0 = não alvo + Classe 1 = alvo/pulverizável + """ + valid = (gt_veg != ignore_index) & (gt_cana != ignore_index) + + gt_target = ((gt_veg == 1) & (gt_cana == 0)).long() + pred_target = ((pred_veg == 1) & (pred_cana == 0)).long() + + gt_target = gt_target[valid] + pred_target = pred_target[valid] + + if gt_target.numel() == 0: + return + + idx = gt_target.reshape(-1) * 2 + pred_target.reshape(-1) + bins = torch.bincount(idx, minlength=4) + cm_target += bins.view(2, 2) + + +def compute_losses_for_batch( + logits_by_head: Dict[str, torch.Tensor], + masks_by_head: Dict[str, torch.Tensor], + heads_config: Dict[str, dict], + criterions: Dict[str, nn.Module], + dice_weight: float = 0.30, + config: Optional[dict] = None, + epoch: Optional[int] = None, +): + total = None + loss_parts = {} + + cfg = config or {} + target_distill_cfg = cfg.get("target_distillation", {}) or {} + target_distill_enabled = bool(target_distill_cfg.get("enabled", False)) + target_distill_ramp = get_target_distill_ramp_factor(cfg, epoch) + target_distill_enabled = target_distill_enabled and target_distill_ramp > 0.0 + + for head_name, logits in logits_by_head.items(): + target = masks_by_head[head_name] + hcfg = heads_config[head_name] + num_classes = int(hcfg["num_classes"]) + ignore_index = int(hcfg.get("ignore_index", 255)) + head_weight = float(hcfg.get("loss_weight_norm", 1.0)) + + if logits.shape[-2:] != target.shape[-2:]: + logits = F.interpolate( + logits, + size=target.shape[-2:], + mode="bilinear", + align_corners=False, + ) + + ce = criterions[head_name](logits, target) + dl = dice_loss( + logits=logits, + target=target, + num_classes=num_classes, + ignore_index=ignore_index, + smooth=1.0, + ) + + head_loss = (1.0 - dice_weight) * ce + dice_weight * dl + + if head_name == "target" and target_distill_enabled: + teacher = build_target_teacher_from_heads( + logits_by_head=logits_by_head, + target_shape=target.shape[-2:], + config=cfg, + erva_id=int( + cfg.get("heads", {}) + .get("semantic", {}) + .get("classes", {}) + .get("erva", 2) + ), + ) + + if teacher is not None: + logits_target = logits + + if logits_target.shape[-2:] != target.shape[-2:]: + logits_target = F.interpolate( + logits_target, + size=target.shape[-2:], + mode="bilinear", + align_corners=False, + ) + + logits_binary = logits_target[:, 1, :, :] - logits_target[:, 0, :, :] + + valid = target != ignore_index + + if valid.any(): + logits_binary_v = logits_binary[valid] + teacher_v = teacher[valid].clamp(0.0, 1.0) + + distill = F.binary_cross_entropy_with_logits( + logits_binary_v, + teacher_v, + ) + + hard_weight = float(target_distill_cfg.get("hard_weight", 0.70)) + distill_weight = float(target_distill_cfg.get("distill_weight", 0.30)) + + # Aplica ramp-up só na parte destilada + distill_weight = distill_weight * float(target_distill_ramp) + + denom = max(hard_weight + distill_weight, 1e-6) + hard_weight = hard_weight / denom + distill_weight = distill_weight / denom + + head_loss = hard_weight * head_loss + distill_weight * distill + + loss_parts[f"{head_name}_distill"] = distill.detach() + loss_parts[f"{head_name}_distill_ramp"] = torch.as_tensor( + target_distill_ramp, + device=distill.device, + dtype=distill.dtype, + ).detach() + + weighted = head_weight * head_loss + + total = weighted if total is None else total + weighted + + loss_parts[f"{head_name}_loss"] = head_loss.detach() + loss_parts[f"{head_name}_ce"] = ce.detach() + loss_parts[f"{head_name}_dice"] = dl.detach() + + return total, loss_parts + + +def resize_logits_to_target_if_needed( + logits: torch.Tensor, + target: torch.Tensor, +) -> torch.Tensor: + """ + Garante que logits estejam no mesmo HxW da máscara. + + Usado para métricas, sem depender do efeito colateral da loss. + """ + if logits.shape[-2:] != target.shape[-2:]: + logits = F.interpolate( + logits, + size=target.shape[-2:], + mode="bilinear", + align_corners=False, + ) + + return logits + + +def get_target_distill_ramp_factor(config: dict, epoch: Optional[int]) -> float: + cfg = (config or {}).get("target_distillation", {}) or {} + + if not bool(cfg.get("enabled", False)): + return 0.0 + + if not bool(cfg.get("rampup_enabled", True)): + return 1.0 + + if epoch is None: + return 1.0 + + start_epoch = int(cfg.get("start_epoch", 8)) + rampup_epochs = int(cfg.get("rampup_epochs", 12)) + + if epoch < start_epoch: + return 0.0 + + if rampup_epochs <= 0: + return 1.0 + + t = (float(epoch) - float(start_epoch) + 1.0) / float(rampup_epochs) + t = max(0.0, min(1.0, t)) + + # Smoothstep: sobe suave, sem tranco + return float(t * t * (3.0 - 2.0 * t)) + + +# ============================================================ +# Train / Val +# ============================================================ + +def run_one_epoch( + model, + loader, + optimizer, + device, + heads_config, + criterions, + amp, + scaler, + train, + grad_accum=1, + normalizer=None, + dice_weight=0.30, + config=None, + epoch: Optional[int] = None, +): + model.train(train) + + total_loss = 0.0 + n_batches = 0 + + loss_sums = {} + + cms = { + head_name: torch.zeros( + (int(hcfg["num_classes"]), int(hcfg["num_classes"])), + dtype=torch.int64, + device=device, + ) + for head_name, hcfg in heads_config.items() + } + + cm_target = torch.zeros((2, 2), dtype=torch.int64, device=device) + + t0 = time.time() + + with torch.set_grad_enabled(train): + if train and optimizer is not None: + optimizer.zero_grad(set_to_none=True) + + for step, (imgs, masks, _meta) in enumerate(loader): + imgs = imgs.to(device, non_blocking=True) + masks = {k: v.to(device, non_blocking=True) for k, v in masks.items()} + + if normalizer is not None: + imgs = normalizer(imgs) + else: + imgs = normalize_per_batch(imgs) + + with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): + logits_by_head = model(pixel_values=imgs) + + loss, loss_parts = compute_losses_for_batch( + logits_by_head=logits_by_head, + masks_by_head=masks, + heads_config=heads_config, + criterions=criterions, + dice_weight=dice_weight, + config=config, + epoch=epoch, + ) + + 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 + 1) % grad_accum) == 0 or (step + 1) == len(loader): + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + else: + loss.backward() + + if ((step + 1) % grad_accum) == 0 or (step + 1) == len(loader): + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + loss_value = float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0) + total_loss += loss_value + n_batches += 1 + + for k, v in loss_parts.items(): + loss_sums[k] = loss_sums.get(k, 0.0) + float(v.item()) + + with torch.no_grad(): + preds_by_head = {} + + for head_name, logits in logits_by_head.items(): + if head_name not in masks: + continue + + target_mask = masks[head_name] + + logits_metric = resize_logits_to_target_if_needed( + logits=logits, + target=target_mask, + ) + + preds = torch.argmax(logits_metric, dim=1) + preds_by_head[head_name] = preds + + update_confusion_matrix( + cms[head_name], + preds, + target_mask, + int(heads_config[head_name]["num_classes"]), + int(heads_config[head_name].get("ignore_index", 255)), + ) + + # Métrica operacional antiga: + # target_op = vegetation == 1 AND cana == 0 + # Continua útil para comparar com a nova head target direta. + if "vegetation" in preds_by_head and "cana" in preds_by_head: + update_operational_target_cm( + cm_target=cm_target, + pred_veg=preds_by_head["vegetation"], + pred_cana=preds_by_head["cana"], + gt_veg=masks["vegetation"], + gt_cana=masks["cana"], + ignore_index=int(heads_config["vegetation"].get("ignore_index", 255)), + ) + + avg_loss = total_loss / max(1, n_batches) + + metrics = { + "loss": avg_loss, + "time_s": time.time() - t0, + "heads": {}, + "loss_parts": {k: v / max(1, n_batches) for k, v in loss_sums.items()}, + } + + for head_name, cm in cms.items(): + miou, iou_per_class = compute_iou_from_cm(cm) + acc = compute_pixel_acc_from_cm(cm) + metrics["heads"][head_name] = { + "miou": miou, + "iou_per_class": iou_per_class, + "acc": acc, + "cm": cm.detach().cpu().tolist(), + } + + target_miou, target_iou = compute_iou_from_cm(cm_target) + target_acc = compute_pixel_acc_from_cm(cm_target) + + metrics["operational_target"] = { + "miou": target_miou, + "iou_background": target_iou[0], + "iou_target": target_iou[1], + "acc": target_acc, + "cm": cm_target.detach().cpu().tolist(), + } + + return metrics + + +# ============================================================ +# Checkpoint / logs +# ============================================================ + +def save_checkpoint(path, model, optimizer, scaler, epoch, best: dict, extra=None): + ckpt = { + "epoch": epoch, + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "best": best, + } + + if scaler is not None: + ckpt["scaler"] = scaler.state_dict() + + if extra: + ckpt["extra"] = extra + + torch.save(ckpt, path) + + +def load_checkpoint(path, model, optimizer=None, scaler=None, map_location="cpu"): + ckpt = torch.load(path, map_location=map_location, weights_only=False) + model.load_state_dict(ckpt["model"], 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 + + +def append_train_log(path: Path, row: dict): + ensure_dir(path.parent) + exists = path.exists() + + with path.open("a", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=list(row.keys())) + if not exists: + w.writeheader() + w.writerow(row) + + +def flatten_epoch_log(epoch: int, lr: float, tr: dict, va: dict, best: dict) -> dict: + row = { + "epoch": epoch, + "lr": lr, + "train_loss": tr["loss"], + "val_loss": va["loss"], + "best_score": best.get("score", -1.0), + "best_target_iou": best.get("target_iou", -1.0), + "best_cana_head_iou": best.get("cana_head_iou", -1.0), + "best_semantic_miou": best.get("semantic_miou", -1.0), + "best_target_head_iou": best.get("target_head_iou", -1.0), + "best_operational_target_iou": best.get("operational_target_iou", -1.0), + } + + for prefix, obj in (("train", tr), ("val", va)): + for head_name, hm in obj.get("heads", {}).items(): + row[f"{prefix}_{head_name}_miou"] = hm.get("miou") + row[f"{prefix}_{head_name}_acc"] = hm.get("acc") + + ious = hm.get("iou_per_class", []) or [] + for i, v in enumerate(ious): + row[f"{prefix}_{head_name}_iou_{i}"] = v + + op = obj.get("operational_target", {}) or {} + row[f"{prefix}_target_miou"] = op.get("miou") + row[f"{prefix}_target_iou"] = op.get("iou_target") + row[f"{prefix}_target_acc"] = op.get("acc") + + for k, v in obj.get("loss_parts", {}).items(): + row[f"{prefix}_{k}"] = v + + return row + + +# ============================================================ +# Pretty print +# ============================================================ + +def pretty_iou(names: Dict[int, str], iou_list: List[float]): + return " | ".join([ + f"{names.get(i, i)}:{v:.3f}" + for i, v in enumerate(iou_list) + ]) + + +def binary_iou_text(head_name: str, iou_list: List[float]): + if head_name == "vegetation": + names = {0: "bg", 1: "veg"} + elif head_name == "cana": + names = {0: "not_cana", 1: "cana"} + elif head_name == "target": + names = {0: "bg", 1: "target"} + else: + names = {0: "0", 1: "1"} + + return pretty_iou(names, iou_list) + + +def compute_selection_score(va: dict) -> dict: + heads = va.get("heads", {}) + op = va.get("operational_target", {}) + + # Preferir target direta, se existir + target_head_iou = 0.0 + if "target" in heads: + vals = heads["target"].get("iou_per_class", []) or [] + if len(vals) > 1: + target_head_iou = safe_float(vals[1], 0.0) + + operational_target_iou = safe_float(op.get("iou_target"), 0.0) + + # Se target head existir, ela manda. Se não existir, usa operacional antigo. + target_iou = target_head_iou if "target" in heads else operational_target_iou + + cana_iou = 0.0 + if "cana" in heads: + vals = heads["cana"].get("iou_per_class", []) or [] + if len(vals) > 1: + cana_iou = safe_float(vals[1], 0.0) + + veg_miou = safe_float(heads.get("vegetation", {}).get("miou"), 0.0) + semantic_miou = safe_float(heads.get("semantic", {}).get("miou"), 0.0) + + score = ( + 0.50 * target_iou + + 0.25 * cana_iou + + 0.15 * veg_miou + + 0.10 * semantic_miou + ) + + return { + "score": float(score), + "target_iou": float(target_iou), + "target_head_iou": float(target_head_iou), + "operational_target_iou": float(operational_target_iou), + "cana_head_iou": float(cana_iou), + "vegetation_miou": float(veg_miou), + "semantic_miou": float(semantic_miou), + } + + +# ============================================================ +# Main +# ============================================================ + +def main(): + parser = argparse.ArgumentParser() + + parser.add_argument("--config", default="config.json") + parser.add_argument("--epochs", type=int, default=80) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--lr", type=float, default=3e-5) + parser.add_argument("--wd", type=float, default=0.01) + parser.add_argument("--num_workers", type=int, default=2) + + parser.add_argument("--amp", action="store_true") + parser.add_argument("--amp_val", action="store_true") + parser.add_argument("--grad_accum", type=int, default=4) + parser.add_argument("--grad_ckpt", action="store_true") + + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--ignore_index", type=int, default=None) + + parser.add_argument("--class_weights", default="auto") + parser.add_argument("--dice_weight", type=float, default=0.30) + + parser.add_argument("--norm_stats", default=None) + parser.add_argument("--src-root", default="dataset/split") + parser.add_argument("--save-every", type=int, default=10) + + parser.add_argument("--resume", action="store_true") + parser.add_argument("--resume-ckpt", default=None) + + parser.add_argument("--early-stop", type=int, default=25) + + args = parser.parse_args() + + set_seed(args.seed) + + config = load_json(args.config) + + W, H = config["resolucao"] + channels = int(config.get("channels", 5)) + backbone = config.get("backbone", "nvidia/mit-b1") + fusion_mode = config.get("fusion_mode", "stacked") + + input_channel_names = get_input_channel_names(config) + input_channel_indices = get_input_channel_indices(config) + channels = len(input_channel_names) + + print(f"Input channels: {input_channel_names} idx={input_channel_indices}") + + if fusion_mode != "stacked": + raise RuntimeError("Este script é para fusion_mode='stacked'.") + + model_family = config.get("modelo", "segformer") + model_name = config.get("model_name", "test") + stats_source_tag = config.get("stats_source_tag", f"stacked_raw{channels}") + + save_dir = Path("backup") / model_family / model_name / f"{fusion_mode}_raw{channels}" + ensure_dir(save_dir) + + labelmap_path = Path("dataset") / "labelmap.txt" + if not labelmap_path.exists(): + raise RuntimeError(f"Labelmap não encontrado: {labelmap_path}") + + semantic_id2label, semantic_label2id, ignore_from_labelmap = load_labelmap(str(labelmap_path)) + ignore_index = int(args.ignore_index if args.ignore_index is not None else ignore_from_labelmap) + + heads_config = build_heads_config(config, ignore_index=ignore_index) + + # Garante que a semantic conhece o num_classes real do labelmap. + heads_config["semantic"]["num_classes"] = int(len(semantic_id2label)) + heads_config["semantic"]["ignore_index"] = int(ignore_index) + + print("==========================================") + print("Train SegFormer OAK-FCC-3 Multi-Head") + print(f"Backbone : {backbone}") + print(f"Save dir : {save_dir}") + print(f"Split root : {args.src_root}") + print(f"Resolution : {W}x{H}") + print(f"Channels : {channels}") + print(f"Semantic : {heads_config['semantic']['num_classes']} -> {semantic_id2label}") + print(f"Ignore index : {ignore_index}") + print("Heads:") + for name, hcfg in heads_config.items(): + print( + f" - {name}: classes={hcfg['num_classes']} " + f"mask_dir={hcfg['mask_dir']} " + f"loss_weight={hcfg['loss_weight']} norm={hcfg['loss_weight_norm']:.3f}" + ) + print("==========================================") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + + train_root = Path(args.src_root) / "train" + val_root = Path(args.src_root) / "val" + + resize_hw = (H, W) + + ds_train = OakFcc3TensorMultiHeadDataset( + train_root, + heads_config=heads_config, + channels=channels, + channel_indices=input_channel_indices, + strict_channels=False, + resize_hw=resize_hw, + ) + + ds_val = OakFcc3TensorMultiHeadDataset( + val_root, + heads_config=heads_config, + channels=channels, + channel_indices=input_channel_indices, + strict_channels=False, + resize_hw=resize_hw, + ) + + print(f"[DATA] train={len(ds_train)} | val={len(ds_val)}") + + dl_train = DataLoader( + ds_train, + batch_size=args.batch, + shuffle=True, + num_workers=args.num_workers, + pin_memory=True, + collate_fn=collate_fn, + drop_last=True if len(ds_train) >= args.batch else False, + ) + + dl_val = DataLoader( + ds_val, + batch_size=1, + shuffle=False, + num_workers=max(0, args.num_workers // 2), + pin_memory=True, + collate_fn=collate_fn, + drop_last=False, + ) + + normalizer, norm_stats_path = build_normalizer(config, args, device) + + model = build_model( + backbone=backbone, + channels=channels, + heads_config=heads_config, + semantic_id2label=semantic_id2label, + semantic_label2id=semantic_label2id, + ) + + if args.grad_ckpt: + try: + model.segformer.gradient_checkpointing_enable() + print("[MODEL] gradient checkpointing enabled") + except Exception as e: + print(f"[WARN] gradient checkpointing não suportado: {e}") + + model.to(device) + + criterions, weight_debug = build_criterions( + ds_train=ds_train, + heads_config=heads_config, + class_weights_mode=args.class_weights, + device=device, + seed=args.seed, + ) + + optimizer = torch.optim.AdamW( + model.parameters(), + lr=args.lr, + weight_decay=args.wd, + ) + + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, + mode="min", + factor=0.5, + patience=6, + threshold=1e-4, + ) + + scaler = GradScaler(enabled=args.amp and device.type == "cuda") + + last_path = save_dir / "last.pt" + best_score_path = save_dir / "best_score.pt" + best_target_path = save_dir / "best_target.pt" + best_cana_path = save_dir / "best_cana_head.pt" + best_semantic_path = save_dir / "best_semantic_miou.pt" + train_log_path = save_dir / "train_log.csv" + + save_json(save_dir / "train_config_snapshot.json", { + "config": config, + "args": vars(args), + "semantic_id2label": semantic_id2label, + "semantic_label2id": semantic_label2id, + "heads_config": heads_config, + "ignore_index": ignore_index, + "norm_stats_path": norm_stats_path, + "class_weight_debug": weight_debug, + "input_channel_names": input_channel_names, + "input_channel_indices": input_channel_indices, + "checkpoint_selection_score": { + "score": "0.45*target_iou + 0.30*cana_head_iou + 0.15*vegetation_miou + 0.10*semantic_miou" + }, + }) + + start_epoch = 1 + best = { + "score": -1.0, + "target_iou": -1.0, + "target_head_iou": -1.0, + "operational_target_iou": -1.0, + "cana_head_iou": -1.0, + "semantic_miou": -1.0, + "vegetation_miou": -1.0, + } + epochs_without_improve = 0 + + resume_path = Path(args.resume_ckpt) if args.resume_ckpt else last_path + + if args.resume and resume_path.exists(): + ckpt = load_checkpoint( + resume_path, + model, + optimizer, + scaler=scaler, + map_location="cpu", + ) + start_epoch = int(ckpt["epoch"]) + 1 + best = ckpt.get("best", best) + print(f"[RESUME] {resume_path} epoch={start_epoch}") + + for epoch in range(start_epoch, args.epochs + 1): + lr_now = optimizer.param_groups[0]["lr"] + print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====") + + if device.type == "cuda": + torch.cuda.empty_cache() + + tr = run_one_epoch( + model=model, + loader=dl_train, + optimizer=optimizer, + device=device, + heads_config=heads_config, + criterions=criterions, + amp=args.amp, + scaler=scaler, + train=True, + grad_accum=max(1, args.grad_accum), + normalizer=normalizer, + dice_weight=args.dice_weight, + config=config, + epoch=epoch, + ) + + if device.type == "cuda": + torch.cuda.empty_cache() + + va = run_one_epoch( + model=model, + loader=dl_val, + optimizer=None, + device=device, + heads_config=heads_config, + criterions=criterions, + amp=args.amp_val, + scaler=None, + train=False, + grad_accum=1, + normalizer=normalizer, + dice_weight=args.dice_weight, + config=config, + epoch=epoch, + ) + + scheduler.step(va["loss"]) + + score_now = compute_selection_score(va) + + print( + f"TRAIN: loss={tr['loss']:.4f} " + f"sem_mIoU={tr['heads']['semantic']['miou']:.4f} " + f"veg_mIoU={tr['heads']['vegetation']['miou']:.4f} " + f"cana_mIoU={tr['heads']['cana']['miou']:.4f} " + f"targetIoU={tr['operational_target']['iou_target']:.4f} " + f"t={tr['time_s']:.1f}s" + ) + target_head_iou_val = 0.0 + if "target" in va["heads"]: + vals = va["heads"]["target"].get("iou_per_class", []) or [] + if len(vals) > 1: + target_head_iou_val = vals[1] + print( + f"VAL : loss={va['loss']:.4f} " + f"sem_mIoU={va['heads']['semantic']['miou']:.4f} " + f"veg_mIoU={va['heads']['vegetation']['miou']:.4f} " + f"cana_mIoU={va['heads']['cana']['miou']:.4f} " + f"targetHeadIoU={target_head_iou_val:.4f} " + f"opTargetIoU={va['operational_target']['iou_target']:.4f} " + f"score={score_now['score']:.4f} " + f"t={va['time_s']:.1f}s" + ) + + print("IoU semantic:", pretty_iou(semantic_id2label, va["heads"]["semantic"]["iou_per_class"])) + print("IoU vegetation:", binary_iou_text("vegetation", va["heads"]["vegetation"]["iou_per_class"])) + print("IoU cana_head:", binary_iou_text("cana", va["heads"]["cana"]["iou_per_class"])) + print("IoU target:", f"bg:{va['operational_target']['iou_background']:.3f} | alvo:{va['operational_target']['iou_target']:.3f}") + if "target" in va["heads"]: + print("IoU target_head:", binary_iou_text("target", va["heads"]["target"]["iou_per_class"])) + + save_checkpoint( + last_path, + model, + optimizer, + scaler=scaler, + epoch=epoch, + best=best, + extra={ + "train": tr, + "val": va, + "score_now": score_now, + }, + ) + + if args.save_every > 0 and epoch % args.save_every == 0: + save_checkpoint( + save_dir / f"epoch_{epoch:04d}.pt", + model, + optimizer, + scaler=scaler, + epoch=epoch, + best=best, + ) + + improved = False + + if score_now["score"] > best.get("score", -1.0): + best["score"] = score_now["score"] + improved = True + save_checkpoint( + best_score_path, + model, + optimizer, + scaler=scaler, + epoch=epoch, + best=best, + extra={"val": va, "score_now": score_now}, + ) + print(f"[BEST SCORE] {best['score']:.4f} -> {best_score_path}") + + if score_now["target_iou"] > best.get("target_iou", -1.0): + best["target_iou"] = score_now["target_iou"] + improved = True + save_checkpoint( + best_target_path, + model, + optimizer, + scaler=scaler, + epoch=epoch, + best=best, + extra={"val": va, "score_now": score_now}, + ) + print(f"[BEST TARGET] {best['target_iou']:.4f} -> {best_target_path}") + + if score_now["cana_head_iou"] > best.get("cana_head_iou", -1.0): + best["cana_head_iou"] = score_now["cana_head_iou"] + improved = True + save_checkpoint( + best_cana_path, + model, + optimizer, + scaler=scaler, + epoch=epoch, + best=best, + extra={"val": va, "score_now": score_now}, + ) + print(f"[BEST CANA HEAD] {best['cana_head_iou']:.4f} -> {best_cana_path}") + + if score_now["semantic_miou"] > best.get("semantic_miou", -1.0): + best["semantic_miou"] = score_now["semantic_miou"] + improved = True + save_checkpoint( + best_semantic_path, + model, + optimizer, + scaler=scaler, + epoch=epoch, + best=best, + extra={"val": va, "score_now": score_now}, + ) + print(f"[BEST SEMANTIC] {best['semantic_miou']:.4f} -> {best_semantic_path}") + + append_train_log( + train_log_path, + flatten_epoch_log( + epoch=epoch, + lr=lr_now, + tr=tr, + va=va, + best=best, + ), + ) + + if improved: + epochs_without_improve = 0 + else: + epochs_without_improve += 1 + + if args.early_stop > 0 and epochs_without_improve >= args.early_stop: + print(f"[EARLY STOP] {epochs_without_improve} épocas sem melhora.") + break + + print("\nTreino finalizado.") + print(f"Best score : {best.get('score', -1.0):.4f}") + print(f"Best target IoU : {best.get('target_iou', -1.0):.4f}") + print(f"Best cana-head IoU: {best.get('cana_head_iou', -1.0):.4f}") + print(f"Best semantic mIoU: {best.get('semantic_miou', -1.0):.4f}") + print(f"Save dir : {save_dir}") + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/datasets/oak-fcc-3/_9_test_multihead.py b/Python/OAK/datasets/oak-fcc-3/_9_test_multihead.py new file mode 100644 index 000000000..ac337a511 --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/_9_test_multihead.py @@ -0,0 +1,1548 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +_9_test_infer_multihead.py + +Teste/visualização do SegFormer OAK-FCC-3 Multi-Head. + +Contrato esperado após normalize + split: + + dataset/split/val/group// + tensors/.npy # CHW float32 [R,G,B,RE,NIR] + masks/.npy # semantic: 0=chao, 1=cana, 2=erva, 255=ignore + masks_vegetation/.npy # vegetation: 0=background, 1=vegetation, 255=ignore + masks_cana/.npy # cana: 0=not_cana, 1=cana, 255=ignore + metas/.json + previews/.png + +Mostra: + - RGB preview do tensor + - GT/pred/overlay semantic + - GT/pred/overlay vegetation + - GT/pred/overlay cana + - alvo operacional = vegetation == 1 AND cana == 0 + - mapas de confiança das heads binárias + +Exemplo: + +python .\_9_test_infer_multihead.py ^ + --config config.json ^ + --split_folder val ^ + --ckpt backup\segformer_b1\test_multi\stacked_raw5_multihead\best_score.pt + +Controles: + D / seta direita : próxima amostra + A / seta esquerda: amostra anterior + S : salvar painel atual em --out_dir + SPACE : alterna modo compacto/detalhado + Q / ESC : sair +""" + +from __future__ import annotations + +import argparse +import copy +import json +import time +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +import cv2 +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import SegformerForSemanticSegmentation + +from core.raw_processor_core import RawProcessorCore + + +# ============================================================ +# Config default +# ============================================================ + +DEFAULT_HEADS = { + "semantic": { + "enabled": True, + "type": "multiclass", + "num_classes": 3, + "mask_dir": "masks", + "classes": {"chao": 0, "cana": 1, "erva": 2}, + "ignore_index": 255, + }, + "vegetation": { + "enabled": True, + "type": "binary", + "num_classes": 2, + "mask_dir": "masks_vegetation", + "classes": {"background": 0, "vegetation": 1}, + "ignore_index": 255, + }, + "cana": { + "enabled": True, + "type": "binary", + "num_classes": 2, + "mask_dir": "masks_cana", + "classes": {"not_cana": 0, "cana": 1}, + "ignore_index": 255, + }, + "target": { + "enabled": True, + "type": "binary", + "num_classes": 2, + "mask_dir": "__derived_target__", + "classes": { + "background": 0, + "target": 1, + }, + "ignore_index": 255, + "derived_from": ["vegetation", "cana"], + }, +} + +SEMANTIC_COLORS_RGB = { + 0: (85, 85, 85), # chao + 1: (0, 190, 0), # cana + 2: (230, 55, 55), # erva +} + +BINARY_COLORS_RGB = { + 0: (30, 30, 30), + 1: (0, 220, 80), +} + +CANA_COLORS_RGB = { + 0: (30, 30, 30), + 1: (40, 210, 255), +} + +TARGET_COLORS_RGB = { + 0: (30, 30, 30), + 1: (255, 70, 30), +} + + +@dataclass +class SampleItem: + group: str + base: str + tensor_path: Optional[Path] + masks: Dict[str, Optional[Path]] + meta_path: Optional[Path] = None + preview_path: Optional[Path] = None + + # novo + source_kind: str = "tensor" # "tensor" ou "raw_native_multi" + raw_group: Optional[dict] = None + + +# ============================================================ +# Util +# ============================================================ + +def load_json(path: str | Path) -> dict: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +def ensure_dir(path: Path): + path.mkdir(parents=True, exist_ok=True) + + +def resolve_path(path_like: Optional[str], base: Optional[Path] = None) -> Optional[Path]: + if path_like is None: + return None + p = Path(path_like) + if p.is_absolute(): + return p + if base is None: + base = Path.cwd() + return (base / p).resolve() + + +def merge_dict(dst: dict, src: dict) -> dict: + out = copy.deepcopy(dst) + + def rec(a, b): + for k, v in b.items(): + if isinstance(v, dict) and isinstance(a.get(k), dict): + rec(a[k], v) + else: + a[k] = v + + if isinstance(src, dict): + rec(out, src) + return out + + +def build_heads_config(config: dict, ignore_index: int) -> Dict[str, dict]: + cfg = merge_dict(DEFAULT_HEADS, config.get("heads", {}) or {}) + active = {} + for name, hcfg in cfg.items(): + if not bool(hcfg.get("enabled", True)): + continue + hcfg.setdefault("ignore_index", ignore_index) + hcfg["ignore_index"] = int(hcfg.get("ignore_index", ignore_index)) + hcfg["num_classes"] = int(hcfg.get("num_classes", 2)) + hcfg["mask_dir"] = str(hcfg.get("mask_dir", "masks")) + active[name] = hcfg + + for required in ("semantic", "vegetation", "cana"): + if required not in active: + raise RuntimeError(f"Head obrigatória ausente no config: {required}") + + return active + + +def load_labelmap(labelmap_path: Path) -> Tuple[Dict[int, str], Dict[str, int], int, Dict[int, Tuple[int, int, int]]]: + try: + from helpers import carregar_labelmap_completo, _infer_ignore_id + + _cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(str(labelmap_path)) + ignore_id = int(_infer_ignore_id(ignore_rgb, 255)) + + id2label = { + int(k): str(v) + for k, v in id_para_nome.items() + if str(v).lower() not in ("ignore", "void", "background_ignore") + } + label2id = {v.lower(): k for k, v in id2label.items()} + + colormap_rgb = {} + if isinstance(_colormap_rgb, dict): + for k, v in _colormap_rgb.items(): + ik = int(k) + if ik in id2label: + colormap_rgb[ik] = tuple(map(int, v[:3])) + elif isinstance(_colormap_rgb, (list, tuple)): + for ik, v in enumerate(_colormap_rgb): + if ik in id2label: + colormap_rgb[ik] = tuple(map(int, v[:3])) + + return id2label, label2id, ignore_id, colormap_rgb + + except Exception as e: + print(f"[WARN] Não consegui usar helpers.carregar_labelmap_completo: {e}") + print("[WARN] Usando parser simples do labelmap.") + + id2label: Dict[int, str] = {} + colormap_rgb: Dict[int, Tuple[int, int, int]] = {} + ignore_id = 255 + next_id = 0 + + with labelmap_path.open("r", encoding="utf-8") as f: + for raw_line in f: + s = raw_line.strip() + if not s or s.startswith("#"): + continue + + name = None + color = None + cid = None + + if ":" in s and not s.split(":", 1)[0].strip().isdigit(): + name_part, rest = s.split(":", 1) + name = name_part.strip() + color_txt = rest.split("::", 1)[0].strip().strip(":") + rgb_parts = [p.strip() for p in color_txt.split(",") if p.strip()] + if len(rgb_parts) >= 3: + color = tuple(int(float(p)) for p in rgb_parts[:3]) + else: + parts = s.replace(",", " ").replace(":", " ").split() + if len(parts) >= 2 and parts[0].isdigit(): + cid = int(parts[0]) + name = parts[1] + if len(parts) >= 5: + color = tuple(int(float(p)) for p in parts[2:5]) + elif len(parts) >= 1: + name = parts[0] + + if not name: + continue + + if name.lower() in ("ignore", "void", "background_ignore"): + ignore_id = 255 + continue + + if cid is None: + cid = next_id + next_id = max(next_id, cid + 1) + + id2label[int(cid)] = str(name) + if color is not None: + colormap_rgb[int(cid)] = color + + if not id2label: + id2label = {0: "chao", 1: "cana", 2: "erva"} + + label2id = {v.lower(): k for k, v in id2label.items()} + for cid in id2label: + colormap_rgb.setdefault(cid, SEMANTIC_COLORS_RGB.get(cid, (255, 255, 255))) + + return id2label, label2id, int(ignore_id), colormap_rgb + + +def find_raw_dataset_layout_root(path: Path) -> Optional[Path]: + """ + Detecta layout bruto: + root/metas + root/previews + root/bins + root/masks opcional + + Aceita root, root/metas, root/bins, root/previews ou arquivo dentro deles. + """ + p = path.resolve() + candidates = [] + + if p.is_file(): + candidates.append(p.parent) + candidates.append(p.parent.parent) + else: + candidates.append(p) + candidates.append(p.parent) + + for c in candidates: + if not c: + continue + + if c.name.lower() in ("metas", "metadata", "jsons", "previews", "bins", "masks"): + root = c.parent + else: + root = c + + if (root / "metas").is_dir() and (root / "previews").is_dir() and (root / "bins").is_dir(): + return root + + return None + + +def resolve_raw_sibling_file(root: Path, subdir: str, stem: str, exts: Tuple[str, ...]) -> Optional[Path]: + folder = root / subdir + if not folder.is_dir(): + return None + + for ext in exts: + p = folder / f"{stem}{ext}" + if p.exists(): + return p + + return None + + +def resolve_raw_capture_group_from_json(json_path: Path, dataset_root: Path) -> dict: + """ + Resolve uma captura bruta: + metas/.json + previews/.png + bins/ + masks/ opcional + """ + json_path = json_path.resolve() + meta = load_json(json_path) + base_name = json_path.stem + bins_dir = dataset_root / "bins" + + preview_path = resolve_raw_sibling_file( + dataset_root, + "previews", + base_name, + (".png", ".jpg", ".jpeg"), + ) + + mask_path = resolve_raw_sibling_file( + dataset_root, + "masks", + base_name, + (".npy", ".png", ".tif", ".tiff"), + ) + + group = { + "json": json_path, + "png": preview_path, + "mask": mask_path, + "final_raw": None, + "cameras": {}, + "dataset_root": dataset_root, + } + + if "saved_payload_paths" in meta: + for cam_id, fname in meta["saved_payload_paths"].items(): + fname_path = Path(fname) + + candidates = [ + bins_dir / fname_path.name, + json_path.parent / fname, + dataset_root / fname, + bins_dir / f"{base_name}_{cam_id}.bin", + bins_dir / f"{base_name}_{cam_id}.raw", + bins_dir / f"{base_name}_{cam_id.lower()}.bin", + bins_dir / f"{base_name}_{cam_id.lower()}.raw", + ] + + found = next((c for c in candidates if c.exists()), None) + if found is not None: + group["cameras"][cam_id] = found + else: + print(f"[WARN] bin não encontrado para {cam_id}: {fname}") + + return group + + +def build_multispec_from_raw_native_multi_for_infer(group: dict, meta: dict) -> Tuple[np.ndarray, dict]: + """ + Gera tensor MULTISPEC [R,G,B,RE,NIR] em CHW float32 0..1 + a partir dos bins RAW_BRUTO salvos. + """ + if meta.get("saved_payload_type") != "raw_native_multi": + raise RuntimeError( + f"Captura bruta não suportada aqui: saved_payload_type={meta.get('saved_payload_type')}" + ) + + stream_meta = meta.get("stream_meta", {}) or {} + saved_dtypes = meta.get("saved_payload_dtypes", {}) or {} + saved_shapes = meta.get("saved_payload_shapes", {}) or {} + + frame = {} + + for cam_id, path in group["cameras"].items(): + saved_dtype = saved_dtypes.get(cam_id) + saved_shape = saved_shapes.get(cam_id) + + if saved_dtype is None or saved_shape is None: + raise RuntimeError(f"Faltam dtype/shape para {cam_id} no JSON: {group['json']}") + + arr = np.fromfile(str(path), dtype=np.dtype(saved_dtype)).reshape(tuple(saved_shape)) + frame[cam_id] = arr + + if not frame: + raise RuntimeError(f"Nenhum bin de câmera encontrado para: {group['json']}") + + sensor_width = int(meta.get("sensor_width", 1280)) + sensor_height = int(meta.get("sensor_height", 800)) + bayer = meta.get("bayer_pattern", "RGGB") + + calib_path = meta.get("camera_params_json") or "calibration/module_params.json" + + if calib_path and not os.path.isfile(calib_path): + json_dir = Path(group["json"]).parent + alt = json_dir / calib_path + + if alt.exists(): + calib_path = str(alt) + else: + raise RuntimeError( + f"module_params.json não encontrado: {calib_path}. " + f"Sem ele eu não gero MULTISPEC, porque faltaria fusion/radiometria/flatfield." + ) + + core = RawProcessorCore( + sensor_width=sensor_width, + sensor_height=sensor_height, + bayer_pattern=bayer, + calibration_json_path=calib_path, + ) + + processing_meta = dict(stream_meta) + + if meta.get("actual_camera_controls") is not None: + processing_meta["actual_camera_controls"] = meta.get("actual_camera_controls") + + if meta.get("startup_camera_controls") is not None: + processing_meta["startup_camera_controls"] = meta.get("startup_camera_controls") + + tensor = core.build_infer_tensor_from_stream(frame, processing_meta, 5) + + if tensor is None: + raise RuntimeError(f"RawProcessorCore retornou tensor None para: {group['json']}") + + tensor = np.asarray(tensor, dtype=np.float32) + tensor = np.nan_to_num(tensor, nan=0.0, posinf=1.0, neginf=0.0) + tensor = np.clip(tensor, 0.0, 1.0) + + processing_info = { + "fusion_result": getattr(core, "last_fusion_result", None), + "fusion_config_used": getattr(core, "fusion_config", None), + "radiometric_normalization_result": getattr(core, "last_radiometric_normalization_result", None), + "patch_normalization_result": getattr(core, "last_patch_normalization_result", None), + "frame_quality": getattr(core, "last_frame_quality_result", None), + } + + return tensor, processing_info + + +DEFAULT_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"] + +def get_input_channel_names(config: dict, channels_override: Optional[int] = None) -> List[str]: + if "input_channels" in config: + names = [str(c).upper() for c in config["input_channels"]] + else: + n = int(channels_override or config.get("channels", 5)) + names = DEFAULT_CHANNEL_ORDER[:n] + + invalid = [c for c in names if c not in DEFAULT_CHANNEL_ORDER] + if invalid: + raise RuntimeError(f"Canais inválidos em input_channels: {invalid}") + + return names + + +def get_input_channel_indices(config: dict, channels_override: Optional[int] = None) -> List[int]: + names = get_input_channel_names(config, channels_override) + return [DEFAULT_CHANNEL_ORDER.index(c) for c in names] + + +# ============================================================ +# Dataset +# ============================================================ + +def collect_samples(root: Path, heads_config: Dict[str, dict], require_masks: bool = False) -> Tuple[List[SampleItem], bool]: + tensor_paths: List[Path] = [] + + direct = root / "tensors" + if direct.is_dir(): + tensor_paths.extend(sorted(direct.glob("*.npy"))) + + group_root = root / "group" + if group_root.is_dir(): + for gdir in sorted(group_root.iterdir()): + tdir = gdir / "tensors" + if tdir.is_dir(): + tensor_paths.extend(sorted(tdir.glob("*.npy"))) + + if not tensor_paths: + tensor_paths.extend(sorted(root.glob("**/tensors/*.npy"))) + + # ============================================================ + # MODO 1: dataset já normalizado com tensors/*.npy + # ============================================================ + if tensor_paths: + samples: List[SampleItem] = [] + has_any_gt = False + + for tp in tensor_paths: + base = tp.stem + group_name = tp.parent.parent.name if tp.parent.name == "tensors" else "default" + group_dir = tp.parent.parent if tp.parent.name == "tensors" else tp.parent + + masks: Dict[str, Optional[Path]] = {} + missing = [] + + for head_name, hcfg in heads_config.items(): + mask_dir = group_dir / str(hcfg.get("mask_dir", "masks")) + mask_path = mask_dir / f"{base}.npy" + + if mask_path.exists(): + masks[head_name] = mask_path + has_any_gt = True + else: + masks[head_name] = None + missing.append(f"{head_name}:{mask_path}") + + if require_masks and missing: + raise RuntimeError(f"Masks ausentes para {tp}: {missing}") + + meta_path = group_dir / "metas" / f"{base}.json" + preview_path = group_dir / "previews" / f"{base}.png" + + samples.append(SampleItem( + group=group_name, + base=base, + tensor_path=tp, + masks=masks, + meta_path=meta_path if meta_path.exists() else None, + preview_path=preview_path if preview_path.exists() else None, + source_kind="tensor", + raw_group=None, + )) + + return samples, has_any_gt + + # ============================================================ + # MODO 2: dataset bruto com bins/metas/previews/masks + # ============================================================ + raw_root = find_raw_dataset_layout_root(root) + + if raw_root is None: + raise RuntimeError( + f"Nenhum tensor .npy encontrado e também não detectei layout bruto " + f"com bins/metas/previews em: {root}" + ) + + meta_paths = sorted((raw_root / "metas").glob("*.json")) + if not meta_paths: + raise RuntimeError(f"Nenhum meta .json encontrado em: {raw_root / 'metas'}") + + samples: List[SampleItem] = [] + has_any_gt = False + + for mp in meta_paths: + base = mp.stem + meta = load_json(mp) + + if meta.get("saved_payload_type") != "raw_native_multi": + print(f"[WARN] pulando {mp.name}: saved_payload_type={meta.get('saved_payload_type')}") + continue + + raw_group = resolve_raw_capture_group_from_json(mp, raw_root) + + masks: Dict[str, Optional[Path]] = {} + + # Tenta masks específicas multi-head primeiro. + for head_name, hcfg in heads_config.items(): + mask_dir = raw_root / str(hcfg.get("mask_dir", "masks")) + mask_path = None + + mask_dir_name = str(hcfg.get("mask_dir", "masks")) + + if mask_dir_name == "__derived_target__" or bool(hcfg.get("derived", False)): + masks[head_name] = None + continue + + for ext in (".npy", ".png", ".tif", ".tiff"): + p = mask_dir / f"{base}{ext}" + if p.exists(): + mask_path = p + break + + masks[head_name] = mask_path + if mask_path is not None: + has_any_gt = True + + # Se só existir masks/.png ou .npy semantic, conecta na head semantic. + if masks.get("semantic") is None: + semantic_mask = resolve_raw_sibling_file( + raw_root, + "masks", + base, + (".npy", ".png", ".tif", ".tiff"), + ) + if semantic_mask is not None: + masks["semantic"] = semantic_mask + has_any_gt = True + + if require_masks and not any(p is not None for p in masks.values()): + raise RuntimeError(f"Mask ausente para captura bruta: {mp}") + + preview_path = resolve_raw_sibling_file( + raw_root, + "previews", + base, + (".png", ".jpg", ".jpeg"), + ) + + samples.append(SampleItem( + group=raw_root.name, + base=base, + tensor_path=None, + masks=masks, + meta_path=mp, + preview_path=preview_path, + source_kind="raw_native_multi", + raw_group=raw_group, + )) + + if not samples: + raise RuntimeError(f"Nenhuma captura raw_native_multi válida encontrada em: {raw_root}") + + return samples, has_any_gt + + +def load_tensor( + path: Path, + channels: int, + channel_indices: Optional[List[int]] = None, +) -> np.ndarray: + arr = np.load(str(path)).astype(np.float32) + + if arr.ndim != 3: + raise RuntimeError(f"Tensor inválido {path}: shape={arr.shape}, esperado 3D") + + # Normaliza para CHW. + if arr.shape[0] in (3, 4, 5): + chw = arr + elif arr.shape[-1] in (3, 4, 5): + chw = np.transpose(arr, (2, 0, 1)) + else: + raise RuntimeError(f"Tensor com layout inesperado: {path} shape={arr.shape}") + + if channel_indices is not None: + max_idx = max(channel_indices) + if chw.shape[0] <= max_idx: + raise RuntimeError( + f"Tensor {path} tem {chw.shape[0]} canais, " + f"mas precisa acessar índice {max_idx}. Shape={chw.shape}" + ) + chw = chw[channel_indices, :, :] + else: + if chw.shape[0] < channels: + raise RuntimeError( + f"Tensor {path} tem {chw.shape[0]} canais, " + f"mas config pediu {channels}." + ) + chw = chw[:channels, :, :] + + finite = np.isfinite(chw) + if finite.any(): + mx = float(np.nanmax(chw[finite])) + if mx > 2.0 and mx <= 255.0: + chw = chw / 255.0 + elif mx > 255.0: + chw = chw / 65535.0 + + chw = np.nan_to_num(chw, nan=0.0, posinf=1.0, neginf=0.0) + return np.clip(chw, 0.0, 1.0).astype(np.float32) + + +def load_sample_tensor( + sample: SampleItem, + channels: int, + channel_indices: Optional[List[int]] = None, +) -> np.ndarray: + """ + Carrega tensor de uma amostra. + - Se for tensor pronto: lê .npy. + - Se for RAW_BRUTO: gera MULTISPEC em tempo real a partir dos bins. + """ + if sample.source_kind == "tensor": + if sample.tensor_path is None: + raise RuntimeError(f"Sample tensor sem tensor_path: {sample.base}") + return load_tensor( + sample.tensor_path, + channels=channels, + channel_indices=channel_indices, + ) + + if sample.source_kind == "raw_native_multi": + if sample.raw_group is None or sample.meta_path is None: + raise RuntimeError(f"Sample raw sem raw_group/meta_path: {sample.base}") + + meta = load_json(sample.meta_path) + tensor, _processing_info = build_multispec_from_raw_native_multi_for_infer(sample.raw_group, meta) + + if tensor.ndim != 3: + raise RuntimeError(f"Tensor RAW gerado inválido: {sample.base} shape={tensor.shape}") + + if channel_indices is not None: + max_idx = max(channel_indices) + if tensor.shape[0] <= max_idx: + raise RuntimeError( + f"Tensor RAW tem {tensor.shape[0]} canais, " + f"mas precisa acessar índice {max_idx}." + ) + tensor = tensor[channel_indices, :, :] + else: + if tensor.shape[0] < channels: + raise RuntimeError( + f"Tensor RAW com canais incompatíveis: {sample.base} " + f"shape={tensor.shape}, esperado pelo menos {channels}" + ) + tensor = tensor[:channels, :, :] + + return tensor.astype(np.float32) + + raise RuntimeError(f"source_kind desconhecido: {sample.source_kind}") + + +def load_mask(path: Optional[Path]) -> Optional[np.ndarray]: + if path is None: + return None + if path.suffix.lower() == ".npy": + mask = np.load(str(path)) + else: + mask = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) + if mask is None: + raise RuntimeError(f"Falha ao ler mask: {path}") + if mask.ndim == 3: + mask = mask[:, :, 0] + return mask.astype(np.int64) + + +# ============================================================ +# Normalização +# ============================================================ + +def load_norm_stats( + path: Optional[Path], + channels: int, + channel_indices: Optional[List[int]] = None, + channel_names: Optional[List[str]] = None, +) -> Tuple[Optional[List[float]], Optional[List[float]]]: + if path is None or not path.is_file(): + if path is not None: + print(f"[NORM] não encontrei norm_stats em {path}. Usando tensor 0..1 sem padronização.") + return None, None + + js = load_json(path) + mean = js.get("mean", None) + std = js.get("std", None) + names = js.get("channels", []) + + if mean is None or std is None: + raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}") + if channel_indices is not None: + max_idx = max(channel_indices) + if len(mean) <= max_idx or len(std) <= max_idx: + raise RuntimeError( + f"norm_stats incompatível: precisa índices={channel_indices}, " + f"mean={len(mean)} std={len(std)}" + ) + + mean = [mean[i] for i in channel_indices] + std = [std[i] for i in channel_indices] + + if names: + names = [names[i] for i in channel_indices] + elif channel_names: + names = channel_names + + else: + if len(mean) != channels or len(std) != channels: + raise RuntimeError( + f"norm_stats incompatível com channels={channels}: " + f"mean={len(mean)} std={len(std)}" + ) + + print(f"[NORM] usando {path}") + print(f"[NORM] channels={names}") + print(f"[NORM] mean={mean}") + print(f"[NORM] std ={std}") + return list(map(float, mean)), list(map(float, std)) + + +# ============================================================ +# Modelo multi-head, igual ao treino +# ============================================================ + +def patch_segformer_encoder_input_channels(segformer_encoder: nn.Module, in_ch: int): + if in_ch == 3: + return segformer_encoder + + proj = segformer_encoder.encoder.patch_embeddings[0].proj + + if proj.in_channels == in_ch: + return segformer_encoder + + old_weight = proj.weight.data.clone() + old_bias = proj.bias.data.clone() if proj.bias is not None else None + + new_proj = nn.Conv2d( + in_channels=in_ch, + out_channels=proj.out_channels, + kernel_size=proj.kernel_size, + stride=proj.stride, + padding=proj.padding, + dilation=proj.dilation, + groups=proj.groups, + bias=proj.bias is not None, + padding_mode=proj.padding_mode, + ) + + with torch.no_grad(): + if in_ch <= old_weight.shape[1]: + new_proj.weight.copy_(old_weight[:, :in_ch, :, :]) + else: + new_proj.weight[:, :old_weight.shape[1], :, :].copy_(old_weight) + extra = in_ch - old_weight.shape[1] + mean_w = old_weight.mean(dim=1, keepdim=True) + new_proj.weight[:, old_weight.shape[1]:, :, :].copy_(mean_w.repeat(1, extra, 1, 1)) + + if old_bias is not None: + new_proj.bias.copy_(old_bias) + + segformer_encoder.encoder.patch_embeddings[0].proj = new_proj + print(f"[MODEL] patch input channels: 3 -> {in_ch}") + return segformer_encoder + + +def replace_segformer_decode_classifier(decode_head: nn.Module, num_classes: int): + old = decode_head.classifier + if not isinstance(old, nn.Conv2d): + raise RuntimeError(f"decode_head.classifier não é Conv2d: {type(old)}") + + new = nn.Conv2d( + in_channels=old.in_channels, + out_channels=int(num_classes), + kernel_size=old.kernel_size, + stride=old.stride, + padding=old.padding, + dilation=old.dilation, + groups=old.groups, + bias=old.bias is not None, + padding_mode=old.padding_mode, + ) + decode_head.classifier = new + return decode_head + + +class MultiHeadSegFormer(nn.Module): + def __init__( + self, + backbone: str, + channels: int, + heads_config: Dict[str, dict], + semantic_id2label: Dict[int, str], + semantic_label2id: Dict[str, int], + ): + super().__init__() + + semantic_classes = int(heads_config["semantic"].get("num_classes", len(semantic_id2label))) + + base = SegformerForSemanticSegmentation.from_pretrained( + backbone, + num_labels=semantic_classes, + id2label={int(k): str(v) for k, v in semantic_id2label.items()}, + label2id={str(k): int(v) for k, v in semantic_label2id.items()}, + ignore_mismatched_sizes=True, + ) + + patch_segformer_encoder_input_channels(base.segformer, channels) + base.config.num_channels = int(channels) + + self.segformer = base.segformer + self.decode_heads = nn.ModuleDict() + self.heads_config = heads_config + + for head_name, hcfg in heads_config.items(): + h = copy.deepcopy(base.decode_head) + h = replace_segformer_decode_classifier(h, int(hcfg["num_classes"])) + self.decode_heads[head_name] = h + + self.config = base.config + + def forward(self, pixel_values: torch.Tensor, head_names=None) -> Dict[str, torch.Tensor]: + outputs = self.segformer( + pixel_values=pixel_values, + output_hidden_states=True, + return_dict=True, + ) + + hidden_states = outputs.hidden_states + + if head_names is None: + selected = list(self.decode_heads.keys()) + else: + selected = [h for h in head_names if h in self.decode_heads] + + return { + head_name: self.decode_heads[head_name](hidden_states) + for head_name in selected + } + + +class MultiHeadTester: + def __init__( + self, + config: dict, + ckpt_path: Path, + device: torch.device, + channels: int, + heads_config: Dict[str, dict], + semantic_id2label: Dict[int, str], + semantic_label2id: Dict[str, int], + mean: Optional[Sequence[float]], + std: Optional[Sequence[float]], + use_amp: bool = True, + ): + self.config = config + self.ckpt_path = ckpt_path + self.device = device + self.channels = channels + self.heads_config = heads_config + self.use_amp = use_amp and device.type == "cuda" + self.runtime_mode = str(config.get("runtime_mode", "all")).lower() + + self.mean = None if mean is None else torch.tensor(mean, dtype=torch.float32).view(1, channels, 1, 1).to(device) + self.std = None if std is None else torch.tensor(std, dtype=torch.float32).view(1, channels, 1, 1).to(device) + + backbone = config.get("backbone", config.get("pretrained_model", "nvidia/mit-b1")) + print(f"[MODEL] backbone={backbone}") + print(f"[MODEL] ckpt={ckpt_path}") + + self.model = MultiHeadSegFormer( + backbone=backbone, + channels=channels, + heads_config=heads_config, + semantic_id2label=semantic_id2label, + semantic_label2id=semantic_label2id, + ) + + self._load_checkpoint(ckpt_path) + self.model.to(device) + self.model.eval() + + def _load_checkpoint(self, ckpt_path: Path): + ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False) + + if isinstance(ckpt, dict): + for key in ("model", "model_state", "model_state_dict", "state_dict"): + if key in ckpt and isinstance(ckpt[key], dict): + state = ckpt[key] + break + else: + state = ckpt + else: + raise RuntimeError(f"Checkpoint em formato inesperado: {type(ckpt)}") + + clean = {} + for k, v in state.items(): + nk = k + for prefix in ("module.", "model."): + if nk.startswith(prefix): + nk = nk[len(prefix):] + clean[nk] = v + + missing, unexpected = self.model.load_state_dict(clean, strict=False) + print(f"[MODEL] load_state_dict strict=False | missing={len(missing)} unexpected={len(unexpected)}") + if missing: + print("[MODEL] primeiros missing:", missing[:8]) + if unexpected: + print("[MODEL] primeiros unexpected:", unexpected[:8]) + + def _normalize(self, x: torch.Tensor) -> torch.Tensor: + if self.mean is not None and self.std is not None: + return (x - self.mean) / torch.clamp(self.std, min=1e-6) + return x + + @torch.inference_mode() + def infer(self, chw_01: np.ndarray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], float]: + x = torch.from_numpy(chw_01).unsqueeze(0).to(self.device, non_blocking=True) + x = self._normalize(x) + + h, w = int(chw_01.shape[1]), int(chw_01.shape[2]) + + if self.device.type == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + + with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp): + head_names = None + if self.runtime_mode in ("target_direct", "target_head"): + head_names = ["target"] + elif self.runtime_mode in ("operational", "target_op"): + head_names = ["vegetation", "cana"] + + logits_by_head = self.model(pixel_values=x, head_names=head_names) + + preds = {} + probs = {} + + for head_name, logits in logits_by_head.items(): + logits = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False) + prob = torch.softmax(logits, dim=1)[0] + pred = torch.argmax(prob, dim=0) + preds[head_name] = pred.detach().cpu().numpy().astype(np.uint8) + probs[head_name] = prob.detach().cpu().numpy().astype(np.float32) + + if self.device.type == "cuda": + torch.cuda.synchronize() + t_ms = (time.perf_counter() - t0) * 1000.0 + + return preds, probs, t_ms + + +# ============================================================ +# Checkpoints / paths +# ============================================================ + +def infer_experiment_tag(config: dict, channels: int) -> str: + fusion_mode = config.get("fusion_mode", "stacked") + return f"{fusion_mode}_raw{channels}" + + +def infer_save_dir(config: dict, config_dir: Path, channels: int) -> Path: + model_name = config.get("model_name", "test_multi") + modelo_folder = config.get("modelo", "segformer_b1") + exp_tag = infer_experiment_tag(config, channels) + return (config_dir / "backup" / modelo_folder / model_name / exp_tag).resolve() + + +def find_checkpoint(save_dir: Path, preferred: Optional[str] = None) -> Path: + if preferred is not None: + ckpt = Path(preferred) + if not ckpt.is_absolute(): + ckpt_cwd = (Path.cwd() / ckpt).resolve() + ckpt_save = (save_dir / ckpt).resolve() + ckpt = ckpt_cwd if ckpt_cwd.is_file() else ckpt_save + if not ckpt.is_file(): + raise FileNotFoundError(f"Checkpoint não encontrado: {ckpt}") + return ckpt + + candidates = [ + save_dir / "best_score.pt", + save_dir / "best_target.pt", + save_dir / "best_cana_head.pt", + save_dir / "best_semantic_miou.pt", + save_dir / "last.pt", + ] + for c in candidates: + if c.is_file(): + return c + + raise FileNotFoundError("Nenhum checkpoint encontrado. Procurei:\n" + "\n".join(str(c) for c in candidates)) + + +# ============================================================ +# Visualização +# ============================================================ + +def tensor_to_preview_rgb(chw: np.ndarray, gamma: float = 0.85) -> np.ndarray: + c, h, w = chw.shape + if c >= 3: + rgb = np.transpose(chw[:3], (1, 2, 0)).copy() + else: + one = chw[0] + rgb = np.stack([one, one, one], axis=-1) + + rgb = np.nan_to_num(rgb, nan=0.0, posinf=1.0, neginf=0.0) + lo = np.percentile(rgb, 1.0) + hi = np.percentile(rgb, 99.0) + if hi > lo: + rgb = (rgb - lo) / (hi - lo) + rgb = np.clip(rgb, 0.0, 1.0) + if gamma and gamma > 0: + rgb = np.power(rgb, gamma) + return (rgb * 255.0).astype(np.uint8) + + +def ids_to_rgb(mask: np.ndarray, colormap_rgb: Dict[int, Tuple[int, int, int]], ignore_id: int = 255) -> np.ndarray: + h, w = mask.shape[:2] + out = np.zeros((h, w, 3), dtype=np.uint8) + for cid, color in colormap_rgb.items(): + out[mask == cid] = color + out[mask == ignore_id] = (0, 0, 0) + return out + + +def prob_to_heat_rgb(prob01: np.ndarray) -> np.ndarray: + p = np.clip(prob01, 0.0, 1.0) + u8 = (p * 255.0).astype(np.uint8) + bgr = cv2.applyColorMap(u8, cv2.COLORMAP_TURBO) + return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + + +def overlay_rgb(base_rgb: np.ndarray, mask_rgb: np.ndarray, alpha: float) -> np.ndarray: + return cv2.addWeighted(base_rgb, 1.0 - alpha, mask_rgb, alpha, 0.0) + + +def put_label(img_rgb: np.ndarray, title: str, subtitle: str = "") -> np.ndarray: + out = img_rgb.copy() + cv2.putText(out, title, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.68, (0, 0, 0), 4, cv2.LINE_AA) + cv2.putText(out, title, (10, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.68, (255, 255, 255), 2, cv2.LINE_AA) + if subtitle: + cv2.putText(out, subtitle, (10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 0, 0), 3, cv2.LINE_AA) + cv2.putText(out, subtitle, (10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (0, 255, 90), 1, cv2.LINE_AA) + return out + + +def resize_panel(img: np.ndarray, size: Tuple[int, int]) -> np.ndarray: + w, h = size + if img.shape[1] == w and img.shape[0] == h: + return img + return cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST) + + +def compose_grid(panels: List[Tuple[str, np.ndarray, str]], cols: int = 3, max_width: int = 1800) -> np.ndarray: + if not panels: + return np.zeros((480, 640, 3), dtype=np.uint8) + + base_h, base_w = panels[0][1].shape[:2] + labeled = [] + for title, img, subtitle in panels: + img = resize_panel(img, (base_w, base_h)) + labeled.append(put_label(img, title, subtitle)) + + rows = [] + blank = np.zeros_like(labeled[0]) + for i in range(0, len(labeled), cols): + row_imgs = labeled[i:i + cols] + while len(row_imgs) < cols: + row_imgs.append(blank.copy()) + rows.append(np.hstack(row_imgs)) + + canvas = np.vstack(rows) + + if canvas.shape[1] > max_width: + scale = max_width / canvas.shape[1] + canvas = cv2.resize(canvas, (int(canvas.shape[1] * scale), int(canvas.shape[0] * scale)), interpolation=cv2.INTER_AREA) + + return canvas + + +def class_percent(mask: np.ndarray, class_id: int, ignore_id: int = 255) -> float: + valid = mask != ignore_id + den = int(valid.sum()) + if den <= 0: + return 0.0 + return float(((mask == class_id) & valid).sum() * 100.0 / den) + + +# ============================================================ +# Métricas numpy +# ============================================================ + +def confusion_matrix_np(pred: np.ndarray, gt: np.ndarray, num_classes: int, ignore_id: int) -> np.ndarray: + if pred.shape != gt.shape: + pred = cv2.resize(pred.astype(np.uint8), (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST) + + valid = gt != ignore_id + valid &= gt >= 0 + valid &= gt < num_classes + gt_v = gt[valid].astype(np.int64) + pred_v = pred[valid].astype(np.int64) + pred_v = np.clip(pred_v, 0, num_classes - 1) + + cm = np.bincount(num_classes * gt_v + pred_v, minlength=num_classes * num_classes) + return cm.reshape(num_classes, num_classes).astype(np.int64) + + +def metrics_from_cm(cm: np.ndarray) -> Tuple[np.ndarray, float, float]: + tp = np.diag(cm).astype(np.float64) + fp = cm.sum(axis=0).astype(np.float64) - tp + fn = cm.sum(axis=1).astype(np.float64) - tp + denom = tp + fp + fn + iou = np.divide(tp, denom, out=np.zeros_like(tp), where=denom > 0) + miou = float(np.mean(iou)) if len(iou) else 0.0 + acc = float(tp.sum() / max(cm.sum(), 1)) + return iou, miou, acc + + +def operational_target_mask(veg_mask: np.ndarray, cana_mask: np.ndarray, ignore_id: int = 255) -> np.ndarray: + out = np.zeros_like(veg_mask, dtype=np.uint8) + ignore = (veg_mask == ignore_id) | (cana_mask == ignore_id) + out[(veg_mask == 1) & (cana_mask == 0)] = 1 + out[ignore] = ignore_id + return out + + +# ============================================================ +# Main +# ============================================================ + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="config.json") + parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"]) + parser.add_argument("--root_override", default=None) + parser.add_argument("--test_folder", default=None) + parser.add_argument("--ckpt", default=None) + parser.add_argument("--norm_stats", default=None) + parser.add_argument("--channels", type=int, default=None) + parser.add_argument("--resize_w", type=int, default=None) + parser.add_argument("--resize_h", type=int, default=None) + parser.add_argument("--alpha", type=float, default=0.45) + parser.add_argument("--ignore_id", type=int, default=None) + parser.add_argument("--no_amp", action="store_true") + parser.add_argument("--require_masks", action="store_true") + parser.add_argument("--out_dir", default="outputs_test_multihead") + parser.add_argument("--start_idx", type=int, default=0) + parser.add_argument("--max_width", type=int, default=1800) + parser.add_argument("--runtime_mode", default="all", choices=["all", "target_direct", "operational"]) + args = parser.parse_args() + + config_path = resolve_path(args.config, Path.cwd()) + if config_path is None or not config_path.is_file(): + raise FileNotFoundError(f"Config não encontrado: {config_path}") + + config_dir = config_path.parent + config = load_json(config_path) + config["runtime_mode"] = args.runtime_mode + + input_channel_names = get_input_channel_names(config, args.channels) + input_channel_indices = get_input_channel_indices(config, args.channels) + channels = len(input_channel_names) + + print(f"Input channels: {input_channel_names} idx={input_channel_indices}") + res = config.get("resolucao", [1024, 640]) + default_w, default_h = int(res[0]), int(res[1]) + target_w = int(args.resize_w or default_w) + target_h = int(args.resize_h or default_h) + + dataset_path = config_dir / "dataset" + labelmap_path = dataset_path / "labelmap.txt" + semantic_id2label, semantic_label2id, labelmap_ignore_id, loaded_colormap_rgb = load_labelmap(labelmap_path) + + ignore_id = int(args.ignore_id if args.ignore_id is not None else labelmap_ignore_id) + heads_config = build_heads_config(config, ignore_index=ignore_id) + heads_config["semantic"]["num_classes"] = int(len(semantic_id2label)) + heads_config["semantic"]["ignore_index"] = int(ignore_id) + + semantic_cmap = dict(SEMANTIC_COLORS_RGB) + semantic_cmap.update({int(k): tuple(map(int, v)) for k, v in loaded_colormap_rgb.items()}) + + save_dir = infer_save_dir(config, config_dir, channels) + ckpt_path = find_checkpoint(save_dir, args.ckpt) + + if args.norm_stats is not None: + norm_stats_path = resolve_path(args.norm_stats, Path.cwd()) + else: + # O treino multihead usa stats do dataset normalizado, mas também tentamos alguns fallbacks. + candidates = [ + dataset_path / f"{default_w}x{default_h}" / "group" / "norm_stats.json", + save_dir / "norm_stats.json", + config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test_multi") / config.get("stats_source_tag", "stacked_raw5") / "norm_stats.json", + ] + norm_stats_path = next((p for p in candidates if p.is_file()), candidates[0]) + + mean, std = load_norm_stats( + norm_stats_path, + channels=channels, + channel_indices=input_channel_indices, + channel_names=input_channel_names, + ) + + if args.test_folder is not None: + root = resolve_path(args.test_folder, Path.cwd()) + elif args.root_override is not None: + root = resolve_path(args.root_override, Path.cwd()) + else: + root = (dataset_path / "split" / args.split_folder).resolve() + + if root is None or not root.is_dir(): + raise FileNotFoundError(f"Root de dados não encontrado: {root}") + + samples, has_gt = collect_samples(root, heads_config=heads_config, require_masks=args.require_masks) + n = len(samples) + + print("==========================================") + print("Teste SegFormer OAK-FCC-3 Multi-Head") + print(f"Root : {root}") + print(f"Samples : {n}") + print(f"GT : {'sim' if has_gt else 'não'}") + print(f"Resolution : {target_w}x{target_h}") + print(f"Channels : {channels}") + print(f"Semantic : {semantic_id2label}") + print(f"Ignore index: {ignore_id}") + print("Heads:") + for name, hcfg in heads_config.items(): + print(f" - {name}: classes={hcfg['num_classes']} mask_dir={hcfg['mask_dir']}") + print("==========================================") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + + tester = MultiHeadTester( + config=config, + ckpt_path=ckpt_path, + device=device, + channels=channels, + heads_config=heads_config, + semantic_id2label=semantic_id2label, + semantic_label2id=semantic_label2id, + mean=mean, + std=std, + use_amp=not args.no_amp, + ) + + out_dir = Path(args.out_dir) + ensure_dir(out_dir) + + idx = max(0, min(args.start_idx, n - 1)) + detailed = False + + cms_total = { + name: np.zeros((int(hcfg["num_classes"]), int(hcfg["num_classes"])), dtype=np.int64) + for name, hcfg in heads_config.items() + } + cm_target_total = np.zeros((2, 2), dtype=np.int64) + visited = set() + + win_name = "OAK-FCC-3 MultiHead Test | D/A navega | S salva | SPACE detalhado | Q sai" + cv2.namedWindow(win_name, cv2.WINDOW_NORMAL) + + while True: + sample = samples[idx] + chw = load_sample_tensor( + sample, + channels=channels, + channel_indices=input_channel_indices, + ) + + if (chw.shape[2], chw.shape[1]) != (target_w, target_h): + hwc = np.transpose(chw, (1, 2, 0)) + hwc = cv2.resize(hwc, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + chw = np.transpose(hwc, (2, 0, 1)).astype(np.float32) + + preds, probs, t_inf = tester.infer(chw) + preview_rgb = tensor_to_preview_rgb(chw) + + first_pred = next(iter(preds.values())) + pred_h, pred_w = first_pred.shape[:2] + + pred_sem = preds.get("semantic") + pred_veg = preds.get("vegetation") + pred_cana = preds.get("cana") + + # Target operacional antigo: vegetation AND not cana + pred_target_op = None + if pred_veg is not None and pred_cana is not None: + pred_target_op = operational_target_mask(pred_veg, pred_cana, ignore_id=ignore_id) + + # Target direta nova: saída da 4ª cabeça + pred_target_head = preds.get("target") + + # Fallback visual: se não tiver head target, mostra operacional + pred_target = pred_target_head if pred_target_head is not None else pred_target_op + + prob_veg = probs["vegetation"][1] if "vegetation" in probs and probs["vegetation"].shape[0] > 1 else None + prob_cana = probs["cana"][1] if "cana" in probs and probs["cana"].shape[0] > 1 else None + + prob_target_op = None + if prob_veg is not None and prob_cana is not None: + prob_target_op = np.clip(prob_veg * (1.0 - prob_cana), 0.0, 1.0) + + prob_target_head = None + if "target" in probs: + prob_target_head = probs["target"][1] if probs["target"].shape[0] > 1 else probs["target"][0] + + prob_target = prob_target_head if prob_target_head is not None else prob_target_op + + gt_masks = {name: load_mask(path) for name, path in sample.masks.items()} + for name, gt in list(gt_masks.items()): + if gt is not None and gt.shape != (pred_h, pred_w): + gt_masks[name] = cv2.resize( + gt.astype(np.uint8), + (pred_w, pred_h), + interpolation=cv2.INTER_NEAREST, + ) + + gt_target = None + if gt_masks.get("target") is not None: + gt_target = gt_masks["target"] + elif gt_masks.get("vegetation") is not None and gt_masks.get("cana") is not None: + gt_target = operational_target_mask( + gt_masks["vegetation"], + gt_masks["cana"], + ignore_id=ignore_id, + ) + gt_masks["target"] = gt_target + + # Métricas da amostra. + metric_lines = [] + sample_metrics = {} + for head_name, pred in preds.items(): + gt = gt_masks.get(head_name) + if gt is not None: + cm = confusion_matrix_np(pred, gt, int(heads_config[head_name]["num_classes"]), ignore_id) + iou, miou, acc = metrics_from_cm(cm) + sample_metrics[head_name] = {"iou": iou, "miou": miou, "acc": acc} + metric_lines.append(f"{head_name}: mIoU={miou:.3f} acc={acc:.3f}") + + if gt_target is not None: + cm_t = confusion_matrix_np(pred_target, gt_target, 2, ignore_id) + iou_t, miou_t, acc_t = metrics_from_cm(cm_t) + sample_metrics["target_op"] = {"iou": iou_t, "miou": miou_t, "acc": acc_t} + metric_lines.append(f"target_op: IoU_alvo={iou_t[1]:.3f} acc={acc_t:.3f}") + if "target" in sample_metrics: + iou_head = sample_metrics["target"]["iou"] + metric_lines.append(f"target_head: IoU_alvo={iou_head[1]:.3f}") + + if idx not in visited: + for head_name, pred in preds.items(): + gt = gt_masks.get(head_name) + if gt is not None: + cms_total[head_name] += confusion_matrix_np(pred, gt, int(heads_config[head_name]["num_classes"]), ignore_id) + if gt_target is not None and pred_target_op is not None: + cm_target_total += confusion_matrix_np(pred_target_op, gt_target, 2, ignore_id) + visited.add(idx) + + # Visuals. + panels: List[Tuple[str, np.ndarray, str]] = [] + if pred_sem is not None: + pred_sem_rgb = ids_to_rgb(pred_sem, semantic_cmap, ignore_id) + if gt_masks.get("semantic") is not None: + panels.append(("GT semantic", ids_to_rgb(gt_masks["semantic"], semantic_cmap, ignore_id), "chao/cana/erva")) + panels.append(("Pred semantic", pred_sem_rgb, f"erva={class_percent(pred_sem, 2, ignore_id):.1f}% cana={class_percent(pred_sem, 1, ignore_id):.1f}%")) + panels.append(("Overlay semantic", overlay_rgb(preview_rgb, pred_sem_rgb, args.alpha), "")) + + if pred_veg is not None: + pred_veg_rgb = ids_to_rgb(pred_veg, BINARY_COLORS_RGB, ignore_id) + if gt_masks.get("vegetation") is not None: + panels.append(("GT vegetation", ids_to_rgb(gt_masks["vegetation"], BINARY_COLORS_RGB, ignore_id), "0=fundo 1=veg")) + panels.append(("Pred vegetation", pred_veg_rgb, f"veg={class_percent(pred_veg, 1, ignore_id):.1f}%")) + if prob_veg is not None: + panels.append(("P vegetation", prob_to_heat_rgb(prob_veg), f"mean={float(prob_veg.mean()):.3f}")) + + if pred_cana is not None: + pred_cana_rgb = ids_to_rgb(pred_cana, CANA_COLORS_RGB, ignore_id) + if gt_masks.get("cana") is not None: + panels.append(("GT cana", ids_to_rgb(gt_masks["cana"], CANA_COLORS_RGB, ignore_id), "0=not_cana 1=cana")) + panels.append(("Pred cana", pred_cana_rgb, f"cana={class_percent(pred_cana, 1, ignore_id):.1f}%")) + if prob_cana is not None: + panels.append(("P cana", prob_to_heat_rgb(prob_cana), f"mean={float(prob_cana.mean()):.3f}")) + + if detailed: + if gt_target is not None: + panels.append(("GT target", ids_to_rgb(gt_target, TARGET_COLORS_RGB, ignore_id), "derivado")) + + if pred_target_head is not None: + panels.append(("Pred target HEAD", ids_to_rgb(pred_target_head, TARGET_COLORS_RGB, ignore_id), f"alvo={class_percent(pred_target_head, 1, ignore_id):.1f}%")) + if prob_target_head is not None: + panels.append(("P target HEAD", prob_to_heat_rgb(prob_target_head), f"mean={float(prob_target_head.mean()):.3f}")) + + if pred_target_op is not None: + panels.append(("Pred target OP", ids_to_rgb(pred_target_op, TARGET_COLORS_RGB, ignore_id), f"alvo={class_percent(pred_target_op, 1, ignore_id):.1f}%")) + if prob_target_op is not None: + panels.append(("P target OP", prob_to_heat_rgb(prob_target_op), f"mean={float(prob_target_op.mean()):.3f}")) + else: + if pred_target_head is not None: + rgb = ids_to_rgb(pred_target_head, TARGET_COLORS_RGB, ignore_id) + panels.append(("Pred target HEAD", rgb, f"alvo={class_percent(pred_target_head, 1, ignore_id):.1f}%")) + panels.append(("Overlay target HEAD", overlay_rgb(preview_rgb, rgb, args.alpha), "head direta")) + elif pred_target_op is not None: + rgb = ids_to_rgb(pred_target_op, TARGET_COLORS_RGB, ignore_id) + panels.append(("Pred target OP", rgb, f"alvo={class_percent(pred_target_op, 1, ignore_id):.1f}%")) + panels.append(("Overlay target OP", overlay_rgb(preview_rgb, rgb, args.alpha), "veg & !cana")) + + canvas = compose_grid(panels, cols=3, max_width=args.max_width) + + header_h = 78 + header = np.zeros((header_h, canvas.shape[1], 3), dtype=np.uint8) + header[:] = (25, 25, 25) + source_name = sample.tensor_path.name if sample.tensor_path is not None else f"{sample.base}.json [{sample.source_kind}]" + h1 = f"idx {idx + 1}/{n} | {source_name} | inf={t_inf:.1f}ms | {'detalhado' if detailed else 'compacto'}" + h2 = " | ".join(metric_lines[:3]) if metric_lines else "sem GT" + cv2.putText(header, h1, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (235, 235, 235), 1, cv2.LINE_AA) + cv2.putText(header, h2, (12, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0, 255, 120), 1, cv2.LINE_AA) + canvas = np.vstack([header, canvas]) + + cv2.imshow(win_name, cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)) + k = cv2.waitKey(0) & 0xFF + + if k in (ord("q"), ord("Q"), 27): + break + elif k in (ord("d"), ord("D"), 83): + idx = (idx + 1) % n + elif k in (ord("a"), ord("A"), 81): + idx = (idx - 1 + n) % n + elif k == ord(" "): + detailed = not detailed + elif k in (ord("s"), ord("S")): + out_path = out_dir / f"multihead_{idx:05d}_{sample.base}.png" + cv2.imwrite(str(out_path), cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)) + print(f"[SAVE] {out_path}") + + cv2.destroyAllWindows() + + if visited: + print("\n========== RESUMO DOS SAMPLES VISITADOS ==========") + print(f"visitados={len(visited)}/{n}") + for head_name, cm in cms_total.items(): + if cm.sum() <= 0: + continue + iou, miou, acc = metrics_from_cm(cm) + print(f"\n[{head_name}] acc={acc:.4f} mIoU={miou:.4f}") + for i, v in enumerate(iou): + print(f" IoU {i} = {v:.4f}") + + if cm_target_total.sum() > 0: + iou, miou, acc = metrics_from_cm(cm_target_total) + print(f"\n[target operacional] acc={acc:.4f} mIoU={miou:.4f}") + print(f" IoU background = {iou[0]:.4f}") + print(f" IoU alvo = {iou[1]:.4f}") + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/datasets/oak-fcc-3/calibration/module_params.json b/Python/OAK/datasets/oak-fcc-3/calibration/module_params.json index 6ad12f326..860780c90 100644 --- a/Python/OAK/datasets/oak-fcc-3/calibration/module_params.json +++ b/Python/OAK/datasets/oak-fcc-3/calibration/module_params.json @@ -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 } } \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/config.json b/Python/OAK/datasets/oak-fcc-3/config.json index 39cdf68ce..4e8a302cc 100644 --- a/Python/OAK/datasets/oak-fcc-3/config.json +++ b/Python/OAK/datasets/oak-fcc-3/config.json @@ -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 + } } \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/core/__init__.py b/Python/OAK/datasets/oak-fcc-3/core/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/Python/OAK/datasets/oak-fcc-3/core/benchmark.py b/Python/OAK/datasets/oak-fcc-3/core/benchmark.py new file mode 100644 index 000000000..755895987 --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/core/benchmark.py @@ -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() \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/core/benchmark_preview.py b/Python/OAK/datasets/oak-fcc-3/core/benchmark_preview.py new file mode 100644 index 000000000..407a3bbbe --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/core/benchmark_preview.py @@ -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 só 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() diff --git a/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_client.py b/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_client.py index 018926041..b9e515dd7 100644 --- a/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_client.py +++ b/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_client.py @@ -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 já 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) diff --git a/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_manager.py b/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_manager.py index 3879e0e4d..4827b9006 100644 --- a/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_manager.py +++ b/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_manager.py @@ -1,10 +1,37 @@ +import json +import os import time from collections import deque +import threading +import copy + +import cv2 import depthai as dai import numpy as np class OakFcc3Manager: + """ + Manager OAK-FFC-3 com dois fluxos principais: + + 1) RAW_BRUTO + - Mantém o comportamento antigo. + - CAM_A/CAM_B/CAM_C enviam RAW10 packed direto para o PC. + - O PC faz decode, flat/radiometric, homografia/fusão/crop/resize. + + 2) MULTISPEC + - Câmeras sempre em 800p nativo. + - OAK aplica homografia/crop/resize via ImageManip. + - PC recebe frames já alinhados: + CAM_A/rgb -> BGR uint8 + CAM_B/re -> GRAY uint8 + CAM_C/nir -> GRAY uint8 + - O Client deve montar o tensor sem reaplicar homografia. + """ + + SENSOR_W = 1280 + SENSOR_H = 800 + def __init__( self, fps=30, @@ -18,14 +45,24 @@ class OakFcc3Manager: sync_mode="best", sync_tolerance_ms=12.0, buffer_size=8, - only_camera=None + only_camera=None, + mx_id=None, + module_calibration_json=None, + module_params=None, ): self.fps = fps - self.width = width - self.height = height - self.size = (width, height) - self.frame_type = frame_type + # Para compatibilidade, mantemos width/height. + # No RAW_BRUTO isso não muda o sensor, pois usamos 800p fixo. + # No MULTISPEC isso representa a saída final alinhada da OAK. + self.width = int(width) + self.height = int(height) + self.size = (self.width, self.height) + + self.sensor_width = self.SENSOR_W + self.sensor_height = self.SENSOR_H + + self.frame_type = str(frame_type).upper() self.output_dtype = output_dtype self.capture_mode = capture_mode self.raw_policy = raw_policy @@ -41,12 +78,18 @@ class OakFcc3Manager: self.sync_tolerance_ms = sync_tolerance_ms self.buffer_size = buffer_size + self.mx_id = str(mx_id) if mx_id else None + self.dev_info = None self.device = None self.pipeline = None self.queues = {} self.buffers = {} self.camera_info = {} + self.has_imu_pipeline = False + self.tem_imu = False + self.q_imu = None + self.running = False self.frame_id = 0 self.control_queues = {} @@ -56,6 +99,31 @@ class OakFcc3Manager: } self._last_raw_dims = {} + self.module_calibration_json = module_calibration_json + self.module_params = module_params if isinstance(module_params, dict) else self._load_module_params(module_calibration_json) + self.fusion_config = (self.module_params or {}).get("fusion_config", {}) or {} + self.aligned_geometry = None + + self.async_capture_enabled = True + self.async_capture_mode = "latest" # latest | queue + self.async_capture_max_queue = 2 + self._capture_thread = None + self._capture_stop_event = threading.Event() + self._capture_lock = threading.RLock() + self._capture_cond = threading.Condition(self._capture_lock) + self._latest_packet = None + self._latest_packet_seq = 0 + self._last_consumed_packet_seq = 0 + self._packet_queue = deque(maxlen=self.async_capture_max_queue) + self._capture_thread_stats = { + "started": False, + "packets": 0, + "dropped_latest": 0, + "dropped_queue": 0, + "last_error": None, + "last_loop_ms": 0.0, + } + def __enter__(self): self.start() return self @@ -63,9 +131,30 @@ class OakFcc3Manager: def __exit__(self, exc_type, exc, tb): self.stop() + # ============================================================ + # Modes + # ============================================================ + def _is_preview_mode(self): return str(self.frame_type).upper() == "PREVIEW" + def _is_multispec_mode(self): + return str(self.frame_type).upper() == "MULTISPEC" + + def _is_raw_mode(self): + return str(self.frame_type).upper() == "RAW_BRUTO" + + # ============================================================ + # Config helpers + # ============================================================ + + def _load_module_params(self, path): + if not path or not os.path.isfile(path): + return {} + + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + def _default_controls_for_role(self, role: str): role = str(role).lower() @@ -86,8 +175,14 @@ class OakFcc3Manager: "colour_gains": None, } + # ============================================================ + # Device discovery + # ============================================================ + def list_cameras(self): - with dai.Device() as dev: + dev_info = self._resolve_device_info() + + with dai.Device(dev_info) as dev: result = [] for f in dev.getConnectedCameraFeatures(): result.append({ @@ -97,80 +192,669 @@ class OakFcc3Manager: }) return result - def _get_available_cam_ids_ordered(self): - role_order = ["rgb", "nir", "re"] - available = list(self.queues.keys()) - - def sort_key(cam_id): - role = str(self.roles.get(cam_id, "unknown")).lower() + def _device_id_from_info(self, dev_info): + for name in ("getMxId", "getDeviceId"): try: - return role_order.index(role) - except ValueError: - return 99 + fn = getattr(dev_info, name, None) + if callable(fn): + value = fn() + if value: + return str(value) + except Exception: + pass - return sorted(available, key=sort_key) + try: + value = getattr(dev_info, "mxid", None) + if value: + return str(value) + except Exception: + pass + + try: + value = getattr(dev_info, "deviceId", None) + if value: + return str(value) + except Exception: + pass + + return None + + def _resolve_device_info(self): + devices = dai.Device.getAllAvailableDevices() + + if not devices: + raise RuntimeError("Nenhum dispositivo DepthAI/OAK encontrado.") + + if self.mx_id is None: + return devices[0] + + target = str(self.mx_id).strip() + + for dev_info in devices: + dev_id = self._device_id_from_info(dev_info) + if dev_id == target: + return dev_info + + disponiveis = [ + self._device_id_from_info(d) or str(getattr(d, "name", "unknown")) + for d in devices + ] + + raise RuntimeError( + f"Dispositivo DepthAI com ID '{target}' não encontrado. " + f"Disponíveis: {disponiveis}" + ) + + def _socket_from_name(self, socket_name): + socket_name = str(socket_name).upper() + if socket_name == "CAM_A": + return dai.CameraBoardSocket.CAM_A + if socket_name == "CAM_B": + return dai.CameraBoardSocket.CAM_B + if socket_name == "CAM_C": + return dai.CameraBoardSocket.CAM_C + if socket_name == "CAM_D": + return dai.CameraBoardSocket.CAM_D + raise ValueError(f"Socket não suportado: {socket_name}") + + # ============================================================ + # Pipeline creation, classic RAW/PREVIEW + # ============================================================ + + def _create_camera_node_classic(self, socket, sensor_name: str, role: str): + """ + Fluxo clássico. + + RGB/OV9782: + ColorCamera raw para RAW_BRUTO. + + MONO/OV9282: + MonoCamera raw quando disponível. + """ + sensor_name_u = str(sensor_name or "").upper() + role_u = str(role or "").lower() + + is_rgb = ( + role_u == "rgb" + or "OV9782" in sensor_name_u + or socket == dai.CameraBoardSocket.CAM_A + ) + + if is_rgb: + cam = self.pipeline.createColorCamera() + cam.setBoardSocket(socket) + + try: + cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_800_P) + except Exception: + try: + cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P) + except Exception: + pass + + cam.setInterleaved(False) + cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB) + cam.setFps(float(self.fps)) + + if self._is_preview_mode(): + try: + cam.setVideoSize(int(self.width), int(self.height)) + return cam, cam.video + except Exception: + return cam, cam.preview + + return cam, cam.raw + + mono = self.pipeline.create(dai.node.MonoCamera) + mono.setBoardSocket(socket) + + try: + mono.setResolution(dai.MonoCameraProperties.SensorResolution.THE_800_P) + except Exception: + try: + mono.setResolution(dai.MonoCameraProperties.SensorResolution.THE_720_P) + except Exception: + try: + mono.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P) + except Exception: + pass + + mono.setFps(float(self.fps)) + + if self._is_preview_mode(): + return mono, mono.out + + if hasattr(mono, "raw"): + return mono, mono.raw + + return mono, mono.out + + def _create_imu_node(self, pipeline): + self.has_imu_pipeline = False + + try: + imu = pipeline.create(dai.node.IMU) + imu.enableIMUSensor(dai.IMUSensor.ACCELEROMETER_RAW, 100) + imu.enableIMUSensor(dai.IMUSensor.GYROSCOPE_RAW, 100) + imu.setBatchReportThreshold(1) + imu.setMaxBatchReports(20) + + xout_imu = pipeline.create(dai.node.XLinkOut) + xout_imu.setStreamName("imu") + imu.out.link(xout_imu.input) + + self.has_imu_pipeline = True + print("[OAK] Pipeline IMU criado") + + except Exception as e: + self.has_imu_pipeline = False + print(f"[WARN] IMU indisponível no pipeline: {e}") + + # ============================================================ + # Pipeline creation, MULTISPEC aligned on OAK + # ============================================================ + + def _create_color_camera_multispec(self, socket): + cam = self.pipeline.create(dai.node.ColorCamera) + cam.setBoardSocket(socket) + cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_800_P) + cam.setFps(float(self.fps)) + cam.setInterleaved(False) + + # Usamos BGR porque getCvFrame/OpenCV lida direto com BGR. + # O Client converte para RGB float no _frame_to_float01. + cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR) + cam.setVideoSize(int(self.sensor_width), int(self.sensor_height)) + return cam + + def _create_mono_camera_multispec(self, socket): + cam = self.pipeline.create(dai.node.MonoCamera) + cam.setBoardSocket(socket) + cam.setResolution(dai.MonoCameraProperties.SensorResolution.THE_800_P) + cam.setFps(float(self.fps)) + return cam + + def _make_xout(self, name): + if hasattr(dai.node, "XLinkOut"): + xout = self.pipeline.create(dai.node.XLinkOut) + xout.setStreamName(name) + return xout + + if hasattr(self.pipeline, "createXLinkOut"): + xout = self.pipeline.createXLinkOut() + xout.setStreamName(name) + return xout + + raise RuntimeError("Não encontrei XLinkOut nesta versão do DepthAI.") + + def _apply_four_point_transform(self, config, src_quad_px, dst_quad_px): + src_pts = [dai.Point2f(float(x), float(y)) for x, y in src_quad_px] + dst_pts = [dai.Point2f(float(x), float(y)) for x, y in dst_quad_px] + + if hasattr(config, "setWarpTransformFourPoints"): + try: + config.setWarpTransformFourPoints(src_pts, dst_pts, False) + return + except TypeError: + config.setWarpTransformFourPoints(src_pts, False) + return + + if hasattr(config, "addTransformFourPoints"): + config.addTransformFourPoints(src_pts, dst_pts, False) + return + + raise RuntimeError( + "ImageManipConfig não tem setWarpTransformFourPoints nem addTransformFourPoints" + ) + + def _create_warp_manip( + self, + name, + out_w, + out_h, + src_quad_px, + frame_type=None, + max_output_frame_size=None, + ): + manip = self.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)), + ] + + self._apply_four_point_transform( + manip.initialConfig, + src_quad_px, + dst_quad_px, + ) + + if hasattr(manip.initialConfig, "setResize"): + try: + manip.initialConfig.setResize(int(out_w), int(out_h)) + except Exception: + pass + + if frame_type is not None: + try: + manip.initialConfig.setFrameType(frame_type) + except Exception: + pass + + if max_output_frame_size is None: + max_output_frame_size = int(out_w * out_h * 3) + + manip.setMaxOutputFrameSize(int(max_output_frame_size)) + + xout = self._make_xout(name) + manip.out.link(xout.input) + + return manip, xout + + def _start_multispec_pipeline(self, features): + """ + Cria pipeline onde a OAK entrega frames já alinhados. + Stream names continuam CAM_A/CAM_B/CAM_C para preservar o contrato. + """ + self.aligned_geometry = self._prepare_multispec_geometry() + + feature_by_socket = {f.socket.name: f for f in features} + + required = ["CAM_A", "CAM_B", "CAM_C"] + missing = [cam_id for cam_id in required if cam_id not in feature_by_socket] + if missing: + raise RuntimeError( + f"MULTISPEC exige CAM_A/CAM_B/CAM_C ativos. Ausentes: {missing}" + ) + + for cam_id in required: + if self.only_camera is not None and cam_id != self.only_camera: + continue + + f = feature_by_socket[cam_id] + role = str(self.roles.get(cam_id, "unknown")).lower() + socket = f.socket + + print( + f"[OAK] Criando câmera MULTISPEC {cam_id} " + f"sensor={f.sensorName} role={role}" + ) + + if role == "rgb": + cam = self._create_color_camera_multispec(socket) + src_output = cam.video + quad = self.aligned_geometry["quad_rgb"] + out_type = dai.ImgFrame.Type.BGR888p + max_size = self.width * self.height * 3 + channels = 3 + bit_depth = 8 + raw_format = "BGR888p" + elif role == "re": + cam = self._create_mono_camera_multispec(socket) + src_output = cam.out + quad = self.aligned_geometry["quad_re"] + out_type = dai.ImgFrame.Type.GRAY8 + max_size = self.width * self.height + channels = 1 + bit_depth = 8 + raw_format = "GRAY8" + elif role == "nir": + cam = self._create_mono_camera_multispec(socket) + src_output = cam.out + quad = self.aligned_geometry["quad_nir"] + out_type = dai.ImgFrame.Type.GRAY8 + max_size = self.width * self.height + channels = 1 + bit_depth = 8 + raw_format = "GRAY8" + else: + raise RuntimeError(f"Role não suportada no MULTISPEC: cam_id={cam_id}, role={role}") + + self.apply_initial_camera_controls_to_node(cam, cam_id) + + xin_ctrl = self.pipeline.create(dai.node.XLinkIn) + xin_ctrl.setStreamName(f"{cam_id}_ctrl") + xin_ctrl.out.link(cam.inputControl) + + manip, _ = self._create_warp_manip( + name=cam_id, + out_w=self.width, + out_h=self.height, + src_quad_px=quad, + frame_type=out_type, + max_output_frame_size=max_size, + ) + src_output.link(manip.inputImage) + + self.queues[cam_id] = None + self.buffers[cam_id] = deque(maxlen=self.buffer_size) + self.control_queues[cam_id] = None + + self.camera_info[cam_id] = { + "id": cam_id, + "socket": cam_id, + "sensor": f.sensorName, + "role": role, + "interface": "OAK_ALIGNED", + "raw_format": raw_format, + "channels": channels, + "bit_depth": bit_depth, + "width": int(self.width), + "height": int(self.height), + "aligned_by_oak": True, + "homography_applied": role in ("re", "nir"), + "crop_resize_applied": True, + } + + # ============================================================ + # Homography helpers for MULTISPEC + # ============================================================ + + def _prepare_multispec_geometry(self): + fusion = self.fusion_config or {} + + if str(fusion.get("alignment_mode", "homography")).lower() != "homography": + raise RuntimeError( + "frame_type=MULTISPEC na OAK exige fusion_config.alignment_mode='homography'." + ) + + 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 " + "para frame_type=MULTISPEC." + ) + + calib_size = fusion.get("homography_calibration_size", None) + runtime_size = (self.sensor_width, self.sensor_height) + + H_re = self._scale_homography_to_runtime(H_re_raw, calib_size, runtime_size) + H_nir = self._scale_homography_to_runtime(H_nir_raw, calib_size, runtime_size) + + crop_box = self._compute_common_crop_box( + self.sensor_width, + self.sensor_height, + H_re, + H_nir, + ) + + quad_rgb = self._identity_quad_for_crop(crop_box) + quad_re = self._quad_for_output_crop_to_input(H_re, crop_box) + quad_nir = self._quad_for_output_crop_to_input(H_nir, crop_box) + + quad_rgb = self._clamp_quad(quad_rgb, self.sensor_width, self.sensor_height) + quad_re = self._clamp_quad(quad_re, self.sensor_width, self.sensor_height) + quad_nir = self._clamp_quad(quad_nir, self.sensor_width, self.sensor_height) + + print("============================================") + print("[OAK MULTISPEC] Geometria alinhada na OAK") + print(f"sensor : {self.sensor_width}x{self.sensor_height}") + print(f"output : {self.width}x{self.height}") + 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("============================================") + + return { + "reference": "rgb", + "mode": "homography", + "sensor_size": [int(self.sensor_width), int(self.sensor_height)], + "output_size": [int(self.width), int(self.height)], + "crop_box": [int(v) for v in crop_box], + "H_re": H_re, + "H_nir": H_nir, + "quad_rgb": quad_rgb, + "quad_re": quad_re, + "quad_nir": quad_nir, + } + + def _scale_homography_to_runtime(self, H, calib_size, runtime_size): + 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) + + if calib_w <= 0 or calib_h <= 0: + return H.astype(np.float32) + + 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(self, 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(self, runtime_w, runtime_h, H_re, H_nir): + base = np.ones((runtime_h, runtime_w), dtype=np.uint8) * 255 + + rgb_mask = base + re_mask = self._warp_mask(base, H_re, runtime_w, runtime_h) + nir_mask = self._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(self, 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(self, H_src_to_rgb, crop_box): + x0, y0, x1, y1 = crop_box + + dst_corners_rgb = [ + (x0, y0), + (x1, y0), + (x1, y1), + (x0, y1), + ] + + H_inv = np.linalg.inv(H_src_to_rgb).astype(np.float32) + + src_quad = [] + for x, y in dst_corners_rgb: + sx, sy = self._apply_H_to_point(H_inv, x, y) + src_quad.append((sx, sy)) + + return src_quad + + def _identity_quad_for_crop(self, 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 _clamp_quad(self, 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 + + # ============================================================ + # Start / stop + # ============================================================ def start(self): if self.running: return - self.device = dai.Device() - self.pipeline = dai.Pipeline(self.device) + self.dev_info = self._resolve_device_info() + self.mx_id = self._device_id_from_info(self.dev_info) or self.mx_id + + self.device = dai.Device(self.dev_info) + self.pipeline = dai.Pipeline() features = self.device.getConnectedCameraFeatures() self.queues.clear() self.buffers.clear() self.camera_info.clear() + self.control_queues.clear() self._last_raw_dims.clear() + self.aligned_geometry = None - for f in features: - socket = f.socket - socket_name = socket.name - if self.only_camera is not None and socket_name != self.only_camera: - continue + if self._is_multispec_mode(): + self._start_multispec_pipeline(features) + else: + for f in features: + socket = f.socket + socket_name = socket.name + if self.only_camera is not None and socket_name != self.only_camera: + continue - role = self.roles.get(socket_name, "unknown") + role = self.roles.get(socket_name, "unknown") - print(f"[OAK] Criando câmera {socket_name} sensor={f.sensorName} role={role}") + print(f"[OAK] Criando câmera {socket_name} sensor={f.sensorName} role={role}") - cam = self.pipeline.create(dai.node.Camera).build(socket) - self.apply_initial_camera_controls_to_node(cam, socket_name) - ctrl_q = cam.inputControl.createInputQueue() - - if self._is_preview_mode(): - out = cam.requestOutput( - self.size, - fps=self.fps + cam, output = self._create_camera_node_classic( + socket=socket, + sensor_name=f.sensorName, + role=role, ) - else: - out = cam.raw - q = out.createOutputQueue() - cam_id = socket_name + self.apply_initial_camera_controls_to_node(cam, socket_name) - self.queues[cam_id] = q - self.buffers[cam_id] = deque(maxlen=self.buffer_size) - self.control_queues[cam_id] = ctrl_q + xin_ctrl = self.pipeline.create(dai.node.XLinkIn) + xin_ctrl.setStreamName(f"{socket_name}_ctrl") + xin_ctrl.out.link(cam.inputControl) - self.camera_info[cam_id] = { - "id": cam_id, - "socket": socket_name, - "sensor": f.sensorName, - "role": role, - } + xout = self.pipeline.create(dai.node.XLinkOut) + xout.setStreamName(socket_name) + output.link(xout.input) + + cam_id = socket_name + + self.queues[cam_id] = None + self.buffers[cam_id] = deque(maxlen=self.buffer_size) + self.control_queues[cam_id] = None + + self.camera_info[cam_id] = { + "id": cam_id, + "socket": socket_name, + "sensor": f.sensorName, + "role": role, + } self._validate_capture_mode() - self.pipeline.start() + self._create_imu_node(self.pipeline) + + self.device.startPipeline(self.pipeline) + + self.q_imu = None + self.tem_imu = False + + if self.has_imu_pipeline: + try: + self.q_imu = self.device.getOutputQueue( + name="imu", + maxSize=50, + blocking=False, + ) + self.tem_imu = True + print("[OAK] Fila IMU criada") + except Exception as e: + self.q_imu = None + self.tem_imu = False + print(f"[WARN] Fila IMU indisponível: {e}") + + for cam_id in self.camera_info.keys(): + self.queues[cam_id] = self.device.getOutputQueue( + name=cam_id, + maxSize=self.buffer_size, + blocking=False, + ) + + self.control_queues[cam_id] = self.device.getInputQueue( + name=f"{cam_id}_ctrl", + maxSize=4, + blocking=False, + ) + self.running = True + if not hasattr(self, "async_capture_enabled"): + self.async_capture_enabled = True + if not hasattr(self, "async_capture_mode"): + self.async_capture_mode = "latest" + if not hasattr(self, "async_capture_max_queue"): + self.async_capture_max_queue = 2 + self._reset_async_capture_state() + self._start_async_capture_thread() + time.sleep(0.05) + def stop(self): if not self.running: return + + self._stop_async_capture_thread() try: - if self.pipeline is not None: + if self.pipeline is not None and hasattr(self.pipeline, "stop"): self.pipeline.stop() except Exception: pass @@ -188,57 +872,165 @@ class OakFcc3Manager: self.camera_info.clear() self.control_queues.clear() self._last_raw_dims.clear() + self.aligned_geometry = None + self.q_imu = None + self.tem_imu = False + self.has_imu_pipeline = False self.running = False + # ============================================================ + # Status + # ============================================================ + def get_status(self): return { + "mx_id": self.mx_id, "backend": "oak_fcc3", "running": self.running, "fps": self.fps, "width": self.width, "height": self.height, + "sensor_width": self.sensor_width, + "sensor_height": self.sensor_height, "frame_type": self.frame_type, "output_dtype": self.output_dtype, "capture_mode": self.capture_mode, "raw_policy": self.raw_policy, "sync_tolerance_ms": self.sync_tolerance_ms, "buffer_size": self.buffer_size, + "geometry_stage": "oak" if self._is_multispec_mode() else "pc", + "aligned_geometry": self._serializable_aligned_geometry(), "cameras": list(self.camera_info.values()), + "async_capture": self.get_async_capture_status() if hasattr(self, "get_async_capture_status") else None, + "tem_imu": bool(getattr(self, "tem_imu", False)), + "has_imu_pipeline": bool(getattr(self, "has_imu_pipeline", False)), } + def _serializable_aligned_geometry(self): + if not isinstance(self.aligned_geometry, dict): + return None + + out = {} + for k, v in self.aligned_geometry.items(): + if isinstance(v, np.ndarray): + out[k] = v.tolist() + else: + out[k] = v + return out + + # ============================================================ + # Frame capture + # ============================================================ + def get_next_frame(self, timeout=1.0): if not self.running: raise RuntimeError("OakFcc3Manager não está rodando. Chame start() primeiro.") - t0 = time.time() + # Fallback síncrono se desligar async. + if not bool(getattr(self, "async_capture_enabled", True)): + return self._get_next_frame_sync_instrumented(timeout=timeout) - while time.time() - t0 < timeout: - self._drain_queues_to_buffers() + t0 = time.perf_counter() + deadline = t0 + float(timeout) - synced = self._try_get_synced_packet() + with self._capture_cond: + while time.perf_counter() < deadline: + packet = None + + mode = str(getattr(self, "async_capture_mode", "latest")).lower() + + if mode == "queue": + if len(self._packet_queue) > 0: + packet = self._packet_queue.popleft() + else: + latest = self._latest_packet + if latest is not None and int(latest.get("seq", 0)) > int(self._last_consumed_packet_seq): + packet = latest + + if packet is not None: + seq = int(packet.get("seq", 0)) + self._last_consumed_packet_seq = seq + + frames = packet["frames"] + meta = dict(packet["meta"]) + + age_ms = (time.perf_counter() - float(packet.get("created_perf_counter", time.perf_counter()))) * 1000.0 + get_wait_ms = (time.perf_counter() - t0) * 1000.0 + + cp = dict(meta.get("capture_perf", {}) or {}) + cp["async_consumer"] = True + cp["async_packet_seq"] = seq + cp["async_packet_age_ms"] = float(age_ms) + cp["async_get_wait_ms"] = float(get_wait_ms) + cp["async_status"] = self.get_async_capture_status() + meta["capture_perf"] = cp + + return frames, meta + + remaining = deadline - time.perf_counter() + if remaining <= 0: + break + + self._capture_cond.wait(timeout=min(0.005, remaining)) + + raise TimeoutError( + f"Timeout aguardando pacote assíncrono do OAK-FFC-3. " + f"status={self.get_async_capture_status()}" + ) + + def _get_next_frame_sync_instrumented(self, timeout=1.0): + if not self.running: + raise RuntimeError("OakFcc3Manager não está rodando. Chame start() primeiro.") + + perf = self._new_capture_perf() if hasattr(self, "_new_capture_perf") else None + t_start_wall = time.time() + t_start = time.perf_counter() + + while time.time() - t_start_wall < timeout: + if perf is not None: + perf["loop_count"] += 1 + perf["buffer_lengths_before_sync"] = self._buffer_lengths_snapshot() + + t0 = time.perf_counter() + self._drain_queues_to_buffers(perf=perf) + if perf is not None: + perf["drain_total_ms"] += (time.perf_counter() - t0) * 1000.0 + perf["buffer_lengths_after_drain"] = self._buffer_lengths_snapshot() + + t0 = time.perf_counter() + synced = self._try_get_synced_packet(perf=perf) + if perf is not None: + perf["sync_select_ms"] += (time.perf_counter() - t0) * 1000.0 if synced is not None: frames, timestamps, sync_dt_ms, sync_ok, frame_controls = synced - self.frame_id += 1 - meta = self._build_meta( - frames, - timestamps, - sync_dt_ms, - sync_ok, - frame_controls=frame_controls, - ) + + t0 = time.perf_counter() + meta = self._build_meta(frames, timestamps, sync_dt_ms, sync_ok, frame_controls=frame_controls) + if perf is not None: + perf["meta_ms"] += (time.perf_counter() - t0) * 1000.0 + perf["wait_total_ms"] = (time.perf_counter() - t_start) * 1000.0 + perf["sync_dt_ms"] = float(sync_dt_ms) + perf["sync_ok"] = bool(sync_ok) + perf["buffer_lengths_after_sync"] = self._buffer_lengths_snapshot() + perf["wait_reason"] = "synced_packet_ready_sync" + meta["capture_perf"] = perf return frames, meta + t0 = time.perf_counter() time.sleep(0.001) + if perf is not None: + perf["sleep_ms"] += (time.perf_counter() - t0) * 1000.0 + perf["sleep_count"] += 1 raise TimeoutError( f"Timeout aguardando pacote sincronizado do OAK-FFC-3. " - f"Tolerância atual={self.sync_tolerance_ms} ms. " - f"Tente aumentar para 25 ou 35 ms para diagnóstico." + f"Tolerância atual={self.sync_tolerance_ms} ms." ) + def _extract_frame_controls(self, msg): controls = { "exposure_time_us": None, @@ -302,25 +1094,69 @@ class OakFcc3Manager: return controls - def _drain_queues_to_buffers(self): + def _drain_queues_to_buffers(self, perf=None): for cam_id, q in self.queues.items(): - while q.has(): - msg = q.get() + if perf is not None: + perf["drained_by_cam"].setdefault(cam_id, 0) + perf["queue_has_true_by_cam"].setdefault(cam_id, 0) + while True: + t0_has = self._cap_now_ms() + has_msg = q.has() + if perf is not None: + perf["drain_has_ms"] += self._cap_now_ms() - t0_has + + if not has_msg: + break + + if perf is not None: + perf["queue_has_true_by_cam"][cam_id] += 1 + + t0_get = self._cap_now_ms() + msg = q.get() + if perf is not None: + perf["drain_get_msg_ms"] += self._cap_now_ms() - t0_get + + t0_ts = self._cap_now_ms() try: ts = msg.getTimestamp().total_seconds() except Exception: ts = time.time() + if perf is not None: + perf["drain_get_timestamp_ms"] += self._cap_now_ms() - t0_ts - if self._is_preview_mode(): + if self._is_preview_mode() or self._is_multispec_mode(): + t0_data = self._cap_now_ms() frame = msg.getCvFrame() - else: - data = msg.getData() - raw = np.frombuffer(data, dtype=np.uint8).copy() + if perf is not None: + perf["drain_get_data_ms"] += self._cap_now_ms() - t0_data + if frame is None: + continue + + if self._is_multispec_mode(): + self._last_raw_dims[cam_id] = { + "sensor_width": int(frame.shape[1]), + "sensor_height": int(frame.shape[0]), + "stride": int(frame.strides[0]) if hasattr(frame, "strides") else int(frame.shape[1]), + "packed_width": int(frame.shape[1]), + } + + else: + t0_data = self._cap_now_ms() + data = msg.getData() + if perf is not None: + perf["drain_get_data_ms"] += self._cap_now_ms() - t0_data + + t0_copy = self._cap_now_ms() + raw = np.frombuffer(data, dtype=np.uint8).copy() + if perf is not None: + perf["drain_frombuffer_copy_ms"] += self._cap_now_ms() - t0_copy + + t0_shape = self._cap_now_ms() h = int(msg.getHeight()) w = int(msg.getWidth()) - stride = int(msg.getStride()) + stride = self._get_imgframe_stride(msg, raw.size, h, w) expected = h * stride @@ -331,6 +1167,8 @@ class OakFcc3Manager: ) frame = raw[:expected].reshape((h, stride)) + if perf is not None: + perf["drain_reshape_ms"] += self._cap_now_ms() - t0_shape self._last_raw_dims[cam_id] = { "sensor_width": w, @@ -338,8 +1176,11 @@ class OakFcc3Manager: "stride": stride, "packed_width": stride, } - + + t0_ctrl = self._cap_now_ms() frame_controls = self._extract_frame_controls(msg) + if perf is not None: + perf["drain_controls_ms"] += self._cap_now_ms() - t0_ctrl self.buffers[cam_id].append({ "frame": frame, @@ -347,14 +1188,33 @@ class OakFcc3Manager: "controls": frame_controls, }) - def _try_get_synced_packet(self): + if perf is not None: + perf["drained_total"] += 1 + perf["drained_by_cam"][cam_id] += 1 + + def _get_imgframe_stride(self, msg, raw_size: int, h: int, w: int) -> int: + try: + return int(msg.getStride()) + except Exception: + pass + + try: + if h > 0 and raw_size % h == 0: + return int(raw_size // h) + except Exception: + pass + + return int(np.ceil(w * 5.0 / 4.0)) + + def _try_get_synced_packet(self, perf=None): required_cam_ids = self._get_required_cam_ids() for cam_id in required_cam_ids: if cam_id not in self.buffers or len(self.buffers[cam_id]) == 0: + if perf is not None: + perf["wait_reason"] = f"empty_buffer:{cam_id}" return None - # Usa o timestamp mais antigo da câmera com menor buffer como referência. ref_cam_id = min(required_cam_ids, key=lambda cid: len(self.buffers[cid])) ref_item = self.buffers[ref_cam_id][0] ref_ts = ref_item["timestamp"] @@ -372,6 +1232,8 @@ class OakFcc3Manager: best_item = item if best_item is None: + if perf is not None: + perf["wait_reason"] = f"no_best_item:{cam_id}" return None selected[cam_id] = best_item @@ -386,6 +1248,13 @@ class OakFcc3Manager: for cam_id, item in selected.items() } + if perf is not None: + perf["selected_ts_by_cam"] = {cam_id: float(ts) for cam_id, ts in timestamps.items()} + perf["selected_seq_by_cam"] = { + cam_id: item.get("controls", {}).get("sequence_num") + for cam_id, item in selected.items() + } + ts_values = list(timestamps.values()) sync_dt_ms = (max(ts_values) - min(ts_values)) * 1000.0 if len(ts_values) >= 2 else 0.0 sync_ok = sync_dt_ms <= self.sync_tolerance_ms @@ -394,6 +1263,10 @@ class OakFcc3Manager: oldest_cam_id = min(timestamps, key=timestamps.get) if len(self.buffers[oldest_cam_id]) > 0: self.buffers[oldest_cam_id].popleft() + if perf is not None: + perf["wait_reason"] = f"strict_drop_oldest:{oldest_cam_id}" + perf["sync_dt_ms"] = float(sync_dt_ms) + perf["sync_ok"] = False return None frames = { @@ -401,18 +1274,39 @@ class OakFcc3Manager: for cam_id, item in selected.items() } - # Remove dos buffers tudo até os frames usados. for cam_id, used_item in selected.items(): while len(self.buffers[cam_id]) > 0: item = self.buffers[cam_id].popleft() if item is used_item: break + if perf is not None: + perf["wait_reason"] = "synced_selected" + perf["sync_dt_ms"] = float(sync_dt_ms) + perf["sync_ok"] = bool(sync_ok) + return frames, timestamps, sync_dt_ms, sync_ok, frame_controls + def _get_available_cam_ids_ordered(self): + role_order = ["rgb", "re", "nir"] + available = list(self.queues.keys()) + + def sort_key(cam_id): + role = str(self.roles.get(cam_id, "unknown")).lower() + try: + return role_order.index(role) + except ValueError: + return 99 + + return sorted(available, key=sort_key) + def _get_required_cam_ids(self): available = self._get_available_cam_ids_ordered() + if self._is_multispec_mode(): + # MULTISPEC precisa sempre do trio completo para montar [R,G,B,RE,NIR]. + return available[:3] + if self.capture_mode == "SINGLE": return available[:1] @@ -429,6 +1323,10 @@ class OakFcc3Manager: return available + # ============================================================ + # Meta + # ============================================================ + def _build_meta(self, frames, timestamps, sync_dt_ms, sync_ok, frame_controls): payload_sources = list(frames.keys()) @@ -451,12 +1349,28 @@ class OakFcc3Manager: item["shape"] = list(arr.shape) item["dtype"] = str(arr.dtype) - if self._is_preview_mode(): + if self._is_multispec_mode(): + role = str(item.get("role", "")).lower() + item["interface"] = "OAK_ALIGNED" + item["raw_format"] = "BGR888p" if role == "rgb" else "GRAY8" + item["channels"] = 3 if arr.ndim == 3 else 1 + item["bit_depth"] = 8 if arr.dtype == np.uint8 else 16 + item["height"] = int(arr.shape[0]) + item["width"] = int(arr.shape[1]) + item["packed"] = False + item["aligned_by_oak"] = True + item["homography_applied"] = role in ("re", "nir") + item["crop_applied"] = True + item["resize_applied"] = True + item["color_order"] = "BGR" if role == "rgb" else None + + elif self._is_preview_mode(): item["interface"] = "OAK" item["channels"] = 3 if arr.ndim == 3 else 1 item["bit_depth"] = 8 if arr.dtype == np.uint8 else 16 item["height"] = int(arr.shape[0]) item["width"] = int(arr.shape[1]) + else: raw_dims = self._last_raw_dims.get(cam_id, {}) item["interface"] = "OAK_RAW" @@ -473,14 +1387,13 @@ class OakFcc3Manager: camera_info[cam_id] = item - return { + meta = { "frame_id": self.frame_id, "backend": "oak_fcc3", "frame_type": self.frame_type, "capture_mode": self.capture_mode, "output_dtype": self.output_dtype, "dtype": self.output_dtype, - "output_layout": "dict_by_camera", "payload_sources": payload_sources, "camera_info": camera_info, "timestamps": timestamps, @@ -496,9 +1409,32 @@ class OakFcc3Manager: "dt_send_payload_prev": 0.0, } + if self._is_multispec_mode(): + meta.update({ + "output_layout": "dict_by_camera_aligned", + "geometry_stage": "oak", + "aligned_by_oak": True, + "fusion_alignment": self._serializable_aligned_geometry(), + }) + else: + meta.update({ + "output_layout": "dict_by_camera", + "geometry_stage": "pc", + "aligned_by_oak": False, + }) + + return meta + + # ============================================================ + # Validation + # ============================================================ + def _validate_capture_mode(self): n = len(self.queues) + if self._is_multispec_mode() and n < 3: + raise RuntimeError(f"frame_type=MULTISPEC exige 3 câmeras, mas detectou {n}.") + if self.capture_mode == "TRIPLE" and n < 3: raise RuntimeError(f"CaptureMode TRIPLE exige 3 câmeras, mas detectou {n}.") @@ -508,6 +1444,10 @@ class OakFcc3Manager: if self.raw_policy == "require_triple" and n < 3: raise RuntimeError(f"raw_policy=require_triple exige 3 câmeras, mas detectou {n}.") + # ============================================================ + # Camera controls + # ============================================================ + def get_camera_controls(self, cam_id): self._validate_cam_id_known(cam_id) return dict(self.camera_controls.get(cam_id, {})) @@ -522,12 +1462,8 @@ class OakFcc3Manager: ctrl = dai.CameraControl() if enable: - # Em alguns builds v3 esse método existe. if hasattr(ctrl, "setAutoExposureEnable"): ctrl.setAutoExposureEnable() - else: - # fallback: deixa AE assumir por região/algoritmo interno quando possível - pass else: exp_us = int(ctrl_state.get("exposure_time_us") or 15000) gain = float(ctrl_state.get("analogue_gain") or 1.0) @@ -604,7 +1540,7 @@ class OakFcc3Manager: ctrl = dai.CameraControl() if hasattr(ctrl, "setManualWhiteBalance"): - # Nem sempre esse método usa red/blue diretamente. Fica como placeholder seguro. + # Placeholder seguro. Algumas versões não expõem red/blue diretamente. pass self._send_control(cam_id, ctrl) @@ -616,16 +1552,12 @@ class OakFcc3Manager: result = dict(self.camera_controls.get(cam_id, {})) - ae_requested = controls.get("ae_enable", None) - if "ae_enable" in controls: result = self.set_ae_enable(cam_id, bool(controls["ae_enable"])) if "awb_enable" in controls: result = self.set_awb_enable(cam_id, bool(controls["awb_enable"])) - # Se AE está ligado, NÃO aplicar exposição/ganho manual. - # exposure_time_us e analogue_gain ficam apenas como referência/snapshot. ae_is_on = bool(self.camera_controls[cam_id].get("ae_enable", False)) if not ae_is_on: @@ -636,7 +1568,6 @@ class OakFcc3Manager: result = self.set_analogue_gain(cam_id, float(controls["analogue_gain"])) else: - # Mantém os valores no estado interno só como referência, sem mandar manual exposure. if "exposure_time_us" in controls and controls["exposure_time_us"] is not None: self.camera_controls[cam_id]["exposure_time_us"] = int(controls["exposure_time_us"]) @@ -657,7 +1588,7 @@ class OakFcc3Manager: if not ae: cam.initialControl.setManualExposure( exp_us, - self._gain_to_iso(gain) + self._gain_to_iso(gain), ) def _send_control(self, cam_id, ctrl): @@ -681,9 +1612,204 @@ class OakFcc3Manager: @staticmethod def _gain_to_iso(gain: float) -> int: - # DepthAI usa ISO no setManualExposure(exposure_us, sensitivity_iso). - # Mantemos analogue_gain estilo Pi e convertemos para ISO aproximado. gain = max(1.0, float(gain)) iso = int(round(gain * 100)) return max(100, min(1600, iso)) - \ No newline at end of file + + + + def _cap_now_ms(self): + return time.perf_counter() * 1000.0 + + def _new_capture_perf(self): + return { + "wait_total_ms": 0.0, + "drain_total_ms": 0.0, + "drain_has_ms": 0.0, + "drain_get_msg_ms": 0.0, + "drain_get_timestamp_ms": 0.0, + "drain_get_data_ms": 0.0, + "drain_frombuffer_copy_ms": 0.0, + "drain_reshape_ms": 0.0, + "drain_controls_ms": 0.0, + "sync_select_ms": 0.0, + "meta_ms": 0.0, + "sleep_ms": 0.0, + "sleep_count": 0, + "loop_count": 0, + "drained_total": 0, + "drained_by_cam": {}, + "queue_has_true_by_cam": {}, + "buffer_lengths_before_sync": {}, + "buffer_lengths_after_drain": {}, + "buffer_lengths_after_sync": {}, + "selected_seq_by_cam": {}, + "selected_ts_by_cam": {}, + "sync_dt_ms": None, + "sync_ok": None, + "wait_reason": None, + } + + def _add_ms(self, perf, key, t0_ms): + if perf is not None: + perf[key] = float(perf.get(key, 0.0) or 0.0) + float(self._cap_now_ms() - t0_ms) + + def _buffer_lengths_snapshot(self): + return { + cam_id: int(len(buf)) + for cam_id, buf in self.buffers.items() + } + + + + def _reset_async_capture_state(self): + self._capture_stop_event = threading.Event() + self._capture_lock = threading.RLock() + self._capture_cond = threading.Condition(self._capture_lock) + self._latest_packet = None + self._latest_packet_seq = 0 + self._last_consumed_packet_seq = 0 + self._packet_queue = deque(maxlen=int(getattr(self, "async_capture_max_queue", 2))) + self._capture_thread_stats = { + "started": False, + "packets": 0, + "dropped_latest": 0, + "dropped_queue": 0, + "last_error": None, + "last_loop_ms": 0.0, + "last_packet_age_ms": None, + } + + def _start_async_capture_thread(self): + if not bool(getattr(self, "async_capture_enabled", True)): + return + + if self._capture_thread is not None and self._capture_thread.is_alive(): + return + + self._capture_stop_event.clear() + + self._capture_thread = threading.Thread( + target=self._async_capture_loop, + name="OakFcc3AsyncCapture", + daemon=True, + ) + self._capture_thread.start() + + def _stop_async_capture_thread(self): + try: + if hasattr(self, "_capture_stop_event") and self._capture_stop_event is not None: + self._capture_stop_event.set() + + if hasattr(self, "_capture_cond") and self._capture_cond is not None: + with self._capture_cond: + self._capture_cond.notify_all() + + th = getattr(self, "_capture_thread", None) + if th is not None and th.is_alive(): + th.join(timeout=1.0) + except Exception: + pass + + self._capture_thread = None + + def _async_capture_loop(self): + if not hasattr(self, "_capture_thread_stats"): + self._reset_async_capture_state() + + self._capture_thread_stats["started"] = True + + while not self._capture_stop_event.is_set(): + t_loop0 = time.perf_counter() + + try: + # Perf interno leve para diagnóstico. Não precisa imprimir todo frame. + perf = self._new_capture_perf() if hasattr(self, "_new_capture_perf") else None + + # Importante: esta thread é a única que mexe nas queues/buffers. + self._drain_queues_to_buffers(perf=perf) + synced = self._try_get_synced_packet(perf=perf) + + if synced is None: + # Dorme curto. Pode testar 0.0005 se quiser reduzir latência. + time.sleep(0.001) + continue + + frames, timestamps, sync_dt_ms, sync_ok, frame_controls = synced + + self.frame_id += 1 + meta = self._build_meta( + frames, + timestamps, + sync_dt_ms, + sync_ok, + frame_controls=frame_controls, + ) + + now = time.perf_counter() + packet = { + "seq": int(self._latest_packet_seq + 1), + "created_perf_counter": float(now), + "frames": frames, + "meta": meta, + } + + # Adiciona perf de captura assíncrona no meta. + if perf is not None: + perf["async_thread"] = True + perf["wait_total_ms"] = (time.perf_counter() - t_loop0) * 1000.0 + perf["sync_dt_ms"] = float(sync_dt_ms) + perf["sync_ok"] = bool(sync_ok) + perf["wait_reason"] = "async_packet_ready" + meta["capture_perf"] = perf + + with self._capture_cond: + self._latest_packet_seq += 1 + packet["seq"] = int(self._latest_packet_seq) + + if str(getattr(self, "async_capture_mode", "latest")).lower() == "queue": + before = len(self._packet_queue) + self._packet_queue.append(packet) + if before == self._packet_queue.maxlen: + self._capture_thread_stats["dropped_queue"] += 1 + else: + # latest mode: substitui pacote antigo se consumidor não pegou. + if self._latest_packet is not None and self._last_consumed_packet_seq < self._latest_packet.get("seq", 0): + self._capture_thread_stats["dropped_latest"] += 1 + self._latest_packet = packet + + self._capture_thread_stats["packets"] += 1 + self._capture_thread_stats["last_error"] = None + self._capture_thread_stats["last_loop_ms"] = (time.perf_counter() - t_loop0) * 1000.0 + + self._capture_cond.notify_all() + + except Exception as e: + try: + self._capture_thread_stats["last_error"] = f"{type(e).__name__}: {e}" + except Exception: + pass + time.sleep(0.005) + + try: + self._capture_thread_stats["started"] = False + except Exception: + pass + + def get_async_capture_status(self): + with self._capture_lock: + latest_age_ms = None + if self._latest_packet is not None: + latest_age_ms = (time.perf_counter() - float(self._latest_packet.get("created_perf_counter", 0.0))) * 1000.0 + + st = dict(getattr(self, "_capture_thread_stats", {}) or {}) + st.update({ + "enabled": bool(getattr(self, "async_capture_enabled", True)), + "mode": str(getattr(self, "async_capture_mode", "latest")), + "thread_alive": bool(self._capture_thread is not None and self._capture_thread.is_alive()), + "latest_seq": int(getattr(self, "_latest_packet_seq", 0)), + "last_consumed_seq": int(getattr(self, "_last_consumed_packet_seq", 0)), + "queue_len": int(len(getattr(self, "_packet_queue", []))), + "latest_age_ms": latest_age_ms, + }) + return st diff --git a/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_service.py b/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_service.py index 119746c52..1fde71780 100644 --- a/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_service.py +++ b/Python/OAK/datasets/oak-fcc-3/core/oak_fcc3_service.py @@ -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, diff --git a/Python/OAK/datasets/oak-fcc-3/core/radiometric_controller_bkp.py b/Python/OAK/datasets/oak-fcc-3/core/radiometric_controller_bkp.py deleted file mode 100644 index ffdfdf4d0..000000000 --- a/Python/OAK/datasets/oak-fcc-3/core/radiometric_controller_bkp.py +++ /dev/null @@ -1,1223 +0,0 @@ -import time -import json -import numpy as np - - -class RadiometricController: - """ - Controlador radiométrico para o módulo RGB/RE/NIR. - - Modos principais: - metering_mode: - - "global": usa uma região grande da cena, robusta por percentis. - - "reference_patches": usa patches/ROIs conhecidos, por exemplo branco/cinza/preto. - - "legacy_patch": compatível com o comportamento antigo: strip_y/patch_x. - - spectral_control_mode: - - "independent": controla rgb, re e nir separadamente. - - "shared": controla rgb separado e aplica uma decisão conjunta para re/nir. - """ - - def __init__( - self, - client, - enabled=True, - config_json_path=None, - interval_s=0.5, - strip_y0_pct=0.95, - strip_y1_pct=1.0, - patch_x0_pct=0.35, - patch_x1_pct=0.75, - target_mean=0.55, - deadband=0.04, - alpha=0.18, - exp_min_us=100, - exp_max_us=80000, - gain_min=1.0, - gain_max=4.0, - exp_step_gain=0.55, - prefer_exposure=True, - verbose=False, - ): - self.client = client - cfg = self._load_config_json(config_json_path) - - self.enabled = bool(cfg.get("enabled", enabled)) - self.interval_s = float(cfg.get("interval_s", interval_s)) - self.verbose = bool(cfg.get("verbose", verbose)) - - self.metering_mode = str(cfg.get("metering_mode", "global")).lower() - self.spectral_control_mode = str(cfg.get("spectral_control_mode", "shared")).lower() - - if self.metering_mode not in ("global", "reference_patches", "legacy_patch"): - self.metering_mode = "global" - - if self.spectral_control_mode not in ("shared", "independent"): - self.spectral_control_mode = "shared" - - self.strip_y0_pct = float(cfg.get("strip_y0_pct", strip_y0_pct)) - self.strip_y1_pct = float(cfg.get("strip_y1_pct", strip_y1_pct)) - self.patch_x0_pct = float(cfg.get("patch_x0_pct", patch_x0_pct)) - self.patch_x1_pct = float(cfg.get("patch_x1_pct", patch_x1_pct)) - - global_roi = cfg.get("global_roi_pct", {}) or {} - self.global_roi_pct = self._safe_roi_pct( - global_roi, - fallback={"x0": 0.08, "y0": 0.08, "x1": 0.92, "y1": 0.92}, - ) - - raw_global_by_role = cfg.get("global_roi_pct_by_role", {}) or {} - self.global_roi_pct_by_role = {} - - for role in self.ROLES: - roi = raw_global_by_role.get(role) - if isinstance(roi, dict) and roi: - self.global_roi_pct_by_role[role] = self._safe_roi_pct( - roi, - fallback=self.global_roi_pct, - ) - else: - self.global_roi_pct_by_role[role] = dict(self.global_roi_pct) - - self.reference_patches = cfg.get("reference_patches", []) or [] - self.patch_aggregation = str(cfg.get("patch_aggregation", "weighted_mean")).lower() - - self.patch_control_mode = str(cfg.get("patch_control_mode", "gray_primary")).lower() - - self.patch_require_order = bool(cfg.get("patch_require_order", True)) - self.patch_min_separation = float(cfg.get("patch_min_separation", 0.08)) - - self.patch_white_sat_limit_pct = float(cfg.get("patch_white_sat_limit_pct", 0.50)) - self.patch_white_p95_limit = float(cfg.get("patch_white_p95_limit", 0.90)) - - self.patch_black_dark_limit_pct = float(cfg.get("patch_black_dark_limit_pct", 80.0)) - self.patch_black_max_p50 = float(cfg.get("patch_black_max_p50", 0.20)) - - self.patch_gray_min_p50 = float(cfg.get("patch_gray_min_p50", 0.08)) - self.patch_gray_max_p50 = float(cfg.get("patch_gray_max_p50", 0.85)) - - self.control_metric = str(cfg.get("control_metric", "p50")).lower() - self.target_value = float(cfg.get("target_value", cfg.get("target_mean", target_mean))) - self.target_mean = self.target_value - self.deadband = float(cfg.get("deadband", deadband)) - self.alpha = float(cfg.get("alpha", alpha)) - - self.p95_limit = float(cfg.get("p95_limit", 0.92)) - self.saturation_limit_pct = float(cfg.get("saturation_limit_pct", 0.50)) - self.dark_limit_pct = float(cfg.get("dark_limit_pct", 35.0)) - - self.reduce_fast_factor = float(cfg.get("reduce_fast_factor", 0.82)) - self.factor_min = float(cfg.get("factor_min", 0.72)) - self.factor_max = float(cfg.get("factor_max", 1.28)) - - self.saturation_hard_pct = float(cfg.get("saturation_hard_pct", 20.0)) - self.saturation_extreme_pct = float(cfg.get("saturation_extreme_pct", 60.0)) - - self.gain_return_enabled = bool(cfg.get("gain_return_enabled", True)) - self.gain_return_factor = float(cfg.get("gain_return_factor", 0.60)) - self.gain_reduce_on_saturation = bool(cfg.get("gain_reduce_on_saturation", True)) - - self.exp_high_ratio_for_gain = float(cfg.get("exp_high_ratio_for_gain", 0.85)) - self.exp_low_ratio_for_gain_return = float(cfg.get("exp_low_ratio_for_gain_return", 0.65)) - - self.gain_increase_required_cycles = int(cfg.get("gain_increase_required_cycles", 5)) - self.gain_decrease_required_cycles = int(cfg.get("gain_decrease_required_cycles", 2)) - - self.gain_step_up = float(cfg.get("gain_step_up", 0.25)) - self.gain_step_down = float(cfg.get("gain_step_down", 0.50)) - - self.gain_hard_reset_on_saturation = bool(cfg.get("gain_hard_reset_on_saturation", False)) - - self._underexposed_cycles = { - "rgb": 0, - "re": 0, - "nir": 0, - "spectral_shared": 0, - } - - self._overexposed_cycles = { - "rgb": 0, - "re": 0, - "nir": 0, - "spectral_shared": 0, - } - - self.control_strategy = str(cfg.get("control_strategy", "ratio")).lower() - - self.ratio_alpha = float(cfg.get("ratio_alpha", 0.55)) - self.ratio_min = float(cfg.get("ratio_min", 0.55)) - self.ratio_max = float(cfg.get("ratio_max", 1.85)) - - self.ready_required_cycles = int(cfg.get("ready_required_cycles", 3)) - self._ready_cycles = { - "rgb": 0, - "re": 0, - "nir": 0, - "spectral_shared": 0, - } - - self.exp_min_us = int(cfg.get("exp_min_us", exp_min_us)) - self.exp_max_us = int(cfg.get("exp_max_us", exp_max_us)) - self.gain_min = float(cfg.get("gain_min", gain_min)) - self.gain_max = float(cfg.get("gain_max", gain_max)) - self.role_limits = cfg.get("role_limits", {}) or {} - - self.exp_step_gain = float(cfg.get("exp_step_gain", exp_step_gain)) - self.prefer_exposure = bool(cfg.get("prefer_exposure", prefer_exposure)) - - self.exp_apply_threshold_us = int(cfg.get("exp_apply_threshold_us", 80)) - self.gain_apply_threshold = float(cfg.get("gain_apply_threshold", 0.05)) - - self.apply_same_spectral_to_both = bool(cfg.get("apply_same_spectral_to_both", True)) - self.spectral_roles = tuple(cfg.get("spectral_roles", ["re", "nir"])) - - self.last_update_ts = 0.0 - self.last_result = {} - - self.state = { - "rgb": {"exp": 15000, "gain": 1.0}, - "nir": {"exp": 15000, "gain": 1.0}, - "re": {"exp": 15000, "gain": 1.0}, - } - - self._ae_disabled = set() - self._last_applied = { - "rgb": {"exp": None, "gain": None}, - "nir": {"exp": None, "gain": None}, - "re": {"exp": None, "gain": None}, - } - - def _load_config_json(self, path): - if not path: - return {} - try: - with open(path, "r", encoding="utf-8") as f: - data = json.load(f) - except Exception: - return {} - cfg = data.get("radiometric_config", {}) - return cfg if isinstance(cfg, dict) else {} - - ROLES = ("rgb", "re", "nir") - - @classmethod - def _normalize_role(cls, role: str) -> str: - role = str(role or "").lower() - return role if role in cls.ROLES else "rgb" - - @staticmethod - def _safe_roi_pct(roi_pct, fallback=None) -> dict: - if fallback is None: - fallback = {"x0": 0.08, "y0": 0.08, "x1": 0.92, "y1": 0.92} - - if not isinstance(roi_pct, dict): - roi_pct = fallback - - return { - "x0": float(roi_pct.get("x0", fallback.get("x0", 0.08))), - "y0": float(roi_pct.get("y0", fallback.get("y0", 0.08))), - "x1": float(roi_pct.get("x1", fallback.get("x1", 0.92))), - "y1": float(roi_pct.get("y1", fallback.get("y1", 0.92))), - } - - def _get_global_roi_pct_for_role(self, role: str) -> dict: - role = self._normalize_role(role) - - by_role = getattr(self, "global_roi_pct_by_role", {}) or {} - if isinstance(by_role, dict): - roi = by_role.get(role) - if isinstance(roi, dict) and roi: - return self._safe_roi_pct(roi, fallback=self.global_roi_pct) - - return self._safe_roi_pct(self.global_roi_pct) - - def _get_patch_roi_pct_for_role(self, patch: dict, role: str): - role = self._normalize_role(role) - - by_role = patch.get("roi_pct_by_role", {}) - if isinstance(by_role, dict): - roi = by_role.get(role) - if isinstance(roi, dict) and roi: - return self._safe_roi_pct(roi), "roi_pct_by_role" - - legacy = patch.get("roi_pct") - if isinstance(legacy, dict) and legacy: - return self._safe_roi_pct(legacy), "roi_pct" - - return None, "missing" - - def get_patch_target_for_role(self, patch, role, fallback): - by_role = patch.get("target_value_by_role", {}) - if isinstance(by_role, dict) and role in by_role: - return float(by_role[role]) - return float(patch.get("target_value", fallback)) - - def sync_from_camera_controls(self, camera_controls: dict | None): - if not isinstance(camera_controls, dict): - return - for role, ctrl in camera_controls.items(): - role = str(role).lower() - if role not in self.state or not isinstance(ctrl, dict): - continue - src = ctrl - if "requested" in ctrl and isinstance(ctrl["requested"], dict): - src = ctrl["requested"] - exp = src.get("exposure_time_us") - gain = src.get("analogue_gain") - if exp is not None: - self.state[role]["exp"] = int(exp) - if gain is not None: - self.state[role]["gain"] = float(gain) - - def sync_from_actual_camera_controls(self): - for role in ("rgb", "re", "nir"): - try: - ctrl = self.client.svc.get_camera_controls(role=role) - exp = ctrl.get("exposure_time_us") - gain = ctrl.get("analogue_gain") - if exp is not None: - self.state[role]["exp"] = int(exp) - if gain is not None: - self.state[role]["gain"] = float(gain) - except Exception: - pass - - def update(self, decoded: dict, meta: dict | None = None): - if not self.enabled: - return None - now = time.perf_counter() - if now - self.last_update_ts < self.interval_s: - return None - self.last_update_ts = now - - results = {} - rgb_result = self._update_single_role(decoded, "rgb") - if rgb_result is not None: - results["rgb"] = rgb_result - - if self.spectral_control_mode == "independent": - for role in self.spectral_roles: - result = self._update_single_role(decoded, role) - if result is not None: - results[role] = result - else: - result = self._update_spectral_shared(decoded) - if result is not None: - results["spectral_shared"] = result - - self.last_result = results - return results - - def _update_single_role(self, decoded: dict, role: str): - cam_id = self._resolve_cam_id(decoded, role) - if cam_id is None: - return None - img = decoded[cam_id].get("image") - if img is None: - return None - metrics = self.measure_image(img, role=role) - decision = self.compute_control(role, metrics) - apply_resp = self.apply_control(role, decision) - result = { - "mode": "single_role", - "cam_id": cam_id, - "role": role, - "metering_mode": self.metering_mode, - "metrics": metrics, - "decision": decision, - "apply": apply_resp, - } - self._print_metrics_debug(role, result) - - return result - - def _update_spectral_shared(self, decoded: dict): - role_items = {} - for role in self.spectral_roles: - cam_id = self._resolve_cam_id(decoded, role) - if cam_id is None: - continue - img = decoded[cam_id].get("image") - if img is None: - continue - role_items[role] = { - "cam_id": cam_id, - "metrics": self.measure_image(img, role=role), - } - if not role_items: - return None - - shared_metrics = self.aggregate_spectral_metrics(role_items) - state_role = "re" if "re" in role_items else list(role_items.keys())[0] - decision = self.compute_control(state_role, shared_metrics, virtual_role="spectral_shared") - - apply_resp = {} - if self.apply_same_spectral_to_both: - for role in role_items.keys(): - apply_resp[role] = self.apply_control(role, decision) - else: - apply_resp[state_role] = self.apply_control(state_role, decision) - - result = { - "mode": "shared_spectral", - "roles": role_items, - "metering_mode": self.metering_mode, - "metrics": shared_metrics, - "decision": decision, - "apply": apply_resp, - } - self._print_metrics_debug("spectral_shared", result) - - return result - - def _update_ready_state( - self, - log_role: str, - action: str, - error: float, - p95: float, - sat_pct: float, - ) -> tuple[bool, int]: - key = str(log_role).lower() - - is_ready_now = ( - action == "hold" - and abs(float(error)) <= self.deadband - and float(p95) <= self.p95_limit - and float(sat_pct) <= self.saturation_limit_pct - ) - - if is_ready_now: - self._ready_cycles[key] = self._ready_cycles.get(key, 0) + 1 - else: - self._ready_cycles[key] = 0 - - cycles = self._ready_cycles.get(key, 0) - return cycles >= self.ready_required_cycles, cycles - - def _update_exposure_pressure_state( - self, - log_role: str, - error: float, - p95: float, - sat_pct: float, - ) -> tuple[int, int]: - key = str(log_role).lower() - - under = ( - error > self.deadband - and p95 < self.p95_limit - and sat_pct <= self.saturation_limit_pct - ) - - over = ( - error < -self.deadband - or p95 > self.p95_limit - or sat_pct > self.saturation_limit_pct - ) - - if under: - self._underexposed_cycles[key] = self._underexposed_cycles.get(key, 0) + 1 - else: - self._underexposed_cycles[key] = 0 - - if over: - self._overexposed_cycles[key] = self._overexposed_cycles.get(key, 0) + 1 - else: - self._overexposed_cycles[key] = 0 - - return ( - self._underexposed_cycles.get(key, 0), - self._overexposed_cycles.get(key, 0), - ) - - def _resolve_cam_id(self, decoded, role): - role = str(role).lower() - for cam_id, data in decoded.items(): - if str(data.get("role", "")).lower() == role: - return cam_id - return None - - def measure_image(self, img01: np.ndarray, role: str) -> dict: - role = self._normalize_role(role) - gray = self.to_luma_or_gray(img01) - - if self.metering_mode == "reference_patches": - return self.measure_reference_patches(gray, role=role) - - if self.metering_mode == "legacy_patch": - return self.measure_legacy_patch(gray) - - return self.measure_global(gray, role=role) - - @staticmethod - def to_luma_or_gray(img01: np.ndarray) -> np.ndarray: - if img01.ndim == 3: - return ( - 0.299 * img01[:, :, 0] + - 0.587 * img01[:, :, 1] + - 0.114 * img01[:, :, 2] - ).astype(np.float32) - return img01.astype(np.float32) - - def measure_global(self, img_gray: np.ndarray, role: str = "rgb") -> dict: - role = self._normalize_role(role) - - h, w = img_gray.shape[:2] - roi_pct = self._get_global_roi_pct_for_role(role) - - roi = self._roi_pct_to_pixels( - h, w, - roi_pct["x0"], - roi_pct["y0"], - roi_pct["x1"], - roi_pct["y1"], - ) - - arr = self._crop_array(img_gray, roi) - stats = self.compute_stats(arr) - - stats["roi"] = list(roi) - stats["roi_pct"] = dict(roi_pct) - stats["roi_source"] = "global_roi_pct_by_role" - stats["role"] = role - stats["source"] = "global" - - return stats - - def measure_legacy_patch(self, img_gray: np.ndarray) -> dict: - h, w = img_gray.shape[:2] - roi = self._roi_pct_to_pixels( - h, w, - self.patch_x0_pct, - self.strip_y0_pct, - self.patch_x1_pct, - self.strip_y1_pct, - ) - arr = self._crop_array(img_gray, roi) - stats = self.compute_stats(arr) - stats["roi"] = list(roi) - stats["source"] = "legacy_patch" - return stats - - def measure_reference_patches(self, img_gray: np.ndarray, role: str) -> dict: - role = self._normalize_role(role) - - h, w = img_gray.shape[:2] - patch_results = [] - - for patch in self.reference_patches: - if not isinstance(patch, dict): - continue - - roles = patch.get("roles", ["rgb", "re", "nir", "all"]) - roles = [str(r).lower() for r in roles] - - if role not in roles and "all" not in roles: - continue - - roi_pct, roi_source = self._get_patch_roi_pct_for_role(patch, role) - - if not isinstance(roi_pct, dict): - continue - - x0 = float(roi_pct.get("x0", 0.0)) - y0 = float(roi_pct.get("y0", 0.0)) - x1 = float(roi_pct.get("x1", 1.0)) - y1 = float(roi_pct.get("y1", 1.0)) - - roi = self._roi_pct_to_pixels(h, w, x0, y0, x1, y1) - arr = self._crop_array(img_gray, roi) - stats = self.compute_stats(arr) - - target = self.get_patch_target_for_role( - patch=patch, - role=role, - fallback=self.target_value, - ) - if target is not None: - target = float(target) - - patch_results.append({ - "name": patch.get("name", f"patch_{len(patch_results) + 1}"), - "type": patch.get("type", "reference"), - "role": role, - "roi": list(roi), - "roi_pct": {"x0": x0, "y0": y0, "x1": x1, "y1": y1}, - "roi_source": roi_source, - "weight": float(patch.get("weight", 1.0)), - "target_value": target, - "stats": stats, - }) - - if not patch_results: - stats = self.measure_global(img_gray, role=role) - stats["source"] = "reference_patches_fallback_global" - stats["patches"] = [] - return stats - - metrics = self.aggregate_patch_metrics(patch_results) - metrics["role"] = role - - return metrics - - def aggregate_patch_metrics(self, patch_results: list[dict]) -> dict: - valid = [p for p in patch_results if p["stats"].get("valid")] - - if not valid: - return { - "valid": False, - "source": "reference_patches", - "patches": patch_results, - "mean": 0.0, - "p50": 0.0, - "p95": 0.0, - "sat_pct": 0.0, - "dark_pct": 0.0, - "control_value": 0.0, - "target_value": self.target_value, - "weighted_error": 0.0, - "patch_quality": { - "valid": False, - "warnings": ["no_valid_patches"], - }, - } - - black = self._find_patch_result(valid, "black") - gray = self._find_patch_result(valid, "gray") - white = self._find_patch_result(valid, "white") - - warnings = [] - - # Stats gerais de proteção. - p95s = np.array([p["stats"]["p95"] for p in valid], dtype=np.float32) - sats = np.array([p["stats"]["sat_pct"] for p in valid], dtype=np.float32) - darks = np.array([p["stats"]["dark_pct"] for p in valid], dtype=np.float32) - means = np.array([p["stats"]["mean"] for p in valid], dtype=np.float32) - p50s = np.array([p["stats"]["p50"] for p in valid], dtype=np.float32) - - p95_max = float(np.max(p95s)) - sat_max = float(np.max(sats)) - dark_mean = float(np.mean(darks)) - mean_mean = float(np.mean(means)) - p50_mean = float(np.mean(p50s)) - - # Valores por patch, quando existem. - black_p50 = float(black["stats"]["p50"]) if black else None - gray_p50 = float(gray["stats"]["p50"]) if gray else None - white_p50 = float(white["stats"]["p50"]) if white else None - - black_target = float(black.get("target_value", 0.06)) if black else 0.06 - gray_target = float(gray.get("target_value", self.target_value)) if gray else self.target_value - white_target = float(white.get("target_value", 0.80)) if white else 0.80 - - # ============================================================ - # Validações de coerência dos cartões - # ============================================================ - - if gray is None: - warnings.append("missing_gray_patch") - - if self.patch_require_order and black and gray and white: - if not (black_p50 < gray_p50 < white_p50): - warnings.append( - f"patch_order_invalid: black={black_p50:.3f}, gray={gray_p50:.3f}, white={white_p50:.3f}" - ) - - if (gray_p50 - black_p50) < self.patch_min_separation: - warnings.append( - f"black_gray_separation_low: diff={gray_p50 - black_p50:.3f}" - ) - - if (white_p50 - gray_p50) < self.patch_min_separation: - warnings.append( - f"gray_white_separation_low: diff={white_p50 - gray_p50:.3f}" - ) - - if white: - white_sat = float(white["stats"]["sat_pct"]) - white_p95 = float(white["stats"]["p95"]) - if white_sat > self.patch_white_sat_limit_pct: - warnings.append(f"white_patch_saturated: sat={white_sat:.2f}%") - if white_p95 > self.patch_white_p95_limit: - warnings.append(f"white_patch_p95_high: p95={white_p95:.3f}") - - if black: - black_dark = float(black["stats"]["dark_pct"]) - if black_dark > self.patch_black_dark_limit_pct: - warnings.append(f"black_patch_too_dark: dark={black_dark:.1f}%") - if black_p50 > self.patch_black_max_p50: - warnings.append(f"black_patch_too_bright: p50={black_p50:.3f}") - - if gray: - if gray_p50 < self.patch_gray_min_p50: - warnings.append(f"gray_patch_too_dark: p50={gray_p50:.3f}") - if gray_p50 > self.patch_gray_max_p50: - warnings.append(f"gray_patch_too_bright: p50={gray_p50:.3f}") - - # ============================================================ - # Modo recomendado: cinza como controle principal - # ============================================================ - if self.patch_control_mode == "gray_primary" and gray is not None: - control_value = gray_p50 - target_value = gray_target - weighted_error = target_value - control_value - - control_source = "gray_primary" - - else: - # Fallback: média ponderada original, mas preservando guardas. - weights = np.array([max(0.0, p.get("weight", 1.0)) for p in valid], dtype=np.float32) - - if float(weights.sum()) <= 1e-9: - weights = np.ones(len(valid), dtype=np.float32) - - weights = weights / weights.sum() - - patch_errors = [] - control_values = [] - - for p in valid: - target = p.get("target_value") - if target is None: - target = self.target_value - - value = p["stats"].get( - self.control_metric, - p["stats"].get("p50", p["stats"].get("mean", 0.0)) - ) - - control_values.append(float(value)) - patch_errors.append(float(target) - float(value)) - - weighted_error = float(np.sum(np.array(patch_errors, dtype=np.float32) * weights)) - control_value = float(np.sum(np.array(control_values, dtype=np.float32) * weights)) - target_value = self.target_value - control_source = "weighted_patches" - - # ============================================================ - # Guardas de saturação e faixa útil - # ============================================================ - # Se o branco saturou, queremos que o compute_control reduza exposição, - # mesmo que o cinza esteja aparentemente bom. - if white: - white_sat = float(white["stats"]["sat_pct"]) - white_p95 = float(white["stats"]["p95"]) - - sat_max = max(sat_max, white_sat) - p95_max = max(p95_max, white_p95) - - # Se o cinza está ausente, a métrica ainda pode funcionar por fallback, - # mas marcamos warning para debug. - quality_valid = gray is not None and len(warnings) == 0 - - return { - "valid": True, - "source": "reference_patches", - "patches": patch_results, - - # Métricas agregadas informativas. - "mean": mean_mean, - "p50": p50_mean, - "p95": p95_max, - "sat_pct": sat_max, - "dark_pct": dark_mean, - - # Métricas usadas pelo controle. - "control_metric": self.control_metric, - "control_value": float(control_value), - "target_value": float(target_value), - "weighted_error": float(weighted_error), - - # Debug/qualidade. - "patch_control_mode": self.patch_control_mode, - "control_source": control_source, - "patch_quality": { - "valid": bool(quality_valid), - "warnings": warnings, - "black_p50": black_p50, - "gray_p50": gray_p50, - "white_p50": white_p50, - "black_target": black_target, - "gray_target": gray_target, - "white_target": white_target, - }, - } - - def aggregate_spectral_metrics(self, role_items: dict) -> dict: - valid_items = {role: item for role, item in role_items.items() if item["metrics"].get("valid")} - if not valid_items: - return { - "valid": False, - "source": "spectral_shared", - "roles": role_items, - "mean": 0.0, - "p50": 0.0, - "p95": 0.0, - "sat_pct": 0.0, - "dark_pct": 0.0, - "control_value": 0.0, - "target_value": self.target_value, - } - - metrics_list = [item["metrics"] for item in valid_items.values()] - p95 = max(float(m.get("p95", 0.0)) for m in metrics_list) - sat_pct = max(float(m.get("sat_pct", 0.0)) for m in metrics_list) - mean = float(np.mean([float(m.get("mean", 0.0)) for m in metrics_list])) - p50 = float(np.mean([float(m.get("p50", m.get("mean", 0.0))) for m in metrics_list])) - dark_pct = float(np.mean([float(m.get("dark_pct", 0.0)) for m in metrics_list])) - control_values = [ - float(m.get("control_value", m.get(self.control_metric, m.get("p50", m.get("mean", 0.0))))) - for m in metrics_list - ] - - target_values = [ - float(m.get("target_value", self.target_value)) - for m in metrics_list - ] - - errors = [ - float(m.get("weighted_error", target - value)) - for m, target, value in zip(metrics_list, target_values, control_values) - ] - - control_value = float(np.mean(control_values)) - target_value = float(np.mean(target_values)) - weighted_error = float(np.mean(errors)) - - return { - "valid": True, - "source": "spectral_shared", - "roles": role_items, - "mean": mean, - "p50": p50, - "p95": p95, - "sat_pct": sat_pct, - "dark_pct": dark_pct, - "control_metric": self.control_metric, - "control_value": control_value, - "target_value": target_value, - "weighted_error": weighted_error, - "control_values_by_role": { - role: float(item["metrics"].get("control_value", item["metrics"].get("p50", 0.0))) - for role, item in valid_items.items() - }, - "targets_by_role": { - role: float(item["metrics"].get("target_value", self.target_value)) - for role, item in valid_items.items() - }, - "patch_quality_by_role": { - role: item["metrics"].get("patch_quality", {}) - for role, item in valid_items.items() - }, - } - - def compute_stats(self, arr: np.ndarray) -> dict: - arr = np.asarray(arr, dtype=np.float32).reshape(-1) - if arr.size == 0: - return { - "valid": False, - "pixels": 0, - "mean": 0.0, - "std": 0.0, - "p05": 0.0, - "p50": 0.0, - "p95": 0.0, - "sat_pct": 0.0, - "dark_pct": 0.0, - } - return { - "valid": True, - "pixels": int(arr.size), - "mean": float(arr.mean()), - "std": float(arr.std()), - "p05": float(np.percentile(arr, 5)), - "p50": float(np.percentile(arr, 50)), - "p95": float(np.percentile(arr, 95)), - "sat_pct": float((arr >= 0.98).mean() * 100.0), - "dark_pct": float((arr <= 0.02).mean() * 100.0), - } - - @staticmethod - def _crop_array(img: np.ndarray, roi: tuple[int, int, int, int]) -> np.ndarray: - x0, y0, x1, y1 = roi - return img[y0:y1, x0:x1].reshape(-1) - - @staticmethod - def _roi_pct_to_pixels(h: int, w: int, x0_pct: float, y0_pct: float, x1_pct: float, y1_pct: float): - x0 = int(w * x0_pct) - x1 = int(w * x1_pct) - y0 = int(h * y0_pct) - y1 = int(h * y1_pct) - x0 = max(0, min(w - 1, x0)) - x1 = max(x0 + 1, min(w, x1)) - y0 = max(0, min(h - 1, y0)) - y1 = max(y0 + 1, min(h, y1)) - return x0, y0, x1, y1 - - @staticmethod - def _find_patch_result(patch_results: list[dict], patch_type: str): - patch_type = str(patch_type).lower() - for p in patch_results: - if str(p.get("type", "")).lower() == patch_type: - return p - return None - - def compute_control(self, role: str, metrics: dict, virtual_role: str | None = None) -> dict: - state_role = str(role).lower() - log_role = virtual_role or state_role - - st = self.state.setdefault(state_role, {"exp": 15000, "gain": 1.0}) - old_exp = int(st["exp"]) - old_gain = float(st["gain"]) - limits = self._limits_for_role(state_role) - - if not metrics.get("valid"): - ready, ready_cycles = self._update_ready_state( - log_role=log_role, - action="hold", - error=999.0, - p95=1.0, - sat_pct=100.0, - ) - return { - "role": log_role, - "state_role": state_role, - "action": "hold", - "reason": "métrica inválida", - "old_exp": old_exp, - "new_exp": old_exp, - "old_gain": old_gain, - "new_gain": old_gain, - "ready": ready, - "ready_cycles": ready_cycles, - "ready_required_cycles": self.ready_required_cycles, - } - - control_value = float(metrics.get( - "control_value", - metrics.get(self.control_metric, metrics.get("p50", metrics.get("mean", 0.0))) - )) - target = float(metrics.get("target_value", self.target_value)) - error = float(metrics.get("weighted_error", target - control_value)) - p95 = float(metrics.get("p95", 0.0)) - sat_pct = float(metrics.get("sat_pct", 0.0)) - - under_cycles, over_cycles = self._update_exposure_pressure_state( - log_role=log_role, - error=error, - p95=p95, - sat_pct=sat_pct, - ) - - exp_min = int(limits["exp_min_us"]) - exp_max = int(limits["exp_max_us"]) - gain_min = float(limits["gain_min"]) - gain_max = float(limits["gain_max"]) - - new_exp = old_exp - new_gain = old_gain - action = "hold" - reason = "dentro da faixa morta" - - ratio = None - factor = 1.0 - gain_policy = "hold" - - # ============================================================ - # 1) Proteção forte contra saturação / p95 alto - # ============================================================ - if sat_pct > self.saturation_limit_pct or p95 > self.p95_limit: - if sat_pct >= self.saturation_extreme_pct: - exp_factor = 0.45 - elif sat_pct >= self.saturation_hard_pct: - exp_factor = 0.32 - elif sat_pct > self.saturation_limit_pct: - exp_factor = 0.55 - else: - exp_factor = self.reduce_fast_factor - - new_exp = int(self._clamp(old_exp * exp_factor, exp_min, exp_max)) - - if self.gain_reduce_on_saturation and old_gain > gain_min: - if self.gain_hard_reset_on_saturation and sat_pct >= self.saturation_extreme_pct: - new_gain = gain_min - gain_policy = "hard_reset_gain_on_extreme_saturation" - else: - desired_gain = old_gain - self.gain_step_down - new_gain = float(self._clamp(desired_gain, gain_min, gain_max)) - gain_policy = "decrease_gain_step_on_saturation" - else: - new_gain = old_gain - gain_policy = "hold_gain" - - action = "decrease_exposure" - reason = ( - f"saturação/p95 alto: sat={sat_pct:.2f}% p95={p95:.3f} " - f"exp_factor={exp_factor:.3f} gain_policy={gain_policy}" - ) - - # ============================================================ - # 2) Fora da faixa morta: controle por ratio/linear - # ============================================================ - elif abs(error) > self.deadband: - if self.control_strategy == "ratio": - safe_value = max(control_value, 1e-6) - ratio = target / safe_value - ratio = self._clamp(ratio, self.ratio_min, self.ratio_max) - factor = 1.0 + self.ratio_alpha * (ratio - 1.0) - else: - factor = 1.0 + self.exp_step_gain * error - factor = max(self.factor_min, min(self.factor_max, factor)) - - if self.prefer_exposure: - # ---------------------------------------------------- - # 2A) Cena escura: subir exposição primeiro. - # Só subir ganho se exposição já estiver perto do máximo. - # ---------------------------------------------------- - if error > 0: - desired_exp = int(self._clamp(old_exp * factor, exp_min, exp_max)) - new_exp = desired_exp - new_gain = old_gain - gain_policy = "hold_gain_prefer_exposure" - - exp_high_threshold = int(exp_max * self.exp_high_ratio_for_gain) - - if ( - desired_exp >= exp_high_threshold - and under_cycles >= self.gain_increase_required_cycles - ): - # Sobe ganho devagar, em degrau fixo. - desired_gain = old_gain + self.gain_step_up - new_gain = float(self._clamp(desired_gain, gain_min, gain_max)) - gain_policy = f"increase_gain_slow_under_cycles_{under_cycles}" - else: - new_gain = old_gain - gain_policy = f"hold_gain_under_cycles_{under_cycles}" - - action = "increase_exposure" - reason = ( - f"subindo exposição por {self.control_metric}: " - f"value={control_value:.3f} target={target:.3f} " - f"error={error:.3f} factor={factor:.3f} gain_policy={gain_policy}" - ) - - # ---------------------------------------------------- - # 2B) Cena clara: se ganho está acima do mínimo, - # reduzir ganho primeiro ou junto. - # ---------------------------------------------------- - else: - desired_exp = int(self._clamp(old_exp * factor, exp_min, exp_max)) - new_exp = desired_exp - - if ( - self.gain_return_enabled - and old_gain > gain_min - and over_cycles >= self.gain_decrease_required_cycles - ): - desired_gain = old_gain - self.gain_step_down - new_gain = float(self._clamp(desired_gain, gain_min, gain_max)) - gain_policy = f"return_gain_step_over_cycles_{over_cycles}" - else: - new_gain = old_gain - gain_policy = f"hold_gain_over_cycles_{over_cycles}" - - action = "decrease_exposure" - reason = ( - f"reduzindo brilho por {self.control_metric}: " - f"value={control_value:.3f} target={target:.3f} " - f"error={error:.3f} factor={factor:.3f} gain_policy={gain_policy}" - ) - - else: - desired_gain = old_gain * factor - new_gain = float(self._clamp(desired_gain, gain_min, gain_max)) - action = "increase_gain" if error > 0 else "decrease_gain" - gain_policy = "direct_gain_control" - reason = ( - f"corrigindo ganho por {self.control_metric}: " - f"value={control_value:.3f} target={target:.3f} " - f"error={error:.3f} factor={factor:.3f}" - ) - - # ============================================================ - # 3) Dentro da faixa morta: opcionalmente devolver ganho - # se ganho alto não é mais necessário. - # ============================================================ - else: - if self.gain_return_enabled and old_gain > gain_min: - exp_low_threshold = int(exp_max * self.exp_low_ratio_for_gain_return) - - if old_exp < exp_low_threshold: - new_gain = float(self._clamp(old_gain * self.gain_return_factor, gain_min, gain_max)) - gain_policy = "return_gain_while_ready" - action = "decrease_gain" - reason = ( - f"dentro da faixa, devolvendo ganho: " - f"value={control_value:.3f} target={target:.3f} " - f"gain={old_gain:.2f}->{new_gain:.2f}" - ) - else: - gain_policy = "hold_gain_high_exp" - else: - gain_policy = "hold_gain" - - new_exp = int(self._clamp(new_exp, exp_min, exp_max)) - new_gain = float(self._clamp(new_gain, gain_min, gain_max)) - - ready, ready_cycles = self._update_ready_state( - log_role=log_role, - action=action, - error=error, - p95=p95, - sat_pct=sat_pct, - ) - - return { - "role": log_role, - "state_role": state_role, - "action": action, - "reason": reason, - "metering_mode": self.metering_mode, - "spectral_control_mode": self.spectral_control_mode, - "control_metric": self.control_metric, - "control_value": control_value, - "target_value": target, - "error": error, - "p95": p95, - "sat_pct": sat_pct, - #"metrics_source": metrics.get("source"), - #"control_source": metrics.get("control_source"), - #"patch_quality": metrics.get("patch_quality"), - #"patches": metrics.get("patches"), - #"control_values_by_role": metrics.get("control_values_by_role"), - #"targets_by_role": metrics.get("targets_by_role"), - #"patch_quality_by_role": metrics.get("patch_quality_by_role"), - "old_exp": old_exp, - "new_exp": new_exp, - "old_gain": old_gain, - "new_gain": new_gain, - "limits": limits, - "ready": ready, - "ready_cycles": ready_cycles, - "ready_required_cycles": self.ready_required_cycles, - "control_strategy": self.control_strategy, - "ratio": ratio, - "factor": float(factor), - "gain_policy": gain_policy, - } - - def apply_control(self, role: str, decision: dict): - role = str(role).lower() - new_exp = int(decision["new_exp"]) - new_gain = float(decision["new_gain"]) - self.state[role]["exp"] = new_exp - self.state[role]["gain"] = new_gain - responses = {} - last = self._last_applied.setdefault(role, {"exp": None, "gain": None}) - try: - if role not in self._ae_disabled: - responses["ae"] = self.client.svc.set_ae_enable(role=role, enable=False) - if role == "rgb": - responses["awb"] = self.client.svc.set_awb_enable(role=role, enable=False) - self._ae_disabled.add(role) - if last["exp"] is None or abs(new_exp - last["exp"]) >= self.exp_apply_threshold_us: - responses["exposure"] = self.client.svc.set_exposure_time(role=role, exposure_time_us=new_exp) - last["exp"] = new_exp - if last["gain"] is None or abs(new_gain - last["gain"]) >= self.gain_apply_threshold: - responses["gain"] = self.client.svc.set_analogue_gain(role=role, analogue_gain=new_gain) - last["gain"] = new_gain - except Exception as e: - responses["error"] = str(e) - if self.verbose: - print( - f"[RAD_APPLY] role={role} " - f"action={decision.get('action')} " - f"exp={decision.get('old_exp')}->{decision.get('new_exp')} " - f"gain={decision.get('old_gain'):.2f}->{decision.get('new_gain'):.2f} " - f"ok={'error' not in responses}" - ) - return responses - - def _limits_for_role(self, role: str) -> dict: - role_cfg = self.role_limits.get(role, {}) or {} - return { - "exp_min_us": int(role_cfg.get("exp_min_us", self.exp_min_us)), - "exp_max_us": int(role_cfg.get("exp_max_us", self.exp_max_us)), - "gain_min": float(role_cfg.get("gain_min", self.gain_min)), - "gain_max": float(role_cfg.get("gain_max", self.gain_max)), - } - - def _smooth_int(self, old, desired): - return int(round((1.0 - self.alpha) * old + self.alpha * desired)) - - def _smooth_float(self, old, desired): - return float((1.0 - self.alpha) * old + self.alpha * desired) - - @staticmethod - def _clamp(v, lo, hi): - return max(lo, min(hi, v)) - - def _print_metrics_debug(self, role: str, result: dict): - if not self.verbose: - return - - metrics = result.get("metrics", {}) - decision = result.get("decision", {}) - - print( - f"[RAD_METRICS] role={role} " - f"mode={result.get('mode')} " - f"metering={result.get('metering_mode')} " - f"action={decision.get('action')} " - f"exp={decision.get('old_exp')}->{decision.get('new_exp')} " - f"gain={decision.get('old_gain')}->{decision.get('new_gain')} " - f"control={decision.get('control_value'):.3f} " - f"target={decision.get('target_value'):.3f} " - f"p95={decision.get('p95'):.3f} " - f"sat={decision.get('sat_pct'):.2f}%" - ) - - # Caso normal: rgb individual - patches = metrics.get("patches", []) - if patches: - for p in patches: - st = p.get("stats", {}) - print( - f" [PATCH] role={p.get('role', role)} " - f"type={p.get('type')} " - f"roi_source={p.get('roi_source')} " - f"roi_pct={p.get('roi_pct')} " - f"p50={st.get('p50', 0):.3f} " - f"p95={st.get('p95', 0):.3f} " - f"sat={st.get('sat_pct', 0):.2f}% " - f"dark={st.get('dark_pct', 0):.1f}%" - ) - - # Caso spectral_shared: RE/NIR agregados - roles = metrics.get("roles", {}) - if roles: - for r, item in roles.items(): - m = item.get("metrics", {}) - print( - f" [ROLE_METRICS] role={r} " - f"cam_id={item.get('cam_id')} " - f"control={m.get('control_value', 0):.3f} " - f"target={m.get('target_value', 0):.3f} " - f"p95={m.get('p95', 0):.3f} " - f"sat={m.get('sat_pct', 0):.2f}% " - f"warnings={m.get('patch_quality', {}).get('warnings', [])}" - ) - - for p in m.get("patches", []): - st = p.get("stats", {}) - print( - f" [PATCH] role={p.get('role', r)} " - f"type={p.get('type')} " - f"roi_source={p.get('roi_source')} " - f"roi_pct={p.get('roi_pct')} " - f"p50={st.get('p50', 0):.3f} " - f"p95={st.get('p95', 0):.3f} " - f"sat={st.get('sat_pct', 0):.2f}% " - f"dark={st.get('dark_pct', 0):.1f}%" - ) diff --git a/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core.py b/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core.py index 8814d1ec5..439adfca7 100644 --- a/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core.py +++ b/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core.py @@ -1,18 +1,242 @@ import json import os +import time import cv2 import numpy as np import math from typing import Optional +try: + import numba as _numba + _HAS_NUMBA = True +except Exception: + _numba = None + _HAS_NUMBA = False + +if _HAS_NUMBA: + + @_numba.njit(cache=True, fastmath=True) + def _raw10_get_pixel_numba(row, x): + group = (x // 4) * 5 + idx = x & 3 + + b = row[group + idx] + b4 = row[group + 4] + + if idx == 0: + low = b4 & 0x03 + elif idx == 1: + low = (b4 >> 2) & 0x03 + elif idx == 2: + low = (b4 >> 4) & 0x03 + else: + low = (b4 >> 6) & 0x03 + + return (int(b) << 2) | int(low) + + + @_numba.njit(cache=True, fastmath=True, parallel=True) + def _raw10_mono_to_float32_numba(packed, out, width, height, scale): + groups_per_row = width // 4 + + for y in _numba.prange(height): + row = packed[y] + for g in range(groups_per_row): + base = g * 5 + x = g * 4 + + b0 = int(row[base + 0]) + b1 = int(row[base + 1]) + b2 = int(row[base + 2]) + b3 = int(row[base + 3]) + b4 = int(row[base + 4]) + + p0 = (b0 << 2) | ((b4 >> 0) & 0x03) + p1 = (b1 << 2) | ((b4 >> 2) & 0x03) + p2 = (b2 << 2) | ((b4 >> 4) & 0x03) + p3 = (b3 << 2) | ((b4 >> 6) & 0x03) + + out[y, x + 0] = p0 * scale + out[y, x + 1] = p1 * scale + out[y, x + 2] = p2 * scale + out[y, x + 3] = p3 * scale + + + @_numba.njit(cache=True, fastmath=True, parallel=True) + def _raw10_bayer_planes_to_rgb_numba( + packed, + out_rgb, + width, + height, + scale, + pattern_code, + gain_r, + gain_g, + gain_b, + ): + """ + packed: 2D uint8 RAW10 packed, shape=(height, packed_width) + out_rgb: 3D float32, shape=(height//2, width//2, 3) + + pattern_code: + 0 = RGGB + 1 = BGGR + 2 = GRBG + 3 = GBRG + """ + out_h = height // 2 + out_w = width // 2 + + for oy in _numba.prange(out_h): + y0 = oy * 2 + y1 = y0 + 1 + row0 = packed[y0] + row1 = packed[y1] + + for ox in range(out_w): + x0 = ox * 2 + x1 = x0 + 1 + + ee = _raw10_get_pixel_numba(row0, x0) * scale + eo = _raw10_get_pixel_numba(row0, x1) * scale + oe = _raw10_get_pixel_numba(row1, x0) * scale + oo = _raw10_get_pixel_numba(row1, x1) * scale + + if pattern_code == 0: # RGGB + r = ee + g = (eo + oe) * 0.5 + b = oo + elif pattern_code == 1: # BGGR + b = ee + g = (eo + oe) * 0.5 + r = oo + elif pattern_code == 2: # GRBG + g = (ee + oo) * 0.5 + r = eo + b = oe + else: # GBRG + g = (ee + oo) * 0.5 + b = eo + r = oe + + r *= gain_r + g *= gain_g + b *= gain_b + + # Clip manual barato. + if r < 0.0: + r = 0.0 + elif r > 1.0: + r = 1.0 + + if g < 0.0: + g = 0.0 + elif g > 1.0: + g = 1.0 + + if b < 0.0: + b = 0.0 + elif b > 1.0: + b = 1.0 + + out_rgb[oy, ox, 0] = r + out_rgb[oy, ox, 1] = g + out_rgb[oy, ox, 2] = b + + + @_numba.njit(cache=True, fastmath=True, parallel=True) + def _raw10_to_raw16_numba(packed, out, width, height): + groups_per_row = width // 4 + + for y in _numba.prange(height): + row = packed[y] + + for g in range(groups_per_row): + base = g * 5 + x = g * 4 + + b0 = int(row[base + 0]) + b1 = int(row[base + 1]) + b2 = int(row[base + 2]) + b3 = int(row[base + 3]) + b4 = int(row[base + 4]) + + out[y, x + 0] = (b0 << 2) | ((b4 >> 0) & 0x03) + out[y, x + 1] = (b1 << 2) | ((b4 >> 2) & 0x03) + out[y, x + 2] = (b2 << 2) | ((b4 >> 4) & 0x03) + out[y, x + 3] = (b3 << 2) | ((b4 >> 6) & 0x03) + + + @_numba.njit(cache=True, fastmath=True, parallel=True) + def _rgb16_to_float32_gain_clip_numba(rgb16, out, height, width, scale, gain_r, gain_g, gain_b): + for y in _numba.prange(height): + for x in range(width): + r = float(rgb16[y, x, 0]) * scale * gain_r + g = float(rgb16[y, x, 1]) * scale * gain_g + b = float(rgb16[y, x, 2]) * scale * gain_b + + if r < 0.0: + r = 0.0 + elif r > 1.0: + r = 1.0 + + if g < 0.0: + g = 0.0 + elif g > 1.0: + g = 1.0 + + if b < 0.0: + b = 0.0 + elif b > 1.0: + b = 1.0 + + out[y, x, 0] = r + out[y, x, 1] = g + out[y, x, 2] = b + + + @_numba.njit(cache=True, fastmath=True, parallel=True) + def _apply_flat_gain_numba(base, gain, out, height, width, clip_output): + for y in _numba.prange(height): + for x in range(width): + v = float(base[y, x]) * float(gain[y, x]) + + if clip_output: + if v < 0.0: + v = 0.0 + elif v > 1.0: + v = 1.0 + + out[y, x] = v + + + @_numba.njit(cache=True, fastmath=True, parallel=True) + def _apply_tensor_flat_gain_chw_numba(tensor, gain, channels, height, width, clip_output): + for c in _numba.prange(channels): + for y in range(height): + for x in range(width): + v = float(tensor[c, y, x]) * float(gain[c, y, x]) + + if clip_output: + if v < 0.0: + v = 0.0 + elif v > 1.0: + v = 1.0 + + tensor[c, y, x] = v + + + + class RawProcessorCore: def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "BGGR", calibration_json_path=None): self.sensor_width = sensor_width self.sensor_height = sensor_height self.bayer_pattern = bayer_pattern.upper() self.rgb_processing_config = { - "mode": "linear_demosaic", # "linear_demosaic" ou "bayer_planes" + "mode": "linear_demosaic", # "linear_demosaic", "linear_demosaic_half" ou "bayer_planes" + "demosaic_algorithm": "ea", # "ea" ou "bilinear" } self.fusion_config = { @@ -161,9 +385,37 @@ class RawProcessorCore: self.last_fusion_result = None self.camera_settings = {} + self._flatfield_runtime_cache = {} + + self.last_decode_perf = {} + self._last_decode_perf_log_ts = 0.0 + if calibration_json_path: self.load_config_json(calibration_json_path) + + def warmup_numba_raw10_decode(self): + if not _HAS_NUMBA: + return {"ok": False, "reason": "numba_not_available"} + + w, h = 1280, 800 + packed_w = int(np.ceil(w * 10 / 8)) + dummy = np.zeros((h, packed_w), dtype=np.uint8) + + _ = self._raw10_mono_to_float01_aggressive(dummy, w, h, 10) + _ = self._raw10_rgb_bayer_planes_to_rgb_float01_aggressive(dummy, w, h, "RGGB", 10) + _ = self._raw10_to_raw16_aggressive(dummy, w, h) + + raw16 = self._raw10_to_raw16_aggressive(dummy, w, h) + rgb16 = cv2.cvtColor(raw16, cv2.COLOR_BayerRG2RGB) + _ = self._rgb16_to_float32_gain_clip_aggressive(rgb16, 10) + + dummy_tensor = np.zeros((5, 640, 1024), dtype=np.float32) + dummy_gain = np.ones((5, 640, 1024), dtype=np.float32) + _apply_tensor_flat_gain_chw_numba(dummy_tensor, dummy_gain, 5, 640, 1024, True) + + return {"ok": True, "backend": "numba", "shape": [h, packed_w]} + def unpack_raw10_packed( self, packed_frame: np.ndarray, @@ -281,31 +533,25 @@ class RawProcessorCore: bit_depth: int = 10, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - code_map = { - "RGGB": cv2.COLOR_BayerRG2BGR_EA, - "BGGR": cv2.COLOR_BayerBG2BGR_EA, - "GRBG": cv2.COLOR_BayerGR2BGR_EA, - "GBRG": cv2.COLOR_BayerGB2BGR_EA, - } + rgb_cfg = getattr(self, "rgb_processing_config", {}) or {} + algorithm = str(rgb_cfg.get("demosaic_algorithm", "ea")).lower() - p = self.bayer_pattern.upper() - if p not in code_map: - raise ValueError(f"Padrão Bayer não suportado para demosaic: {p}") + cv_code, _ = self._get_bayer_cv2_code( + bayer_pattern=self.bayer_pattern, + algorithm=algorithm, + ) raw16 = np.asarray(raw16) if raw16.dtype != np.uint16: - raw16 = raw16.astype(np.uint16) + raw16 = raw16.astype(np.uint16, copy=False) - out = cv2.cvtColor(raw16, code_map[p]) + rgb16 = cv2.cvtColor(raw16, cv_code) - max_val = float((1 << bit_depth) - 1) - out = np.clip(out.astype(np.float32) / max_val, 0.0, 1.0) + rgb = rgb16.astype(np.float32) + rgb *= np.float32(1.0 / float((1 << bit_depth) - 1)) + np.clip(rgb, 0.0, 1.0, out=rgb) - r = out[:, :, 0] - g = out[:, :, 1] - b = out[:, :, 2] - - return r, g, b + return rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2] def build_training_rgb( self, @@ -496,6 +742,10 @@ class RawProcessorCore: decoded = {} + rgb_mode = str( + (getattr(self, "rgb_processing_config", {}) or {}).get("mode", "linear_demosaic") + ).lower() + for cam_id, data in frame.items(): cam_meta = camera_frames.get(cam_id) or camera_info.get(cam_id) or {} @@ -505,27 +755,48 @@ class RawProcessorCore: bit_depth = int(cam_meta.get("bit_depth", 10)) raw_format = str(cam_meta.get("raw_format", "")).upper() - is_raw10 = raw_format == "RAW10_PACKED" or bit_depth == 10 + packed = bool(cam_meta.get("packed", False)) + is_raw10 = raw_format == "RAW10_PACKED" or packed or bit_depth == 10 if role == "rgb": - # Caso OAK RAW real: câmera RGB também vem RAW10 packed if is_raw10 and data.ndim == 2: sensor_w = int(cam_meta.get("width", self.sensor_width)) sensor_h = int(cam_meta.get("height", self.sensor_height)) - raw16 = self.unpack_raw10_packed( - data, - sensor_width=sensor_w, - sensor_height=sensor_h, + bayer = ( + cam_meta.get("bayer_pattern") + or cam_meta.get("bayer") + or self.bayer_pattern + or "RGGB" ) - rgb_chw = self.build_training_rgb( - raw16, - output_dtype="float32", - bit_depth=bit_depth, - ) + if rgb_mode in ("bayer_planes", "bayer", "half_res"): + rgb_hwc = self._raw10_rgb_bayer_planes_to_rgb_float01_aggressive( + data, + width=sensor_w, + height=sensor_h, + bayer_pattern=bayer, + bit_depth=bit_depth, + ) - rgb_hwc = np.transpose(rgb_chw, (1, 2, 0)) + elif rgb_mode in ( + "linear_demosaic", + "demosaic", + "full_res", + "linear_demosaic_half", + "demosaic_half", + "full_demosaic_half", + ): + rgb_hwc = self._raw10_rgb_linear_demosaic_to_rgb_float01_fast( + data, + width=sensor_w, + height=sensor_h, + bayer_pattern=bayer, + bit_depth=bit_depth, + ) + + else: + raise ValueError(f"rgb_processing.mode inválido: {rgb_mode}") decoded[cam_id] = { "name": "RGB", @@ -535,7 +806,7 @@ class RawProcessorCore: } else: - # Caso preview/processado antigo: BGR HWC uint8 + # Caso preview/processado antigo: BGR HWC uint8. if data.ndim != 3 or data.shape[2] != 3: raise RuntimeError(f"{cam_id} RGB inválida: shape={data.shape}") @@ -564,6 +835,9 @@ class RawProcessorCore: "meta": cam_meta, } + else: + raise RuntimeError(f"Role não suportada em {cam_id}: {role}") + return decoded def _decode_spectral_frame_to_float01(self, data, cam_meta): @@ -595,14 +869,28 @@ class RawProcessorCore: ) if looks_like_raw10_packed: - raw16 = self.unpack_raw10_packed( + t_total0 = time.perf_counter() + + out = self._raw10_mono_to_float01_aggressive( arr, - sensor_width=sensor_width, - sensor_height=sensor_height, + width=sensor_width, + height=sensor_height, + bit_depth=bit_depth, ) - max_val = float((1 << bit_depth) - 1) - return np.clip(raw16.astype(np.float32) / max_val, 0.0, 1.0) + role = str(cam_meta.get("role", "spectral")).lower() + + self._set_decode_perf(role, { + "mode": "mono_raw10_aggressive", + "width": int(sensor_width), + "height": int(sensor_height), + "numba": bool(_HAS_NUMBA), + "total_ms": (time.perf_counter() - t_total0) * 1000.0, + "out_shape": list(out.shape), + "out_dtype": str(out.dtype), + }) + + return out # Caso preview/processado: mono já vem uint8 normal. if arr.ndim == 2 and arr.dtype == np.uint8: @@ -619,73 +907,128 @@ class RawProcessorCore: return np.clip(arr01, 0.0, 1.0) def fuse_multispec_cameras(self, decoded, meta, channels_expected): + t_total0 = time.perf_counter() + + t0 = time.perf_counter() decoded = self.apply_dark_to_decoded(decoded) + t_dark_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() decoded = self.normalize_decoded_by_capture_controls(decoded, meta) - decoded = self.apply_flat_gain_to_decoded(decoded) + t_radnorm_ms = (time.perf_counter() - t0) * 1000.0 - rgb_cam_id = self._find_cam_by_role(decoded, "rgb") - if rgb_cam_id is None: - raise RuntimeError("Fusão requer câmera com role='rgb' como referência") + flat_cfg = self.flatfield_config or {} + flat_apply_space = str(flat_cfg.get("apply_space", "native_camera_space")).lower() - rgb = decoded[rgb_cam_id]["image"] - h, w = rgb.shape[:2] + apply_native_flat = ( + bool(flat_cfg.get("enabled", False)) + and flat_apply_space != "final_tensor_space" + ) - self.last_fusion_result = { - "ref_shape": [int(h), int(w)], - "crop_valid_common": bool(self.fusion_config.get("crop_valid_common", False)), - "resize_after_crop": bool(self.fusion_config.get("resize_after_crop", False)), - "crop_box": None, - "crop_applied": False, - "roles": [], - } + t0 = time.perf_counter() - rgb_chw = np.transpose(rgb, (2, 0, 1)) - channels = [rgb_chw] - names = ["R", "G", "B"] + if apply_native_flat: + decoded = self.apply_flat_gain_to_decoded(decoded) - valid_masks = [np.ones((h, w), dtype=np.uint8)] + t_flat_native_ms = (time.perf_counter() - t0) * 1000.0 - role_to_cam = { - item.get("role", data.get("meta", {}).get("role")): cam_id - for cam_id, data in decoded.items() - for item in [data] - } + # Caminho espacial otimizado: direto para o target final. + tensor, direct_perf = self._fuse_multispec_direct_to_target_fast( + decoded, + meta, + channels_expected, + ) - for role, ch_name in (("re", "RE"), ("nir", "NIR")): - cam_id = role_to_cam.get(role) - if cam_id is None: - continue - - img = decoded[cam_id]["image"] - aligned, valid_mask = self._warp_with_valid_mask(img, role, (h, w), meta) - - channels.append(aligned[None, :, :]) - names.append(ch_name) - valid_masks.append(valid_mask) - - cfg = self.fusion_config - if cfg.get("crop_valid_common", False): - crop_box = self._compute_common_crop_box(valid_masks) - if crop_box is not None: - x0, y0, x1, y1 = crop_box - - self.last_fusion_result["crop_box"] = [int(x0), int(y0), int(x1), int(y1)] - self.last_fusion_result["crop_applied"] = True - - channels = self._crop_and_resize_channels(channels, crop_box, (h, w)) - - tensor = np.concatenate(channels, axis=0) + t_flat_ms = float(t_flat_native_ms + float(direct_perf.get("final_flat_ms", 0.0))) + t0 = time.perf_counter() if tensor.shape[0] != channels_expected: raise RuntimeError( f"Tensor fundido com canais inesperados: {tensor.shape[0]} | " - f"esperado={channels_expected} | got={names}" + f"esperado={channels_expected}" ) - - self.last_fusion_result["roles"] = list(names) - self.last_fusion_result["output_shape"] = list(tensor.shape) - return tensor.astype(np.float32, copy=False) + out = tensor.astype(np.float32, copy=False) + t_final_ms = (time.perf_counter() - t0) * 1000.0 + + t_total_ms = (time.perf_counter() - t_total0) * 1000.0 + + # Mantém contrato de perf do benchmark. + t_prepare_ms = float(direct_perf.get("prepare_ms", 0.0)) + t_rgb_ms = float(direct_perf.get("rgb_crop_resize_ms", 0.0)) + t_warp_total_ms = float(direct_perf.get("warp_total_ms", 0.0)) + warp_details = direct_perf.get("warp_details_ms", {}) or {} + + # Aqui crop_resize_ms representa somente RGB crop/resize no modo direto. + # O custo espacial total útil para analisar é: + # spatial_direct_ms = rgb_crop_resize_ms + warp_total_ms + t_crop_resize_ms = float(direct_perf.get("crop_resize_ms", t_rgb_ms)) + t_concat_ms = 0.0 + + if self.last_fusion_result is None: + self.last_fusion_result = {} + + self.last_fusion_result["roles"] = ["R", "G", "B", "RE", "NIR"] + self.last_fusion_result["output_shape"] = list(out.shape) + self.last_fusion_result["perf"] = { + "dark_ms": t_dark_ms, + "radnorm_ms": t_radnorm_ms, + "flat_ms": t_flat_ms, + "prepare_ms": t_prepare_ms, + "rgb_crop_resize_ms": t_rgb_ms, + "warp_total_ms": t_warp_total_ms, + "warp_details_ms": warp_details, + "crop_resize_ms": t_crop_resize_ms, + "concat_ms": t_concat_ms, + "spatial_direct_ms": float(t_rgb_ms + t_warp_total_ms), + "final_ms": t_final_ms, + "total_ms": t_total_ms, + "decode_perf": getattr(self, "last_decode_perf", {}), + } + + #if not hasattr(self, "_last_perf_log_ts"): + # self._last_perf_log_ts = 0.0 + #now = time.time() + #if now - self._last_perf_log_ts >= 1.0: + # self._last_perf_log_ts = now + # print( + # "[PERF][CORE_FUSE] " + # f"dark={t_dark_ms:.1f}ms " + # f"radnorm={t_radnorm_ms:.1f}ms " + # f"flat={t_flat_ms:.1f}ms " + # f"prepare={t_prepare_ms:.1f}ms " + # f"rgb_resize={t_rgb_ms:.1f}ms " + # f"warp={t_warp_total_ms:.1f}ms " + # f"warp_re={warp_details.get('re', -1):.1f}ms " + # f"warp_nir={warp_details.get('nir', -1):.1f}ms " + # f"spatial={t_rgb_ms + t_warp_total_ms:.1f}ms " + # f"crop_resize={t_crop_resize_ms:.1f}ms " + # f"concat={t_concat_ms:.1f}ms " + # f"final={t_final_ms:.1f}ms " + # f"total={t_total_ms:.1f}ms " + # f"shape={out.shape}" + # ) + + #if not hasattr(self, "_last_remap_perf_log_ts"): + # self._last_remap_perf_log_ts = 0.0 + #now = time.time() + #if now - self._last_remap_perf_log_ts >= 1.0: + # self._last_remap_perf_log_ts = now + # print( + # "[PERF][REMAP] " + # f"enabled={direct_perf.get('remap_enabled')} " + # f"hit={direct_perf.get('remap_cache_hit')} " + # f"cache={direct_perf.get('remap_cache_ms', -1):.2f}ms " + # f"rgb={direct_perf.get('rgb_crop_resize_ms', -1):.2f}ms " + # f"warp={direct_perf.get('warp_total_ms', -1):.2f}ms " + # f"re={(direct_perf.get('warp_details_ms') or {}).get('re', -1):.2f}ms " + # f"nir={(direct_perf.get('warp_details_ms') or {}).get('nir', -1):.2f}ms " + # f"spatial={direct_perf.get('spatial_direct_ms', -1):.2f}ms " + # f"hits={direct_perf.get('remap_cache_hits')} " + # f"misses={direct_perf.get('remap_cache_misses')}" + # ) + + return out def _shift_image(self, img, dx, dy): h, w = img.shape[:2] @@ -750,27 +1093,36 @@ class RawProcessorCore: H = cfg.get("homographies", {}).get(f"{role}_to_rgb") if H is None: - warped = img - warped_mask = mask - else: - H = np.asarray(H, dtype=np.float32) - - if H.shape != (3, 3): - raise RuntimeError(f"Homografia inválida para {role}: shape={H.shape}") - - warped = cv2.warpPerspective( - img, H, (ref_w, ref_h), - flags=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0 + raise RuntimeError( + f"fusion_config.alignment_mode='homography', " + f"mas homografia '{role}_to_rgb' está ausente. " + f"Isso deixaria o canal {role.upper()} sem alinhamento." ) - warped_mask = cv2.warpPerspective( - mask, H, (ref_w, ref_h), - flags=cv2.INTER_NEAREST, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0 - ) + calib_size = cfg.get("homography_calibration_size", None) + + H = self._scale_homography_to_runtime( + H, + calib_size=calib_size, + runtime_size=(ref_w, ref_h), + ) + + if H.shape != (3, 3): + raise RuntimeError(f"Homografia inválida para {role}: shape={H.shape}") + + warped = cv2.warpPerspective( + img, H, (ref_w, ref_h), + flags=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0 + ) + + warped_mask = cv2.warpPerspective( + mask, H, (ref_w, ref_h), + flags=cv2.INTER_NEAREST, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0 + ) else: raise RuntimeError(f"alignment_mode inválido: {mode}") @@ -850,6 +1202,58 @@ class RawProcessorCore: return np.stack(chans, axis=0) + def _scale_homography_to_runtime(self, H, calib_size, runtime_size): + """ + Ajusta uma homografia calculada em calib_size para ser aplicada em runtime_size. + + H original: + ponto_spec_calib -> ponto_rgb_calib + + H runtime: + ponto_spec_runtime -> ponto_rgb_runtime + """ + if H is None: + return None + + H = np.asarray(H, dtype=np.float32) + + if calib_size is None: + return H + + 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) + + if calib_w <= 0 or calib_h <= 0: + return 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 apply_patch_normalization_to_tensor(self, tensor: np.ndarray) -> np.ndarray: self.last_patch_normalization_result = None @@ -2016,84 +2420,78 @@ class RawProcessorCore: out = np.maximum(base - dark, 0.0) return out.astype(np.float32, copy=False) - def _apply_flat_gain_single_channel_bkp( + def _apply_flat_gain_single_channel( self, img: np.ndarray, channel_name: str, clip_output: bool = True, + saturation_mask: np.ndarray | None = None, ) -> np.ndarray: + t_total0 = time.perf_counter() + t_cache_ms = 0.0 + t_gain_eff_ms = 0.0 + t_mul_clip_ms = 0.0 + ch = str(channel_name).upper() - entry = self.flatfield_maps.get(ch) - - if not entry: - return img.astype(np.float32, copy=False) - - gain = entry.get("gain") - if gain is None: - return img.astype(np.float32, copy=False) - - base = img.astype(np.float32) - - gain = gain.astype(np.float32) - if gain.shape[:2] != base.shape[:2]: - gain = cv2.resize( - gain, - (base.shape[1], base.shape[0]), - interpolation=cv2.INTER_LINEAR, - ) - - out = base * gain - - if clip_output: - out = np.clip(out, 0.0, 1.0) - - return out.astype(np.float32, copy=False) - - def _apply_flat_gain_single_channel( - self, - img: np.ndarray, - channel_name: str, - clip_output: bool = True, - saturation_mask: np.ndarray | None = None, - ) -> np.ndarray: - ch = str(channel_name).upper() - entry = self.flatfield_maps.get(ch) - - if not entry: - return img.astype(np.float32, copy=False) - - gain = entry.get("gain") - if gain is None: - return img.astype(np.float32, copy=False) - cfg = self.flatfield_config or {} - base = img.astype(np.float32) + base = img.astype(np.float32, copy=False) - gain = gain.astype(np.float32) - if gain.shape[:2] != base.shape[:2]: - gain = cv2.resize( - gain, - (base.shape[1], base.shape[0]), - interpolation=cv2.INTER_LINEAR, + fast_runtime = bool(cfg.get("fast_runtime", False)) + sat_guard_enabled = bool(cfg.get("saturation_guard_enabled", True)) + + if fast_runtime and not sat_guard_enabled: + t0 = time.perf_counter() + + gain_eff = self._get_runtime_gain_eff_map(ch, base.shape, cfg) + + if gain_eff is None: + return base + + out = self._apply_flat_gain_simple_aggressive( + base, + gain_eff, + clip_output=clip_output, ) - # Suavização extra em runtime. - # Útil para corrigir apenas o borrão grande, não microtextura/ruído. + t_total_ms = (time.perf_counter() - t_total0) * 1000.0 + if not hasattr(self, "_last_flat_ch_perf_log_ts"): + self._last_flat_ch_perf_log_ts = 0.0 + now = time.time() + if now - self._last_flat_ch_perf_log_ts >= 1.0: + self._last_flat_ch_perf_log_ts = now + print( + "[PERF][FLAT_CH_FAST] " + f"ch={ch} " + f"shape={base.shape} " + f"total={t_total_ms:.2f}ms " + f"clip={clip_output} " + f"strength={cfg.get('strength_by_channel', {}).get(ch, cfg.get('strength', 1.0))}" + ) + + return out.astype(np.float32, copy=False) + + # ============================================================ + # Gain map runtime cacheado + # ============================================================ + t0 = time.perf_counter() + gain = self._get_runtime_gain_map(ch, base.shape, cfg) + t_cache_ms = (time.perf_counter() - t0) * 1000.0 + + if gain is None: + return base + runtime_smooth_ksize = int(cfg.get("runtime_smooth_ksize", 0) or 0) - if runtime_smooth_ksize >= 3: - if runtime_smooth_ksize % 2 == 0: - runtime_smooth_ksize += 1 + if runtime_smooth_ksize >= 3 and runtime_smooth_ksize % 2 == 0: + runtime_smooth_ksize += 1 - gain = cv2.GaussianBlur( - gain, - (runtime_smooth_ksize, runtime_smooth_ksize), - 0, - ) + # Importante: + # NÃO aplicar GaussianBlur aqui. + # O gain já veio pronto/cacheado de _get_runtime_gain_map(). - # Intensidade global e por canal: - # strength=0.0 -> não aplica flatfield - # strength=1.0 -> aplica mapa integral + # ============================================================ + # Intensidade global e por canal + # ============================================================ strength = float(cfg.get("strength", 1.0)) strength_by_channel = cfg.get("strength_by_channel", {}) or {} if ch in strength_by_channel: @@ -2105,14 +2503,15 @@ class RawProcessorCore: gain_min_runtime = float(cfg.get("gain_min_runtime", 0.0)) gain_max_runtime = float(cfg.get("gain_max_runtime", 999.0)) - # Guarda de saturação. + # ============================================================ + # Guarda de saturação + # ============================================================ sat_guard_enabled = bool(cfg.get("saturation_guard_enabled", True)) sat_mode = str(cfg.get("saturation_guard_mode", "fade_strength")).lower() sat_soft_start = float(cfg.get("saturation_guard_soft_start", 0.90)) sat_hard = float(cfg.get("saturation_guard_hard", 0.98)) - # Se não veio uma máscara RGB comum, usa máscara do próprio canal. if saturation_mask is None and sat_guard_enabled: sat_threshold = float(cfg.get("saturation_guard_threshold", 0.97)) saturation_mask = base >= sat_threshold @@ -2120,45 +2519,66 @@ class RawProcessorCore: # ============================================================ # Calcula ganho efetivo # ============================================================ + t0 = time.perf_counter() if sat_guard_enabled and sat_mode == "fade_strength": - # Reduz gradualmente a força do flatfield conforme aproxima saturação. - # A força cai de 1.0 para 0.0 entre soft_start e hard. denom = max(sat_hard - sat_soft_start, 1e-6) + t = (base - sat_soft_start) / denom t = np.clip(t, 0.0, 1.0) - # strength_mask: - # 1.0 longe da saturação - # 0.0 perto/acima de sat_hard strength_mask = 1.0 - t - # Se uma máscara comum RGB foi passada, zera força nesses pixels. - # Isso mantém a neutralidade entre R/G/B em pixels suspeitos. if saturation_mask is not None: strength_mask = np.where(saturation_mask, 0.0, strength_mask) gain_eff = 1.0 + (strength * strength_mask) * (gain - 1.0) else: - # Modo normal com força fixa. gain_eff = 1.0 + strength * (gain - 1.0) gain_eff = np.clip(gain_eff, gain_min_runtime, gain_max_runtime) + t_gain_eff_ms = (time.perf_counter() - t0) * 1000.0 + + # ============================================================ + # Aplica ganho + # ============================================================ + t0 = time.perf_counter() + out = base * gain_eff - # Modo skip: onde saturou, não aplica flatfield. - # O pixel continua queimado, mas não vira magenta/roxo artificial. if sat_guard_enabled and sat_mode == "skip" and saturation_mask is not None: out = np.where(saturation_mask, base, out) if clip_output: out = np.clip(out, 0.0, 1.0) + t_mul_clip_ms = (time.perf_counter() - t0) * 1000.0 + + t_total_ms = (time.perf_counter() - t_total0) * 1000.0 + + if not hasattr(self, "_last_flat_ch_perf_log_ts"): + self._last_flat_ch_perf_log_ts = 0.0 + + now = time.time() + if now - self._last_flat_ch_perf_log_ts >= 1.0: + self._last_flat_ch_perf_log_ts = now + print( + "[PERF][FLAT_CH] " + f"ch={ch} " + f"shape={base.shape} " + f"cache={t_cache_ms:.1f}ms " + f"gain_eff={t_gain_eff_ms:.1f}ms " + f"mul_clip={t_mul_clip_ms:.1f}ms " + f"total={t_total_ms:.1f}ms " + f"ksize={runtime_smooth_ksize} " + f"strength={strength:.2f}" + ) + return out.astype(np.float32, copy=False) - + def normalize_decoded_by_capture_controls(self, decoded: dict, meta: dict | None = None) -> dict: cfg = self.radiometric_normalization_config or {} self.last_radiometric_normalization_result = None @@ -2198,6 +2618,10 @@ class RawProcessorCore: reference_controls = cfg.get("reference_controls", {}) or {} clip_output = bool(cfg.get("clip_output", False)) + # Cache por role para não recalcular factor/scale 3x quando RGB tem 3 canais no mesmo item. + scale_by_role = {} + debug_by_role = {} + normalized = {} for cam_id, item in decoded.items(): @@ -2209,78 +2633,84 @@ class RawProcessorCore: result["warnings"].append(f"{cam_id}:missing_image_or_role") continue - actual_ctrl = controls_by_role.get(role, {}) or {} - ref_ctrl = reference_controls.get(role, {}) or {} + if role in scale_by_role: + scale = scale_by_role[role] + debug = dict(debug_by_role[role]) + debug["camera_id"] = cam_id + else: + actual_ctrl = controls_by_role.get(role, {}) or {} + ref_ctrl = reference_controls.get(role, {}) or {} - actual_factor = self._radiometric_factor_from_controls(actual_ctrl, cfg) - ref_factor = self._radiometric_factor_from_controls(ref_ctrl, cfg) + actual_factor = self._radiometric_factor_from_controls(actual_ctrl, cfg) + ref_factor = self._radiometric_factor_from_controls(ref_ctrl, cfg) - if actual_factor <= 0 or ref_factor <= 0: - normalized[cam_id] = item - result["warnings"].append( - f"{role}:invalid_factor actual={actual_factor:.6g} ref={ref_factor:.6g}" + if actual_factor <= 0 or ref_factor <= 0: + normalized[cam_id] = item + result["warnings"].append( + f"{role}:invalid_factor actual={actual_factor:.6g} ref={ref_factor:.6g}" + ) + + if str(cfg.get("invalid_controls_policy", "skip")).lower() == "raise": + raise RuntimeError(f"Controles radiométricos inválidos para role={role}: {actual_ctrl}") + + continue + + raw_scale = float(ref_factor / actual_factor) + scale = self._clip_radiometric_scale(raw_scale, role, cfg) + + debug = self._build_radiometric_debug( + method=method, + role=role, + cam_id=cam_id, + actual_ctrl=actual_ctrl, + ref_ctrl=ref_ctrl, + actual_factor=actual_factor, + ref_factor=ref_factor, + raw_scale=raw_scale, + scale=scale, + clip_output=clip_output, ) - if str(cfg.get("invalid_controls_policy", "skip")).lower() == "raise": - raise RuntimeError(f"Controles radiométricos inválidos para role={role}: {actual_ctrl}") + scale_by_role[role] = scale + debug_by_role[role] = dict(debug) - continue - - raw_scale = float(ref_factor / actual_factor) - scale = self._clip_radiometric_scale(raw_scale, role, cfg) - - out = img.astype(np.float32) * scale - - if clip_output: - out = np.clip(out, 0.0, 1.0) + out, reused = self._apply_radiometric_scale_inplace( + img, + scale=scale, + clip_output=clip_output, + ) new_item = dict(item) new_meta = dict(item.get("meta", {}) or {}) - debug = { - "applied": True, - "method": method, - "role": role, - "camera_id": cam_id, - "actual_controls": dict(actual_ctrl), - "reference_controls": dict(ref_ctrl), - "actual_factor": float(actual_factor), - "reference_factor": float(ref_factor), - "scale_raw": float(raw_scale), - "scale_applied": float(scale), - "clip_output": clip_output, - } + debug["inplace_reused_input"] = bool(reused) + debug["image_shape"] = list(out.shape) if hasattr(out, "shape") else None + debug["image_dtype"] = str(out.dtype) if hasattr(out, "dtype") else None new_meta["radiometric_normalization"] = debug new_meta["radiometric_normalization_applied"] = True new_meta["radiometric_normalization_scale"] = float(scale) - new_item["image"] = out.astype(np.float32, copy=False) + new_item["image"] = out new_item["meta"] = new_meta normalized[cam_id] = new_item - result["by_role"][role] = debug + result["applied"] = True result["by_camera"][cam_id] = debug + result["by_role"][role] = debug - result["applied"] = any( - bool(v.get("applied", False)) - for v in result["by_camera"].values() - if isinstance(v, dict) - ) + if result["applied"]: + scales = [v.get("scale_applied") for v in result["by_camera"].values() if isinstance(v, dict)] + scales = [float(s) for s in scales if s is not None] - scales = [ - float(v.get("scale_applied", 1.0)) - for v in result["by_camera"].values() - if isinstance(v, dict) and v.get("applied", False) - ] - - result["summary"] = { - "applied_count": int(len(scales)), - "scale_min": float(min(scales)) if scales else None, - "scale_max": float(max(scales)) if scales else None, - "scale_mean": float(np.mean(scales)) if scales else None, - "warning_count": int(len(result["warnings"])), - } + if scales: + result["summary"] = { + "scale_min": float(np.min(scales)), + "scale_max": float(np.max(scales)), + "scale_mean": float(np.mean(scales)), + "num_normalized": int(len(scales)), + "inplace_fast_path": True, + } self.last_radiometric_normalization_result = result return normalized @@ -2471,4 +2901,1908 @@ class RawProcessorCore: def get_last_radiometric_normalization_result(self): return self.last_radiometric_normalization_result - \ No newline at end of file + + def _get_runtime_gain_map(self, channel_name: str, base_shape: tuple, cfg: dict): + ch = str(channel_name).upper() + entry = self.flatfield_maps.get(ch) + + if not entry: + return None + + gain = entry.get("gain") + if gain is None: + return None + + h, w = base_shape[:2] + runtime_smooth_ksize = int(cfg.get("runtime_smooth_ksize", 0) or 0) + + if runtime_smooth_ksize >= 3 and runtime_smooth_ksize % 2 == 0: + runtime_smooth_ksize += 1 + + cache_key = ( + ch, + h, + w, + runtime_smooth_ksize, + ) + + cached = self._flatfield_runtime_cache.get(cache_key) + if cached is not None: + return cached + + print(f"[FLAT_CACHE] criando gain cache ch={ch} shape=({h},{w}) ksize={runtime_smooth_ksize}") + + gain_rt = gain.astype(np.float32, copy=False) + + if gain_rt.shape[:2] != (h, w): + gain_rt = cv2.resize( + gain_rt, + (w, h), + interpolation=cv2.INTER_LINEAR, + ) + + if runtime_smooth_ksize >= 3: + gain_rt = cv2.GaussianBlur( + gain_rt, + (runtime_smooth_ksize, runtime_smooth_ksize), + 0, + ) + + gain_rt = gain_rt.astype(np.float32, copy=False) + self._flatfield_runtime_cache[cache_key] = gain_rt + + return gain_rt + + + + def _raw10_expected_packed_width_fast(self, width: int) -> int: + """ + RAW10 packed: 4 pixels em 5 bytes. + Para OV9782/OV9282 em 1280 px, packed_width = 1600. + """ + return int(math.ceil(int(width) * 10 / 8)) + + def _prepare_raw10_packed_view_fast(self, packed_frame: np.ndarray, width: int, height: int) -> np.ndarray: + """ + Valida e retorna uma view útil do RAW10 packed, cortando padding se existir. + Não copia quando não precisa. + """ + arr = packed_frame + + if arr.ndim == 3 and arr.shape[2] == 1: + arr = arr[:, :, 0] + + if arr.ndim != 2: + raise ValueError(f"RAW10 packed esperado 2D. Veio shape={arr.shape}") + + width = int(width) + height = int(height) + + if width % 4 != 0: + raise ValueError(f"Largura {width} não é múltipla de 4 para RAW10 packed") + + expected_packed_width = self._raw10_expected_packed_width_fast(width) + + actual_h, actual_w = arr.shape[:2] + padding = actual_w - expected_packed_width + + if actual_h != height: + raise ValueError( + f"[ERRO FRAME] Altura RAW10 packed inesperada: {arr.shape}, esperado altura={height}" + ) + + if actual_w < expected_packed_width: + raise ValueError( + f"[ERRO FRAME] Largura RAW10 packed menor que esperada: {arr.shape}, " + f"esperado pelo menos ({height}, {expected_packed_width})" + ) + + if padding > 64: + raise ValueError( + f"[ERRO FRAME] Padding excessivo no RAW10 packed: {arr.shape}, " + f"esperado útil ({height}, {expected_packed_width}), padding={padding}" + ) + + return arr[:, :expected_packed_width] + + def _raw10_groups_view_fast(self, packed_frame: np.ndarray, width: int, height: int) -> np.ndarray: + """ + Retorna view em grupos RAW10: + shape = (height, width//4, 5) + sem converter tudo para uint16 de uma vez. + """ + arr = self._prepare_raw10_packed_view_fast(packed_frame, width, height) + return arr.reshape(int(height), int(width) // 4, 5) + + def _raw10_unpack_mono_to_float01_fast( + self, + packed_frame: np.ndarray, + width: int, + height: int, + bit_depth: int = 10, + ) -> np.ndarray: + """ + RAW10 packed mono -> float32 0..1. + + Otimização em relação ao caminho antigo: + antigo: RAW10 packed -> raw16 uint16 -> float32 / max + novo : RAW10 packed -> float32 0..1 direto + + Ainda gera imagem full-res, porque RE/NIR continuam full-res antes da fusão. + """ + width = int(width) + height = int(height) + bit_depth = int(bit_depth) + + groups = self._raw10_groups_view_fast(packed_frame, width, height) + + # Conversões por byte. Evita groups.astype(uint16) completo. + b0 = groups[:, :, 0].astype(np.uint16, copy=False) + b1 = groups[:, :, 1].astype(np.uint16, copy=False) + b2 = groups[:, :, 2].astype(np.uint16, copy=False) + b3 = groups[:, :, 3].astype(np.uint16, copy=False) + b4 = groups[:, :, 4].astype(np.uint16, copy=False) + + max_val = float((1 << bit_depth) - 1) + scale = np.float32(1.0 / max_val) + + out = np.empty((height, width), dtype=np.float32) + + # 4 pixels por grupo de 5 bytes. + out[:, 0::4] = ((b0 << 2) | ((b4 >> 0) & 0x03)).astype(np.float32) * scale + out[:, 1::4] = ((b1 << 2) | ((b4 >> 2) & 0x03)).astype(np.float32) * scale + out[:, 2::4] = ((b2 << 2) | ((b4 >> 4) & 0x03)).astype(np.float32) * scale + out[:, 3::4] = ((b3 << 2) | ((b4 >> 6) & 0x03)).astype(np.float32) * scale + + # RAW10 já está no intervalo, mas mantém defesa contra metadado errado. + return np.clip(out, 0.0, 1.0) + + def _raw10_unpack_cols_parity_to_float01_fast( + self, + groups_rows: np.ndarray, + width: int, + parity: int, + scale: np.float32, + ) -> np.ndarray: + """ + Desempacota somente colunas pares ou ímpares de um conjunto de linhas RAW10. + + groups_rows: + shape = (n_rows, width//4, 5) + + parity: + 0 -> colunas x = 0,2,4,6,... + 1 -> colunas x = 1,3,5,7,... + + Retorna: + shape = (n_rows, width//2), float32 0..1 + + Isso é a chave para o RGB bayer_planes sem raw16 full-res. + """ + width = int(width) + n_rows = groups_rows.shape[0] + + b0 = groups_rows[:, :, 0].astype(np.uint16, copy=False) + b1 = groups_rows[:, :, 1].astype(np.uint16, copy=False) + b2 = groups_rows[:, :, 2].astype(np.uint16, copy=False) + b3 = groups_rows[:, :, 3].astype(np.uint16, copy=False) + b4 = groups_rows[:, :, 4].astype(np.uint16, copy=False) + + out = np.empty((n_rows, width // 2), dtype=np.float32) + + if int(parity) == 0: + # Colunas pares: p0, p2 em cada grupo. + out[:, 0::2] = ((b0 << 2) | ((b4 >> 0) & 0x03)).astype(np.float32) * scale + out[:, 1::2] = ((b2 << 2) | ((b4 >> 4) & 0x03)).astype(np.float32) * scale + else: + # Colunas ímpares: p1, p3 em cada grupo. + out[:, 0::2] = ((b1 << 2) | ((b4 >> 2) & 0x03)).astype(np.float32) * scale + out[:, 1::2] = ((b3 << 2) | ((b4 >> 6) & 0x03)).astype(np.float32) * scale + + return out + + def _raw10_rgb_bayer_planes_to_rgb_float01_fast( + self, + packed_frame: np.ndarray, + width: int, + height: int, + bayer_pattern: str | None = None, + bit_depth: int = 10, + ) -> np.ndarray: + """ + RGB Bayer RAW10 packed -> RGB HWC float32 0..1 usando bayer_planes. + + Retorna half-res, exatamente como bayer_planes_to_rgb_linear() fazia: + shape = (height//2, width//2, 3) + canais = RGB + + Vantagem: + evita criar raw16 full-res e evita fatiar raw16 depois. + + Padrões suportados: + RGGB, BGGR, GRBG, GBRG + """ + width = int(width) + height = int(height) + bit_depth = int(bit_depth) + p = str(bayer_pattern or self.bayer_pattern or "RGGB").upper() + + if height % 2 != 0 or width % 2 != 0: + raise ValueError(f"Bayer planes exige width/height pares. Veio {width}x{height}") + + groups = self._raw10_groups_view_fast(packed_frame, width, height) + + even_rows = groups[0::2] + odd_rows = groups[1::2] + + max_val = float((1 << bit_depth) - 1) + scale = np.float32(1.0 / max_val) + + even_even = self._raw10_unpack_cols_parity_to_float01_fast(even_rows, width, parity=0, scale=scale) + even_odd = self._raw10_unpack_cols_parity_to_float01_fast(even_rows, width, parity=1, scale=scale) + odd_even = self._raw10_unpack_cols_parity_to_float01_fast(odd_rows, width, parity=0, scale=scale) + odd_odd = self._raw10_unpack_cols_parity_to_float01_fast(odd_rows, width, parity=1, scale=scale) + + if p == "RGGB": + r = even_even + g1 = even_odd + g2 = odd_even + b = odd_odd + elif p == "BGGR": + b = even_even + g1 = even_odd + g2 = odd_even + r = odd_odd + elif p == "GRBG": + g1 = even_even + r = even_odd + b = odd_even + g2 = odd_odd + elif p == "GBRG": + g1 = even_even + b = even_odd + r = odd_even + g2 = odd_odd + else: + raise ValueError(f"Padrão Bayer não suportado: {p}") + + g = (g1 + g2) * np.float32(0.5) + + rgb = np.empty((height // 2, width // 2, 3), dtype=np.float32) + rgb[:, :, 0] = r + rgb[:, :, 1] = g + rgb[:, :, 2] = b + + rgb_cal = getattr(self, "rgb_calibration", {}) or {} + if rgb_cal.get("enabled", False): + gains = rgb_cal.get("gains", {}) or {} + rgb[:, :, 0] *= float(gains.get("R", 1.0)) + rgb[:, :, 1] *= float(gains.get("G", 1.0)) + rgb[:, :, 2] *= float(gains.get("B", 1.0)) + + return np.clip(rgb, 0.0, 1.0).astype(np.float32, copy=False) + + + + def _radiometric_get_writable_float32_image(self, img): + """ + Garante uma imagem float32 gravável. + + Se já for float32 e writeable, usa a própria referência. + Se não, faz uma única cópia/conversão. + """ + if img is None: + return None, False + + if img.dtype == np.float32 and img.flags.writeable: + return img, True + + return img.astype(np.float32, copy=True), False + + def _apply_radiometric_scale_inplace(self, img, scale: float, clip_output: bool): + """ + Aplica escala radiométrica com o mínimo possível de alocação. + + Retorna: + out, reused_input + """ + out, reused = self._radiometric_get_writable_float32_image(img) + + if out is None: + return img, False + + scale = float(scale) + + # Se a escala é praticamente 1 e não precisa clipar, não faz nada. + if abs(scale - 1.0) <= 1e-6 and not clip_output: + return out, reused + + # Multiplicação in-place. + if abs(scale - 1.0) > 1e-6: + np.multiply(out, np.float32(scale), out=out, casting="unsafe") + + # Clip in-place, se configurado. + if clip_output: + np.clip(out, 0.0, 1.0, out=out) + + return out, reused + + def _build_radiometric_debug(self, method, role, cam_id, actual_ctrl, ref_ctrl, actual_factor, ref_factor, raw_scale, scale, clip_output): + return { + "applied": True, + "method": method, + "role": role, + "camera_id": cam_id, + "actual_controls": dict(actual_ctrl), + "reference_controls": dict(ref_ctrl), + "actual_factor": float(actual_factor), + "reference_factor": float(ref_factor), + "scale_raw": float(raw_scale), + "scale_applied": float(scale), + "clip_output": bool(clip_output), + } + + + + def _bayer_pattern_to_code_fast(self, bayer_pattern: str | None) -> int: + p = str(bayer_pattern or self.bayer_pattern or "RGGB").upper() + if p == "RGGB": + return 0 + if p == "BGGR": + return 1 + if p == "GRBG": + return 2 + if p == "GBRG": + return 3 + raise ValueError(f"Padrão Bayer não suportado: {p}") + + def _get_rgb_calibration_gains_fast(self): + rgb_cal = getattr(self, "rgb_calibration", {}) or {} + if not rgb_cal.get("enabled", False): + return 1.0, 1.0, 1.0 + + gains = rgb_cal.get("gains", {}) or {} + return ( + float(gains.get("R", 1.0)), + float(gains.get("G", 1.0)), + float(gains.get("B", 1.0)), + ) + + def _prepare_raw10_packed_view_numba(self, packed_frame: np.ndarray, width: int, height: int) -> np.ndarray: + """ + Retorna uma view/cópia contígua do RAW10 packed útil. + O Numba gosta de array C-contiguous. + """ + arr = self._prepare_raw10_packed_view_fast(packed_frame, width, height) + + if arr.dtype != np.uint8: + arr = arr.astype(np.uint8, copy=False) + + if not arr.flags.c_contiguous: + arr = np.ascontiguousarray(arr) + + return arr + + def _raw10_mono_to_float01_aggressive( + self, + packed_frame: np.ndarray, + width: int, + height: int, + bit_depth: int = 10, + ) -> np.ndarray: + """ + RAW10 mono -> float32 0..1. + Usa Numba se disponível, fallback para NumPy fast. + """ + if not _HAS_NUMBA: + return self._raw10_unpack_mono_to_float01_fast( + packed_frame, + width=width, + height=height, + bit_depth=bit_depth, + ) + + width = int(width) + height = int(height) + bit_depth = int(bit_depth) + + packed = self._prepare_raw10_packed_view_numba(packed_frame, width, height) + + max_val = float((1 << bit_depth) - 1) + scale = np.float32(1.0 / max_val) + + out = np.empty((height, width), dtype=np.float32) + _raw10_mono_to_float32_numba(packed, out, width, height, scale) + + return out + + def _raw10_rgb_bayer_planes_to_rgb_float01_aggressive( + self, + packed_frame: np.ndarray, + width: int, + height: int, + bayer_pattern: str | None = None, + bit_depth: int = 10, + ) -> np.ndarray: + """ + RAW10 RGB Bayer -> RGB HWC float32 0..1 usando bayer_planes. + Usa Numba se disponível, fallback para NumPy fast. + """ + if not _HAS_NUMBA: + return self._raw10_rgb_bayer_planes_to_rgb_float01_fast( + packed_frame, + width=width, + height=height, + bayer_pattern=bayer_pattern, + bit_depth=bit_depth, + ) + + width = int(width) + height = int(height) + bit_depth = int(bit_depth) + + if width % 2 != 0 or height % 2 != 0: + raise ValueError(f"Bayer planes exige width/height pares. Veio {width}x{height}") + + packed = self._prepare_raw10_packed_view_numba(packed_frame, width, height) + + max_val = float((1 << bit_depth) - 1) + scale = np.float32(1.0 / max_val) + + pattern_code = self._bayer_pattern_to_code_fast(bayer_pattern) + gain_r, gain_g, gain_b = self._get_rgb_calibration_gains_fast() + + out = np.empty((height // 2, width // 2, 3), dtype=np.float32) + + _raw10_bayer_planes_to_rgb_numba( + packed, + out, + width, + height, + scale, + int(pattern_code), + float(gain_r), + float(gain_g), + float(gain_b), + ) + + return out + + def _raw10_to_raw16_aggressive( + self, + packed_frame: np.ndarray, + width: int, + height: int, + ) -> np.ndarray: + """ + RAW10 packed -> raw16 uint16. + Usa Numba quando disponível. + É pensado para o linear_demosaic, porque cv2.cvtColor precisa de raw16. + """ + width = int(width) + height = int(height) + + if not _HAS_NUMBA: + return self.unpack_raw10_packed( + packed_frame, + sensor_width=width, + sensor_height=height, + ) + + packed = self._prepare_raw10_packed_view_numba( + packed_frame, + width, + height, + ) + + out = np.empty((height, width), dtype=np.uint16) + _raw10_to_raw16_numba(packed, out, width, height) + + return out + + def _rgb16_to_float32_gain_clip_aggressive( + self, + rgb16: np.ndarray, + bit_depth: int, + ) -> np.ndarray: + """ + RGB uint16 -> RGB float32 0..1 com calibração e clip em uma passada. + """ + h, w = rgb16.shape[:2] + + gain_r, gain_g, gain_b = self._get_rgb_calibration_gains_fast() + scale = np.float32(1.0 / float((1 << int(bit_depth)) - 1)) + + if not _HAS_NUMBA: + rgb = rgb16.astype(np.float32) + rgb[:, :, 0] *= np.float32(scale * gain_r) + rgb[:, :, 1] *= np.float32(scale * gain_g) + rgb[:, :, 2] *= np.float32(scale * gain_b) + np.clip(rgb, 0.0, 1.0, out=rgb) + return rgb.astype(np.float32, copy=False) + + rgb16_c = rgb16 + if not rgb16_c.flags.c_contiguous: + rgb16_c = np.ascontiguousarray(rgb16_c) + + out = np.empty((h, w, 3), dtype=np.float32) + + _rgb16_to_float32_gain_clip_numba( + rgb16_c, + out, + int(h), + int(w), + np.float32(scale), + float(gain_r), + float(gain_g), + float(gain_b), + ) + + return out + + def _apply_flat_gain_simple_aggressive( + self, + img: np.ndarray, + gain_eff: np.ndarray, + clip_output: bool = True, + ) -> np.ndarray: + base = img.astype(np.float32, copy=False) + + if gain_eff.shape[:2] != base.shape[:2]: + gain_eff = cv2.resize( + gain_eff.astype(np.float32, copy=False), + (base.shape[1], base.shape[0]), + interpolation=cv2.INTER_LINEAR, + ) + + if not _HAS_NUMBA: + out = base * gain_eff + if clip_output: + np.clip(out, 0.0, 1.0, out=out) + return out.astype(np.float32, copy=False) + + if not base.flags.c_contiguous: + base = np.ascontiguousarray(base) + + gain_c = gain_eff.astype(np.float32, copy=False) + if not gain_c.flags.c_contiguous: + gain_c = np.ascontiguousarray(gain_c) + + h, w = base.shape[:2] + out = np.empty((h, w), dtype=np.float32) + + _apply_flat_gain_numba( + base, + gain_c, + out, + int(h), + int(w), + bool(clip_output), + ) + + return out + + + + def _direct_fusion_get_target_size_fast(self, ref_size): + """ + Resolve target_size final como (target_w, target_h). + Usa fusion_config.target_size se existir, senão usa tamanho de referência. + """ + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + fusion = getattr(self, "fusion_config", {}) or {} + + target_size = fusion.get("target_size", None) + if isinstance(target_size, (list, tuple)) and len(target_size) == 2: + return int(target_size[0]), int(target_size[1]) + + return int(ref_w), int(ref_h) + + def _direct_fusion_get_crop_box_fast(self, valid_masks, ref_size): + """ + Calcula crop_box comum se crop_valid_common=true. + Se não houver crop, usa frame inteiro da referência. + """ + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + cfg = getattr(self, "fusion_config", {}) or {} + + if cfg.get("crop_valid_common", False): + crop_box = self._compute_common_crop_box(valid_masks) + if crop_box is not None: + return tuple(int(v) for v in crop_box), True + + return (0, 0, ref_w, ref_h), False + + def _direct_fusion_crop_to_target_matrix_fast(self, crop_box, target_size): + """ + Matriz C que leva coordenadas do espaço RGB/ref para a saída final. + + crop_box está no espaço da imagem de referência: + x0,y0,x1,y1 + + Queremos: + x=x0 -> 0 + x=x1 -> target_w + y=y0 -> 0 + y=y1 -> target_h + + Retorna C_ref_to_target. + """ + x0, y0, x1, y1 = [float(v) for v in crop_box] + target_w, target_h = int(target_size[0]), int(target_size[1]) + + crop_w = max(1.0, x1 - x0) + crop_h = max(1.0, y1 - y0) + + sx = float(target_w) / crop_w + sy = float(target_h) / crop_h + + C = np.array( + [ + [sx, 0.0, -x0 * sx], + [0.0, sy, -y0 * sy], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + + return C + + def _direct_fusion_scale_homography_for_ref_fast(self, H, meta, ref_size): + """ + Escala a homografia calibrada para o runtime da referência RGB. + + Usa a própria função existente _scale_homography_to_runtime() se existir. + Isso mantém compatibilidade com o contrato atual do core. + """ + fusion = getattr(self, "fusion_config", {}) or {} + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + + calib_size = ( + fusion.get("homography_calibration_size") + or fusion.get("calibration_size") + or fusion.get("source_size") + or None + ) + + H = np.asarray(H, dtype=np.float32) + + if hasattr(self, "_scale_homography_to_runtime"): + try: + return self._scale_homography_to_runtime( + H, + calib_size=calib_size, + runtime_size=(ref_w, ref_h), + ).astype(np.float32) + except TypeError: + try: + return self._scale_homography_to_runtime( + H, + calib_size, + (ref_w, ref_h), + ).astype(np.float32) + except Exception: + pass + except Exception: + pass + + # Fallback local. + if calib_size is None: + if abs(float(H[2, 2])) > 1e-9: + H = H / H[2, 2] + return H.astype(np.float32) + + calib_w, calib_h = float(calib_size[0]), float(calib_size[1]) + if calib_w <= 0 or calib_h <= 0: + return H.astype(np.float32) + + sx = float(ref_w) / calib_w + sy = float(ref_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(float(H_runtime[2, 2])) > 1e-9: + H_runtime = H_runtime / H_runtime[2, 2] + + return H_runtime.astype(np.float32) + + def _direct_fusion_get_role_homography_fast(self, role, meta, ref_size): + """ + Retorna H_role_to_rgb escalada para o espaço da referência RGB. + """ + role = str(role).lower() + fusion = getattr(self, "fusion_config", {}) or {} + homographies = fusion.get("homographies", {}) or {} + + key = f"{role}_to_rgb" + H = homographies.get(key) + + if H is None: + # Fallbacks para contratos diferentes. + H = homographies.get(role) + + if H is None: + raise RuntimeError(f"Homografia ausente para role={role}. Esperado fusion_config.homographies.{key}") + + return self._direct_fusion_scale_homography_for_ref_fast(H, meta, ref_size) + + def _direct_fusion_resize_spec_to_ref_if_needed_fast(self, img, ref_size): + """ + Mantém compatibilidade com o fluxo atual: + se RE/NIR não estão no mesmo shape do RGB de referência, redimensiona para ref. + """ + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + + if img.shape[:2] == (ref_h, ref_w): + return img + + return cv2.resize( + img.astype(np.float32, copy=False), + (ref_w, ref_h), + interpolation=cv2.INTER_LINEAR, + ) + + def _direct_fusion_compute_valid_masks_fast(self, decoded, role_to_cam, ref_size, meta): + """ + Calcula máscaras válidas no espaço RGB/ref para crop comum. + Usa warpPerspective apenas em máscara uint8, que costuma ser barato. + """ + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + + valid_masks = [np.ones((ref_h, ref_w), dtype=np.uint8)] + + base = np.ones((ref_h, ref_w), dtype=np.uint8) * 255 + + for role in ("re", "nir"): + if role not in role_to_cam: + continue + + H = self._direct_fusion_get_role_homography_fast(role, meta, ref_size) + mask = cv2.warpPerspective( + base, + H, + (ref_w, ref_h), + flags=cv2.INTER_NEAREST, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0, + ) + valid_masks.append(mask) + + return valid_masks + + def _direct_fusion_write_rgb_fast(self, tensor, rgb, crop_box, target_size): + """ + Escreve RGB no tensor final. + RGB entra HWC float32, sai CHW no tensor[0:3]. + """ + target_w, target_h = int(target_size[0]), int(target_size[1]) + x0, y0, x1, y1 = [int(v) for v in crop_box] + + rgb_crop = rgb[y0:y1, x0:x1, :].astype(np.float32, copy=False) + + if rgb_crop.shape[1] != target_w or rgb_crop.shape[0] != target_h: + rgb_out = cv2.resize( + rgb_crop, + (target_w, target_h), + interpolation=cv2.INTER_LINEAR, + ) + else: + rgb_out = rgb_crop + + tensor[0] = rgb_out[:, :, 0] + tensor[1] = rgb_out[:, :, 1] + tensor[2] = rgb_out[:, :, 2] + + def _direct_fusion_write_spec_fast(self, tensor, channel_index, img, role, C_ref_to_target, ref_size, target_size, meta): + """ + Escreve RE ou NIR direto no tensor final, compondo: + M = C_ref_to_target @ H_role_to_rgb + + img é primeiro redimensionada para ref_size se necessário, para manter o mesmo + comportamento geométrico do fluxo atual. + """ + target_w, target_h = int(target_size[0]), int(target_size[1]) + + img_ref = self._direct_fusion_resize_spec_to_ref_if_needed_fast(img, ref_size) + + H_role_to_rgb = self._direct_fusion_get_role_homography_fast(role, meta, ref_size) + M_role_to_target = (C_ref_to_target @ H_role_to_rgb).astype(np.float32) + + out = cv2.warpPerspective( + img_ref.astype(np.float32, copy=False), + M_role_to_target, + (target_w, target_h), + flags=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0.0, + ) + + tensor[int(channel_index)] = out + + def _fuse_multispec_direct_to_target_fast(self, decoded, meta, channels_expected): + """ + Fusão direta otimizada com cache de geometria fixa. + """ + t0 = time.perf_counter() + + rgb_cam_id = self._find_cam_by_role(decoded, "rgb") + if rgb_cam_id is None: + raise RuntimeError("Fusão direta requer câmera com role='rgb' como referência") + + rgb = decoded[rgb_cam_id]["image"] + if rgb.ndim != 3 or rgb.shape[2] != 3: + raise RuntimeError(f"RGB inválido para fusão direta: shape={rgb.shape}") + + ref_h, ref_w = rgb.shape[:2] + ref_size = (int(ref_h), int(ref_w)) + target_size = self._direct_fusion_get_target_size_fast(ref_size) + target_w, target_h = int(target_size[0]), int(target_size[1]) + + role_to_cam = { + item.get("role", data.get("meta", {}).get("role")): cam_id + for cam_id, data in decoded.items() + for item in [data] + } + + # Cache da geometria fixa. + geom = self._direct_fusion_get_geometry_cached_fast( + decoded=decoded, + role_to_cam=role_to_cam, + ref_size=ref_size, + target_size=target_size, + meta=meta, + ) + + crop_box = geom["crop_box"] + crop_applied = bool(geom["crop_applied"]) + + self.last_fusion_result = { + "ref_shape": [int(ref_h), int(ref_w)], + "target_size": [int(target_w), int(target_h)], + "crop_valid_common": bool((self.fusion_config or {}).get("crop_valid_common", False)), + "resize_after_crop": bool((self.fusion_config or {}).get("resize_after_crop", False)), + "crop_box": [int(v) for v in crop_box], + "crop_applied": bool(crop_applied), + "roles": ["R", "G", "B", "RE", "NIR"], + "direct_fusion_fast": True, + "geometry_cache_hit": bool(geom.get("prepare_cache_hit", False)), + "geometry_cache_hits": int(geom.get("cache_hits", 0)), + "geometry_cache_misses": int(geom.get("cache_misses", 0)), + } + + tensor = np.empty((int(channels_expected), target_h, target_w), dtype=np.float32) + + t_prepare_ms = (time.perf_counter() - t0) * 1000.0 + + # ------------------------------------------------------------ + # Cache de remap + # ------------------------------------------------------------ + use_remap_cache = bool((self.fusion_config or {}).get("use_remap_cache", True)) + use_remap_for_rgb = bool((self.fusion_config or {}).get("use_remap_for_rgb", False)) + use_remap_for_spec = bool((self.fusion_config or {}).get("use_remap_for_spec", True)) + + t0_remap_cache = time.perf_counter() + + if use_remap_cache and (use_remap_for_rgb or use_remap_for_spec): + remap_cache = self._direct_fusion_get_remap_cached_fast( + decoded=decoded, + role_to_cam=role_to_cam, + geom=geom, + ref_size=ref_size, + target_size=target_size, + ) + else: + remap_cache = { + "maps": {}, + "cache_hit": False, + "cache_hits": 0, + "cache_misses": 0, + } + + t_remap_cache_ms = (time.perf_counter() - t0_remap_cache) * 1000.0 + + # ------------------------------------------------------------ + # RGB via remap + # ------------------------------------------------------------ + t0_rgb = time.perf_counter() + + rgb_maps = remap_cache["maps"].get("rgb") + if use_remap_cache and use_remap_for_rgb and rgb_maps is not None: + self._direct_fusion_write_rgb_remap_fast( + tensor=tensor, + rgb=rgb, + remap_entry=rgb_maps, + ) + else: + self._direct_fusion_write_rgb_fast( + tensor, + rgb, + crop_box, + target_size, + ) + + t_rgb_ms = (time.perf_counter() - t0_rgb) * 1000.0 + + warp_details = {} + t_warp_total_ms = 0.0 + + # ------------------------------------------------------------ + # RE via remap + # ------------------------------------------------------------ + if "re" in role_to_cam: + t0w = time.perf_counter() + + re_img = decoded[role_to_cam["re"]]["image"] + re_maps = remap_cache["maps"].get("re") + + if use_remap_cache and use_remap_for_spec and re_maps is not None: + self._direct_fusion_write_spec_remap_fast( + tensor=tensor, + channel_index=3, + img=re_img, + remap_entry=re_maps, + ) + else: + self._direct_fusion_write_spec_cached_fast( + tensor=tensor, + channel_index=3, + img=re_img, + role="re", + geom=geom, + ref_size=ref_size, + target_size=target_size, + ) + + warp_details["re"] = (time.perf_counter() - t0w) * 1000.0 + t_warp_total_ms += warp_details["re"] + + # ------------------------------------------------------------ + # NIR via remap + # ------------------------------------------------------------ + if "nir" in role_to_cam: + t0w = time.perf_counter() + + nir_img = decoded[role_to_cam["nir"]]["image"] + nir_maps = remap_cache["maps"].get("nir") + + if use_remap_cache and use_remap_for_spec and nir_maps is not None: + self._direct_fusion_write_spec_remap_fast( + tensor=tensor, + channel_index=4, + img=nir_img, + remap_entry=nir_maps, + ) + else: + self._direct_fusion_write_spec_cached_fast( + tensor=tensor, + channel_index=4, + img=nir_img, + role="nir", + geom=geom, + ref_size=ref_size, + target_size=target_size, + ) + + warp_details["nir"] = (time.perf_counter() - t0w) * 1000.0 + t_warp_total_ms += warp_details["nir"] + + # ------------------------------------------------------------ + # Flat-field no espaço final do tensor + # ------------------------------------------------------------ + t0_final_flat = time.perf_counter() + final_flat_ms = 0.0 + final_flat_enabled = False + final_flat_cache_available = False + + flat_cfg = self.flatfield_config or {} + apply_final_flat = ( + bool(flat_cfg.get("enabled", False)) + and str(flat_cfg.get("apply_space", "native_camera_space")).lower() == "final_tensor_space" + ) + + if apply_final_flat: + final_flat_enabled = True + + gain_tensor = self._get_final_flat_gain_tensor_cached_fast( + decoded=decoded, + role_to_cam=role_to_cam, + geom=geom, + remap_cache=remap_cache, + crop_box=crop_box, + target_size=target_size, + ) + + if gain_tensor is not None: + final_flat_cache_available = True + tensor = self._apply_final_flat_gain_tensor_inplace( + tensor=tensor, + gain_tensor=gain_tensor, + clip_output=bool(flat_cfg.get("clip_output", True)), + ) + + final_flat_ms = (time.perf_counter() - t0_final_flat) * 1000.0 + + if tensor.shape[0] != channels_expected: + raise RuntimeError( + f"Tensor direto com canais inesperados: {tensor.shape[0]} | esperado={channels_expected}" + ) + + self.last_fusion_result["output_shape"] = list(tensor.shape) + + perf = { + "prepare_ms": float(t_prepare_ms), + "rgb_crop_resize_ms": float(t_rgb_ms), + "warp_total_ms": float(t_warp_total_ms), + "warp_details_ms": warp_details, + "crop_resize_ms": float(t_rgb_ms), + "concat_ms": 0.0, + "spatial_direct_ms": float(t_rgb_ms + t_warp_total_ms), + "geometry_cache_hit": bool(geom.get("prepare_cache_hit", False)), + "geometry_cache_hits": int(geom.get("cache_hits", 0)), + "geometry_cache_misses": int(geom.get("cache_misses", 0)), + "remap_cache_ms": float(t_remap_cache_ms), + "remap_cache_hit": bool(remap_cache.get("cache_hit", False)), + "remap_cache_hits": int(remap_cache.get("cache_hits", 0)), + "remap_cache_misses": int(remap_cache.get("cache_misses", 0)), + "remap_enabled": bool(use_remap_cache), + "remap_rgb_enabled": bool(use_remap_cache and use_remap_for_rgb), + "remap_spec_enabled": bool(use_remap_cache and use_remap_for_spec), + "final_flat_enabled": bool(final_flat_enabled), + "final_flat_cache_available": bool(final_flat_cache_available), + "final_flat_ms": float(final_flat_ms), + } + + return tensor, perf + + + + def _direct_fusion_get_geometry_cache_key_fast(self, ref_size, target_size, role_to_cam): + """ + Chave simples e estável para cache da geometria. + + A geometria depende de: + - tamanho do RGB de referência + - target final + - roles presentes + - crop_valid_common / resize_after_crop + - homografias e calibration_size + + Para evitar custo de serializar o JSON todo por frame, usamos uma versão + simples. Se você editar module_params em runtime, chame + clear_direct_fusion_geometry_cache(). + """ + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + target_w, target_h = int(target_size[0]), int(target_size[1]) + + fusion = getattr(self, "fusion_config", {}) or {} + homographies = fusion.get("homographies", {}) or {} + + # Pequena assinatura numérica das homografias. + def h_sig(key): + H = homographies.get(key) + if H is None: + return None + arr = np.asarray(H, dtype=np.float32).reshape(-1) + # arredonda para evitar ruído float/json, mas detecta mudança real. + return tuple(np.round(arr, 8).tolist()) + + roles = tuple(sorted([str(r).lower() for r in role_to_cam.keys()])) + + return ( + ref_w, + ref_h, + target_w, + target_h, + roles, + bool(fusion.get("crop_valid_common", False)), + bool(fusion.get("resize_after_crop", False)), + tuple(fusion.get("homography_calibration_size") or fusion.get("calibration_size") or []), + h_sig("re_to_rgb"), + h_sig("nir_to_rgb"), + ) + + def clear_direct_fusion_geometry_cache(self): + """ + Chame se mudar fusion_config/module_params em runtime. + """ + self._direct_fusion_geometry_cache = {} + self._direct_fusion_geometry_cache_hits = 0 + self._direct_fusion_geometry_cache_misses = 0 + + self._direct_fusion_remap_cache = {} + self._direct_fusion_remap_cache_hits = 0 + self._direct_fusion_remap_cache_misses = 0 + + def _direct_fusion_get_geometry_cached_fast(self, decoded, role_to_cam, ref_size, target_size, meta): + """ + Retorna geometria cacheada para a fusão direta. + + Saída: + geom = { + key, + crop_box, + crop_applied, + C_ref_to_target, + H_role_to_rgb: {re,nir}, + M_role_to_target: {re,nir}, + valid_masks, # opcional/debug + prepare_cache_hit, + } + """ + if not hasattr(self, "_direct_fusion_geometry_cache"): + self.clear_direct_fusion_geometry_cache() + + key = self._direct_fusion_get_geometry_cache_key_fast(ref_size, target_size, role_to_cam) + cache = self._direct_fusion_geometry_cache + + if key in cache: + self._direct_fusion_geometry_cache_hits += 1 + geom = cache[key] + geom["prepare_cache_hit"] = True + geom["cache_hits"] = int(self._direct_fusion_geometry_cache_hits) + geom["cache_misses"] = int(self._direct_fusion_geometry_cache_misses) + return geom + + self._direct_fusion_geometry_cache_misses += 1 + + ref_h, ref_w = int(ref_size[0]), int(ref_size[1]) + + # ------------------------------------------------------------ + # Homografias escaladas para runtime. + # ------------------------------------------------------------ + H_role_to_rgb = {} + for role in ("re", "nir"): + if role in role_to_cam: + H_role_to_rgb[role] = self._direct_fusion_get_role_homography_fast(role, meta, ref_size) + + # ------------------------------------------------------------ + # Máscaras válidas e crop comum. + # Essa era uma das partes caras e totalmente fixa. + # ------------------------------------------------------------ + valid_masks = [np.ones((ref_h, ref_w), dtype=np.uint8)] + base = np.ones((ref_h, ref_w), dtype=np.uint8) * 255 + + for role in ("re", "nir"): + if role not in H_role_to_rgb: + continue + + mask = cv2.warpPerspective( + base, + H_role_to_rgb[role], + (ref_w, ref_h), + flags=cv2.INTER_NEAREST, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0, + ) + valid_masks.append(mask) + + crop_box, crop_applied = self._direct_fusion_get_crop_box_fast(valid_masks, ref_size) + + # ------------------------------------------------------------ + # Matriz crop RGB/ref -> target final. + # ------------------------------------------------------------ + C_ref_to_target = self._direct_fusion_crop_to_target_matrix_fast(crop_box, target_size) + + # ------------------------------------------------------------ + # Matrizes compostas role original/ref -> target. + # ------------------------------------------------------------ + M_role_to_target = {} + for role, H in H_role_to_rgb.items(): + M_role_to_target[role] = (C_ref_to_target @ H).astype(np.float32) + + geom = { + "key": key, + "crop_box": tuple(int(v) for v in crop_box), + "crop_applied": bool(crop_applied), + "C_ref_to_target": C_ref_to_target.astype(np.float32), + "H_role_to_rgb": H_role_to_rgb, + "M_role_to_target": M_role_to_target, + "valid_masks": valid_masks, + "prepare_cache_hit": False, + "cache_hits": int(self._direct_fusion_geometry_cache_hits), + "cache_misses": int(self._direct_fusion_geometry_cache_misses), + } + + # Cache pequeno: normalmente só uma geometria. Se mudar resolução/config, + # evita crescimento infinito. + if len(cache) > 4: + cache.clear() + + cache[key] = geom + return geom + + def _direct_fusion_write_spec_cached_fast(self, tensor, channel_index, img, role, geom, ref_size, target_size): + """ + Escreve RE/NIR usando matriz composta cacheada. + """ + target_w, target_h = int(target_size[0]), int(target_size[1]) + + img_ref = self._direct_fusion_resize_spec_to_ref_if_needed_fast(img, ref_size) + + M_role_to_target = geom["M_role_to_target"].get(str(role).lower()) + if M_role_to_target is None: + raise RuntimeError(f"Matriz composta ausente para role={role}") + + out = cv2.warpPerspective( + img_ref.astype(np.float32, copy=False), + M_role_to_target, + (target_w, target_h), + flags=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0.0, + ) + + tensor[int(channel_index)] = out + + def _direct_fusion_build_remap_from_src_to_dst_fast( + self, + M_src_to_dst: np.ndarray, + src_shape: tuple, + dst_size: tuple, + ref_shape_for_scaled_src: tuple | None = None, + ): + """ + Cria mapas cacheáveis para cv2.remap. + + M_src_to_dst: + matriz 3x3 que leva coordenadas da imagem fonte para o destino final. + + src_shape: + shape real da imagem fonte, ex: img.shape[:2] + + dst_size: + (target_w, target_h) + + ref_shape_for_scaled_src: + usado quando a matriz foi calculada no espaço ref, mas a imagem fonte + real tem outro tamanho. Ex: spec original 1280x800, ref RGB half 640x400. + """ + src_h, src_w = int(src_shape[0]), int(src_shape[1]) + target_w, target_h = int(dst_size[0]), int(dst_size[1]) + + M = np.asarray(M_src_to_dst, dtype=np.float32) + if M.shape != (3, 3): + raise RuntimeError(f"M_src_to_dst inválida: shape={M.shape}") + + M_inv = np.linalg.inv(M).astype(np.float32) + + xs, ys = np.meshgrid( + np.arange(target_w, dtype=np.float32), + np.arange(target_h, dtype=np.float32), + ) + + den = M_inv[2, 0] * xs + M_inv[2, 1] * ys + M_inv[2, 2] + den = np.where(np.abs(den) < 1e-9, 1e-9, den) + + map_x = (M_inv[0, 0] * xs + M_inv[0, 1] * ys + M_inv[0, 2]) / den + map_y = (M_inv[1, 0] * xs + M_inv[1, 1] * ys + M_inv[1, 2]) / den + + # Se a matriz foi calculada no espaço de referência, mas a imagem fonte + # real tem outro tamanho, converte coordenada ref -> coordenada fonte real. + if ref_shape_for_scaled_src is not None: + ref_h, ref_w = int(ref_shape_for_scaled_src[0]), int(ref_shape_for_scaled_src[1]) + + if (src_w, src_h) != (ref_w, ref_h): + sx = float(src_w) / float(ref_w) + sy = float(src_h) / float(ref_h) + + # Aproxima a convenção de resize do OpenCV: + # x_src = (x_ref + 0.5) * scale - 0.5 + map_x = (map_x + 0.5) * sx - 0.5 + map_y = (map_y + 0.5) * sy - 0.5 + + map_x = map_x.astype(np.float32, copy=False) + map_y = map_y.astype(np.float32, copy=False) + + # convertMaps deixa o remap mais barato em muitos casos. + map1, map2 = cv2.convertMaps(map_x, map_y, cv2.CV_16SC2) + + return { + "map_x": map_x, + "map_y": map_y, + "map1": map1, + "map2": map2, + "src_shape": [src_h, src_w], + "dst_size": [target_w, target_h], + } + + def _direct_fusion_get_remap_cache_key_fast( + self, + geom: dict, + decoded: dict, + role_to_cam: dict, + ref_size: tuple, + target_size: tuple, + ): + """ + Chave do cache de remap. + + Precisa considerar: + - geometria base; + - crop/homografia; + - tamanho real das imagens fonte; + - target final. + """ + role_shapes = {} + + for role, cam_id in role_to_cam.items(): + img = decoded[cam_id]["image"] + role_shapes[str(role).lower()] = tuple(int(v) for v in img.shape[:2]) + + return ( + geom.get("key"), + tuple(int(v) for v in ref_size), + tuple(int(v) for v in target_size), + tuple(sorted(role_shapes.items())), + ) + + def _direct_fusion_get_remap_cached_fast( + self, + decoded: dict, + role_to_cam: dict, + geom: dict, + ref_size: tuple, + target_size: tuple, + ): + """ + Retorna mapas de remap cacheados para RGB, RE e NIR. + """ + if not hasattr(self, "_direct_fusion_remap_cache"): + self._direct_fusion_remap_cache = {} + self._direct_fusion_remap_cache_hits = 0 + self._direct_fusion_remap_cache_misses = 0 + + key = self._direct_fusion_get_remap_cache_key_fast( + geom=geom, + decoded=decoded, + role_to_cam=role_to_cam, + ref_size=ref_size, + target_size=target_size, + ) + + cache = self._direct_fusion_remap_cache + + if key in cache: + self._direct_fusion_remap_cache_hits += 1 + remap = cache[key] + remap["cache_hit"] = True + remap["cache_hits"] = int(self._direct_fusion_remap_cache_hits) + remap["cache_misses"] = int(self._direct_fusion_remap_cache_misses) + return remap + + self._direct_fusion_remap_cache_misses += 1 + + target_w, target_h = int(target_size[0]), int(target_size[1]) + + remap = { + "key": key, + "maps": {}, + "cache_hit": False, + "cache_hits": int(self._direct_fusion_remap_cache_hits), + "cache_misses": int(self._direct_fusion_remap_cache_misses), + } + + # ------------------------------------------------------------ + # RGB: matriz C_ref_to_target leva RGB/ref -> target. + # ------------------------------------------------------------ + rgb_cam_id = role_to_cam.get("rgb") + if rgb_cam_id is not None: + rgb_img = decoded[rgb_cam_id]["image"] + C_ref_to_target = geom["C_ref_to_target"] + + remap["maps"]["rgb"] = self._direct_fusion_build_remap_from_src_to_dst_fast( + M_src_to_dst=C_ref_to_target, + src_shape=rgb_img.shape[:2], + dst_size=(target_w, target_h), + ref_shape_for_scaled_src=None, + ) + + # ------------------------------------------------------------ + # RE/NIR: matriz composta M_role_to_target leva role/ref -> target. + # Se a imagem fonte não tiver o mesmo tamanho do ref, escalamos o mapa. + # ------------------------------------------------------------ + for role in ("re", "nir"): + cam_id = role_to_cam.get(role) + if cam_id is None: + continue + + img = decoded[cam_id]["image"] + M_role_to_target = geom["M_role_to_target"].get(role) + + if M_role_to_target is None: + continue + + remap["maps"][role] = self._direct_fusion_build_remap_from_src_to_dst_fast( + M_src_to_dst=M_role_to_target, + src_shape=img.shape[:2], + dst_size=(target_w, target_h), + ref_shape_for_scaled_src=ref_size, + ) + + if len(cache) > 4: + cache.clear() + + cache[key] = remap + return remap + + def _direct_fusion_write_rgb_remap_fast( + self, + tensor: np.ndarray, + rgb: np.ndarray, + remap_entry: dict, + ): + """ + RGB HWC -> tensor CHW usando cv2.remap direto para target final. + """ + map1 = remap_entry["map1"] + map2 = remap_entry["map2"] + + rgb_out = cv2.remap( + rgb.astype(np.float32, copy=False), + map1, + map2, + interpolation=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0.0, + ) + + tensor[0] = rgb_out[:, :, 0] + tensor[1] = rgb_out[:, :, 1] + tensor[2] = rgb_out[:, :, 2] + + def _direct_fusion_write_spec_remap_fast( + self, + tensor: np.ndarray, + channel_index: int, + img: np.ndarray, + remap_entry: dict, + ): + """ + RE/NIR -> tensor usando cv2.remap direto para target final. + """ + map1 = remap_entry["map1"] + map2 = remap_entry["map2"] + + out = cv2.remap( + img.astype(np.float32, copy=False), + map1, + map2, + interpolation=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=0.0, + ) + + tensor[int(channel_index)] = out + + + + def _raw10_rgb_linear_demosaic_to_rgb_float01_fast( + self, + packed_frame: np.ndarray, + width: int, + height: int, + bayer_pattern: str | None = None, + bit_depth: int = 10, + ) -> np.ndarray: + """ + RAW10 RGB Bayer packed -> RGB HWC float32 0..1 usando demosaic OpenCV. + + Agora com telemetria interna: + - unpack_ms + - cvtColor_ms + - float_ms + - calibration_ms + - clip_ms + - resize_half_ms + - total_ms + """ + t_total0 = time.perf_counter() + + width = int(width) + height = int(height) + bit_depth = int(bit_depth) + + rgb_cfg = getattr(self, "rgb_processing_config", {}) or {} + rgb_mode = str(rgb_cfg.get("mode", "linear_demosaic")).lower() + algorithm = str(rgb_cfg.get("demosaic_algorithm", "ea")).lower() + + perf = { + "mode": rgb_mode, + "algorithm": algorithm, + "width": width, + "height": height, + } + + # ------------------------------------------------------------ + # 1) Resolve código Bayer OpenCV + # ------------------------------------------------------------ + t0 = time.perf_counter() + + cv_code, algorithm_resolved = self._get_bayer_cv2_code( + bayer_pattern=bayer_pattern, + algorithm=algorithm, + ) + + perf["resolve_code_ms"] = (time.perf_counter() - t0) * 1000.0 + perf["algorithm_resolved"] = algorithm_resolved + + # ------------------------------------------------------------ + # 2) RAW10 packed -> raw16 full-res + # ------------------------------------------------------------ + t0 = time.perf_counter() + + raw16 = self._raw10_to_raw16_aggressive( + packed_frame, + width=width, + height=height, + ) + + perf["unpack_ms"] = (time.perf_counter() - t0) * 1000.0 + + # ------------------------------------------------------------ + # 3) Demosaic OpenCV + # ------------------------------------------------------------ + t0 = time.perf_counter() + + rgb16 = cv2.cvtColor(raw16, cv_code) + + perf["cvtColor_ms"] = (time.perf_counter() - t0) * 1000.0 + + # ------------------------------------------------------------ + # 4) uint16 -> float32 0..1 + # ------------------------------------------------------------ + t0 = time.perf_counter() + + rgb = self._rgb16_to_float32_gain_clip_aggressive( + rgb16, + bit_depth=bit_depth, + ) + + perf["float_calib_clip_ms"] = (time.perf_counter() - t0) * 1000.0 + perf["calibration_enabled"] = bool((getattr(self, "rgb_calibration", {}) or {}).get("enabled", False)) + perf["float_ms"] = 0.0 + perf["calibration_ms"] = 0.0 + perf["clip_ms"] = 0.0 + + # ------------------------------------------------------------ + # 7) Half mode, se habilitado + # ------------------------------------------------------------ + t0 = time.perf_counter() + + if rgb_mode in ("linear_demosaic_half", "demosaic_half", "full_demosaic_half"): + rgb = cv2.resize( + rgb, + (width // 2, height // 2), + interpolation=cv2.INTER_AREA, + ).astype(np.float32, copy=False) + perf["resize_half_applied"] = True + else: + perf["resize_half_applied"] = False + + perf["resize_half_ms"] = (time.perf_counter() - t0) * 1000.0 + + # ------------------------------------------------------------ + # Total + # ------------------------------------------------------------ + perf["total_ms"] = (time.perf_counter() - t_total0) * 1000.0 + perf["out_shape"] = list(rgb.shape) + perf["out_dtype"] = str(rgb.dtype) + + self._set_decode_perf("rgb", perf) + + return rgb.astype(np.float32, copy=False) + + def demosaic_raw16_to_rgb_linear_hwc_fast( + self, + raw16: np.ndarray, + bit_depth: int = 10, + ) -> np.ndarray: + rgb_cfg = getattr(self, "rgb_processing_config", {}) or {} + algorithm = str(rgb_cfg.get("demosaic_algorithm", "ea")).lower() + + cv_code, _ = self._get_bayer_cv2_code( + bayer_pattern=self.bayer_pattern, + algorithm=algorithm, + ) + + if raw16.dtype != np.uint16: + raw16 = raw16.astype(np.uint16, copy=False) + + rgb16 = cv2.cvtColor(raw16, cv_code) + + rgb = rgb16.astype(np.float32) + rgb *= np.float32(1.0 / float((1 << bit_depth) - 1)) + + np.clip(rgb, 0.0, 1.0, out=rgb) + + return rgb.astype(np.float32, copy=False) + + def _get_bayer_cv2_code(self, bayer_pattern: str | None, algorithm: str = "ea"): + """ + Retorna o código OpenCV para demosaic Bayer. + + algorithm: + - "ea": Edge-Aware, melhor qualidade, mais pesado + - "bilinear": mais rápido, menor custo + """ + p = str(bayer_pattern or self.bayer_pattern or "RGGB").upper() + algo = str(algorithm or "ea").lower() + + if algo in ("bilinear", "linear", "fast", "normal"): + code_map = { + "BGGR": cv2.COLOR_BayerRG2RGB, + "RGGB": cv2.COLOR_BayerBG2RGB, + "GRBG": cv2.COLOR_BayerGR2RGB, + "GBRG": cv2.COLOR_BayerGB2RGB, + } + elif algo in ("ea", "edge_aware", "edge-aware"): + code_map = { + "BGGR": cv2.COLOR_BayerRG2RGB_EA, + "RGGB": cv2.COLOR_BayerBG2RGB_EA, + "GRBG": cv2.COLOR_BayerGR2RGB_EA, + "GBRG": cv2.COLOR_BayerGB2RGB_EA, + } + else: + raise ValueError(f"demosaic_algorithm inválido: {algorithm}") + + if p not in code_map: + raise ValueError(f"Padrão Bayer não suportado para demosaic: {p}") + + return code_map[p], algo + + + + def _set_decode_perf(self, role: str, perf: dict, log_interval_s: float = 1.0): + """ + Guarda telemetria do decode por role e loga no máximo 1x por segundo. + """ + role = str(role or "unknown").lower() + + if not hasattr(self, "last_decode_perf") or self.last_decode_perf is None: + self.last_decode_perf = {} + + self.last_decode_perf[role] = dict(perf) + + now = time.time() + if not hasattr(self, "_last_decode_perf_log_ts"): + self._last_decode_perf_log_ts = 0.0 + + if now - self._last_decode_perf_log_ts < log_interval_s: + return + + self._last_decode_perf_log_ts = now + + #parts = [] + #for k, v in perf.items(): + # if isinstance(v, (int, float)): + # parts.append(f"{k}={float(v):.2f}ms") + # else: + # parts.append(f"{k}={v}") + #print(f"[PERF][DECODE][{role.upper()}] " + " ".join(parts)) + + + + def _get_runtime_gain_eff_map(self, channel_name: str, base_shape: tuple, cfg: dict): + ch = str(channel_name).upper() + + h, w = base_shape[:2] + + strength = float(cfg.get("strength", 1.0)) + strength_by_channel = cfg.get("strength_by_channel", {}) or {} + if ch in strength_by_channel: + strength = float(strength_by_channel[ch]) + + gain_min_runtime = float(cfg.get("gain_min_runtime", 0.0)) + gain_max_runtime = float(cfg.get("gain_max_runtime", 999.0)) + + runtime_smooth_ksize = int(cfg.get("runtime_smooth_ksize", 0) or 0) + if runtime_smooth_ksize >= 3 and runtime_smooth_ksize % 2 == 0: + runtime_smooth_ksize += 1 + + cache_key = ( + "gain_eff", + ch, + int(h), + int(w), + int(runtime_smooth_ksize), + round(float(strength), 6), + round(float(gain_min_runtime), 6), + round(float(gain_max_runtime), 6), + ) + + cached = self._flatfield_runtime_cache.get(cache_key) + if cached is not None: + return cached + + gain = self._get_runtime_gain_map(ch, base_shape, cfg) + if gain is None: + return None + + gain_eff = 1.0 + np.float32(strength) * (gain.astype(np.float32, copy=False) - 1.0) + gain_eff = np.clip(gain_eff, gain_min_runtime, gain_max_runtime).astype(np.float32, copy=False) + + self._flatfield_runtime_cache[cache_key] = gain_eff + return gain_eff + + + + def _apply_final_flat_gain_tensor_inplace( + self, + tensor: np.ndarray, + gain_tensor: np.ndarray, + clip_output: bool = True, + ) -> np.ndarray: + if tensor is None or gain_tensor is None: + return tensor + + if tensor.ndim != 3: + raise RuntimeError(f"Tensor CHW esperado. Veio shape={tensor.shape}") + + if gain_tensor.shape != tensor.shape: + raise RuntimeError( + f"Gain tensor shape inválido: gain={gain_tensor.shape} tensor={tensor.shape}" + ) + + if not tensor.flags.c_contiguous: + tensor = np.ascontiguousarray(tensor) + + gain_tensor = gain_tensor.astype(np.float32, copy=False) + if not gain_tensor.flags.c_contiguous: + gain_tensor = np.ascontiguousarray(gain_tensor) + + c, h, w = tensor.shape + + if _HAS_NUMBA: + _apply_tensor_flat_gain_chw_numba( + tensor, + gain_tensor, + int(c), + int(h), + int(w), + bool(clip_output), + ) + return tensor + + np.multiply(tensor, gain_tensor, out=tensor) + if clip_output: + np.clip(tensor, 0.0, 1.0, out=tensor) + + return tensor + + def _get_final_flat_gain_tensor_cached_fast( + self, + decoded: dict, + role_to_cam: dict, + geom: dict, + remap_cache: dict, + crop_box: tuple, + target_size: tuple, + ): + cfg = self.flatfield_config or {} + + if not cfg.get("enabled", False): + return None + + if str(cfg.get("apply_space", "native_camera_space")).lower() != "final_tensor_space": + return None + + if not self.flatfield_loaded: + self.load_flatfield_maps() + + if not self.flatfield_loaded: + return None + + target_w, target_h = int(target_size[0]), int(target_size[1]) + + strength = float(cfg.get("strength", 1.0)) + strength_by_channel = cfg.get("strength_by_channel", {}) or {} + + gain_min_runtime = float(cfg.get("gain_min_runtime", 0.0)) + gain_max_runtime = float(cfg.get("gain_max_runtime", 999.0)) + + runtime_smooth_ksize = int(cfg.get("runtime_smooth_ksize", 0) or 0) + if runtime_smooth_ksize >= 3 and runtime_smooth_ksize % 2 == 0: + runtime_smooth_ksize += 1 + + key = ( + "final_flat_gain_tensor", + geom.get("key"), + tuple(int(v) for v in crop_box), + int(target_w), + int(target_h), + int(runtime_smooth_ksize), + round(float(strength), 6), + tuple(sorted((str(k), round(float(v), 6)) for k, v in strength_by_channel.items())), + round(float(gain_min_runtime), 6), + round(float(gain_max_runtime), 6), + ) + + cached = self._flatfield_runtime_cache.get(key) + if cached is not None: + return cached + + gain_tensor = np.ones((5, target_h, target_w), dtype=np.float32) + + # ------------------------------------------------------------ + # RGB: usa crop + resize, igual ao RGB real. + # ------------------------------------------------------------ + rgb_cam = role_to_cam.get("rgb") + if rgb_cam is not None: + rgb_img = decoded[rgb_cam]["image"] + rgb_shape = rgb_img.shape[:2] + + x0, y0, x1, y1 = [int(v) for v in crop_box] + + for ci, ch in enumerate(("R", "G", "B")): + gain_eff = self._get_runtime_gain_eff_map(ch, rgb_shape, cfg) + + if gain_eff is None: + continue + + gain_crop = gain_eff[y0:y1, x0:x1].astype(np.float32, copy=False) + + if gain_crop.shape[1] != target_w or gain_crop.shape[0] != target_h: + gain_out = cv2.resize( + gain_crop, + (target_w, target_h), + interpolation=cv2.INTER_LINEAR, + ) + else: + gain_out = gain_crop + + gain_tensor[ci] = gain_out.astype(np.float32, copy=False) + + # ------------------------------------------------------------ + # RE/NIR: usa o mesmo remap cacheado da imagem real. + # ------------------------------------------------------------ + maps = (remap_cache or {}).get("maps", {}) or {} + + spec_map = { + "re": ("RE", 3), + "nir": ("NIR", 4), + } + + for role, (ch, ci) in spec_map.items(): + cam_id = role_to_cam.get(role) + if cam_id is None: + continue + + img = decoded[cam_id]["image"] + gain_eff = self._get_runtime_gain_eff_map(ch, img.shape[:2], cfg) + + if gain_eff is None: + continue + + remap_entry = maps.get(role) + + if remap_entry is not None: + gain_out = cv2.remap( + gain_eff.astype(np.float32, copy=False), + remap_entry["map1"], + remap_entry["map2"], + interpolation=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=1.0, + ) + else: + M_role_to_target = geom["M_role_to_target"].get(role) + + if M_role_to_target is None: + continue + + gain_out = cv2.warpPerspective( + gain_eff.astype(np.float32, copy=False), + M_role_to_target, + (target_w, target_h), + flags=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + borderValue=1.0, + ) + + gain_tensor[ci] = gain_out.astype(np.float32, copy=False) + + self._flatfield_runtime_cache[key] = gain_tensor + return gain_tensor + diff --git a/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core_bkp.py b/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core_bkp.py deleted file mode 100644 index 10db99e20..000000000 --- a/Python/OAK/datasets/oak-fcc-3/core/raw_processor_core_bkp.py +++ /dev/null @@ -1,1489 +0,0 @@ -import json -import os -import cv2 -import numpy as np -import math -from typing import Optional - - -class RawProcessorCore: - def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "GBRG", calibration_json_path=None): - self.sensor_width = sensor_width - self.sensor_height = sensor_height - self.bayer_pattern = bayer_pattern.upper() - - self.fusion_config = { - "alignment_mode": "manual_affine", - "baseline_mm": 75.0, - "manual_offsets": { - "nir": {"dx": 0, "dy": 0, "theta_deg": 0.0}, - "re": {"dx": 0, "dy": 0, "theta_deg": 0.0}, - }, - "homographies": { - "nir_to_rgb": None, - "re_to_rgb": None, - }, - "crop_valid_common": True, - "resize_after_crop": True, - "target_size": None, - } - - self.rgb_calibration = { - "enabled": False, - "gains": { - "R": 1.0, - "G": 1.0, - "B": 1.0 - } - } - - self.calibration_json_path = calibration_json_path - self.calibration_base_dir = os.path.dirname(os.path.abspath(calibration_json_path)) if calibration_json_path else os.getcwd() - - self.flatfield_config = { - "enabled": False, - "npz_file": None, - "apply_before_fusion": True, - "apply_after_decode": True, - "apply_space": "native_camera_space", - "map_type": "gain", - "channels": ["R", "G", "B", "RE", "NIR"], - "channel_maps": {}, - "subtract_dark": False, - "clip_output": True, - } - self.flatfield_maps = {} - self.flatfield_loaded = False - - self.radiometric_normalization_config = { - "enabled": False, - "method": "exposure_gain_reference", - "reference_controls": {}, - "clip_output": False, - } - self.radiometric_config = {} - self.patch_normalization_config = { - "enabled": False, - "apply_when_metering_mode": "reference_patches", - "apply_stage": "after_fusion", - "method": "gray_scale_with_white_guard", - "space": "multispec_tensor", - "targets": { - "black": 0.06, - "gray": 0.40, - "white": 0.78, - }, - "white_guard_max": 0.92, - "scale_min": 0.35, - "scale_max": 2.50, - "clip_output": True, - "require_valid_gray": True, - "use_black_for_offset": False, - "save_patch_stats": True, - } - self.last_patch_normalization_result = None - self.camera_settings = {} - - if calibration_json_path: - self.load_config_json(calibration_json_path) - - def unpack_raw10_packed( - self, - packed_frame: np.ndarray, - sensor_width: Optional[int] = None, - sensor_height: Optional[int] = None - ): - if packed_frame.ndim == 3 and packed_frame.shape[2] == 1: - packed_frame = packed_frame[:, :, 0] - - width = sensor_width if sensor_width is not None else self.sensor_width - height = sensor_height if sensor_height is not None else self.sensor_height - - if width % 4 != 0: - raise ValueError(f"Largura {width} não é múltipla de 4 para RAW10 packed") - - expected_packed_width = math.ceil(width * 10 / 8) - - actual_h, actual_w = packed_frame.shape[:2] - padding = actual_w - expected_packed_width - - if actual_h != height: - raise ValueError( - f"[ERRO FRAME] Altura packed inesperada: {packed_frame.shape}, " - f"esperado altura={height}" - ) - - if actual_w < expected_packed_width: - raise ValueError( - f"[ERRO FRAME] Largura packed menor que a útil esperada: {packed_frame.shape}, " - f"esperado pelo menos ({height}, {expected_packed_width})" - ) - - if padding > 64: - raise ValueError( - f"[ERRO FRAME] Padding excessivo no packed: {packed_frame.shape}, " - f"esperado útil ({height}, {expected_packed_width}), padding={padding}" - ) - - packed_frame = packed_frame[:, :expected_packed_width] - groups = packed_frame.reshape(height, width // 4, 5).astype(np.uint16) - - b0 = groups[:, :, 0] - b1 = groups[:, :, 1] - b2 = groups[:, :, 2] - b3 = groups[:, :, 3] - b4 = groups[:, :, 4] - - p0 = (b0 << 2) | ((b4 >> 0) & 0x03) - p1 = (b1 << 2) | ((b4 >> 2) & 0x03) - p2 = (b2 << 2) | ((b4 >> 4) & 0x03) - p3 = (b3 << 2) | ((b4 >> 6) & 0x03) - - raw16 = np.empty((height, width), dtype=np.uint16) - raw16[:, 0::4] = p0 - raw16[:, 1::4] = p1 - raw16[:, 2::4] = p2 - raw16[:, 3::4] = p3 - - return raw16 - - def extract_bayer_channels(self, raw16: np.ndarray) -> dict: - p = self.bayer_pattern - - if p == "GBRG": - g1 = raw16[0::2, 0::2] - b = raw16[0::2, 1::2] - r = raw16[1::2, 0::2] - g2 = raw16[1::2, 1::2] - elif p == "GRBG": - g1 = raw16[0::2, 0::2] - r = raw16[0::2, 1::2] - b = raw16[1::2, 0::2] - g2 = raw16[1::2, 1::2] - elif p == "RGGB": - b = raw16[0::2, 0::2] - g1 = raw16[0::2, 1::2] - g2 = raw16[1::2, 0::2] - r = raw16[1::2, 1::2] - elif p == "BGGR": - b = raw16[0::2, 0::2] - g1 = raw16[0::2, 1::2] - g2 = raw16[1::2, 0::2] - r = raw16[1::2, 1::2] - else: - raise ValueError(f"Padrão Bayer não suportado: {p}") - - return {"R": r, "G1": g1, "G2": g2, "B": b} - - def build_training_rgb( - self, - raw16: np.ndarray, - output_dtype: str = "float32", - bit_depth: int = 10, - ) -> np.ndarray: - ch = self.extract_bayer_channels(raw16) - - max_val = float((1 << bit_depth) - 1) - - r = ch["R"].astype(np.float32) / max_val - g = ((ch["G1"].astype(np.float32) + ch["G2"].astype(np.float32)) * 0.5) / max_val - b = ch["B"].astype(np.float32) / max_val - - rgb_cal = getattr(self, "rgb_calibration", {}) or {} - if rgb_cal.get("enabled", False): - gains = rgb_cal.get("gains", {}) or {} - r *= float(gains.get("R", 1.0)) - g *= float(gains.get("G", 1.0)) - b *= float(gains.get("B", 1.0)) - - chw = np.stack([r, g, b], axis=0).astype(np.float32) - chw = np.clip(chw, 0.0, 1.0) - - if output_dtype == "float32": - return chw - - if output_dtype == "uint8": - return (chw * 255.0).clip(0, 255).astype(np.uint8) - - if output_dtype == "uint16": - return (chw * 65535.0).clip(0, 65535).astype(np.uint16) - - raise ValueError(f"output_dtype não suportado: {output_dtype}") - - def _channel_names_from_decoded(self, decoded): - names = ["R", "G", "B"] - - roles = { - data.get("role") or data.get("meta", {}).get("role"): cam_id - for cam_id, data in decoded.items() - } - - if "re" in roles: - names.append("RE") - if "nir" in roles: - names.append("NIR") - - return names - - def _find_cam_by_role(self, decoded, role): - role = str(role).lower() - - for cam_id, data in decoded.items(): - data_role = ( - data.get("role") or - data.get("meta", {}).get("role") or - "" - ) - if str(data_role).lower() == role: - return cam_id - - return None - - def decode_bins_cameras(self, bins_data, bins_meta): - decoded = {} - - for data, meta in zip(bins_data, bins_meta): - role = (meta.get("role") or "").strip().lower() - bit_depth = int(meta.get("bit_depth", 10)) - max_val = float((1 << bit_depth) - 1) - - cam_id = meta.get("cam_id") or meta.get("camera_id") or meta.get("id") or role - - if role == "rgb": - decoded[cam_id] = { - "name": "RGB", - "image": data.astype(np.float32) / max_val, - "meta": meta, - } - - elif role == "re": - decoded[cam_id] = { - "name": "RE", - "image": data.astype(np.float32) / max_val, - "meta": meta, - } - - elif role == "nir": - decoded[cam_id] = { - "name": "NIR", - "image": data.astype(np.float32) / max_val, - "meta": meta, - } - decoded[cam_id]["role"] = role - - return decoded - - def build_multispectral_tensor(self, bins_data, bins_meta, target_size=None): - decoded = self.decode_bins_cameras(bins_data, bins_meta) - - rgb_cam_id = self._find_cam_by_role(decoded, "rgb") - if rgb_cam_id is None: - raise RuntimeError("RGB obrigatório") - - channel_names = self._channel_names_from_decoded(decoded) - tensor = self.fuse_multispec_cameras(decoded, meta=None, channels_expected=len(channel_names)) - tensor = self.resize_tensor_chw(tensor, target_size=target_size) - tensor = self.apply_patch_normalization_to_tensor(tensor) - - return tensor, channel_names - - def build_infer_tensor_from_stream(self, frame, meta, channels_expected, target_size=None): - frame_type = meta.get("frame_type") - - if frame_type == "RAW_BRUTO": - decoded = self.decode_stream_cameras(frame, meta) - tensor = self.fuse_multispec_cameras(decoded, meta, channels_expected) - tensor = self.resize_tensor_chw(tensor, target_size=target_size) - tensor = self.apply_patch_normalization_to_tensor(tensor) - return tensor - - if frame_type in ("RGB", "MULTISPEC"): - if not isinstance(frame, np.ndarray): - raise RuntimeError(f"Frame {frame_type} esperado como ndarray") - - if frame.ndim != 3: - raise RuntimeError(f"Frame {frame_type} inválido: shape={frame.shape}") - - if dtype_str == "uint8": - raw_np = frame.astype(np.float32) / 255.0 - elif dtype_str == "float32": - raw_np = frame.astype(np.float32) - elif dtype_str == "uint16": - raw_np = frame.astype(np.float32) / 65535.0 - else: - raise RuntimeError(f"dtype {frame_type} não suportado: {dtype_str}") - - if raw_np.shape[0] != channels_expected: - raise RuntimeError(f"Frame {frame_type} com canais inesperados: {raw_np.shape[0]} | esperado={channels_expected}") - tensor = raw_np - - tensor = self.resize_tensor_chw(tensor, target_size=target_size) - return tensor - - raise RuntimeError(f"frame_type não suportado para inferência: {frame_type}") - - def decode_stream_cameras(self, frame, meta): - if not isinstance(frame, dict): - raise RuntimeError("RAW_BRUTO esperado como dict de câmeras no modo multi") - - camera_frames = meta.get("camera_frames", {}) or {} - camera_info = meta.get("camera_info", {}) or {} - - decoded = {} - - for cam_id, data in frame.items(): - cam_meta = camera_frames.get(cam_id) or camera_info.get(cam_id) or {} - - role = str(cam_meta.get("role", "")).lower() - if not role: - raise RuntimeError(f"Meta da câmera {cam_id} sem role. Esperado role='rgb', 'nir' ou 're'.") - - bit_depth = int(cam_meta.get("bit_depth", 10)) - raw_format = str(cam_meta.get("raw_format", "")).upper() - is_raw10 = raw_format == "RAW10_PACKED" or bit_depth == 10 - - if role == "rgb": - # Caso OAK RAW real: câmera RGB também vem RAW10 packed - if is_raw10 and data.ndim == 2: - sensor_w = int(cam_meta.get("width", self.sensor_width)) - sensor_h = int(cam_meta.get("height", self.sensor_height)) - - raw16 = self.unpack_raw10_packed( - data, - sensor_width=sensor_w, - sensor_height=sensor_h, - ) - - rgb_chw = self.build_training_rgb( - raw16, - output_dtype="float32", - bit_depth=bit_depth, - ) - - rgb_hwc = np.transpose(rgb_chw, (1, 2, 0)) - - decoded[cam_id] = { - "name": "RGB", - "role": "rgb", - "image": rgb_hwc, - "meta": cam_meta, - } - - else: - # Caso preview/processado antigo: BGR HWC uint8 - if data.ndim != 3 or data.shape[2] != 3: - raise RuntimeError(f"{cam_id} RGB inválida: shape={data.shape}") - - rgb = data[:, :, ::-1].astype(np.float32) / 255.0 - - decoded[cam_id] = { - "name": "RGB", - "role": "rgb", - "image": np.clip(rgb, 0.0, 1.0), - "meta": cam_meta, - } - - elif role == "re": - decoded[cam_id] = { - "name": "RE", - "role": "re", - "image": self._decode_spectral_frame_to_float01(data, cam_meta), - "meta": cam_meta, - } - - elif role == "nir": - decoded[cam_id] = { - "name": "NIR", - "role": "nir", - "image": self._decode_spectral_frame_to_float01(data, cam_meta), - "meta": cam_meta, - } - - return decoded - - def _decode_spectral_frame_to_float01(self, data, cam_meta): - arr = data - - if arr.ndim == 3 and arr.shape[2] == 1: - arr = arr[:, :, 0] - - bit_depth = int(cam_meta.get("bit_depth", 8)) - raw_format = str(cam_meta.get("raw_format", "")).upper() - packed = bool(cam_meta.get("packed", False)) - channels = int(cam_meta.get("channels", 1)) if cam_meta.get("channels") is not None else 1 - - sensor_width = int(cam_meta.get("width", self.sensor_width)) - sensor_height = int(cam_meta.get("height", arr.shape[0])) - packed_width = int(cam_meta.get("packed_width", 0) or 0) - - looks_like_raw10_packed = ( - arr.ndim == 2 - and arr.dtype == np.uint8 - and arr.shape[0] == sensor_height - and ( - raw_format == "RAW10_PACKED" - or packed - or bit_depth == 10 - or (packed_width > 0 and arr.shape[1] == packed_width and packed_width != sensor_width) - or arr.shape[1] == int(sensor_width * 10 / 8) - ) - ) - - if looks_like_raw10_packed: - raw16 = self.unpack_raw10_packed( - arr, - sensor_width=sensor_width, - sensor_height=sensor_height, - ) - - max_val = float((1 << bit_depth) - 1) - return np.clip(raw16.astype(np.float32) / max_val, 0.0, 1.0) - - # Caso preview/processado: mono já vem uint8 normal. - if arr.ndim == 2 and arr.dtype == np.uint8: - return np.clip(arr.astype(np.float32) / 255.0, 0.0, 1.0) - - if arr.ndim == 2 and arr.dtype == np.uint16: - max_val = float((1 << bit_depth) - 1) if bit_depth > 0 and bit_depth <= 16 else 65535.0 - return np.clip(arr.astype(np.float32) / max_val, 0.0, 1.0) - - arr01 = arr.astype(np.float32) - if arr01.max() > 1.5: - arr01 /= 255.0 - - return np.clip(arr01, 0.0, 1.0) - - def fuse_multispec_cameras(self, decoded, meta, channels_expected): - decoded = self.apply_dark_to_decoded(decoded) - decoded = self.normalize_decoded_by_capture_controls(decoded, meta) - decoded = self.apply_flat_gain_to_decoded(decoded) - - rgb_cam_id = self._find_cam_by_role(decoded, "rgb") - if rgb_cam_id is None: - raise RuntimeError("Fusão requer câmera com role='rgb' como referência") - - rgb = decoded[rgb_cam_id]["image"] - h, w = rgb.shape[:2] - - rgb_chw = np.transpose(rgb, (2, 0, 1)) - channels = [rgb_chw] - names = ["R", "G", "B"] - - valid_masks = [np.ones((h, w), dtype=np.uint8)] - - role_to_cam = { - item.get("role", data.get("meta", {}).get("role")): cam_id - for cam_id, data in decoded.items() - for item in [data] - } - - for role, ch_name in (("re", "RE"), ("nir", "NIR")): - cam_id = role_to_cam.get(role) - if cam_id is None: - continue - - img = decoded[cam_id]["image"] - aligned, valid_mask = self._warp_with_valid_mask(img, role, (h, w), meta) - - channels.append(aligned[None, :, :]) - names.append(ch_name) - valid_masks.append(valid_mask) - - cfg = self.fusion_config - if cfg.get("crop_valid_common", False): - crop_box = self._compute_common_crop_box(valid_masks) - if crop_box is not None: - channels = self._crop_and_resize_channels(channels, crop_box, (h, w)) - - tensor = np.concatenate(channels, axis=0) - - if tensor.shape[0] != channels_expected: - raise RuntimeError( - f"Tensor fundido com canais inesperados: {tensor.shape[0]} | " - f"esperado={channels_expected} | got={names}" - ) - - return tensor.astype(np.float32, copy=False) - - def _shift_image(self, img, dx, dy): - h, w = img.shape[:2] - M = np.float32([[1, 0, dx], [0, 1, dy]]) - return cv2.warpAffine( - img, M, (w, h), - flags=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0 - ) - - def _affine_image(self, img, dx, dy, theta_deg): - h, w = img.shape[:2] - center = (w * 0.5, h * 0.5) - - M = cv2.getRotationMatrix2D(center, theta_deg, 1.0) - M[0, 2] += dx - M[1, 2] += dy - - return cv2.warpAffine( - img, - M, - (w, h), - flags=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0 - ) - - def _warp_with_valid_mask(self, img, role, ref_shape, meta): - ref_h, ref_w = ref_shape - - if img.shape[:2] != (ref_h, ref_w): - img = cv2.resize(img, (ref_w, ref_h), interpolation=cv2.INTER_LINEAR) - - cfg = self.fusion_config - mode = cfg.get("alignment_mode", "identity") - - mask = np.ones((ref_h, ref_w), dtype=np.uint8) * 255 - - if mode == "identity": - warped = img - warped_mask = mask - - elif mode == "manual_offset": - offs = cfg.get("manual_offsets", {}).get(role, {}) - dx = int(offs.get("dx", 0)) - dy = int(offs.get("dy", 0)) - - warped = self._shift_image(img, dx, dy) - warped_mask = self._shift_image(mask, dx, dy) - - elif mode == "manual_affine": - offs = cfg.get("manual_offsets", {}).get(role, {}) - dx = int(offs.get("dx", 0)) - dy = int(offs.get("dy", 0)) - theta_deg = float(offs.get("theta_deg", 0.0)) - - warped = self._affine_image(img, dx, dy, theta_deg) - warped_mask = self._affine_image(mask, dx, dy, theta_deg) - - elif mode == "homography": - H = cfg.get("homographies", {}).get(f"{role}_to_rgb") - - if H is None: - warped = img - warped_mask = mask - else: - H = np.asarray(H, dtype=np.float32) - - if H.shape != (3, 3): - raise RuntimeError(f"Homografia inválida para {role}: shape={H.shape}") - - warped = cv2.warpPerspective( - img, H, (ref_w, ref_h), - flags=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0 - ) - - warped_mask = cv2.warpPerspective( - mask, H, (ref_w, ref_h), - flags=cv2.INTER_NEAREST, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0 - ) - - else: - raise RuntimeError(f"alignment_mode inválido: {mode}") - - warped_mask = (warped_mask > 0).astype(np.uint8) - return warped, warped_mask - - def _compute_common_crop_box(self, masks): - if not masks: - return None - - common = masks[0].copy() - for m in masks[1:]: - common = np.logical_and(common > 0, m > 0) - - ys, xs = np.where(common) - if len(xs) == 0 or len(ys) == 0: - return None - - x0 = int(xs.min()) - x1 = int(xs.max()) + 1 - y0 = int(ys.min()) - y1 = int(ys.max()) + 1 - - return x0, y0, x1, y1 - - def _crop_and_resize_channels(self, channels, crop_box, ref_shape): - x0, y0, x1, y1 = crop_box - ref_h, ref_w = ref_shape - - cropped = [ch[:, y0:y1, x0:x1] for ch in channels] - - cfg = self.fusion_config - if not cfg.get("resize_after_crop", False): - return cropped - - target_size = cfg.get("target_size", None) - if target_size is None: - target_w, target_h = ref_w, ref_h - else: - target_w, target_h = target_size - - resized = [] - for ch in cropped: - ch_resized = np.stack([ - cv2.resize( - ch_i, - (target_w, target_h), - interpolation=cv2.INTER_LINEAR - ) - for ch_i in ch - ], axis=0) - resized.append(ch_resized) - - return resized - - def resize_tensor_chw(self, tensor, target_size=None): - if target_size is None: - return tensor - - target_w, target_h = target_size - - if tensor.ndim != 3: - raise RuntimeError(f"Tensor esperado em CHW. Veio shape={tensor.shape}") - - _, h, w = tensor.shape - - if (w, h) == (target_w, target_h): - return tensor.astype(np.float32, copy=False) - - interp = cv2.INTER_AREA if target_w < w or target_h < h else cv2.INTER_LINEAR - - chans = [] - for ch in tensor: - ch_res = cv2.resize(ch, (target_w, target_h), interpolation=interp) - chans.append(ch_res.astype(np.float32)) - - return np.stack(chans, axis=0) - - def apply_patch_normalization_to_tensor(self, tensor: np.ndarray) -> np.ndarray: - self.last_patch_normalization_result = None - - cfg = self.patch_normalization_config or {} - - result = { - "enabled": bool(cfg.get("enabled", False)), - "applied": False, - "method": cfg.get("method", "gray_scale_with_white_guard"), - "space": cfg.get("space", "multispec_tensor"), - "warnings": [], - "scales": {}, - "patch_stats": {}, - } - - if not cfg.get("enabled", False): - result["warnings"].append("patch_normalization_disabled") - self.last_patch_normalization_result = result - return tensor - - rad_cfg = self.radiometric_config or {} - - if cfg.get("apply_when_metering_mode") == "reference_patches": - if rad_cfg.get("metering_mode") != "reference_patches": - result["warnings"].append( - f"metering_mode_not_reference_patches: {rad_cfg.get('metering_mode')}" - ) - self.last_patch_normalization_result = result - return tensor - - if tensor is None or tensor.ndim != 3 or tensor.shape[0] < 5: - result["warnings"].append(f"invalid_tensor_shape: {None if tensor is None else tensor.shape}") - self.last_patch_normalization_result = result - return tensor - - patches = rad_cfg.get("reference_patches", []) or [] - patch_by_type = { - str(p.get("type", "")).lower(): p - for p in patches - if isinstance(p, dict) - } - - gray = patch_by_type.get("gray") - white = patch_by_type.get("white") - black = patch_by_type.get("black") - - if gray is None: - result["warnings"].append("missing_gray_patch") - if cfg.get("require_valid_gray", True): - self.last_patch_normalization_result = result - return tensor - - targets = cfg.get("targets", {}) or {} - gray_target = float(targets.get("gray", 0.40)) - - scale_min = float(cfg.get("scale_min", 0.35)) - scale_max = float(cfg.get("scale_max", 2.50)) - white_guard_max = float(cfg.get("white_guard_max", 0.92)) - clip_output = bool(cfg.get("clip_output", True)) - - channel_names = ["R", "G", "B", "RE", "NIR"] - - out = tensor.astype(np.float32).copy() - h, w = out.shape[1], out.shape[2] - - def roi_from_patch(patch): - if not patch: - return None - return self._roi_pct_to_pixels_from_patch(patch.get("roi_pct", {}) or {}, w, h) - - gray_roi = roi_from_patch(gray) - white_roi = roi_from_patch(white) - black_roi = roi_from_patch(black) - - if gray_roi is None: - result["warnings"].append("invalid_gray_roi") - self.last_patch_normalization_result = result - return tensor - - for ci, ch_name in enumerate(channel_names): - ch = out[ci] - - # ----------------------------- - # Stats do gray - # ----------------------------- - gx0, gy0, gx1, gy1 = gray_roi - gray_vals = ch[gy0:gy1, gx0:gx1].reshape(-1) - - if gray_vals.size <= 0: - result["warnings"].append(f"{ch_name}: empty_gray_roi") - continue - - gray_p50 = float(np.percentile(gray_vals, 50)) - gray_p05 = float(np.percentile(gray_vals, 5)) - gray_p95 = float(np.percentile(gray_vals, 95)) - gray_sat = float((gray_vals >= 0.98).mean() * 100.0) - gray_dark = float((gray_vals <= 0.02).mean() * 100.0) - - result["patch_stats"].setdefault("gray", {})[ch_name] = { - "p05": gray_p05, - "p50": gray_p50, - "p95": gray_p95, - "sat_pct": gray_sat, - "dark_pct": gray_dark, - "roi_px": list(gray_roi), - } - - if gray_p50 <= 1e-6: - result["warnings"].append(f"{ch_name}: gray_p50_too_low") - continue - - scale = gray_target / gray_p50 - - # ----------------------------- - # Stats do white + guarda - # ----------------------------- - if white_roi is not None: - wx0, wy0, wx1, wy1 = white_roi - white_vals = ch[wy0:wy1, wx0:wx1].reshape(-1) - - if white_vals.size > 0: - white_p50 = float(np.percentile(white_vals, 50)) - white_p05 = float(np.percentile(white_vals, 5)) - white_p95 = float(np.percentile(white_vals, 95)) - white_sat = float((white_vals >= 0.98).mean() * 100.0) - white_dark = float((white_vals <= 0.02).mean() * 100.0) - - result["patch_stats"].setdefault("white", {})[ch_name] = { - "p05": white_p05, - "p50": white_p50, - "p95": white_p95, - "sat_pct": white_sat, - "dark_pct": white_dark, - "roi_px": list(white_roi), - } - - if white_sat > 0.5: - result["warnings"].append(f"{ch_name}: white_patch_saturated_{white_sat:.2f}%") - - if white_p50 > 1e-6: - max_scale_by_white = white_guard_max / white_p50 - if scale > max_scale_by_white: - result["warnings"].append( - f"{ch_name}: scale_limited_by_white_guard " - f"{scale:.3f}->{max_scale_by_white:.3f}" - ) - scale = min(scale, max_scale_by_white) - - # ----------------------------- - # Stats do black, só diagnóstico - # ----------------------------- - if black_roi is not None: - bx0, by0, bx1, by1 = black_roi - black_vals = ch[by0:by1, bx0:bx1].reshape(-1) - - if black_vals.size > 0: - black_p50 = float(np.percentile(black_vals, 50)) - black_p05 = float(np.percentile(black_vals, 5)) - black_p95 = float(np.percentile(black_vals, 95)) - black_sat = float((black_vals >= 0.98).mean() * 100.0) - black_dark = float((black_vals <= 0.02).mean() * 100.0) - - result["patch_stats"].setdefault("black", {})[ch_name] = { - "p05": black_p05, - "p50": black_p50, - "p95": black_p95, - "sat_pct": black_sat, - "dark_pct": black_dark, - "roi_px": list(black_roi), - } - - scale_before_clip = float(scale) - scale = float(np.clip(scale, scale_min, scale_max)) - - if abs(scale - scale_before_clip) > 1e-6: - result["warnings"].append( - f"{ch_name}: scale_clipped {scale_before_clip:.3f}->{scale:.3f}" - ) - - out[ci] = ch * scale - - result["scales"][ch_name] = { - "scale": scale, - "gray_target": gray_target, - "gray_measured_p50": gray_p50, - } - - if clip_output: - out = np.clip(out, 0.0, 1.0) - - result["applied"] = True - result["valid"] = bool(len(result["scales"]) == len(channel_names)) - result["clip_output"] = clip_output - result["shape"] = list(out.shape) - result["channel_names"] = channel_names - - self.last_patch_normalization_result = result - return out.astype(np.float32, copy=False) - - def _roi_pct_to_pixels_from_patch(self, roi_pct: dict, w: int, h: int): - x0 = int(float(roi_pct.get("x0", 0.0)) * w) - y0 = int(float(roi_pct.get("y0", 0.0)) * h) - x1 = int(float(roi_pct.get("x1", 1.0)) * w) - y1 = int(float(roi_pct.get("y1", 1.0)) * h) - - x0 = max(0, min(w - 1, x0)) - x1 = max(x0 + 1, min(w, x1)) - y0 = max(0, min(h - 1, y0)) - y1 = max(y0 + 1, min(h, y1)) - - return x0, y0, x1, y1 - - - def extract_camera_meta(self, meta_json: dict, cam_id: str) -> dict: - cam_frames = meta_json.get("camera_frames", {}) or meta_json.get("stream_meta", {}).get("camera_frames", {}) - - cam = cam_frames.get(cam_id) - if not cam: - raise ValueError(f"Camera {cam_id} não encontrada no meta") - - # Detecta RAW10 packed mono - if int(cam.get("channels", 1)) == 1 and int(cam.get("bit_depth", 10)) == 10: - packed_width = int(cam.get("width")) - height = int(cam.get("height")) - - # 🔥 converte packed → real - real_width = int((packed_width * 8) / 10) - - return { - "camera_id": cam_id, - "width": real_width, - "height": height, - "channels": 1, - "bit_depth": 10, - "shape": [height, packed_width], # packed shape - "role": cam.get("role") - } - - # RGB - else: - width = int(cam.get("width")) - height = int(cam.get("height")) - channels = int(cam.get("channels", 3)) - - return { - "camera_id": cam_id, - "width": width, - "height": height, - "channels": channels, - "bit_depth": int(cam.get("bit_depth", 8)), - "shape": [height, width, channels], # 🔥 AQUI está a correção - "role": cam.get("role") - } - - # ============================================================ - # RAW10 PACKED - # ============================================================ - - def packed_width_for_raw10(self, sensor_width: int = None) -> int: - width = sensor_width if sensor_width is not None else self.sensor_width - return math.ceil(width * 10 / 8) - - def pack_raw10_packed(self, raw16: np.ndarray) -> np.ndarray: - h, w = raw16.shape - - if w % 4 != 0: - raise ValueError(f"Width precisa ser múltiplo de 4 para pack otimizado. Veio {w}") - - raw16 = np.clip(raw16, 0, 1023).astype(np.uint16) - - p0 = raw16[:, 0::4] - p1 = raw16[:, 1::4] - p2 = raw16[:, 2::4] - p3 = raw16[:, 3::4] - - b0 = (p0 >> 2).astype(np.uint8) - b1 = (p1 >> 2).astype(np.uint8) - b2 = (p2 >> 2).astype(np.uint8) - b3 = (p3 >> 2).astype(np.uint8) - - b4 = ( - ((p0 & 0x03) << 0) | - ((p1 & 0x03) << 2) | - ((p2 & 0x03) << 4) | - ((p3 & 0x03) << 6) - ).astype(np.uint8) - - packed = np.empty((h, w // 4, 5), dtype=np.uint8) - packed[:, :, 0] = b0 - packed[:, :, 1] = b1 - packed[:, :, 2] = b2 - packed[:, :, 3] = b3 - packed[:, :, 4] = b4 - - return packed.reshape(h, w // 4 * 5) - - def load_raw10_packed_file(self, path: str, width: int, height: int) -> np.ndarray: - packed_width = self.packed_width_for_raw10(width) - - expected_size = height * packed_width - actual_size = os.path.getsize(path) - - if actual_size != expected_size: - raise ValueError( - f"Tamanho inválido RAW10: {actual_size}, esperado {expected_size} em {path}" - ) - - packed = np.fromfile(path, dtype=np.uint8).reshape(height, packed_width) - return self.unpack_raw10_packed(packed, sensor_width=width, sensor_height=height) - - def save_raw10_packed_file(self, path: str, raw16: np.ndarray): - packed = self.pack_raw10_packed(raw16) - packed.tofile(path) - - # ============================================================ - # RGB UINT8 - # ============================================================ - - def load_rgb_u8_file(self, path: str, shape) -> np.ndarray: - arr = np.fromfile(path, dtype=np.uint8) - - expected = np.prod(shape) - if arr.size != expected: - raise ValueError( - f"Tamanho inválido RGB: {arr.size}, esperado {expected} em {path}" - ) - - return arr.reshape(shape) - - def save_rgb_u8_file(self, path: str, arr: np.ndarray): - arr.astype(np.uint8).tofile(path) - - - # ============================================================ - # DISPATCHER (O MAIS IMPORTANTE) - # ============================================================ - - def load_native_bin(self, path: str, cam_meta: dict) -> np.ndarray: - """ - Decide automaticamente como carregar o .bin baseado no meta. - """ - - channels = int(cam_meta.get("channels", 1)) - bit_depth = int(cam_meta.get("bit_depth", 10)) - shape = cam_meta.get("shape") - - if channels == 1 and bit_depth == 10: - width = int(cam_meta["width"]) - height = int(cam_meta["height"]) - return self.load_raw10_packed_file(path, width, height) - - elif channels == 3 and bit_depth == 8: - return self.load_rgb_u8_file(path, shape) - - else: - raise ValueError( - f"Formato não suportado: channels={channels}, bit_depth={bit_depth}" - ) - - - def save_native_bin(self, path: str, arr: np.ndarray, cam_meta: dict): - """ - Salva no formato correto baseado no meta. - """ - - channels = int(cam_meta.get("channels", 1)) - bit_depth = int(cam_meta.get("bit_depth", 10)) - - if channels == 1 and bit_depth == 10: - self.save_raw10_packed_file(path, arr) - - elif channels == 3 and bit_depth == 8: - self.save_rgb_u8_file(path, arr) - - else: - raise ValueError( - f"Formato não suportado para salvar: channels={channels}, bit_depth={bit_depth}" - ) - - - def load_config_json(self, path: str): - if not path or not os.path.isfile(path): - raise FileNotFoundError(f"Arquivo de calibração não encontrado: {path}") - - with open(path, "r", encoding="utf-8") as f: - data = json.load(f) - - fusion = data.get("fusion_config") - if isinstance(fusion, dict): - self.fusion_config = self._merge_config(self.fusion_config, fusion) - else: - print("[WARN] JSON sem fusion_config. Mantendo config padrão.") - - rgb_cal = data.get("rgb_calibration") - if isinstance(rgb_cal, dict): - self.rgb_calibration = self._merge_config(self.rgb_calibration, rgb_cal) - - flatfield = data.get("flatfield_config") - if isinstance(flatfield, dict): - self.flatfield_config = self._merge_config(self.flatfield_config, flatfield) - self.load_flatfield_maps() - else: - self.flatfield_config["enabled"] = False - self.flatfield_maps = {} - self.flatfield_loaded = False - - radiometric = data.get("radiometric_config") - if isinstance(radiometric, dict): - self.radiometric_config = self._merge_config(self.radiometric_config, radiometric) - - rad_norm_config = data.get("radiometric_normalization") - if isinstance(rad_norm_config, dict): - self.radiometric_normalization_config = self._merge_config(self.radiometric_normalization_config, rad_norm_config) - - patch_norm = data.get("patch_normalization") - if isinstance(patch_norm, dict): - self.patch_normalization_config = self._merge_config(self.patch_normalization_config, patch_norm) - - cam_set = data.get("camera_settings") - if isinstance(cam_set, dict): - self.camera_settings = self._merge_config(self.camera_settings, cam_set) - - def _merge_config(self, default_cfg: dict, loaded_cfg: dict) -> dict: - cfg = json.loads(json.dumps(default_cfg)) - - def merge(dst: dict, src: dict): - for key, value in src.items(): - if isinstance(value, dict) and isinstance(dst.get(key), dict): - merge(dst[key], value) - else: - dst[key] = value - - if isinstance(loaded_cfg, dict): - merge(cfg, loaded_cfg) - - return cfg - - - def _resolve_calibration_path(self, path: str) -> str: - if not path: - return "" - - path = str(path).replace("\\", "/") - - if os.path.isabs(path): - return path - - # Primeiro tenta relativo ao diretório de execução. - if os.path.isfile(path): - return path - - # Depois tenta relativo ao diretório do module_params.json. - candidate = os.path.join(self.calibration_base_dir, path) - if os.path.isfile(candidate): - return candidate - - # Por fim, se o path já começa com "calibration/", tenta relativo ao pai da pasta calibration. - base_parent = os.path.dirname(self.calibration_base_dir) - candidate = os.path.join(base_parent, path) - if os.path.isfile(candidate): - return candidate - - return path - - def load_flatfield_maps(self): - cfg = self.flatfield_config or {} - - if not cfg.get("enabled", False): - self.flatfield_maps = {} - self.flatfield_loaded = False - return False - - npz_file = cfg.get("npz_file") - if not npz_file: - print("[WARN] flatfield_config habilitado, mas sem npz_file.") - self.flatfield_maps = {} - self.flatfield_loaded = False - return False - - npz_path = self._resolve_calibration_path(npz_file) - - if not os.path.isfile(npz_path): - print(f"[WARN] Arquivo flat-field não encontrado: {npz_file} -> {npz_path}") - self.flatfield_maps = {} - self.flatfield_loaded = False - return False - - data = np.load(npz_path) - - maps = {} - channel_maps = cfg.get("channel_maps", {}) or {} - channels = cfg.get("channels", ["R", "G", "B", "RE", "NIR"]) - - for ch in channels: - ch = str(ch).upper() - ch_cfg = channel_maps.get(ch, {}) or {} - - gain_key = ch_cfg.get("gain_key", f"gain_{ch}") - dark_key = ch_cfg.get("dark_median_key", f"dark_median_{ch}") - - if gain_key not in data: - print(f"[WARN] Flat-field sem chave {gain_key} para canal {ch}.") - continue - - entry = { - "gain": data[gain_key].astype(np.float32), - "gain_key": gain_key, - } - - if dark_key and dark_key in data: - entry["dark"] = data[dark_key].astype(np.float32) - entry["dark_key"] = dark_key - - maps[ch] = entry - - self.flatfield_maps = maps - self.flatfield_loaded = len(maps) > 0 - - if self.flatfield_loaded: - print(f"[OK] Flat-field carregado: {npz_path} | canais={list(maps.keys())}") - else: - print(f"[WARN] Flat-field habilitado, mas nenhum mapa foi carregado: {npz_path}") - - return self.flatfield_loaded - - def apply_dark_to_decoded(self, decoded: dict) -> dict: - cfg = self.flatfield_config or {} - - if not cfg.get("enabled", False): - return decoded - - subtract_dark = bool(cfg.get("subtract_dark", True)) - if not subtract_dark: - return decoded - - if not self.flatfield_loaded: - self.load_flatfield_maps() - - if not self.flatfield_loaded: - return decoded - - corrected = {} - - for cam_id, item in decoded.items(): - role = str(item.get("role") or item.get("meta", {}).get("role") or "").lower() - img = item.get("image") - - if img is None: - corrected[cam_id] = item - continue - - new_item = dict(item) - new_meta = dict(item.get("meta", {}) or {}) - - if role == "rgb": - if img.ndim != 3 or img.shape[2] < 3: - corrected[cam_id] = item - continue - - out = img.astype(np.float32).copy() - - for idx, ch in enumerate(("R", "G", "B")): - out[:, :, idx] = self._subtract_dark_single_channel( - out[:, :, idx], - ch, - ) - - new_item["image"] = out - - elif role in ("re", "nir"): - ch = "RE" if role == "re" else "NIR" - - new_item["image"] = self._subtract_dark_single_channel( - img.astype(np.float32), - ch, - ) - - else: - corrected[cam_id] = item - continue - - new_meta["dark_applied"] = True - new_item["meta"] = new_meta - corrected[cam_id] = new_item - - return corrected - - def apply_flat_gain_to_decoded(self, decoded: dict) -> dict: - cfg = self.flatfield_config or {} - - if not cfg.get("enabled", False): - return decoded - - if not self.flatfield_loaded: - self.load_flatfield_maps() - - if not self.flatfield_loaded: - return decoded - - clip_output = bool(cfg.get("clip_output", True)) - corrected = {} - - for cam_id, item in decoded.items(): - role = str(item.get("role") or item.get("meta", {}).get("role") or "").lower() - img = item.get("image") - - if img is None: - corrected[cam_id] = item - continue - - new_item = dict(item) - new_meta = dict(item.get("meta", {}) or {}) - - if role == "rgb": - if img.ndim != 3 or img.shape[2] < 3: - corrected[cam_id] = item - continue - - out = img.astype(np.float32).copy() - - for idx, ch in enumerate(("R", "G", "B")): - out[:, :, idx] = self._apply_flat_gain_single_channel( - out[:, :, idx], - ch, - clip_output=clip_output, - ) - - new_item["image"] = out - - elif role in ("re", "nir"): - ch = "RE" if role == "re" else "NIR" - - new_item["image"] = self._apply_flat_gain_single_channel( - img.astype(np.float32), - ch, - clip_output=clip_output, - ) - - else: - corrected[cam_id] = item - continue - - new_meta["flatfield_applied"] = True - new_meta["flatfield_map_type"] = cfg.get("map_type", "gain") - new_item["meta"] = new_meta - corrected[cam_id] = new_item - - return corrected - - def _subtract_dark_single_channel( - self, - img: np.ndarray, - channel_name: str, - ) -> np.ndarray: - ch = str(channel_name).upper() - entry = self.flatfield_maps.get(ch) - - if not entry: - return img.astype(np.float32, copy=False) - - dark = entry.get("dark") - if dark is None: - return img.astype(np.float32, copy=False) - - base = img.astype(np.float32) - - dark = dark.astype(np.float32) - if dark.shape[:2] != base.shape[:2]: - dark = cv2.resize( - dark, - (base.shape[1], base.shape[0]), - interpolation=cv2.INTER_LINEAR, - ) - - out = np.maximum(base - dark, 0.0) - return out.astype(np.float32, copy=False) - - def _apply_flat_gain_single_channel( - self, - img: np.ndarray, - channel_name: str, - clip_output: bool = True, - ) -> np.ndarray: - ch = str(channel_name).upper() - entry = self.flatfield_maps.get(ch) - - if not entry: - return img.astype(np.float32, copy=False) - - gain = entry.get("gain") - if gain is None: - return img.astype(np.float32, copy=False) - - base = img.astype(np.float32) - - gain = gain.astype(np.float32) - if gain.shape[:2] != base.shape[:2]: - gain = cv2.resize( - gain, - (base.shape[1], base.shape[0]), - interpolation=cv2.INTER_LINEAR, - ) - - out = base * gain - - if clip_output: - out = np.clip(out, 0.0, 1.0) - - return out.astype(np.float32, copy=False) - - - def normalize_decoded_by_capture_controls(self, decoded: dict, meta: dict | None = None) -> dict: - cfg = self.radiometric_normalization_config or {} - - if not cfg.get("enabled", False): - return decoded - - method = str(cfg.get("method", "exposure_gain_reference")).lower() - if method != "exposure_gain_reference": - return decoded - - controls = self._extract_actual_controls_from_meta(meta) - if not controls: - return decoded - - reference_controls = cfg.get("reference_controls", {}) or {} - clip_output = bool(cfg.get("clip_output", False)) - - normalized = {} - - for cam_id, item in decoded.items(): - role = str(item.get("role") or item.get("meta", {}).get("role") or "").lower() - img = item.get("image") - - if img is None or not role: - normalized[cam_id] = item - continue - - actual_ctrl = controls.get(role, {}) or {} - - ref_ctrl = ( - reference_controls.get(role) - or self.camera_settings.get(role) - or actual_ctrl - or {} - ) - - actual_factor = self._exposure_gain_factor(actual_ctrl) - ref_factor = self._exposure_gain_factor(ref_ctrl) - - if actual_factor <= 0 or ref_factor <= 0: - normalized[cam_id] = item - continue - - scale = ref_factor / actual_factor - - new_item = dict(item) - new_meta = dict(item.get("meta", {}) or {}) - - out = img.astype(np.float32) * float(scale) - - if clip_output: - out = np.clip(out, 0.0, 1.0) - - new_meta["radiometric_normalization_applied"] = True - new_meta["radiometric_normalization_method"] = method - new_meta["radiometric_normalization_scale"] = float(scale) - new_meta["radiometric_actual_factor"] = float(actual_factor) - new_meta["radiometric_reference_factor"] = float(ref_factor) - - new_item["image"] = out.astype(np.float32, copy=False) - new_item["meta"] = new_meta - - normalized[cam_id] = new_item - - return normalized - - def _extract_actual_controls_from_meta(self, meta: dict | None) -> dict: - if not meta: - return {} - - for key in ("actual_camera_controls", "camera_controls", "startup_camera_controls"): - controls = meta.get(key) - if isinstance(controls, dict) and controls: - return controls - - stream_meta = meta.get("stream_meta") - if isinstance(stream_meta, dict): - for key in ("actual_camera_controls", "camera_controls", "startup_camera_controls"): - controls = stream_meta.get(key) - if isinstance(controls, dict) and controls: - return controls - - return {} - - def _exposure_gain_factor(self, ctrl: dict) -> float: - if not isinstance(ctrl, dict): - return 0.0 - - exp = ctrl.get("exposure_time_us", None) - gain = ctrl.get("analogue_gain", None) - - try: - exp = float(exp) - except Exception: - exp = 0.0 - - try: - gain = float(gain) - except Exception: - gain = 1.0 - - if exp <= 0: - return 0.0 - - if gain <= 0: - gain = 1.0 - - return float(exp * gain) diff --git a/Python/OAK/datasets/oak-fcc-3/core/segformer_service.py b/Python/OAK/datasets/oak-fcc-3/core/segformer_service.py new file mode 100644 index 000000000..d5de9f8c5 --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/core/segformer_service.py @@ -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 + diff --git a/Python/OAK/datasets/oak-fcc-3/core/test.json b/Python/OAK/datasets/oak-fcc-3/core/test.json new file mode 100644 index 000000000..9acd6786c --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/core/test.json @@ -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 + } + } + + } \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/dataset/labelmap.txt b/Python/OAK/datasets/oak-fcc-3/dataset/labelmap.txt new file mode 100644 index 000000000..ffe55badf --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/dataset/labelmap.txt @@ -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:: \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/tests/test.py b/Python/OAK/datasets/oak-fcc-3/tests/test.py index 77a9ad01c..6b5535602 100644 --- a/Python/OAK/datasets/oak-fcc-3/tests/test.py +++ b/Python/OAK/datasets/oak-fcc-3/tests/test.py @@ -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() \ No newline at end of file +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() \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/tests/test_oak_fcc3_aligned_geometry.py b/Python/OAK/datasets/oak-fcc-3/tests/test_oak_fcc3_aligned_geometry.py new file mode 100644 index 000000000..261405df0 --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/tests/test_oak_fcc3_aligned_geometry.py @@ -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 só 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() \ No newline at end of file diff --git a/Python/OAK/datasets/oak-fcc-3/utils/check_saved_files.py b/Python/OAK/datasets/oak-fcc-3/utils/check_saved_files.py index 28e4566db..6a8ca1c7e 100644 --- a/Python/OAK/datasets/oak-fcc-3/utils/check_saved_files.py +++ b/Python/OAK/datasets/oak-fcc-3/utils/check_saved_files.py @@ -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", diff --git a/Python/OAK/datasets/oak-fcc-3/utils/flat_rgb_calibration_tuner.py b/Python/OAK/datasets/oak-fcc-3/utils/flat_rgb_calibration_tuner.py new file mode 100644 index 000000000..67e874d2e --- /dev/null +++ b/Python/OAK/datasets/oak-fcc-3/utils/flat_rgb_calibration_tuner.py @@ -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 lê: + /bins/_CAM_A.bin + /bins/_CAM_B.bin + /bins/_CAM_C.bin + /metas/.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() diff --git a/Python/OAK/datasets/oak-fcc-3/utils/manual_fusion_calibrator.py b/Python/OAK/datasets/oak-fcc-3/utils/manual_fusion_calibrator.py index 3731bdce3..10bb99124 100644 --- a/Python/OAK/datasets/oak-fcc-3/utils/manual_fusion_calibrator.py +++ b/Python/OAK/datasets/oak-fcc-3/utils/manual_fusion_calibrator.py @@ -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}" diff --git a/Python/raspi/Instalar Sistema.txt b/Python/raspi/Instalar Sistema.txt new file mode 100644 index 000000000..51abc76b1 --- /dev/null +++ b/Python/raspi/Instalar Sistema.txt @@ -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 + diff --git a/Python/raspi/cam_3/multispectral_client.py b/Python/raspi/cam_3/multispectral_client.py index ea1304f87..ae4cc4346 100644 --- a/Python/raspi/cam_3/multispectral_client.py +++ b/Python/raspi/cam_3/multispectral_client.py @@ -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)