python
This commit is contained in:
parent
eeb92bbbc0
commit
26aebaf070
|
|
@ -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()
|
||||
|
|
@ -0,0 +1,653 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
corridor_state_labeler.py
|
||||
-------------------------------------------------
|
||||
Rotulador rápido de estado global do corredor para dataset de segmentação.
|
||||
|
||||
Objetivo
|
||||
- Percorrer uma raiz de dataset no formato:
|
||||
|
||||
dataset_root/
|
||||
├── grupo_1/
|
||||
│ ├── images/
|
||||
│ ├── masks/
|
||||
│ └── labels/ # criado automaticamente
|
||||
├── grupo_2/
|
||||
│ ├── images/
|
||||
│ ├── masks/
|
||||
│ └── labels/
|
||||
└── grupo_3/
|
||||
├── images/
|
||||
├── masks/
|
||||
└── labels/
|
||||
|
||||
- Mostrar imagem original e máscara lado a lado.
|
||||
- Permitir classificar cada frame com teclas numéricas.
|
||||
- Salvar um JSON por imagem dentro de labels/.
|
||||
- Suportar resume, skip, voltar, avançar e undo.
|
||||
|
||||
Exemplo de labels globais:
|
||||
Direcionando,EntrandoRua,CaminhandoRua,SaindoRua
|
||||
|
||||
Dependências:
|
||||
pip install pillow
|
||||
|
||||
Uso com GUI:
|
||||
python corridor_state_labeler.py
|
||||
|
||||
Uso por linha de comando:
|
||||
python corridor_state_labeler.py --dataset-root "C:/.../oak-d/dataset/original/group" --states Direcionando EntrandoRua CaminhandoRua SaindoRua --resume
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import tkinter as tk
|
||||
from tkinter import filedialog, messagebox
|
||||
|
||||
try:
|
||||
from PIL import Image, ImageTk, ImageOps
|
||||
except Exception:
|
||||
print("ERRO: Pillow não encontrado. Instale com: pip install pillow", file=sys.stderr)
|
||||
raise
|
||||
|
||||
|
||||
IMG_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
|
||||
MASK_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
|
||||
|
||||
IMAGE_DIR_NAMES = ["images", "image", "imgs", "rgb"]
|
||||
MASK_DIR_NAMES = ["masks", "mask", "segmentacao", "seg", "segmentations"]
|
||||
LABEL_DIR_NAME = "labels"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SampleItem:
|
||||
group: str
|
||||
image_path: Path
|
||||
mask_path: Path
|
||||
label_path: Path
|
||||
base: str
|
||||
|
||||
|
||||
def normalize_stem(stem: str) -> str:
|
||||
"""
|
||||
Remove sufixos comuns para tentar casar image e mask mesmo quando nomes diferem.
|
||||
Exemplos:
|
||||
001_rgb -> 001
|
||||
001_image -> 001
|
||||
001_mask -> 001
|
||||
001_segmentacao -> 001
|
||||
"""
|
||||
suffixes = [
|
||||
"_rgb", "_RGB", "_Rgb",
|
||||
"_image", "_img", "_frame",
|
||||
"_mask", "_masks",
|
||||
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
|
||||
]
|
||||
out = stem
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for sfx in suffixes:
|
||||
if out.endswith(sfx):
|
||||
out = out[: -len(sfx)]
|
||||
changed = True
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def natural_key(text: str):
|
||||
"""Ordenação natural simples: 2 antes de 10."""
|
||||
import re
|
||||
parts = re.split(r"(\d+)", text)
|
||||
return [int(p) if p.isdigit() else p.lower() for p in parts]
|
||||
|
||||
|
||||
def find_first_existing_dir(group_dir: Path, candidates: List[str]) -> Optional[Path]:
|
||||
for name in candidates:
|
||||
p = group_dir / name
|
||||
if p.is_dir():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def list_image_files(folder: Path, exts: List[str]) -> List[Path]:
|
||||
files = []
|
||||
for p in folder.iterdir():
|
||||
if p.is_file() and p.suffix.lower() in exts:
|
||||
files.append(p)
|
||||
files.sort(key=lambda x: natural_key(x.name))
|
||||
return files
|
||||
|
||||
|
||||
def collect_samples(dataset_root: Path) -> List[SampleItem]:
|
||||
"""
|
||||
Procura grupos dentro de dataset_root.
|
||||
Cada grupo precisa ter uma pasta images/ e masks/.
|
||||
Cria labels/ automaticamente.
|
||||
"""
|
||||
if not dataset_root.exists() or not dataset_root.is_dir():
|
||||
raise FileNotFoundError(f"Raiz inválida: {dataset_root}")
|
||||
|
||||
samples: List[SampleItem] = []
|
||||
|
||||
group_dirs = [p for p in dataset_root.iterdir() if p.is_dir()]
|
||||
group_dirs.sort(key=lambda x: natural_key(x.name))
|
||||
|
||||
for group_dir in group_dirs:
|
||||
image_dir = find_first_existing_dir(group_dir, IMAGE_DIR_NAMES)
|
||||
mask_dir = find_first_existing_dir(group_dir, MASK_DIR_NAMES)
|
||||
|
||||
if image_dir is None or mask_dir is None:
|
||||
continue
|
||||
|
||||
label_dir = group_dir / LABEL_DIR_NAME
|
||||
label_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
image_files = list_image_files(image_dir, IMG_EXTS)
|
||||
mask_files = list_image_files(mask_dir, MASK_EXTS)
|
||||
|
||||
mask_by_base: Dict[str, Path] = {}
|
||||
for m in mask_files:
|
||||
mask_by_base[normalize_stem(m.stem)] = m
|
||||
|
||||
for img in image_files:
|
||||
base = normalize_stem(img.stem)
|
||||
mask = mask_by_base.get(base)
|
||||
if mask is None:
|
||||
# fallback: tenta mesmo stem exato
|
||||
exact_candidates = [m for m in mask_files if m.stem == img.stem]
|
||||
mask = exact_candidates[0] if exact_candidates else None
|
||||
|
||||
if mask is None:
|
||||
continue
|
||||
|
||||
label_path = label_dir / f"{base}.json"
|
||||
samples.append(
|
||||
SampleItem(
|
||||
group=group_dir.name,
|
||||
image_path=img,
|
||||
mask_path=mask,
|
||||
label_path=label_path,
|
||||
base=base,
|
||||
)
|
||||
)
|
||||
|
||||
samples.sort(key=lambda s: (natural_key(s.group), natural_key(s.base)))
|
||||
return samples
|
||||
|
||||
|
||||
def read_existing_label(label_path: Path) -> Optional[Dict]:
|
||||
if not label_path.exists():
|
||||
return None
|
||||
try:
|
||||
with label_path.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def safe_relative(path: Path, root: Path) -> str:
|
||||
try:
|
||||
return str(path.resolve().relative_to(root.resolve())).replace("\\", "/")
|
||||
except Exception:
|
||||
return str(path).replace("\\", "/")
|
||||
|
||||
|
||||
class CorridorStateLabelerApp:
|
||||
def __init__(
|
||||
self,
|
||||
dataset_root: Path,
|
||||
samples: List[SampleItem],
|
||||
states: List[str],
|
||||
resume: bool = True,
|
||||
display_height: int = 520,
|
||||
show_existing: bool = False,
|
||||
):
|
||||
self.dataset_root = dataset_root
|
||||
self.states = states
|
||||
self.display_height = int(display_height)
|
||||
self.show_existing = show_existing
|
||||
|
||||
if resume:
|
||||
self.samples = [s for s in samples if not s.label_path.exists()]
|
||||
else:
|
||||
self.samples = samples
|
||||
|
||||
self.idx = 0
|
||||
self.history: List[Dict] = []
|
||||
|
||||
self.root = tk.Tk()
|
||||
self.root.title("Agrobot - Corridor State Labeler")
|
||||
self.root.geometry("1600x880")
|
||||
self.root.bind("<Key>", self.on_key)
|
||||
|
||||
self.top_frame = tk.Frame(self.root)
|
||||
self.top_frame.pack(side=tk.TOP, fill=tk.X)
|
||||
|
||||
self.info_label = tk.Label(self.top_frame, text="", font=("Segoe UI", 11), anchor="w")
|
||||
self.info_label.pack(side=tk.LEFT, padx=10, pady=6, fill=tk.X, expand=True)
|
||||
|
||||
self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10), anchor="e")
|
||||
self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6)
|
||||
|
||||
self.image_frame = tk.Frame(self.root)
|
||||
self.image_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
|
||||
|
||||
self.left_frame = tk.LabelFrame(self.image_frame, text="Imagem")
|
||||
self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)
|
||||
|
||||
self.right_frame = tk.LabelFrame(self.image_frame, text="Máscara")
|
||||
self.right_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)
|
||||
|
||||
self.image_label = tk.Label(self.left_frame)
|
||||
self.image_label.pack(fill=tk.BOTH, expand=True)
|
||||
|
||||
self.mask_label = tk.Label(self.right_frame)
|
||||
self.mask_label.pack(fill=tk.BOTH, expand=True)
|
||||
|
||||
self.current_label_var = tk.StringVar(value="")
|
||||
self.current_label = tk.Label(self.root, textvariable=self.current_label_var, font=("Segoe UI", 11), anchor="w")
|
||||
self.current_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=2)
|
||||
|
||||
self.status_var = tk.StringVar(value="Pronto.")
|
||||
self.status_label = tk.Label(self.root, textvariable=self.status_var, font=("Segoe UI", 10), anchor="w")
|
||||
self.status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=4)
|
||||
|
||||
self.footer = tk.Label(
|
||||
self.root,
|
||||
text="1..9/0=classificar | Espaço=saltar | →/n=próxima | ←/p=anterior | b=desfazer | r=recarregar | q/Esc=sair",
|
||||
font=("Segoe UI", 10),
|
||||
)
|
||||
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
|
||||
|
||||
self.render()
|
||||
|
||||
def build_legend_text(self) -> str:
|
||||
parts = []
|
||||
for i, state in enumerate(self.states, start=1):
|
||||
key = i if i <= 9 else 0
|
||||
parts.append(f"[{key}] {state}")
|
||||
return " | ".join(parts)
|
||||
|
||||
def load_and_resize(self, path: Path) -> ImageTk.PhotoImage:
|
||||
img = Image.open(path).convert("RGB")
|
||||
img = ImageOps.exif_transpose(img)
|
||||
|
||||
w, h = img.size
|
||||
if h <= 0:
|
||||
raise ValueError(f"Imagem inválida: {path}")
|
||||
|
||||
new_h = self.display_height
|
||||
new_w = max(1, int(w * (new_h / h)))
|
||||
img = img.resize((new_w, new_h), Image.BILINEAR)
|
||||
return ImageTk.PhotoImage(img)
|
||||
|
||||
def render(self):
|
||||
if not self.samples:
|
||||
messagebox.showinfo("Fim", "Nenhum item para rotular. Talvez tudo já esteja rotulado com resume ligado.")
|
||||
self.root.destroy()
|
||||
return
|
||||
|
||||
self.idx = max(0, min(self.idx, len(self.samples) - 1))
|
||||
item = self.samples[self.idx]
|
||||
|
||||
try:
|
||||
self.tk_img = self.load_and_resize(item.image_path)
|
||||
self.tk_mask = self.load_and_resize(item.mask_path)
|
||||
self.image_label.configure(image=self.tk_img)
|
||||
self.mask_label.configure(image=self.tk_mask)
|
||||
except Exception as e:
|
||||
self.status_var.set(f"ERRO ao abrir imagem/máscara: {e}")
|
||||
return
|
||||
|
||||
existing = read_existing_label(item.label_path)
|
||||
existing_txt = ""
|
||||
if existing:
|
||||
existing_txt = f" | Já rotulado: {existing.get('estado_corredor')}"
|
||||
|
||||
self.info_label.configure(
|
||||
text=(
|
||||
f"{self.idx + 1}/{len(self.samples)} | Grupo: {item.group} | Base: {item.base} "
|
||||
f"| Imagem: {item.image_path.name} | Máscara: {item.mask_path.name}{existing_txt}"
|
||||
)
|
||||
)
|
||||
|
||||
self.current_label_var.set(
|
||||
f"Arquivo de label: {safe_relative(item.label_path, self.dataset_root)}"
|
||||
)
|
||||
|
||||
def save_label(self, state_index: int):
|
||||
if state_index < 0 or state_index >= len(self.states):
|
||||
return
|
||||
|
||||
item = self.samples[self.idx]
|
||||
state = self.states[state_index]
|
||||
|
||||
previous_content = None
|
||||
previous_existed = item.label_path.exists()
|
||||
if previous_existed:
|
||||
try:
|
||||
previous_content = item.label_path.read_text(encoding="utf-8")
|
||||
except Exception:
|
||||
previous_content = None
|
||||
|
||||
rec = {
|
||||
"estado_corredor": state,
|
||||
"label_id": state_index,
|
||||
"states": self.states,
|
||||
"group": item.group,
|
||||
"base": item.base,
|
||||
"image": safe_relative(item.image_path, self.dataset_root),
|
||||
"mask": safe_relative(item.mask_path, self.dataset_root),
|
||||
"label": safe_relative(item.label_path, self.dataset_root),
|
||||
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
||||
"source": "corridor_state_labeler",
|
||||
}
|
||||
|
||||
try:
|
||||
item.label_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with item.label_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(rec, f, ensure_ascii=False, indent=2)
|
||||
|
||||
self.history.append(
|
||||
{
|
||||
"action": "label",
|
||||
"index": self.idx,
|
||||
"label_path": item.label_path,
|
||||
"previous_existed": previous_existed,
|
||||
"previous_content": previous_content,
|
||||
"new_state": state,
|
||||
}
|
||||
)
|
||||
|
||||
self.status_var.set(f"Salvo: {item.group}/{item.base} → {state}")
|
||||
self.idx += 1
|
||||
if self.idx >= len(self.samples):
|
||||
messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
|
||||
self.root.destroy()
|
||||
return
|
||||
self.render()
|
||||
|
||||
except Exception as e:
|
||||
messagebox.showerror("Erro", f"Falha ao salvar label: {e}")
|
||||
self.status_var.set(f"ERRO ao salvar label: {e}")
|
||||
|
||||
def skip(self):
|
||||
item = self.samples[self.idx]
|
||||
self.history.append({"action": "skip", "index": self.idx})
|
||||
self.status_var.set(f"Pulou: {item.group}/{item.base}")
|
||||
self.idx += 1
|
||||
if self.idx >= len(self.samples):
|
||||
messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
|
||||
self.root.destroy()
|
||||
return
|
||||
self.render()
|
||||
|
||||
def undo(self):
|
||||
if not self.history:
|
||||
self.status_var.set("Nada para desfazer.")
|
||||
return
|
||||
|
||||
last = self.history.pop()
|
||||
action = last.get("action")
|
||||
|
||||
if action == "label":
|
||||
label_path: Path = last["label_path"]
|
||||
try:
|
||||
if last.get("previous_existed"):
|
||||
previous_content = last.get("previous_content")
|
||||
if previous_content is not None:
|
||||
label_path.write_text(previous_content, encoding="utf-8")
|
||||
else:
|
||||
if label_path.exists():
|
||||
label_path.unlink()
|
||||
|
||||
self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
|
||||
self.status_var.set(f"Desfeito label: {label_path.name}")
|
||||
self.render()
|
||||
except Exception as e:
|
||||
messagebox.showerror("Erro", f"Falha ao desfazer: {e}")
|
||||
self.status_var.set(f"ERRO no undo: {e}")
|
||||
|
||||
elif action == "skip":
|
||||
self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
|
||||
self.status_var.set("Desfeito pulo.")
|
||||
self.render()
|
||||
|
||||
def next_item(self):
|
||||
self.idx = min(len(self.samples) - 1, self.idx + 1)
|
||||
self.status_var.set("Avançou uma imagem.")
|
||||
self.render()
|
||||
|
||||
def prev_item(self):
|
||||
self.idx = max(0, self.idx - 1)
|
||||
self.status_var.set("Voltou uma imagem.")
|
||||
self.render()
|
||||
|
||||
def on_key(self, event):
|
||||
ch = event.keysym.lower()
|
||||
|
||||
if ch in [str(i) for i in range(1, 10)] or ch == "0":
|
||||
idx = 9 if ch == "0" else int(ch) - 1
|
||||
self.save_label(idx)
|
||||
return
|
||||
|
||||
if ch in ("space", "s"):
|
||||
self.skip()
|
||||
return
|
||||
|
||||
if ch in ("right", "n"):
|
||||
self.next_item()
|
||||
return
|
||||
|
||||
if ch in ("left", "p"):
|
||||
self.prev_item()
|
||||
return
|
||||
|
||||
if ch == "b":
|
||||
self.undo()
|
||||
return
|
||||
|
||||
if ch == "r":
|
||||
self.status_var.set("Recarregado.")
|
||||
self.render()
|
||||
return
|
||||
|
||||
if ch in ("q", "escape"):
|
||||
self.root.destroy()
|
||||
return
|
||||
|
||||
def run(self):
|
||||
self.root.mainloop()
|
||||
|
||||
|
||||
class SetupWindow:
|
||||
def __init__(self):
|
||||
self.root = tk.Tk()
|
||||
self.root.title("Configurar - Corridor State Labeler")
|
||||
self.root.geometry("760x430")
|
||||
|
||||
frm_root = tk.LabelFrame(self.root, text="Raiz do dataset")
|
||||
frm_root.pack(fill=tk.X, padx=10, pady=8)
|
||||
|
||||
self.dataset_root_var = tk.StringVar(value="")
|
||||
tk.Entry(frm_root, textvariable=self.dataset_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6)
|
||||
tk.Button(frm_root, text="Escolher...", command=self.choose_dataset_root).pack(side=tk.RIGHT, padx=6, pady=6)
|
||||
|
||||
frm_states = tk.LabelFrame(self.root, text="Estados do corredor separados por vírgula")
|
||||
frm_states.pack(fill=tk.X, padx=10, pady=8)
|
||||
|
||||
self.states_var = tk.StringVar(value="Parado,EntrandoRua,CaminhandoRua,SaindoRua,Manobrando,Direcionando,RetornandoBase,Indefinido")
|
||||
tk.Entry(frm_states, textvariable=self.states_var).pack(fill=tk.X, padx=6, pady=6)
|
||||
|
||||
frm_opts = tk.LabelFrame(self.root, text="Opções")
|
||||
frm_opts.pack(fill=tk.X, padx=10, pady=8)
|
||||
|
||||
self.resume_var = tk.BooleanVar(value=True)
|
||||
self.height_var = tk.IntVar(value=520)
|
||||
|
||||
tk.Checkbutton(frm_opts, text="Retomar: pular imagens que já possuem labels/*.json", variable=self.resume_var).pack(anchor="w", padx=6, pady=4)
|
||||
|
||||
frm_height = tk.Frame(frm_opts)
|
||||
frm_height.pack(fill=tk.X, padx=6, pady=4)
|
||||
tk.Label(frm_height, text="Altura de exibição em px:").pack(side=tk.LEFT)
|
||||
tk.Entry(frm_height, textvariable=self.height_var, width=8).pack(side=tk.LEFT, padx=6)
|
||||
|
||||
self.preview_var = tk.StringVar(value="")
|
||||
tk.Label(self.root, textvariable=self.preview_var, fg="#555", anchor="w", justify="left").pack(fill=tk.X, padx=12, pady=4)
|
||||
|
||||
btn_frame = tk.Frame(self.root)
|
||||
btn_frame.pack(fill=tk.X, padx=10, pady=10)
|
||||
tk.Button(btn_frame, text="Verificar dataset", command=self.preview_dataset).pack(side=tk.LEFT, padx=4)
|
||||
tk.Button(btn_frame, text="Iniciar classificação", command=self.start).pack(side=tk.RIGHT, padx=4)
|
||||
|
||||
tk.Label(
|
||||
self.root,
|
||||
text="Teclas: 1..9/0=classificar | Espaço=saltar | b=desfazer | ←/→ navegar | q/Esc=sair",
|
||||
fg="#555",
|
||||
).pack(pady=4)
|
||||
|
||||
self.result = None
|
||||
|
||||
def choose_dataset_root(self):
|
||||
p = filedialog.askdirectory(title="Selecione a raiz do dataset")
|
||||
if p:
|
||||
self.dataset_root_var.set(p)
|
||||
self.preview_dataset()
|
||||
|
||||
def parse_states(self) -> List[str]:
|
||||
return [s.strip() for s in self.states_var.get().split(",") if s.strip()]
|
||||
|
||||
def preview_dataset(self):
|
||||
root_raw = self.dataset_root_var.get().strip()
|
||||
if not root_raw:
|
||||
self.preview_var.set("Selecione uma raiz de dataset.")
|
||||
return
|
||||
root = Path(root_raw)
|
||||
try:
|
||||
samples = collect_samples(root)
|
||||
groups = sorted(set(s.group for s in samples), key=natural_key)
|
||||
labeled = sum(1 for s in samples if s.label_path.exists())
|
||||
self.preview_var.set(
|
||||
f"Amostras encontradas: {len(samples)} | Já rotuladas: {labeled} | Grupos: {', '.join(groups) if groups else '-'}"
|
||||
)
|
||||
except Exception as e:
|
||||
self.preview_var.set(f"Erro ao verificar dataset: {e}")
|
||||
|
||||
def start(self):
|
||||
root_raw = self.dataset_root_var.get().strip()
|
||||
states = self.parse_states()
|
||||
|
||||
if not root_raw:
|
||||
messagebox.showwarning("Faltando raiz", "Selecione a raiz do dataset.")
|
||||
return
|
||||
if not states:
|
||||
messagebox.showwarning("Faltando estados", "Informe pelo menos um estado.")
|
||||
return
|
||||
if len(states) > 10:
|
||||
messagebox.showwarning("Muitos estados", "Este rotulador suporta até 10 estados nas teclas 1..9 e 0.")
|
||||
return
|
||||
|
||||
root = Path(root_raw)
|
||||
try:
|
||||
samples = collect_samples(root)
|
||||
except Exception as e:
|
||||
messagebox.showerror("Erro", f"Falha ao ler dataset: {e}")
|
||||
return
|
||||
|
||||
if not samples:
|
||||
messagebox.showinfo(
|
||||
"Sem amostras",
|
||||
"Nenhum par image/mask encontrado. Verifique se cada grupo possui images/ e masks/.",
|
||||
)
|
||||
return
|
||||
|
||||
self.result = {
|
||||
"dataset_root": root,
|
||||
"states": states,
|
||||
"resume": self.resume_var.get(),
|
||||
"display_height": self.height_var.get(),
|
||||
}
|
||||
self.root.destroy()
|
||||
|
||||
def run(self):
|
||||
self.root.mainloop()
|
||||
return self.result
|
||||
|
||||
|
||||
def run_gui_setup():
|
||||
setup = SetupWindow()
|
||||
res = setup.run()
|
||||
if not res:
|
||||
return
|
||||
|
||||
dataset_root = res["dataset_root"]
|
||||
samples = collect_samples(dataset_root)
|
||||
|
||||
app = CorridorStateLabelerApp(
|
||||
dataset_root=dataset_root,
|
||||
samples=samples,
|
||||
states=res["states"],
|
||||
resume=res["resume"],
|
||||
display_height=res["display_height"],
|
||||
)
|
||||
app.run()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Rotulador de estado global do corredor para dataset de segmentação.")
|
||||
parser.add_argument("--dataset-root", help="Raiz do dataset contendo grupos com images/ e masks/.")
|
||||
parser.add_argument("--states", nargs="+", help="Estados do corredor. Ex: Direcionando EntrandoRua CaminhandoRua SaindoRua")
|
||||
parser.add_argument("--resume", action="store_true", help="Pular itens que já possuem labels/*.json.")
|
||||
parser.add_argument("--no-resume", action="store_true", help="Não pular itens já rotulados.")
|
||||
parser.add_argument("--display-height", type=int, default=520, help="Altura de exibição das imagens em px.")
|
||||
parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir janela de configuração.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
need_gui = not args.no_gui_setup and (not args.dataset_root or not args.states)
|
||||
if need_gui:
|
||||
run_gui_setup()
|
||||
return
|
||||
|
||||
if not args.dataset_root or not args.states:
|
||||
print("ERRO: informe --dataset-root e --states, ou rode sem argumentos para abrir a GUI.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
dataset_root = Path(args.dataset_root)
|
||||
states = args.states
|
||||
if len(states) > 10:
|
||||
print("ERRO: máximo de 10 estados: teclas 1..9 e 0.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
samples = collect_samples(dataset_root)
|
||||
if not samples:
|
||||
print("ERRO: nenhum par image/mask encontrado.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
resume = True
|
||||
if args.no_resume:
|
||||
resume = False
|
||||
elif args.resume:
|
||||
resume = True
|
||||
|
||||
app = CorridorStateLabelerApp(
|
||||
dataset_root=dataset_root,
|
||||
samples=samples,
|
||||
states=states,
|
||||
resume=resume,
|
||||
display_height=args.display_height,
|
||||
)
|
||||
app.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -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"]
|
||||
|
||||
|
|
|
|||
|
|
@ -1,22 +1,34 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Agrupa imagens e máscaras novas em subpastas por combinação de classes presentes.
|
||||
Agrupa imagens, máscaras e labels globais em subpastas por combinação de classes presentes.
|
||||
|
||||
Estrutura lida (via config.json -> MODELO):
|
||||
MODELO/dataset/original/new_images/
|
||||
MODELO/dataset/original/new_masks/
|
||||
Estrutura lida via config.json -> camera:
|
||||
|
||||
MODELO/dataset/original/images/
|
||||
MODELO/dataset/original/masks/
|
||||
MODELO/dataset/original/masks2/ opcional, se dual_head_mask=true
|
||||
MODELO/dataset/original/labels/ opcional/necessário, se dual_head_label=true
|
||||
|
||||
Saída:
|
||||
MODELO/dataset/original/group/<grupo>/images
|
||||
MODELO/dataset/original/group/<grupo>/masks
|
||||
|
||||
Onde <grupo> é os nomes das classes presentes unidos por "_", ex:
|
||||
chao, erva, cana, chao_erva, erva_cana, chao_erva_cana, etc.
|
||||
MODELO/dataset/original/group/<grupo>/images/
|
||||
MODELO/dataset/original/group/<grupo>/masks/
|
||||
MODELO/dataset/original/group/<grupo>/masks2/ se dual_head_mask=true
|
||||
MODELO/dataset/original/group/<grupo>/labels/ se dual_head_label=true
|
||||
|
||||
Requer: utils.carregar_labelmap_completo(labelmap_path)
|
||||
O labelmap define mapeamento de cores/ids/nomes das classes.
|
||||
Onde <grupo> é o nome das classes presentes na máscara unidos por "_", ex:
|
||||
navegavel
|
||||
naonavegavel
|
||||
naonavegavel_navegavel
|
||||
|
||||
Observação importante:
|
||||
- dual_head_mask = segunda cabeça também é máscara pixel-a-pixel.
|
||||
- dual_head_label = segunda cabeça é label global por frame, vindo de JSON.
|
||||
|
||||
O JSON de label é copiado para labels/ com o mesmo base name final da imagem/máscara.
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import csv
|
||||
|
|
@ -27,42 +39,60 @@ import numpy as np
|
|||
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
|
||||
# ====================== Configurações base ======================
|
||||
|
||||
def carregar_config_e_paths():
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
|
||||
EXT_MASKS = (".png", ".jpg", ".jpeg")
|
||||
EXT_MASKS2 = (".png", ".jpg", ".jpeg")
|
||||
EXT_LABELS = (".json", ".txt")
|
||||
|
||||
MANIFESTO_DEFAULT = "manifest.csv"
|
||||
|
||||
|
||||
def carregar_config_e_paths(modelo_cli=None):
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
||||
MODELO = modelo_cli or config.get("camera")
|
||||
USE_MASKS2 = bool(config.get("dual_head_mask", False))
|
||||
USE_LABELS = bool(config.get("dual_head_label", False))
|
||||
|
||||
pasta_base = os.path.join(MODELO, "dataset")
|
||||
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
|
||||
|
||||
# Pastas origem/destino
|
||||
PASTA_NEW_IMAGES = os.path.join(pasta_base, "original", "images")
|
||||
PASTA_NEW_MASKS = os.path.join(pasta_base, "original", "masks")
|
||||
PASTA_NEW_MASKS2 = os.path.join(pasta_base, "original", "masks2")
|
||||
PASTA_NEW_LABELS = os.path.join(pasta_base, "original", "labels")
|
||||
PASTA_FINAL = os.path.join(pasta_base, "original", "group")
|
||||
|
||||
return MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_NEW_MASKS2, USE_MASKS2, PASTA_FINAL
|
||||
return {
|
||||
"MODELO": MODELO,
|
||||
"pasta_base": pasta_base,
|
||||
"labelmap_path": labelmap_path,
|
||||
"PASTA_NEW_IMAGES": PASTA_NEW_IMAGES,
|
||||
"PASTA_NEW_MASKS": PASTA_NEW_MASKS,
|
||||
"PASTA_NEW_MASKS2": PASTA_NEW_MASKS2,
|
||||
"PASTA_NEW_LABELS": PASTA_NEW_LABELS,
|
||||
"PASTA_FINAL": PASTA_FINAL,
|
||||
"USE_MASKS2": USE_MASKS2,
|
||||
"USE_LABELS": USE_LABELS,
|
||||
}
|
||||
|
||||
# Extensões aceitas
|
||||
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
|
||||
EXT_MASKS = (".png", ".jpg", ".jpeg") # prioridade será .png quando houver
|
||||
EXT_MASKS2 = (".png", ".jpg", ".jpeg") # idem
|
||||
|
||||
# Manifesto padrão
|
||||
MANIFESTO_DEFAULT = "manifest.csv"
|
||||
|
||||
# ====================== Utilitários ======================
|
||||
|
||||
|
||||
def garantir_pasta(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
|
||||
def nome_disponivel(dest_dir, base_name, ext):
|
||||
"""Gera nome único em dest_dir com sufixos _001, _002, ... se necessário."""
|
||||
cand = os.path.join(dest_dir, base_name + ext)
|
||||
if not os.path.exists(cand):
|
||||
return cand
|
||||
|
||||
i = 1
|
||||
while True:
|
||||
cand = os.path.join(dest_dir, f"{base_name}_{i:03d}{ext}")
|
||||
|
|
@ -70,41 +100,69 @@ def nome_disponivel(dest_dir, base_name, ext):
|
|||
return cand
|
||||
i += 1
|
||||
|
||||
def mapear_masks_por_base(pasta_masks):
|
||||
"""Retorna {base: caminho_mask}, priorizando .png se houver múltiplas por base."""
|
||||
|
||||
def normalizar_base(stem):
|
||||
"""
|
||||
Tenta remover sufixos comuns para casar image/mask/label.
|
||||
Ex:
|
||||
001_rgb -> 001
|
||||
001_segmentacao -> 001
|
||||
001_mask -> 001
|
||||
"""
|
||||
sufixos = [
|
||||
"_rgb", "_RGB", "_Rgb",
|
||||
"_image", "_img", "_frame",
|
||||
"_mask", "_masks",
|
||||
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
|
||||
"_label", "_labels",
|
||||
]
|
||||
|
||||
out = stem
|
||||
mudou = True
|
||||
while mudou:
|
||||
mudou = False
|
||||
for sfx in sufixos:
|
||||
if out.endswith(sfx):
|
||||
out = out[: -len(sfx)]
|
||||
mudou = True
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def mapear_arquivos_por_base(pasta, extensoes, normalizar=True):
|
||||
"""
|
||||
Retorna {base: caminho}, priorizando .png para máscaras/imagens e .json para labels.
|
||||
"""
|
||||
if not pasta or not os.path.isdir(pasta):
|
||||
return {}
|
||||
|
||||
mapa = {}
|
||||
for nome in os.listdir(pasta_masks):
|
||||
prioridade = {
|
||||
".json": 0,
|
||||
".png": 1,
|
||||
".jpg": 2,
|
||||
".jpeg": 3,
|
||||
".txt": 4,
|
||||
}
|
||||
|
||||
for nome in os.listdir(pasta):
|
||||
lower = nome.lower()
|
||||
if not lower.endswith(EXT_MASKS):
|
||||
if not lower.endswith(extensoes):
|
||||
continue
|
||||
base, ext = os.path.splitext(nome)
|
||||
cam = os.path.join(pasta_masks, nome)
|
||||
|
||||
stem, ext = os.path.splitext(nome)
|
||||
base = normalizar_base(stem) if normalizar else stem
|
||||
caminho = os.path.join(pasta, nome)
|
||||
|
||||
if base not in mapa:
|
||||
mapa[base] = cam
|
||||
mapa[base] = caminho
|
||||
else:
|
||||
atual_ext = os.path.splitext(mapa[base])[1].lower()
|
||||
if atual_ext != ".png" and ext.lower() == ".png":
|
||||
mapa[base] = cam
|
||||
if prioridade.get(ext.lower(), 99) < prioridade.get(atual_ext, 99):
|
||||
mapa[base] = caminho
|
||||
|
||||
return mapa
|
||||
|
||||
def mapear_masks2_por_base(pasta_masks2):
|
||||
"""Retorna {base: caminho_mask2}, priorizando .png se houver múltiplas por base."""
|
||||
if not pasta_masks2 or not os.path.isdir(pasta_masks2):
|
||||
return {}
|
||||
mapa = {}
|
||||
for nome in os.listdir(pasta_masks2):
|
||||
lower = nome.lower()
|
||||
if not lower.endswith(EXT_MASKS2):
|
||||
continue
|
||||
base, ext = os.path.splitext(nome)
|
||||
cam = os.path.join(pasta_masks2, nome)
|
||||
if base not in mapa:
|
||||
mapa[base] = cam
|
||||
else:
|
||||
atual_ext = os.path.splitext(mapa[base])[1].lower()
|
||||
if atual_ext != ".png" and ext.lower() == ".png":
|
||||
mapa[base] = cam
|
||||
return mapa
|
||||
|
||||
def localizar_imagem_por_base(pasta_imgs, base):
|
||||
"""Retorna caminho da imagem correspondente ao base se existir."""
|
||||
|
|
@ -112,119 +170,123 @@ def localizar_imagem_por_base(pasta_imgs, base):
|
|||
p = os.path.join(pasta_imgs, base + ext)
|
||||
if os.path.isfile(p):
|
||||
return p
|
||||
return None
|
||||
|
||||
# fallback com normalização
|
||||
mapa = mapear_arquivos_por_base(pasta_imgs, EXT_IMAGENS, normalizar=True)
|
||||
return mapa.get(base)
|
||||
|
||||
|
||||
def inferir_ignore_id(ignore_rgb, cor_para_id):
|
||||
"""
|
||||
Tenta inferir o ID da classe ignorada a partir do retorno ignore_rgb e do mapa cor->id.
|
||||
- Se ignore_rgb for [id] ou (id,), retorna esse id.
|
||||
- Se ignore_rgb parecer uma cor RGB (len==3), usa cor_para_id[(R,G,B)] se existir.
|
||||
- Caso não consiga, retorna None.
|
||||
"""
|
||||
if ignore_rgb is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
# caso [id]
|
||||
if isinstance(ignore_rgb, (list, tuple)) and len(ignore_rgb) == 1:
|
||||
return int(ignore_rgb[0])
|
||||
# caso [R,G,B]
|
||||
|
||||
if isinstance(ignore_rgb, (list, tuple)) and len(ignore_rgb) == 3:
|
||||
key = tuple(int(v) for v in ignore_rgb)
|
||||
return cor_para_id.get(key)
|
||||
except Exception:
|
||||
pass
|
||||
# pode já ser um inteiro simples
|
||||
|
||||
if isinstance(ignore_rgb, (int, np.integer)):
|
||||
return int(ignore_rgb)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
|
||||
"""
|
||||
Versão corrigida e otimizada:
|
||||
|
||||
- Se máscara for 1 canal: np.unique direto -> IDs.
|
||||
- Se for 3 canais:
|
||||
* se assume_rgb=True: labelmap está em RGB,
|
||||
mas OpenCV 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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,113 +1,185 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Augmenta imagens e máscaras *por grupo*.
|
||||
Augmenta imagens, máscaras e opcionalmente labels globais por grupo.
|
||||
|
||||
Entrada (via config.json -> MODELO):
|
||||
MODELO/dataset/original/group/<grupo>/images
|
||||
MODELO/dataset/original/group/<grupo>/masks
|
||||
Entrada via config_oak.json -> camera:
|
||||
|
||||
MODELO/dataset/original/group/<grupo>/images/
|
||||
MODELO/dataset/original/group/<grupo>/masks/
|
||||
MODELO/dataset/original/group/<grupo>/masks2/ opcional, se dual_head_mask=true
|
||||
MODELO/dataset/original/group/<grupo>/labels/ opcional, se dual_head_label=true
|
||||
|
||||
Saída:
|
||||
MODELO/dataset/augmented/group/<grupo>/images
|
||||
MODELO/dataset/augmented/group/<grupo>/masks
|
||||
|
||||
Se "original/group" não existir, faz fallback para:
|
||||
MODELO/dataset/original/{images,masks}
|
||||
MODELO/dataset/augmented/{images,masks}
|
||||
MODELO/dataset/augmented/group/<grupo>/images/
|
||||
MODELO/dataset/augmented/group/<grupo>/masks/
|
||||
MODELO/dataset/augmented/group/<grupo>/masks2/ se dual_head_mask=true
|
||||
MODELO/dataset/augmented/group/<grupo>/labels/ se dual_head_label=true
|
||||
|
||||
Transf. geométricas (aplicam a img e máscara) e fotométricas (apenas imagem).
|
||||
Observação:
|
||||
- A máscara e a mask2 recebem as mesmas transformações geométricas da imagem.
|
||||
- O label global NÃO é transformado visualmente. Ele é copiado e atualizado para apontar para o novo base augmentado.
|
||||
|
||||
Uso:
|
||||
python _3_augmentation_grouped.py --copies 5
|
||||
python _3_augmentation_grouped.py --copies 5 --groups chao,chao_erva,cana
|
||||
python _3_augmentation_grouped_with_labels.py --copies 5
|
||||
python _3_augmentation_grouped_with_labels.py --copies 5 --groups navegavel,naonavegavel_navegavel
|
||||
python _3_augmentation_grouped_with_labels.py --copies 5 --strict-label
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
from PIL import Image
|
||||
import albumentations as A
|
||||
import argparse
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
|
||||
# ====================== Configurações ======================
|
||||
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera", ".")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
||||
# Pastas base
|
||||
MODELO = config.get("camera", ".")
|
||||
USE_MASKS2 = bool(config.get("dual_head_mask", False))
|
||||
USE_LABELS = bool(config.get("dual_head_label", False))
|
||||
|
||||
DATASET_BASE = os.path.join(MODELO, "dataset")
|
||||
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
|
||||
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
|
||||
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
|
||||
|
||||
# Fallback (modo antigo, sem grupos)
|
||||
# Fallback legacy, sem grupos
|
||||
ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images")
|
||||
ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks")
|
||||
ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2")
|
||||
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
|
||||
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
|
||||
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
|
||||
ORIG_OLD_LABELS = os.path.join(DATASET_BASE, "original", "labels")
|
||||
|
||||
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
|
||||
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
|
||||
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
|
||||
AUG_OLD_LABELS = os.path.join(DATASET_BASE, "augmented", "labels")
|
||||
|
||||
# Extensões aceitas
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXTS = (".png", ".jpg", ".jpeg") # manter prioridade PNG quando possível
|
||||
MSK_EXTS = (".png", ".jpg", ".jpeg")
|
||||
MSK2_EXTS = (".png", ".jpg", ".jpeg")
|
||||
LABEL_EXTS = (".json", ".txt")
|
||||
|
||||
|
||||
# ====================== Augmentation ======================
|
||||
|
||||
train_tf = A.Compose(
|
||||
[
|
||||
A.HorizontalFlip(p=0.5),
|
||||
|
||||
A.ShiftScaleRotate(
|
||||
shift_limit=0.01,
|
||||
scale_limit=0.10,
|
||||
rotate_limit=5,
|
||||
border_mode=cv2.BORDER_REFLECT_101,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
p=0.30,
|
||||
),
|
||||
|
||||
A.OneOf(
|
||||
[
|
||||
A.RandomBrightnessContrast(0.2, 0.2, p=1.0),
|
||||
A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1.0),
|
||||
A.RandomGamma(gamma_limit=(90, 110), p=1.0),
|
||||
],
|
||||
p=0.70,
|
||||
),
|
||||
|
||||
A.OneOf(
|
||||
[
|
||||
A.MotionBlur(blur_limit=3, p=1.0),
|
||||
A.GaussianBlur(blur_limit=3, p=1.0),
|
||||
],
|
||||
p=0.20,
|
||||
),
|
||||
|
||||
A.OneOf(
|
||||
[
|
||||
A.GaussNoise(var_limit=(5.0, 15.0), p=1.0),
|
||||
A.ImageCompression(quality_lower=50, quality_upper=85, p=1.0),
|
||||
],
|
||||
p=0.20,
|
||||
),
|
||||
|
||||
A.RandomShadow(p=0.10),
|
||||
A.RandomSunFlare(p=0.10),
|
||||
A.ChannelShuffle(p=0.05),
|
||||
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10),
|
||||
],
|
||||
additional_targets={"mask2": "mask"},
|
||||
)
|
||||
|
||||
|
||||
# ====================== Utilitários ======================
|
||||
|
||||
|
||||
def garantir_dir(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
# Pipeline de augmentations
|
||||
train_tf = A.Compose([
|
||||
A.HorizontalFlip(p=0.5),
|
||||
|
||||
# Geométricas (aplicam em imagem e máscara)
|
||||
A.ShiftScaleRotate(
|
||||
shift_limit=0.01,
|
||||
scale_limit=0.10,
|
||||
rotate_limit=5,
|
||||
border_mode=cv2.BORDER_REFLECT_101,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
p=0.30
|
||||
),
|
||||
def normalizar_base(stem: str) -> str:
|
||||
sufixos = [
|
||||
"_rgb", "_RGB", "_Rgb",
|
||||
"_image", "_img", "_frame",
|
||||
"_mask", "_masks",
|
||||
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
|
||||
"_label", "_labels",
|
||||
]
|
||||
|
||||
# Fotométricas (somente imagem)
|
||||
A.OneOf([
|
||||
A.RandomBrightnessContrast(0.2, 0.2, p=1.0),
|
||||
A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1.0),
|
||||
A.RandomGamma(gamma_limit=(90, 110), p=1.0),
|
||||
], p=0.70),
|
||||
out = stem
|
||||
mudou = True
|
||||
while mudou:
|
||||
mudou = False
|
||||
for sfx in sufixos:
|
||||
if out.endswith(sfx):
|
||||
out = out[: -len(sfx)]
|
||||
mudou = True
|
||||
break
|
||||
return out
|
||||
|
||||
A.OneOf([
|
||||
A.MotionBlur(blur_limit=3, p=1.0),
|
||||
A.GaussianBlur(blur_limit=3, p=1.0),
|
||||
], p=0.20),
|
||||
|
||||
A.OneOf([
|
||||
A.GaussNoise(var_limit=(5.0, 15.0), p=1.0),
|
||||
A.ImageCompression(quality_lower=50, quality_upper=85, p=1.0),
|
||||
], p=0.20),
|
||||
def map_files_by_base(folder, exts):
|
||||
by_base = {}
|
||||
if not os.path.isdir(folder):
|
||||
return by_base
|
||||
|
||||
A.RandomShadow(p=0.10),
|
||||
A.RandomSunFlare(p=0.10),
|
||||
A.ChannelShuffle(p=0.05),
|
||||
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10),
|
||||
], additional_targets={
|
||||
'mask2': 'mask'
|
||||
})
|
||||
prioridade = {
|
||||
".json": 0,
|
||||
".png": 1,
|
||||
".jpg": 2,
|
||||
".jpeg": 3,
|
||||
".txt": 4,
|
||||
}
|
||||
|
||||
def load_rgb(path):
|
||||
# cv2 lê BGR → converte para RGB
|
||||
im = cv2.imread(path, cv2.IMREAD_COLOR)
|
||||
if im is None:
|
||||
raise FileNotFoundError(path)
|
||||
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
||||
for fname in os.listdir(folder):
|
||||
lower = fname.lower()
|
||||
if not lower.endswith(exts):
|
||||
continue
|
||||
|
||||
stem, ext = os.path.splitext(fname)
|
||||
base = normalizar_base(stem)
|
||||
cand = os.path.join(folder, fname)
|
||||
|
||||
if base not in by_base:
|
||||
by_base[base] = cand
|
||||
else:
|
||||
cur_ext = os.path.splitext(by_base[base])[1].lower()
|
||||
if prioridade.get(ext.lower(), 99) < prioridade.get(cur_ext, 99):
|
||||
by_base[base] = cand
|
||||
|
||||
return by_base
|
||||
|
||||
def save_rgb(path, arr_rgb):
|
||||
Image.fromarray(arr_rgb).save(path)
|
||||
|
||||
def list_groups(root):
|
||||
"""Lista grupos válidos (que contêm subpastas images e masks)."""
|
||||
if not os.path.isdir(root):
|
||||
return []
|
||||
|
||||
grupos = []
|
||||
for name in sorted(os.listdir(root)):
|
||||
gdir = os.path.join(root, name)
|
||||
|
|
@ -117,68 +189,107 @@ def list_groups(root):
|
|||
grupos.append(name)
|
||||
return grupos
|
||||
|
||||
def map_masks_by_base(msk_dir):
|
||||
"""Mapeia máscaras por base (prioriza .png)."""
|
||||
by_base = {}
|
||||
if not os.path.isdir(msk_dir):
|
||||
return by_base
|
||||
for fname in os.listdir(msk_dir):
|
||||
f_lower = fname.lower()
|
||||
if not f_lower.endswith(MSK_EXTS):
|
||||
continue
|
||||
base, ext = os.path.splitext(fname)
|
||||
cand = os.path.join(msk_dir, fname)
|
||||
if base not in by_base:
|
||||
by_base[base] = cand
|
||||
else:
|
||||
# mantém .png se disponível
|
||||
cur_ext = os.path.splitext(by_base[base])[1].lower()
|
||||
if cur_ext != ".png" and ext.lower() == ".png":
|
||||
by_base[base] = cand
|
||||
return by_base
|
||||
|
||||
def map_masks2_by_base(msk2_dir):
|
||||
"""Mapeia máscaras2 por base (prioriza .png)."""
|
||||
by_base = {}
|
||||
if not os.path.isdir(msk2_dir):
|
||||
return by_base
|
||||
for fname in os.listdir(msk2_dir):
|
||||
f_lower = fname.lower()
|
||||
if not f_lower.endswith(MSK2_EXTS):
|
||||
continue
|
||||
base, ext = os.path.splitext(fname)
|
||||
cand = os.path.join(msk2_dir, fname)
|
||||
if base not in by_base:
|
||||
by_base[base] = cand
|
||||
else:
|
||||
cur_ext = os.path.splitext(by_base[base])[1].lower()
|
||||
if cur_ext != ".png" and ext.lower() == ".png":
|
||||
by_base[base] = cand
|
||||
return by_base
|
||||
def load_rgb(path):
|
||||
im = cv2.imread(path, cv2.IMREAD_COLOR)
|
||||
if im is None:
|
||||
raise FileNotFoundError(path)
|
||||
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
||||
|
||||
def ensure_aug_dirs(group_name=None, use_masks2=False):
|
||||
"""Cria diretórios de saída para o grupo ou modo antigo. Se use_masks2, cria masks2."""
|
||||
|
||||
def save_rgb(path, arr_rgb):
|
||||
garantir_dir(os.path.dirname(path))
|
||||
Image.fromarray(arr_rgb).save(path)
|
||||
|
||||
|
||||
def ensure_aug_dirs(group_name=None, use_masks2=False, use_labels=False):
|
||||
if group_name:
|
||||
img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images")
|
||||
msk_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks")
|
||||
msk2_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks2") if use_masks2 else None
|
||||
labels_out = os.path.join(AUG_GROUP_ROOT, group_name, "labels") if use_labels else None
|
||||
else:
|
||||
img_out = AUG_OLD_IMG
|
||||
msk_out = AUG_OLD_MSK
|
||||
msk2_out = AUG_OLD_MSK2 if use_masks2 else None
|
||||
labels_out = AUG_OLD_LABELS if use_labels else None
|
||||
|
||||
garantir_dir(img_out)
|
||||
garantir_dir(msk_out)
|
||||
|
||||
if use_masks2 and msk2_out:
|
||||
garantir_dir(msk2_out)
|
||||
return img_out, msk_out, msk2_out
|
||||
|
||||
def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies, msk2_path=None, msk2_out_dir=None):
|
||||
if use_labels and labels_out:
|
||||
garantir_dir(labels_out)
|
||||
|
||||
return img_out, msk_out, msk2_out, labels_out
|
||||
|
||||
|
||||
def safe_rel(path, root):
|
||||
try:
|
||||
return str(Path(path).resolve().relative_to(Path(root).resolve())).replace("\\", "/")
|
||||
except Exception:
|
||||
return str(path).replace("\\", "/")
|
||||
|
||||
|
||||
def copiar_label_aug(label_path, out_label_path, new_base, out_img_path, out_msk_path, out_msk2_path=None, group_name=None):
|
||||
"""
|
||||
Copia label global para a amostra augmentada.
|
||||
Se for JSON, atualiza campos úteis.
|
||||
Se for TXT, copia o conteúdo como está.
|
||||
"""
|
||||
if not label_path or not out_label_path:
|
||||
return
|
||||
|
||||
garantir_dir(os.path.dirname(out_label_path))
|
||||
ext = os.path.splitext(label_path)[1].lower()
|
||||
|
||||
if ext == ".json":
|
||||
try:
|
||||
with open(label_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except Exception:
|
||||
data = {}
|
||||
|
||||
data["base"] = new_base
|
||||
data["group"] = group_name if group_name is not None else data.get("group")
|
||||
data["image"] = safe_rel(out_img_path, DATASET_BASE)
|
||||
data["mask"] = safe_rel(out_msk_path, DATASET_BASE)
|
||||
data["label"] = safe_rel(out_label_path, DATASET_BASE)
|
||||
data["is_augmented"] = True
|
||||
data["source_label"] = safe_rel(label_path, DATASET_BASE)
|
||||
|
||||
if out_msk2_path:
|
||||
data["mask2"] = safe_rel(out_msk2_path, DATASET_BASE)
|
||||
|
||||
with open(out_label_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
else:
|
||||
with open(label_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
with open(out_label_path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
|
||||
def augment_pair(
|
||||
img_path,
|
||||
msk_path,
|
||||
img_out_dir,
|
||||
msk_out_dir,
|
||||
copies,
|
||||
msk2_path=None,
|
||||
msk2_out_dir=None,
|
||||
label_path=None,
|
||||
label_out_dir=None,
|
||||
group_name=None,
|
||||
):
|
||||
base_img, img_ext = os.path.splitext(os.path.basename(img_path))
|
||||
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
|
||||
_, msk_ext = os.path.splitext(os.path.basename(msk_path))
|
||||
msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
|
||||
label_ext = os.path.splitext(os.path.basename(label_path))[1] if label_path else None
|
||||
|
||||
# padroniza pelo base da imagem
|
||||
base = base_img
|
||||
base = normalizar_base(base_img)
|
||||
|
||||
img = load_rgb(img_path)
|
||||
msk = load_rgb(msk_path)
|
||||
|
|
@ -194,122 +305,223 @@ def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies, msk2_path
|
|||
img_aug = aug["image"]
|
||||
msk_aug = aug["mask"]
|
||||
|
||||
out_img = os.path.join(img_out_dir, f"{base}_aug_{i:02d}{img_ext}")
|
||||
out_msk = os.path.join(msk_out_dir, f"{base}_aug_{i:02d}{msk_ext}")
|
||||
new_base = f"{base}_aug_{i:02d}"
|
||||
out_img = os.path.join(img_out_dir, f"{new_base}{img_ext}")
|
||||
out_msk = os.path.join(msk_out_dir, f"{new_base}{msk_ext}")
|
||||
|
||||
save_rgb(out_img, img_aug)
|
||||
save_rgb(out_msk, msk_aug)
|
||||
|
||||
out_msk2 = None
|
||||
if msk2 is not None and msk2_out_dir:
|
||||
msk2_aug = aug["mask2"]
|
||||
out_msk2 = os.path.join(msk2_out_dir, f"{base}_aug_{i:02d}{msk2_ext}")
|
||||
out_msk2 = os.path.join(msk2_out_dir, f"{new_base}{msk2_ext}")
|
||||
save_rgb(out_msk2, msk2_aug)
|
||||
|
||||
if label_path and label_out_dir:
|
||||
out_label = os.path.join(label_out_dir, f"{new_base}{label_ext}")
|
||||
copiar_label_aug(
|
||||
label_path=label_path,
|
||||
out_label_path=out_label,
|
||||
new_base=new_base,
|
||||
out_img_path=out_img,
|
||||
out_msk_path=out_msk,
|
||||
out_msk2_path=out_msk2,
|
||||
group_name=group_name,
|
||||
)
|
||||
|
||||
gen += 1
|
||||
|
||||
return gen
|
||||
|
||||
def process_group(group_name, copies):
|
||||
"""Processa um grupo único (images/masks dentro de ORIG_GROUP_ROOT/<group_name>/)."""
|
||||
|
||||
# ====================== Processamento ======================
|
||||
|
||||
|
||||
def process_group(group_name, copies, strict_label=False):
|
||||
img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images")
|
||||
msk_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks")
|
||||
msk2_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks2")
|
||||
labels_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "labels")
|
||||
|
||||
if not (os.path.isdir(img_dir) and os.path.isdir(msk_dir)):
|
||||
print(f"[WARN] Grupo '{group_name}' inválido (sem images/masks). Pulando.")
|
||||
print(f"[WARN] Grupo '{group_name}' inválido, sem images/masks. Pulando.")
|
||||
return 0
|
||||
|
||||
imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||||
msk_map = map_masks_by_base(msk_dir)
|
||||
|
||||
msk_map = map_files_by_base(msk_dir, MSK_EXTS)
|
||||
|
||||
use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
|
||||
msk2_map = map_masks2_by_base(msk2_dir) if use_masks2 else {}
|
||||
img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name, use_masks2=use_masks2)
|
||||
msk2_map = map_files_by_base(msk2_dir, MSK2_EXTS) if use_masks2 else {}
|
||||
|
||||
use_labels = USE_LABELS and os.path.isdir(labels_dir)
|
||||
label_map = map_files_by_base(labels_dir, LABEL_EXTS) if use_labels else {}
|
||||
|
||||
if USE_MASKS2 and not use_masks2:
|
||||
print(f"[WARN] [{group_name}] dual_head_mask=true, mas pasta masks2 não existe.")
|
||||
|
||||
if USE_LABELS and not use_labels:
|
||||
msg = f"[WARN] [{group_name}] dual_head_label=true, mas pasta labels não existe."
|
||||
if strict_label:
|
||||
print(msg + " Pulando grupo.")
|
||||
return 0
|
||||
print(msg + " Gerando sem labels.")
|
||||
|
||||
img_out_dir, msk_out_dir, msk2_out_dir, label_out_dir = ensure_aug_dirs(
|
||||
group_name,
|
||||
use_masks2=use_masks2,
|
||||
use_labels=use_labels,
|
||||
)
|
||||
|
||||
count = 0
|
||||
sem_mask = 0
|
||||
sem_mask2 = 0
|
||||
sem_label = 0
|
||||
|
||||
for img_file in sorted(imgs):
|
||||
base, _ = os.path.splitext(img_file)
|
||||
msk_file = msk_map.get(base)
|
||||
base_norm = normalizar_base(base)
|
||||
|
||||
msk_file = msk_map.get(base_norm)
|
||||
if not msk_file:
|
||||
sem_mask += 1
|
||||
print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.")
|
||||
continue
|
||||
msk2_file = msk2_map.get(base) if use_masks2 else None
|
||||
|
||||
msk2_file = msk2_map.get(base_norm) if use_masks2 else None
|
||||
if use_masks2 and not msk2_file:
|
||||
sem_mask2 += 1
|
||||
print(f"[WARN] [{group_name}] mask2 não encontrada para {img_file}, gerando só img+mask.")
|
||||
|
||||
label_file = label_map.get(base_norm) if use_labels else None
|
||||
if use_labels and not label_file:
|
||||
sem_label += 1
|
||||
msg = f"[WARN] [{group_name}] label não encontrado para {img_file}."
|
||||
if strict_label:
|
||||
print(msg + " Pulando item.")
|
||||
continue
|
||||
print(msg + " Gerando augmentation sem label.")
|
||||
|
||||
try:
|
||||
count += augment_pair(
|
||||
os.path.join(img_dir, img_file),
|
||||
msk_file,
|
||||
img_out_dir,
|
||||
msk_out_dir,
|
||||
img_path=os.path.join(img_dir, img_file),
|
||||
msk_path=msk_file,
|
||||
img_out_dir=img_out_dir,
|
||||
msk_out_dir=msk_out_dir,
|
||||
copies=copies,
|
||||
msk2_path=msk2_file,
|
||||
msk2_out_dir=msk2_out_dir
|
||||
msk2_out_dir=msk2_out_dir,
|
||||
label_path=label_file,
|
||||
label_out_dir=label_out_dir,
|
||||
group_name=group_name,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[ERRO] [{group_name}] {img_file}: {e}")
|
||||
print(f"[OK] Grupo '{group_name}' → {count} pares gerados.")
|
||||
|
||||
print(
|
||||
f"[OK] Grupo '{group_name}' → {count} pares gerados. "
|
||||
f"sem_mask={sem_mask} | sem_mask2={sem_mask2} | sem_label={sem_label}"
|
||||
)
|
||||
return count
|
||||
|
||||
def process_legacy(copies):
|
||||
"""Fallback: modo sem grupos (original/images e original/masks)."""
|
||||
|
||||
def process_legacy(copies, strict_label=False):
|
||||
if not (os.path.isdir(ORIG_OLD_IMG) and os.path.isdir(ORIG_OLD_MSK)):
|
||||
print("[WARN] Modo legacy não encontrado. Nada a fazer.")
|
||||
return 0
|
||||
|
||||
imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||||
msk_map = map_masks_by_base(ORIG_OLD_MSK)
|
||||
|
||||
msk_map = map_files_by_base(ORIG_OLD_MSK, MSK_EXTS)
|
||||
|
||||
use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
|
||||
msk2_map = map_masks2_by_base(ORIG_OLD_MSK2) if use_masks2 else {}
|
||||
img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name=None, use_masks2=use_masks2)
|
||||
msk2_map = map_files_by_base(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}
|
||||
|
||||
use_labels = USE_LABELS and os.path.isdir(ORIG_OLD_LABELS)
|
||||
label_map = map_files_by_base(ORIG_OLD_LABELS, LABEL_EXTS) if use_labels else {}
|
||||
|
||||
if USE_LABELS and not use_labels and strict_label:
|
||||
print("[WARN] Legacy com dual_head_label=true, mas sem pasta labels. Pulando.")
|
||||
return 0
|
||||
|
||||
img_out_dir, msk_out_dir, msk2_out_dir, label_out_dir = ensure_aug_dirs(
|
||||
group_name=None,
|
||||
use_masks2=use_masks2,
|
||||
use_labels=use_labels,
|
||||
)
|
||||
|
||||
count = 0
|
||||
for img_file in sorted(imgs):
|
||||
base, _ = os.path.splitext(img_file)
|
||||
msk_file = msk_map.get(base)
|
||||
base_norm = normalizar_base(base)
|
||||
|
||||
msk_file = msk_map.get(base_norm)
|
||||
if not msk_file:
|
||||
print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
|
||||
continue
|
||||
msk2_file = msk2_map.get(base) if use_masks2 else None
|
||||
if use_masks2 and not msk2_file:
|
||||
print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.")
|
||||
|
||||
msk2_file = msk2_map.get(base_norm) if use_masks2 else None
|
||||
label_file = label_map.get(base_norm) if use_labels else None
|
||||
|
||||
if use_labels and not label_file and strict_label:
|
||||
print(f"[WARN] (legacy) Label não encontrado para {img_file}, pulando.")
|
||||
continue
|
||||
|
||||
try:
|
||||
count += augment_pair(
|
||||
os.path.join(ORIG_OLD_IMG, img_file),
|
||||
msk_file,
|
||||
img_out_dir,
|
||||
msk_out_dir,
|
||||
img_path=os.path.join(ORIG_OLD_IMG, img_file),
|
||||
msk_path=msk_file,
|
||||
img_out_dir=img_out_dir,
|
||||
msk_out_dir=msk_out_dir,
|
||||
copies=copies,
|
||||
msk2_path=msk2_file,
|
||||
msk2_out_dir=msk2_out_dir
|
||||
msk2_out_dir=msk2_out_dir,
|
||||
label_path=label_file,
|
||||
label_out_dir=label_out_dir,
|
||||
group_name=None,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[ERRO] (legacy) {img_file}: {e}")
|
||||
|
||||
print(f"[OK] Legacy → {count} pares gerados.")
|
||||
return count
|
||||
|
||||
def main(copies=5, groups_csv=None):
|
||||
|
||||
def main(copies=5, groups_csv=None, strict_label=False):
|
||||
total = 0
|
||||
|
||||
print(f"[INFO] MODELO={MODELO}")
|
||||
print(f"[INFO] dual_head_mask={USE_MASKS2}")
|
||||
print(f"[INFO] dual_head_label={USE_LABELS}")
|
||||
|
||||
if os.path.isdir(ORIG_GROUP_ROOT):
|
||||
grupos = list_groups(ORIG_GROUP_ROOT)
|
||||
|
||||
if groups_csv:
|
||||
# filtra pelos grupos desejados
|
||||
want = {g.strip() for g in groups_csv.split(",") if g.strip()}
|
||||
grupos = [g for g in grupos if g in want]
|
||||
if not grupos:
|
||||
print("[WARN] Nenhum grupo válido encontrado após filtro.")
|
||||
|
||||
if not grupos:
|
||||
print("[WARN] Nenhum grupo encontrado em original/group. Tentando modo legacy...")
|
||||
total += process_legacy(copies)
|
||||
total += process_legacy(copies, strict_label=strict_label)
|
||||
else:
|
||||
print(f"Grupos encontrados: {', '.join(grupos)}")
|
||||
for g in grupos:
|
||||
total += process_group(g, copies)
|
||||
total += process_group(g, copies, strict_label=strict_label)
|
||||
else:
|
||||
# sem estrutura de grupos
|
||||
total += process_legacy(copies)
|
||||
total += process_legacy(copies, strict_label=strict_label)
|
||||
|
||||
print(f"\nAugmentation completed! Total: {total} pares gerados.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser(description="Augmentação por grupos (images/masks)")
|
||||
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem (default=5).")
|
||||
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula (ex: chao,erva_cana).")
|
||||
ap = argparse.ArgumentParser(description="Augmentação por grupos com suporte a images/masks/masks2/labels.")
|
||||
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem.")
|
||||
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula.")
|
||||
ap.add_argument("--strict-label", action="store_true", help="Se dual_head_label=true e faltar label, pula o item/grupo.")
|
||||
args = ap.parse_args()
|
||||
main(copies=args.copies, groups_csv=args.groups)
|
||||
|
||||
main(copies=args.copies, groups_csv=args.groups, strict_label=args.strict_label)
|
||||
|
|
|
|||
|
|
@ -1,67 +1,86 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Normaliza/redimensiona imagens e máscaras mantendo a ESTRUTURA POR GRUPO.
|
||||
Normaliza/redimensiona imagens, máscaras, masks2 e labels mantendo a ESTRUTURA POR GRUPO.
|
||||
|
||||
Entradas (via config.json -> MODELO, RESOLUCAO):
|
||||
- MODELO/dataset/original/group/<grupo>/{images,masks}
|
||||
- MODELO/dataset/augmented/group/<grupo>/{images,masks}
|
||||
Entradas via config.json -> camera, resolucao:
|
||||
- MODELO/dataset/original/group/<grupo>/{images,masks,(masks2),(labels)}
|
||||
- MODELO/dataset/augmented/group/<grupo>/{images,masks,(masks2),(labels)}
|
||||
|
||||
Saídas (por resolução):
|
||||
- MODELO/dataset/<WxH>/group/<grupo>/{images,masks}
|
||||
Saídas por resolução:
|
||||
- MODELO/dataset/<WxH>/group/<grupo>/{images,masks,(masks2),(labels)}
|
||||
|
||||
Fallback (modo legado, se não houver "group/"):
|
||||
- original/{images,masks} e augmented/{images,masks} -> <WxH>/{images,masks}
|
||||
Fallback legado, se não houver group/:
|
||||
- original/{images,masks,(masks2),(labels)}
|
||||
- augmented/{images,masks,(masks2),(labels)}
|
||||
- saída: <WxH>/{images,masks,(masks2),(labels)}
|
||||
|
||||
Conversão de máscara:
|
||||
- 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 .../<fonte>/group/ agrupando por cada subpasta de grupo."""
|
||||
|
||||
# ===================== PROCESSAMENTO =====================
|
||||
|
||||
def process_group_root(
|
||||
fonte_root: str,
|
||||
fonte_nome: str,
|
||||
cor_para_id,
|
||||
ignore_id: int,
|
||||
groups_except: str = "",
|
||||
strict_label: bool = False,
|
||||
):
|
||||
total = 0
|
||||
grupos = list_groups(fonte_root)
|
||||
if not grupos:
|
||||
return 0
|
||||
|
||||
|
||||
not_want = {g.strip() for g in groups_except.split(",") if g.strip()}
|
||||
grupos_desconsiderar = [g for g in grupos if g in not_want]
|
||||
|
||||
for nome_res, dim in RESOLUCOES.items():
|
||||
out_root = os.path.join(pasta_base, nome_res, "group")
|
||||
|
||||
for grupo in grupos:
|
||||
if grupo in grupos_desconsiderar:
|
||||
print(f"[WARN] Grupo desconsiderado nao sera processado: {grupo}")
|
||||
print(f"[WARN] Grupo desconsiderado não será processado: {grupo}")
|
||||
continue
|
||||
|
||||
in_img_dir = os.path.join(fonte_root, grupo, "images")
|
||||
in_msk_dir = os.path.join(fonte_root, grupo, "masks")
|
||||
in_msk2_dir = os.path.join(fonte_root, grupo, "masks2")
|
||||
in_label_dir = os.path.join(fonte_root, grupo, "labels")
|
||||
|
||||
if not (os.path.isdir(in_img_dir) and os.path.isdir(in_msk_dir)):
|
||||
print(f"[WARN] Grupo inválido (sem images/masks): {grupo}")
|
||||
print(f"[WARN] Grupo inválido sem images/masks: {grupo}")
|
||||
continue
|
||||
|
||||
out_img_dir = os.path.join(out_root, grupo, "images")
|
||||
out_msk_dir = os.path.join(out_root, grupo, "masks")
|
||||
|
||||
usar_masks2 = USE_MASKS2 and os.path.isdir(in_msk2_dir)
|
||||
usar_labels = USE_LABELS and os.path.isdir(in_label_dir)
|
||||
|
||||
out_msk2_dir = os.path.join(out_root, grupo, "masks2") if usar_masks2 else None
|
||||
out_label_dir = os.path.join(out_root, grupo, "labels") if usar_labels else None
|
||||
|
||||
msk_map = map_masks_by_base(in_msk_dir)
|
||||
msk2_map = map_masks2_by_base(in_msk2_dir) if usar_masks2 else {}
|
||||
label_map = map_labels_by_base(in_label_dir) if usar_labels else {}
|
||||
|
||||
if USE_LABELS and not usar_labels:
|
||||
msg = f"[WARN] [{fonte_nome} | {grupo}] dual_head_label=true, mas labels/ não existe."
|
||||
if strict_label:
|
||||
print(msg + " Pulando grupo.")
|
||||
continue
|
||||
print(msg + " Seguindo sem labels.")
|
||||
|
||||
imgs = [f for f in os.listdir(in_img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||||
n = len(imgs)
|
||||
|
||||
for i, fname in enumerate(sorted(imgs), 1):
|
||||
base, _ = os.path.splitext(fname)
|
||||
base_norm = normalizar_base(base)
|
||||
|
||||
caminho_rgb = os.path.join(in_img_dir, fname)
|
||||
caminho_mask = msk_map.get(base)
|
||||
caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
|
||||
ok = normalize_pair(
|
||||
caminho_rgb, caminho_mask, cor_para_id, ignore_id,
|
||||
out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_"
|
||||
caminho_mask = msk_map.get(base_norm)
|
||||
caminho_mask2 = msk2_map.get(base_norm) if usar_masks2 else None
|
||||
caminho_label = label_map.get(base_norm) if usar_labels else None
|
||||
|
||||
if usar_labels and not caminho_label:
|
||||
msg = f"[WARN] [{fonte_nome} | {grupo}] label não encontrado para {fname}."
|
||||
if strict_label:
|
||||
print(msg + " Pulando item.")
|
||||
continue
|
||||
print(msg + " Seguindo sem label.")
|
||||
|
||||
ok, out_img_path, out_msk_path = normalize_pair(
|
||||
caminho_rgb,
|
||||
caminho_mask,
|
||||
cor_para_id,
|
||||
ignore_id,
|
||||
out_img_dir,
|
||||
out_msk_dir,
|
||||
dim,
|
||||
prefix=f"{fonte_nome}_",
|
||||
)
|
||||
|
||||
out_msk2_path = None
|
||||
if usar_masks2 and out_msk2_dir:
|
||||
if not caminho_mask2:
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname} (vou seguir).")
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname}.")
|
||||
else:
|
||||
normalize_pair_mask2(
|
||||
caminho_rgb, caminho_mask2,
|
||||
out_msk2_dir, dim, prefix=f"{fonte_nome}_"
|
||||
_, out_msk2_path = normalize_pair_mask2(
|
||||
caminho_rgb,
|
||||
caminho_mask2,
|
||||
out_msk2_dir,
|
||||
dim,
|
||||
prefix=f"{fonte_nome}_",
|
||||
)
|
||||
|
||||
if usar_labels and caminho_label and out_label_dir:
|
||||
normalize_label(
|
||||
caminho_rgb=caminho_rgb,
|
||||
caminho_label=caminho_label,
|
||||
out_label_dir=out_label_dir,
|
||||
out_img_path=out_img_path,
|
||||
out_msk_path=out_msk_path,
|
||||
out_msk2_path=out_msk2_path,
|
||||
source_root=fonte_root,
|
||||
output_root=os.path.join(pasta_base, nome_res),
|
||||
grupo=grupo,
|
||||
fonte_nome=fonte_nome,
|
||||
prefix=f"{fonte_nome}_",
|
||||
)
|
||||
|
||||
if ok:
|
||||
total += 1
|
||||
|
||||
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n} → {fname}")
|
||||
|
||||
return total
|
||||
|
||||
def process_legacy_root(legacy_img: str, legacy_msk: str, fonte_nome: str, cor_para_id, ignore_id: int):
|
||||
"""Processa estrutura legado (sem grupos)."""
|
||||
|
||||
def process_legacy_root(
|
||||
legacy_img: str,
|
||||
legacy_msk: str,
|
||||
fonte_nome: str,
|
||||
cor_para_id,
|
||||
ignore_id: int,
|
||||
strict_label: bool = False,
|
||||
):
|
||||
if not (os.path.isdir(legacy_img) and os.path.isdir(legacy_msk)):
|
||||
return 0
|
||||
|
||||
total = 0
|
||||
|
||||
for nome_res, dim in RESOLUCOES.items():
|
||||
out_img_dir = os.path.join(pasta_base, nome_res, "images")
|
||||
out_msk_dir = os.path.join(pasta_base, nome_res, "masks")
|
||||
legacy_msk2 = os.path.join(os.path.dirname(legacy_msk), "masks2")
|
||||
|
||||
legacy_root = os.path.dirname(legacy_msk)
|
||||
legacy_msk2 = os.path.join(legacy_root, "masks2")
|
||||
legacy_labels = os.path.join(legacy_root, "labels")
|
||||
|
||||
usar_masks2 = USE_MASKS2 and os.path.isdir(legacy_msk2)
|
||||
usar_labels = USE_LABELS and os.path.isdir(legacy_labels)
|
||||
|
||||
out_msk2_dir = os.path.join(pasta_base, nome_res, "masks2") if usar_masks2 else None
|
||||
out_label_dir = os.path.join(pasta_base, nome_res, "labels") if usar_labels else None
|
||||
|
||||
msk_map = map_masks_by_base(legacy_msk)
|
||||
msk2_map = map_masks2_by_base(legacy_msk2) if usar_masks2 else {}
|
||||
label_map = map_labels_by_base(legacy_labels) if usar_labels else {}
|
||||
|
||||
if USE_LABELS and not usar_labels:
|
||||
msg = f"[WARN] [{fonte_nome} | legacy] dual_head_label=true, mas labels/ não existe."
|
||||
if strict_label:
|
||||
print(msg + " Pulando legacy.")
|
||||
return total
|
||||
print(msg + " Seguindo sem labels.")
|
||||
|
||||
imgs = [f for f in os.listdir(legacy_img) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||||
n = len(imgs)
|
||||
|
||||
for i, fname in enumerate(sorted(imgs), 1):
|
||||
base, _ = os.path.splitext(fname)
|
||||
base_norm = normalizar_base(base)
|
||||
|
||||
caminho_rgb = os.path.join(legacy_img, fname)
|
||||
caminho_mask = msk_map.get(base)
|
||||
caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
|
||||
ok = normalize_pair(
|
||||
caminho_rgb, caminho_mask, cor_para_id, ignore_id,
|
||||
out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_"
|
||||
caminho_mask = msk_map.get(base_norm)
|
||||
caminho_mask2 = msk2_map.get(base_norm) if usar_masks2 else None
|
||||
caminho_label = label_map.get(base_norm) if usar_labels else None
|
||||
|
||||
if usar_labels and not caminho_label:
|
||||
msg = f"[WARN] [{fonte_nome} | legacy] label não encontrado para {fname}."
|
||||
if strict_label:
|
||||
print(msg + " Pulando item.")
|
||||
continue
|
||||
print(msg + " Seguindo sem label.")
|
||||
|
||||
ok, out_img_path, out_msk_path = normalize_pair(
|
||||
caminho_rgb,
|
||||
caminho_mask,
|
||||
cor_para_id,
|
||||
ignore_id,
|
||||
out_img_dir,
|
||||
out_msk_dir,
|
||||
dim,
|
||||
prefix=f"{fonte_nome}_",
|
||||
)
|
||||
|
||||
out_msk2_path = None
|
||||
if usar_masks2 and out_msk2_dir:
|
||||
if not caminho_mask2:
|
||||
print(f"[WARN] [{fonte_nome} | legacy] masks2 existe, mas não achei mask2 p/ {fname} (vou seguir).")
|
||||
print(f"[WARN] [{fonte_nome} | legacy] masks2 existe, mas não achei mask2 p/ {fname}.")
|
||||
else:
|
||||
normalize_pair_mask2(
|
||||
caminho_rgb, caminho_mask2,
|
||||
out_msk2_dir, dim, prefix=f"{fonte_nome}_"
|
||||
_, out_msk2_path = normalize_pair_mask2(
|
||||
caminho_rgb,
|
||||
caminho_mask2,
|
||||
out_msk2_dir,
|
||||
dim,
|
||||
prefix=f"{fonte_nome}_",
|
||||
)
|
||||
|
||||
if usar_labels and caminho_label and out_label_dir:
|
||||
normalize_label(
|
||||
caminho_rgb=caminho_rgb,
|
||||
caminho_label=caminho_label,
|
||||
out_label_dir=out_label_dir,
|
||||
out_img_path=out_img_path,
|
||||
out_msk_path=out_msk_path,
|
||||
out_msk2_path=out_msk2_path,
|
||||
source_root=legacy_root,
|
||||
output_root=os.path.join(pasta_base, nome_res),
|
||||
grupo=None,
|
||||
fonte_nome=fonte_nome,
|
||||
prefix=f"{fonte_nome}_",
|
||||
)
|
||||
|
||||
if ok:
|
||||
total += 1
|
||||
|
||||
print(f"[{fonte_nome} | legacy | {nome_res}] {i}/{n} → {fname}")
|
||||
|
||||
return total
|
||||
|
||||
|
||||
# ===================== MAIN =====================
|
||||
|
||||
def main(args):
|
||||
# === Labelmap ===
|
||||
# Espera tupla na ordem: (cor_para_id, colormap_rgb, id_para_nome, ignore_rgb)
|
||||
cor_para_id, _colormap_rgb, _id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
ignore_id = infer_ignore_id(ignore_rgb, default_id=255)
|
||||
print(cor_para_id, _colormap_rgb, _id_para_nome)
|
||||
|
||||
print(f"[INFO] cor_para_id: {cor_para_id}")
|
||||
print(f"[INFO] classes: {_id_para_nome}")
|
||||
print(f"[INFO] dual_head_mask/masks2: {USE_MASKS2}")
|
||||
print(f"[INFO] dual_head_label/labels: {USE_LABELS}")
|
||||
|
||||
total_geral = 0
|
||||
# === ORIGINAL ===
|
||||
|
||||
orig_group_root = os.path.join(pasta_base, "original", "group")
|
||||
if os.path.isdir(orig_group_root):
|
||||
total_geral += process_group_root(orig_group_root, "original", cor_para_id, ignore_id, groups_except=args.groups_except)
|
||||
total_geral += process_group_root(
|
||||
orig_group_root,
|
||||
"original",
|
||||
cor_para_id,
|
||||
ignore_id,
|
||||
groups_except=args.groups_except,
|
||||
strict_label=args.strict_label,
|
||||
)
|
||||
else:
|
||||
legacy_img = os.path.join(pasta_base, "original", "images")
|
||||
legacy_msk = os.path.join(pasta_base, "original", "masks")
|
||||
total_geral += process_legacy_root(legacy_img, legacy_msk, "original", cor_para_id, ignore_id)
|
||||
total_geral += process_legacy_root(
|
||||
legacy_img,
|
||||
legacy_msk,
|
||||
"original",
|
||||
cor_para_id,
|
||||
ignore_id,
|
||||
strict_label=args.strict_label,
|
||||
)
|
||||
|
||||
# === AUGMENTED ===
|
||||
aug_group_root = os.path.join(pasta_base, "augmented", "group")
|
||||
if os.path.isdir(aug_group_root):
|
||||
total_geral += process_group_root(aug_group_root, "augmented", cor_para_id, ignore_id, groups_except=args.groups_except)
|
||||
total_geral += process_group_root(
|
||||
aug_group_root,
|
||||
"augmented",
|
||||
cor_para_id,
|
||||
ignore_id,
|
||||
groups_except=args.groups_except,
|
||||
strict_label=args.strict_label,
|
||||
)
|
||||
else:
|
||||
legacy_img = os.path.join(pasta_base, "augmented", "images")
|
||||
legacy_msk = os.path.join(pasta_base, "augmented", "masks")
|
||||
total_geral += process_legacy_root(legacy_img, legacy_msk, "augmented", cor_para_id, ignore_id)
|
||||
total_geral += process_legacy_root(
|
||||
legacy_img,
|
||||
legacy_msk,
|
||||
"augmented",
|
||||
cor_para_id,
|
||||
ignore_id,
|
||||
strict_label=args.strict_label,
|
||||
)
|
||||
|
||||
print(f"\n✅ Concluído! Total normalizados: {total_geral}")
|
||||
|
||||
# === calcula mean/std globais e salva em JSON ===
|
||||
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
|
||||
if GLOBAL_SUM is not None and GLOBAL_PIXELS > 0:
|
||||
# média e variância por canal
|
||||
mean = (GLOBAL_SUM / GLOBAL_PIXELS)
|
||||
var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean**2
|
||||
std = np.sqrt(np.maximum(var, 1e-6))
|
||||
mean = GLOBAL_SUM / GLOBAL_PIXELS
|
||||
var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean ** 2
|
||||
std = np.sqrt(np.maximum(var, 1e-6))
|
||||
|
||||
# Converte para list pra salvar em JSON
|
||||
mean_list = mean.tolist()
|
||||
std_list = std.tolist()
|
||||
|
||||
# Se quiser, você pode nomear os canais explicitamente
|
||||
# dependendo da convenção do raw4:
|
||||
std_list = std.tolist()
|
||||
channel_names = ["R", "G", "B"]
|
||||
|
||||
stats = {
|
||||
|
|
@ -369,10 +664,21 @@ def main(args):
|
|||
print(f" mean: {mean_list}")
|
||||
print(f" std : {std_list}")
|
||||
else:
|
||||
print("⚠️ Nenhum RAW processado, não há stats para salvar.")
|
||||
print("⚠️ Nenhuma imagem processada, não há stats para salvar.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser(description="Augmentação por grupos (images/masks)")
|
||||
ap.add_argument("--groups-except", type=str, default="", help="Lista de grupos para nao usar, separados por vírgula (ex: chao,erva_cana).")
|
||||
ap = argparse.ArgumentParser(description="Normalização por grupos images/masks/masks2/labels.")
|
||||
ap.add_argument(
|
||||
"--groups-except",
|
||||
type=str,
|
||||
default="",
|
||||
help="Lista de grupos para não usar, separados por vírgula. Ex: chao,erva_cana",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--strict-label",
|
||||
action="store_true",
|
||||
help="Se dual_head_label=true e faltar label, pula item/grupo.",
|
||||
)
|
||||
args = ap.parse_args()
|
||||
main(args)
|
||||
|
|
|
|||
|
|
@ -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/<WxH>/group/<grupo>/{images,masks}
|
||||
MODELO/dataset/<WxH>/group/<grupo>/{images,masks,(masks2),(labels)}
|
||||
|
||||
Escreve em:
|
||||
MODELO/dataset/split/<split>/group/<grupo>/{images,masks}
|
||||
MODELO/dataset/split/<split>/group/<grupo>/{images,masks,(masks2),(labels)}
|
||||
|
||||
Definições:
|
||||
- "Família" = todas as variações da MESMA base original:
|
||||
original_<base>.* e augmented_<base>_aug_XX.*
|
||||
- Val/Test: só **original_<base>** (sem augmented)
|
||||
- Train: original_<base> **e** todos augmented_<base>_aug_XX
|
||||
- Família = todas as variações da mesma base original:
|
||||
original_<base>.*
|
||||
augmented_<base>_aug_XX.*
|
||||
- Val/Test: somente original_<base>.
|
||||
- Train: original_<base> + todos augmented_<base>_aug_XX.
|
||||
|
||||
Se não houver prefixos (legado), cai para o comportamento antigo (sem família),
|
||||
mas ainda evita colocar augmented em val/test se detectar sufixo "_aug_XX".
|
||||
Labels:
|
||||
- Ativados por config['dual_head_label'].
|
||||
- Copia labels .json/.txt e .npy quando existirem.
|
||||
- O pareamento é feito pelo mesmo base name da imagem.
|
||||
|
||||
Uso:
|
||||
python _7_split_grouped_noleak.py
|
||||
python _7_split_grouped_noleak.py --train 0.7 --val 0.29 --test 0.01 --seed 42
|
||||
python _7_split_grouped_noleak.py --min-train 1 --min-val 1 --min-test 0
|
||||
python _7_split_grouped_noleak.py --modelo OAK-1-Lite-W --resolucao 640x384
|
||||
python _7_split_grouped_noleak_with_labels.py
|
||||
python _7_split_grouped_noleak_with_labels.py --train 0.7 --val 0.29 --test 0.01 --seed 42
|
||||
python _7_split_grouped_noleak_with_labels.py --strict-label
|
||||
python _7_split_grouped_noleak_with_labels.py --cap-train-families "navegavel:300,naonavegavel_navegavel:800"
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import shutil
|
||||
import random
|
||||
import argparse
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
|
||||
# ===================== CONFIG =====================
|
||||
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config.get("camera")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
USE_MASKS2 = bool(config.get("dual_head_mask", config.get("dual_head", False)))
|
||||
USE_LABELS = bool(config.get("dual_head_label", False))
|
||||
RESOLUCAO = tuple(config.get("resolucao"))
|
||||
|
||||
# Pastas
|
||||
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
|
||||
pasta_destino = os.path.join(MODELO, "dataset", "split")
|
||||
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXT = ".png" # máscaras normalizadas em PNG (recomendado)
|
||||
MSK_EXT = ".png"
|
||||
LABEL_EXTS = (".json", ".txt", ".npy")
|
||||
|
||||
# Regex para identificar famílias
|
||||
RE_ORIGINAL_PREFIX = re.compile(r'^original_(.+)$', re.IGNORECASE)
|
||||
RE_AUGMENTED_FAMILY = re.compile(r'^augmented_(.+?)(?:_aug_\d+)?$', re.IGNORECASE)
|
||||
RE_AUG_SUFFIX = re.compile(r'_aug_\d+$', re.IGNORECASE)
|
||||
RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
|
||||
RE_AUGMENTED_FAMILY = re.compile(r"^augmented_(.+?)(?:_aug_\d+)?$", re.IGNORECASE)
|
||||
RE_AUG_SUFFIX = re.compile(r"_aug_\d+$", re.IGNORECASE)
|
||||
|
||||
|
||||
# ===================== HELPERS =====================
|
||||
|
||||
def garantir(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
|
||||
def lista_grupos(root):
|
||||
if not os.path.isdir(root):
|
||||
return []
|
||||
|
||||
out = []
|
||||
for g in sorted(os.listdir(root)):
|
||||
gdir = os.path.join(root, g)
|
||||
if not os.path.isdir(gdir): continue
|
||||
if not os.path.isdir(gdir):
|
||||
continue
|
||||
if os.path.isdir(os.path.join(gdir, "images")) and os.path.isdir(os.path.join(gdir, "masks")):
|
||||
out.append(g)
|
||||
return out
|
||||
|
||||
|
||||
def listar_imagens(img_dir):
|
||||
if not os.path.isdir(img_dir): return []
|
||||
if not os.path.isdir(img_dir):
|
||||
return []
|
||||
|
||||
fs = []
|
||||
for f in os.listdir(img_dir):
|
||||
ext = os.path.splitext(f.lower())[1]
|
||||
|
|
@ -74,19 +90,27 @@ def listar_imagens(img_dir):
|
|||
fs.append(f)
|
||||
return sorted(fs)
|
||||
|
||||
|
||||
def mask_from_image_name(img_name):
|
||||
base, _ = os.path.splitext(img_name)
|
||||
return base + MSK_EXT
|
||||
|
||||
|
||||
def mask2_from_image_name(img_name):
|
||||
base, _ = os.path.splitext(img_name)
|
||||
return base + MSK_EXT # masks2 também normalizadas em PNG no normalize
|
||||
return base + MSK_EXT
|
||||
|
||||
|
||||
def label_candidates_from_image_name(img_name):
|
||||
base, _ = os.path.splitext(img_name)
|
||||
return [base + ext for ext in LABEL_EXTS]
|
||||
|
||||
|
||||
def classify_source_and_family(filename_no_ext):
|
||||
"""
|
||||
Retorna (source, family_key)
|
||||
source ∈ {"original", "augmented", "unknown"}
|
||||
family_key = base associada ao original (sem prefixo/sufixos), ex: "foo_001"
|
||||
source ∈ {original, augmented, unknown}
|
||||
family_key = base original sem prefixo/sufixo.
|
||||
"""
|
||||
m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
|
||||
if m:
|
||||
|
|
@ -96,46 +120,66 @@ def classify_source_and_family(filename_no_ext):
|
|||
if m:
|
||||
return "augmented", m.group(1)
|
||||
|
||||
# legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
|
||||
if RE_AUG_SUFFIX.search(filename_no_ext):
|
||||
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
|
||||
return "augmented", fam
|
||||
|
||||
return "unknown", filename_no_ext
|
||||
|
||||
def build_family_index(img_dir, msk_dir):
|
||||
|
||||
def has_any_label(label_dir: str, img_name: str) -> bool:
|
||||
if not label_dir or not os.path.isdir(label_dir):
|
||||
return False
|
||||
return any(os.path.exists(os.path.join(label_dir, cand)) for cand in label_candidates_from_image_name(img_name))
|
||||
|
||||
|
||||
def build_family_index(img_dir, msk_dir, label_dir=None, require_label=False):
|
||||
"""
|
||||
Constroi índice de famílias a partir de img_dir/msk_dir.
|
||||
Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]}
|
||||
(strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe)
|
||||
Constrói índice de famílias a partir de img_dir/msk_dir/labels.
|
||||
|
||||
Retorna:
|
||||
dict family -> {
|
||||
original: str|None,
|
||||
augmented: [str],
|
||||
all: [str]
|
||||
}
|
||||
|
||||
Os nomes são arquivos de imagem.
|
||||
"""
|
||||
familias = {}
|
||||
imgs = listar_imagens(img_dir)
|
||||
|
||||
for img_name in imgs:
|
||||
base_no_ext, ext = os.path.splitext(img_name)
|
||||
base_no_ext, _ = os.path.splitext(img_name)
|
||||
mask_name = mask_from_image_name(img_name)
|
||||
|
||||
if not os.path.exists(os.path.join(msk_dir, mask_name)):
|
||||
continue # garante pareamento
|
||||
continue
|
||||
|
||||
if require_label and not has_any_label(label_dir, img_name):
|
||||
continue
|
||||
|
||||
source, fam = classify_source_and_family(base_no_ext)
|
||||
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
|
||||
d["all"].append(img_name)
|
||||
|
||||
if source == "original":
|
||||
d["original"] = img_name
|
||||
elif source == "augmented":
|
||||
d["augmented"].append(img_name)
|
||||
else:
|
||||
# trata como original desconhecido para não perder dado
|
||||
if d["original"] is None:
|
||||
d["original"] = img_name
|
||||
else:
|
||||
d["augmented"].append(img_name)
|
||||
|
||||
return familias
|
||||
|
||||
|
||||
def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
||||
n_train = int(round(n * p_train))
|
||||
n_val = int(round(n * p_val))
|
||||
n_test = n - n_train - n_val
|
||||
n_val = int(round(n * p_val))
|
||||
n_test = n - n_train - n_val
|
||||
|
||||
if n_test < 0:
|
||||
excesso = -n_test
|
||||
|
|
@ -151,8 +195,8 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
min_sum = min_train + min_val + min_test
|
||||
if n >= min_sum:
|
||||
n_train = max(n_train, min_train)
|
||||
n_val = max(n_val, min_val)
|
||||
n_test = max(n_test, min_test)
|
||||
n_val = max(n_val, min_val)
|
||||
n_test = max(n_test, min_test)
|
||||
|
||||
total = n_train + n_val + n_test
|
||||
while total > n:
|
||||
|
|
@ -165,6 +209,7 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
else:
|
||||
break
|
||||
total = n_train + n_val + n_test
|
||||
|
||||
while total < n:
|
||||
if n_train - min_train <= n_val - min_val:
|
||||
n_train += 1
|
||||
|
|
@ -177,11 +222,9 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
n_val = max(0, min(resto, min_val))
|
||||
n_test = max(0, resto - n_val)
|
||||
|
||||
# ajuste final
|
||||
diff = n - (n_train + n_val + n_test)
|
||||
if diff != 0:
|
||||
if diff > 0:
|
||||
# adiciona em train, depois val
|
||||
take = min(diff, n - n_train)
|
||||
n_train += take
|
||||
diff -= take
|
||||
|
|
@ -189,7 +232,6 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
n_val += diff
|
||||
else:
|
||||
diff = -diff
|
||||
# tira de test, depois val
|
||||
take = min(diff, n_test)
|
||||
n_test -= take
|
||||
diff -= take
|
||||
|
|
@ -198,38 +240,113 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
|
||||
return n_train, n_val, n_test
|
||||
|
||||
def copiar(nomes, src_img_dir, src_msk_dir, dst_img_dir, dst_msk_dir, src_msk2_dir=None, dst_msk2_dir=None):
|
||||
garantir(dst_img_dir); garantir(dst_msk_dir)
|
||||
|
||||
def copiar_labels_para_item(nome, src_label_dir, dst_label_dir):
|
||||
if not src_label_dir or not dst_label_dir or not os.path.isdir(src_label_dir):
|
||||
return 0
|
||||
|
||||
garantir(dst_label_dir)
|
||||
copied = 0
|
||||
|
||||
for cand in label_candidates_from_image_name(nome):
|
||||
src = os.path.join(src_label_dir, cand)
|
||||
if not os.path.exists(src):
|
||||
continue
|
||||
shutil.copy2(src, os.path.join(dst_label_dir, cand))
|
||||
copied += 1
|
||||
|
||||
return copied
|
||||
|
||||
|
||||
def copiar(
|
||||
nomes,
|
||||
src_img_dir,
|
||||
src_msk_dir,
|
||||
dst_img_dir,
|
||||
dst_msk_dir,
|
||||
src_msk2_dir=None,
|
||||
dst_msk2_dir=None,
|
||||
src_label_dir=None,
|
||||
dst_label_dir=None,
|
||||
strict_label=False,
|
||||
):
|
||||
garantir(dst_img_dir)
|
||||
garantir(dst_msk_dir)
|
||||
|
||||
use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir))
|
||||
use_labels = bool(src_label_dir and dst_label_dir and os.path.isdir(src_label_dir))
|
||||
|
||||
if use_msk2:
|
||||
garantir(dst_msk2_dir)
|
||||
if use_labels:
|
||||
garantir(dst_label_dir)
|
||||
|
||||
moved = 0
|
||||
skipped_no_label = 0
|
||||
|
||||
for nome in nomes:
|
||||
mask_name = mask_from_image_name(nome)
|
||||
src_img = os.path.join(src_img_dir, nome)
|
||||
src_msk = os.path.join(src_msk_dir, mask_name)
|
||||
|
||||
if not (os.path.exists(src_img) and os.path.exists(src_msk)):
|
||||
continue
|
||||
|
||||
if strict_label and use_labels and not has_any_label(src_label_dir, nome):
|
||||
skipped_no_label += 1
|
||||
continue
|
||||
|
||||
shutil.copy2(src_img, os.path.join(dst_img_dir, nome))
|
||||
shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name))
|
||||
|
||||
if use_msk2:
|
||||
m2_name = mask2_from_image_name(nome)
|
||||
src_m2 = os.path.join(src_msk2_dir, m2_name)
|
||||
if os.path.exists(src_m2):
|
||||
shutil.copy2(src_m2, os.path.join(dst_msk2_dir, m2_name))
|
||||
|
||||
if use_labels:
|
||||
copied = copiar_labels_para_item(nome, src_label_dir, dst_label_dir)
|
||||
if strict_label and copied == 0:
|
||||
skipped_no_label += 1
|
||||
continue
|
||||
|
||||
moved += 1
|
||||
|
||||
if skipped_no_label > 0:
|
||||
print(f"[WARN] {skipped_no_label} itens pulados por falta de label.")
|
||||
|
||||
return moved
|
||||
|
||||
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
|
||||
|
||||
# ===================== SPLIT =====================
|
||||
|
||||
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None, strict_label=False):
|
||||
src_img_dir = os.path.join(pasta_origem, group_name, "images")
|
||||
src_msk_dir = os.path.join(pasta_origem, group_name, "masks")
|
||||
src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2")
|
||||
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
|
||||
src_label_dir = os.path.join(pasta_origem, group_name, "labels")
|
||||
|
||||
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
|
||||
use_labels = USE_LABELS and os.path.isdir(src_label_dir)
|
||||
|
||||
if USE_LABELS and not use_labels:
|
||||
msg = f"[{group_name}] dual_head_label=true, mas labels/ não existe."
|
||||
if strict_label:
|
||||
print(f"[WARN] {msg} Pulando grupo.")
|
||||
return {"train": 0, "val": 0, "test": 0, "familias": 0}
|
||||
print(f"[WARN] {msg} Split seguirá sem copiar labels.")
|
||||
|
||||
familias = build_family_index(
|
||||
src_img_dir,
|
||||
src_msk_dir,
|
||||
label_dir=src_label_dir if use_labels else None,
|
||||
require_label=bool(strict_label and use_labels),
|
||||
)
|
||||
|
||||
familias = build_family_index(src_img_dir, src_msk_dir)
|
||||
# apenas famílias que têm ORIGINAL para participar de val/test
|
||||
familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
|
||||
total_familias = len(familias_originais)
|
||||
|
||||
if total_familias == 0:
|
||||
print(f"[{group_name}] 0 famílias com original, pulando.")
|
||||
return {"train": 0, "val": 0, "test": 0, "familias": 0}
|
||||
|
|
@ -238,81 +355,145 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
|
|||
rng.shuffle(familias_originais)
|
||||
|
||||
n_tr, n_va, n_te = allocate_counts(
|
||||
total_familias, p_train, p_val, p_test,
|
||||
mins["train"], mins["val"], mins["test"]
|
||||
total_familias,
|
||||
p_train,
|
||||
p_val,
|
||||
p_test,
|
||||
mins["train"],
|
||||
mins["val"],
|
||||
mins["test"],
|
||||
)
|
||||
|
||||
fam_train = set(familias_originais[:n_tr])
|
||||
fam_val = set(familias_originais[n_tr:n_tr+n_va])
|
||||
fam_test = set(familias_originais[n_tr+n_va: n_tr+n_va+n_te])
|
||||
fam_val = set(familias_originais[n_tr : n_tr + n_va])
|
||||
fam_test = set(familias_originais[n_tr + n_va : n_tr + n_va + n_te])
|
||||
|
||||
# --- CAP por grupo (apenas no TRAIN) ---
|
||||
if caps_map and group_name in caps_map:
|
||||
cap = caps_map[group_name]
|
||||
if len(fam_train) > cap:
|
||||
fam_list = list(fam_train)
|
||||
rng.shuffle(fam_list) # usa o rng já criado com seed
|
||||
kept = set(fam_list[:cap])
|
||||
dropped = set(fam_list[cap:])
|
||||
rng.shuffle(fam_list)
|
||||
kept = set(fam_list[:cap])
|
||||
dropped = set(fam_list[cap:])
|
||||
fam_train = kept
|
||||
print(f"[{group_name}] cap-train-families={cap} → mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN")
|
||||
print(
|
||||
f"[{group_name}] cap-train-families={cap} → "
|
||||
f"mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN"
|
||||
)
|
||||
|
||||
# listas de nomes por split (imagens)
|
||||
nomes_train, nomes_val, nomes_test = [], [], []
|
||||
|
||||
for fam, d in familias.items():
|
||||
if fam in fam_train:
|
||||
# train recebe original + todos augmented
|
||||
if d["original"]:
|
||||
nomes_train.append(d["original"])
|
||||
if d["augmented"]:
|
||||
nomes_train.extend(d["augmented"])
|
||||
elif fam in fam_val:
|
||||
# val recebe somente original
|
||||
if d["original"]:
|
||||
nomes_val.append(d["original"])
|
||||
elif fam in fam_test:
|
||||
# test recebe somente original
|
||||
if d["original"]:
|
||||
nomes_test.append(d["original"])
|
||||
else:
|
||||
# famílias sem original (não devem cair aqui) ficam fora
|
||||
pass
|
||||
|
||||
# dest dirs
|
||||
dest_train_img = os.path.join(pasta_destino, "train", "group", group_name, "images")
|
||||
dest_train_msk = os.path.join(pasta_destino, "train", "group", group_name, "masks")
|
||||
dest_val_img = os.path.join(pasta_destino, "val", "group", group_name, "images")
|
||||
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
|
||||
dest_test_img = os.path.join(pasta_destino, "test", "group", group_name, "images")
|
||||
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks")
|
||||
dest_val_img = os.path.join(pasta_destino, "val", "group", group_name, "images")
|
||||
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
|
||||
dest_test_img = os.path.join(pasta_destino, "test", "group", group_name, "images")
|
||||
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks")
|
||||
|
||||
dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None
|
||||
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
|
||||
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
|
||||
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
|
||||
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
|
||||
|
||||
m_train = copiar(nomes_train, src_img_dir, src_msk_dir, dest_train_img, dest_train_msk, src_msk2_dir, dest_train_msk2)
|
||||
m_val = copiar(nomes_val, src_img_dir, src_msk_dir, dest_val_img, dest_val_msk, src_msk2_dir, dest_val_msk2)
|
||||
m_test = copiar(nomes_test, src_img_dir, src_msk_dir, dest_test_img, dest_test_msk, src_msk2_dir, dest_test_msk2)
|
||||
dest_train_label = os.path.join(pasta_destino, "train", "group", group_name, "labels") if use_labels else None
|
||||
dest_val_label = os.path.join(pasta_destino, "val", "group", group_name, "labels") if use_labels else None
|
||||
dest_test_label = os.path.join(pasta_destino, "test", "group", group_name, "labels") if use_labels else None
|
||||
|
||||
print(f"[{group_name}] famílias={total_familias} → train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}")
|
||||
m_train = copiar(
|
||||
nomes_train,
|
||||
src_img_dir,
|
||||
src_msk_dir,
|
||||
dest_train_img,
|
||||
dest_train_msk,
|
||||
src_msk2_dir,
|
||||
dest_train_msk2,
|
||||
src_label_dir,
|
||||
dest_train_label,
|
||||
strict_label=strict_label,
|
||||
)
|
||||
m_val = copiar(
|
||||
nomes_val,
|
||||
src_img_dir,
|
||||
src_msk_dir,
|
||||
dest_val_img,
|
||||
dest_val_msk,
|
||||
src_msk2_dir,
|
||||
dest_val_msk2,
|
||||
src_label_dir,
|
||||
dest_val_label,
|
||||
strict_label=strict_label,
|
||||
)
|
||||
m_test = copiar(
|
||||
nomes_test,
|
||||
src_img_dir,
|
||||
src_msk_dir,
|
||||
dest_test_img,
|
||||
dest_test_msk,
|
||||
src_msk2_dir,
|
||||
dest_test_msk2,
|
||||
src_label_dir,
|
||||
dest_test_label,
|
||||
strict_label=strict_label,
|
||||
)
|
||||
|
||||
print(
|
||||
f"[{group_name}] famílias={total_familias} → "
|
||||
f"train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}"
|
||||
)
|
||||
return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (val/test só original).")
|
||||
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).")
|
||||
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação (default=0.29).")
|
||||
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste (default=0.01).")
|
||||
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento (default=42).")
|
||||
|
||||
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
|
||||
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
|
||||
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
|
||||
# ===================== MAIN =====================
|
||||
|
||||
def parse_cap_map(s):
|
||||
caps = {}
|
||||
if not s:
|
||||
return caps
|
||||
for item in s.split(","):
|
||||
if not item.strip():
|
||||
continue
|
||||
k, v = item.strip().split(":")
|
||||
caps[k.strip()] = int(v)
|
||||
return caps
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="Split estratificado por grupo sem vazamento, com suporte a labels.")
|
||||
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino.")
|
||||
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação.")
|
||||
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste.")
|
||||
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento.")
|
||||
|
||||
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de famílias por grupo em train.")
|
||||
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de famílias por grupo em val.")
|
||||
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de famílias por grupo em test.")
|
||||
|
||||
ap.add_argument("--modelo", type=str, default=None, help="Sobrescreve MODELO do config.json.")
|
||||
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 640x480).")
|
||||
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH.")
|
||||
|
||||
ap.add_argument("--cap-train-families", type=str, default="", help="Mapa 'grupo:cap,...' p/ limitar número de FAMÍLIAS no TRAIN. Ex.: 'chao:350'")
|
||||
ap.add_argument(
|
||||
"--cap-train-families",
|
||||
type=str,
|
||||
default="",
|
||||
help="Mapa grupo:cap para limitar famílias no TRAIN. Ex: 'navegavel:350'",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--strict-label",
|
||||
action="store_true",
|
||||
help="Se dual_head_label=true e faltar label, pula item/grupo.",
|
||||
)
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
|
|
@ -326,14 +507,6 @@ def main():
|
|||
else:
|
||||
resolucao = RESOLUCAO
|
||||
|
||||
def parse_cap_map(s):
|
||||
caps = {}
|
||||
if not s: return caps
|
||||
for item in s.split(","):
|
||||
k,v = item.strip().split(":")
|
||||
caps[k.strip()] = int(v)
|
||||
return caps
|
||||
|
||||
caps_map = parse_cap_map(args.cap_train_families)
|
||||
|
||||
global pasta_origem, pasta_destino
|
||||
|
|
@ -341,12 +514,18 @@ def main():
|
|||
pasta_destino = os.path.join(modelo, "dataset", "split")
|
||||
|
||||
soma = args.train + args.val + args.test
|
||||
if soma <= 0: raise ValueError("Soma de proporções deve ser > 0.")
|
||||
p_train = args.train / soma
|
||||
p_val = args.val / soma
|
||||
p_test = args.test / soma
|
||||
if soma <= 0:
|
||||
raise ValueError("Soma de proporções deve ser > 0.")
|
||||
|
||||
mins = {"train": max(0, args.min_train), "val": max(0, args.min_val), "test": max(0, args.min_test)}
|
||||
p_train = args.train / soma
|
||||
p_val = args.val / soma
|
||||
p_test = args.test / soma
|
||||
|
||||
mins = {
|
||||
"train": max(0, args.min_train),
|
||||
"val": max(0, args.min_val),
|
||||
"test": max(0, args.min_test),
|
||||
}
|
||||
|
||||
garantir(pasta_destino)
|
||||
|
||||
|
|
@ -357,22 +536,28 @@ def main():
|
|||
|
||||
random.seed(args.seed)
|
||||
|
||||
total_global = {"train":0, "val":0, "test":0, "familias":0}
|
||||
total_global = {"train": 0, "val": 0, "test": 0, "familias": 0}
|
||||
|
||||
print(f"[INFO] Origem: {pasta_origem}")
|
||||
print(f"[INFO] Destino: {pasta_destino}")
|
||||
print(f"[INFO] dual_head_mask/masks2: {USE_MASKS2}")
|
||||
print(f"[INFO] dual_head_label/labels: {USE_LABELS}")
|
||||
print(f"Grupos: {', '.join(grupos)}")
|
||||
print(f"Proporções normalizadas: train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
|
||||
print(f"Mínimos por grupo (famílias): train={mins['train']} val={mins['val']} test={mins['test']}")
|
||||
print(f"Mínimos por grupo: train={mins['train']} val={mins['val']} test={mins['test']}")
|
||||
|
||||
for g in grupos:
|
||||
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map)
|
||||
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map, strict_label=args.strict_label)
|
||||
for k in total_global.keys():
|
||||
total_global[k] += res.get(k, 0)
|
||||
|
||||
print("\nResumo global (imagens copiadas):")
|
||||
print("\nResumo global, imagens copiadas:")
|
||||
print(f" train: {total_global['train']}")
|
||||
print(f" val: {total_global['val']}")
|
||||
print(f" test: {total_global['test']}")
|
||||
print(f" famílias (total): {total_global['familias']}")
|
||||
print(f" famílias total: {total_global['familias']}")
|
||||
print("\n✅ Split sem vazamento concluído!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -1,965 +0,0 @@
|
|||
# _8_train_segformer_b3_dual_v2.py
|
||||
# Treino SegFormer-B3 com 2 cabeças:
|
||||
# - Head 1: segmentação semântica (3 classes + ignore)
|
||||
# - Head 2: corredor binário (0/1 + ignore)
|
||||
#
|
||||
# Features:
|
||||
# - class weights (opcional, estimado do train)
|
||||
# - pos_weight (opcional, estimado do train p/ BCE do corredor)
|
||||
# - warmup + cosine (opcional) + ReduceLROnPlateau
|
||||
# - early stopping com métrica agregada
|
||||
# - resume completo (model + corridor_head + optimizer + scaler)
|
||||
#
|
||||
# Observação: este script assume que seu ROISegDataset retorna dict com:
|
||||
# {"image": <tensor/ndarray>, "mask": <tensor/ndarray>, ...}
|
||||
# e que você já acertou para também expor "image_path"/"mask_path" OU
|
||||
# que seu DualMaskROISegDataset consiga derivar paths corretamente.
|
||||
|
||||
#python _8_train_segformer_b3_dual.py --epochs 80 --batch 2 --num_workers 2 --amp --amp_val --use_weights --auto_pos_weight --warmup_epochs 2 --cosine_epochs 15 --lambda_corr 0.5 --corr_dice_max 0.10 --corr_dice_ramp_epochs 10 --es_metric harmonic --es_patience 12
|
||||
|
||||
#python _8_train_segformer_b3_dual.py --epochs 120 --batch 2 --num_workers 4 --amp --amp_val --grad_accum 1 --use_weights --auto_pos_weight --warmup_epochs 2 --cosine_epochs 20 --lambda_corr 0.7 --corr_dice_max 0.15 --corr_dice_ramp_epochs 10 --es_metric harmonic --es_patience 15
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import math
|
||||
import json
|
||||
import random
|
||||
import argparse
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from torch.amp import autocast
|
||||
from torch.cuda.amp import GradScaler
|
||||
|
||||
# HuggingFace transformers
|
||||
from transformers import SegformerForSemanticSegmentation
|
||||
|
||||
# Seu dataset (o mesmo que você já usa no projeto)
|
||||
from roi_seg_dataset import ROISegDataset
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Utils
|
||||
# -----------------------------
|
||||
def seed_everything(seed: int = 42):
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
|
||||
|
||||
def normalize_img(img: torch.Tensor) -> torch.Tensor:
|
||||
return (img - IMAGENET_MEAN.to(img.device)) / IMAGENET_STD.to(img.device)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def update_confusion_matrix(cm: torch.Tensor, preds: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = 255):
|
||||
preds = preds.view(-1)
|
||||
labels = labels.view(-1)
|
||||
|
||||
valid = labels != ignore_index
|
||||
preds = preds[valid]
|
||||
labels = labels[valid]
|
||||
|
||||
idx = labels * num_classes + preds
|
||||
bins = torch.bincount(idx, minlength=num_classes * num_classes)
|
||||
cm += bins.view(num_classes, num_classes)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_iou_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Tuple[float, List[float]]:
|
||||
cm = cm.float()
|
||||
tp = torch.diag(cm)
|
||||
fp = cm.sum(0) - tp
|
||||
fn = cm.sum(1) - tp
|
||||
denom = tp + fp + fn + eps
|
||||
iou = (tp / denom).cpu().tolist()
|
||||
miou = float(np.mean(iou))
|
||||
return miou, iou
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float:
|
||||
cm = cm.float()
|
||||
acc = (torch.diag(cm).sum() / (cm.sum() + eps)).item()
|
||||
return acc
|
||||
|
||||
|
||||
def default_collate_dual(batch):
|
||||
imgs, m1s, m2s = [], [], []
|
||||
for item in batch:
|
||||
# item pode vir como (img, mask1, mask2) ou dict etc.
|
||||
if isinstance(item, dict):
|
||||
img = item["image"]
|
||||
mask1 = item["mask1"]
|
||||
mask2 = item["mask2"]
|
||||
else:
|
||||
img, mask1, mask2 = item
|
||||
|
||||
if isinstance(img, np.ndarray):
|
||||
img = torch.from_numpy(img)
|
||||
if isinstance(mask1, np.ndarray):
|
||||
mask1 = torch.from_numpy(mask1)
|
||||
if isinstance(mask2, np.ndarray):
|
||||
mask2 = torch.from_numpy(mask2)
|
||||
|
||||
if img.ndim == 3 and img.shape[-1] == 3:
|
||||
img = img.permute(2, 0, 1)
|
||||
|
||||
if img.dtype != torch.float32:
|
||||
img = img.float()
|
||||
if img.max() > 1.5:
|
||||
img = img / 255.0
|
||||
|
||||
imgs.append(img)
|
||||
m1s.append(mask1.long())
|
||||
m2s.append(mask2.long())
|
||||
|
||||
return torch.stack(imgs, 0), torch.stack(m1s, 0), torch.stack(m2s, 0)
|
||||
|
||||
|
||||
def estimate_class_weights_from_ds(ds: ROISegDataset, num_classes: int, ignore_index: int = 255, max_samples: int = 800) -> torch.Tensor:
|
||||
n = min(len(ds), max_samples)
|
||||
idxs = np.random.choice(len(ds), size=n, replace=False)
|
||||
|
||||
counts = np.zeros(num_classes, dtype=np.float64)
|
||||
for i in idxs:
|
||||
item = ds[i]
|
||||
mask = item["mask"] if isinstance(item, dict) else item[1]
|
||||
m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
|
||||
|
||||
m = m.reshape(-1)
|
||||
m = m[m != ignore_index]
|
||||
if m.size == 0:
|
||||
continue
|
||||
counts += np.bincount(m, minlength=num_classes)[:num_classes]
|
||||
|
||||
freq = counts / (counts.sum() + 1e-12)
|
||||
freq = np.clip(freq, 1e-12, 1.0)
|
||||
weights = 1.0 / np.log(1.02 + freq)
|
||||
weights = weights / weights.mean()
|
||||
return torch.tensor(weights, dtype=torch.float32)
|
||||
|
||||
|
||||
def estimate_pos_weight_from_masks2(ds_dual, ignore_index2: int = 255, max_samples: int = 800) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Estima pos_weight = neg/pos para BCEWithLogitsLoss.
|
||||
Considera targets binários {0,1} ignorando ignore_index2.
|
||||
"""
|
||||
n = min(len(ds_dual), max_samples)
|
||||
if n <= 0:
|
||||
return None
|
||||
|
||||
idxs = np.random.choice(len(ds_dual), size=n, replace=False)
|
||||
pos = 0
|
||||
neg = 0
|
||||
for i in idxs:
|
||||
it = ds_dual[i]
|
||||
m2 = it["mask2"] if isinstance(it, dict) else it[2]
|
||||
|
||||
if isinstance(m2, torch.Tensor):
|
||||
m2 = m2.cpu().numpy()
|
||||
m2 = np.array(m2).reshape(-1)
|
||||
m2 = m2[m2 != ignore_index2]
|
||||
if m2.size == 0:
|
||||
continue
|
||||
pos += int((m2 == 1).sum())
|
||||
neg += int((m2 == 0).sum())
|
||||
|
||||
if pos <= 0:
|
||||
return None
|
||||
pw = float(neg) / float(pos)
|
||||
# clamp leve pra não explodir treino
|
||||
pw = float(np.clip(pw, 1.0, 50.0))
|
||||
return torch.tensor([pw], dtype=torch.float32)
|
||||
|
||||
|
||||
def dice_loss_with_logits(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = 255, eps: float = 1e-6) -> torch.Tensor:
|
||||
"""
|
||||
logits: [B,1,H,W]
|
||||
targets: [B,H,W] com {0,1} e possivelmente ignore_index
|
||||
"""
|
||||
probs = torch.sigmoid(logits)
|
||||
t = targets.float().unsqueeze(1)
|
||||
|
||||
if ignore_index is not None:
|
||||
valid = (targets != ignore_index).unsqueeze(1)
|
||||
probs = probs[valid]
|
||||
t = t[valid]
|
||||
|
||||
probs = probs.reshape(-1)
|
||||
t = t.reshape(-1)
|
||||
|
||||
inter = (probs * t).sum()
|
||||
denom = probs.sum() + t.sum() + eps
|
||||
dice = (2.0 * inter + eps) / denom
|
||||
return 1.0 - dice
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def binary_iou_and_acc_from_logits(logits: torch.Tensor, targets: torch.Tensor, thr: float = 0.5, ignore_index: int = 255, eps: float = 1e-6) -> Tuple[float, float]:
|
||||
probs = torch.sigmoid(logits)
|
||||
preds = (probs >= thr).long().squeeze(1)
|
||||
|
||||
if logits.shape[-2:] != targets.shape[-2:]:
|
||||
preds = F.interpolate(preds.unsqueeze(1).float(), size=targets.shape[-2:], mode="nearest").long().squeeze(1)
|
||||
|
||||
valid = (targets != ignore_index)
|
||||
if valid.sum().item() == 0:
|
||||
return 0.0, 0.0
|
||||
|
||||
p = preds[valid]
|
||||
t = targets[valid].long()
|
||||
|
||||
tp = ((p == 1) & (t == 1)).sum().float()
|
||||
fp = ((p == 1) & (t == 0)).sum().float()
|
||||
fn = ((p == 0) & (t == 1)).sum().float()
|
||||
|
||||
iou = (tp / (tp + fp + fn + eps)).item()
|
||||
acc = ((p == t).sum().float() / (t.numel() + eps)).item()
|
||||
return iou, acc
|
||||
|
||||
|
||||
def _pretty_iou(iou_list: List[float], class_name_by_id: Dict[int, str]) -> str:
|
||||
parts = []
|
||||
for cid, v in enumerate(iou_list):
|
||||
parts.append(f"{class_name_by_id.get(cid, str(cid))}:{v:.3f}")
|
||||
return " | ".join(parts)
|
||||
|
||||
|
||||
def save_checkpoint(path: str,
|
||||
base_model: nn.Module,
|
||||
corridor_head: nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
scaler: Optional[GradScaler],
|
||||
epoch: int,
|
||||
best_miou: float,
|
||||
best_main_iou: float,
|
||||
best_corr_iou: float,
|
||||
extra: Optional[Dict[str, Any]] = None):
|
||||
ckpt = {
|
||||
"epoch": epoch,
|
||||
"model": base_model.state_dict(),
|
||||
"corridor_head": corridor_head.state_dict(),
|
||||
"optimizer": optimizer.state_dict(),
|
||||
"best_miou": best_miou,
|
||||
"best_main_iou": best_main_iou,
|
||||
"best_corr_iou": best_corr_iou,
|
||||
}
|
||||
if scaler is not None:
|
||||
ckpt["scaler"] = scaler.state_dict()
|
||||
if extra:
|
||||
ckpt["extra"] = extra
|
||||
torch.save(ckpt, path)
|
||||
|
||||
|
||||
def load_checkpoint(path: str,
|
||||
base_model: nn.Module,
|
||||
corridor_head: nn.Module,
|
||||
optimizer: Optional[torch.optim.Optimizer] = None,
|
||||
scaler: Optional[GradScaler] = None,
|
||||
map_location: str = "cpu") -> Dict[str, Any]:
|
||||
ckpt = torch.load(path, map_location=map_location, weights_only=False)
|
||||
base_model.load_state_dict(ckpt["model"], strict=True)
|
||||
corridor_head.load_state_dict(ckpt["corridor_head"], strict=True)
|
||||
if optimizer is not None and "optimizer" in ckpt:
|
||||
optimizer.load_state_dict(ckpt["optimizer"])
|
||||
if scaler is not None and "scaler" in ckpt:
|
||||
scaler.load_state_dict(ckpt["scaler"])
|
||||
return ckpt
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Dual Dataset Wrapper
|
||||
# -----------------------------
|
||||
class DualMaskROISegDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, base_ds: ROISegDataset, split_root: str, ignore_index2: int = 255):
|
||||
self.base_ds = base_ds
|
||||
self.split_root = split_root
|
||||
self.ignore_index2 = ignore_index2
|
||||
|
||||
# NÃO falha aqui por pasta. Vamos resolver por caminho da mask1.
|
||||
self.masks2_dir = os.path.join(split_root, "masks2")
|
||||
self.has_flat_masks2 = os.path.isdir(self.masks2_dir)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.base_ds)
|
||||
|
||||
def _get_mask1_path(self, item: Dict[str, Any], idx: int) -> str:
|
||||
# 1) se vier no dict, usa
|
||||
for k in ("mask_path", "mask_file", "maskname", "mask_name"):
|
||||
if k in item and item[k]:
|
||||
return str(item[k])
|
||||
|
||||
# 2) fallback: usa o índice no ROISegDataset (ele tem msk_paths)
|
||||
if hasattr(self.base_ds, "msk_paths"):
|
||||
return str(self.base_ds.msk_paths[idx])
|
||||
|
||||
raise RuntimeError(
|
||||
"ROISegDataset não retornou mask_path/mask_file e não possui base_ds.msk_paths. "
|
||||
"Não dá pra localizar a máscara para buscar a masks2."
|
||||
)
|
||||
|
||||
def _load_mask2_from_mask1_path(self, mask1_path: str, target_shape_hw: Tuple[int, int]) -> torch.Tensor:
|
||||
# 1) modo robusto: espelha estrutura trocando /masks/ -> /masks2/
|
||||
norm = os.path.normpath(mask1_path)
|
||||
parts = norm.split(os.sep)
|
||||
try:
|
||||
i = parts.index("masks")
|
||||
parts[i] = "masks2"
|
||||
path2 = os.sep.join(parts)
|
||||
except ValueError:
|
||||
path2 = ""
|
||||
|
||||
# 2) fallback: modo “flat” (split_root/masks2/<filename>)
|
||||
if (not path2) or (not os.path.exists(path2)):
|
||||
if not self.has_flat_masks2:
|
||||
raise FileNotFoundError(
|
||||
f"Mask2 não encontrada espelhando a mask1.\n"
|
||||
f"mask1: {mask1_path}\n"
|
||||
f"tentativa: {path2}\n"
|
||||
f"E também não existe: {self.masks2_dir}"
|
||||
)
|
||||
filename = os.path.basename(mask1_path)
|
||||
path2 = os.path.join(self.masks2_dir, filename)
|
||||
if not os.path.exists(path2):
|
||||
raise FileNotFoundError(f"Mask2 não encontrada: {path2}")
|
||||
|
||||
m = np.array(Image.open(path2))
|
||||
if m.ndim == 3:
|
||||
m = m[..., 0]
|
||||
|
||||
if m.max() > 1:
|
||||
m = (m >= 128).astype(np.uint8)
|
||||
|
||||
H, W = target_shape_hw
|
||||
if m.shape[0] != H or m.shape[1] != W:
|
||||
m = np.array(Image.fromarray(m).resize((W, H), resample=Image.NEAREST))
|
||||
|
||||
return torch.from_numpy(m).long()
|
||||
|
||||
def __getitem__(self, idx):
|
||||
item = self.base_ds[idx]
|
||||
# caso venha (dict,) ou [dict]
|
||||
if isinstance(item, (list, tuple)) and len(item) == 1 and isinstance(item[0], dict):
|
||||
item = item[0]
|
||||
|
||||
if isinstance(item, dict):
|
||||
pass
|
||||
elif isinstance(item, (list, tuple)) and len(item) >= 2:
|
||||
img, mask = item[0], item[1]
|
||||
item = {"image": img, "mask": mask}
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"ROISegDataset retornou formato inesperado: type={type(item)} value={item}"
|
||||
)
|
||||
|
||||
img = item["image"]
|
||||
mask1 = item["mask"]
|
||||
|
||||
if isinstance(mask1, torch.Tensor):
|
||||
h, w = int(mask1.shape[-2]), int(mask1.shape[-1])
|
||||
else:
|
||||
m = np.array(mask1)
|
||||
h, w = m.shape[0], m.shape[1]
|
||||
|
||||
mask1_path = self._get_mask1_path(item, idx)
|
||||
mask2 = self._load_mask2_from_mask1_path(mask1_path, (h, w))
|
||||
|
||||
item["mask1"] = mask1
|
||||
item["mask2"] = mask2
|
||||
return item
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Corridor head
|
||||
# -----------------------------
|
||||
class CorridorHead(nn.Module):
|
||||
"""
|
||||
Head simples:
|
||||
feat (B,C,H,W) + logits_seg (B,K,H,W) -> concat -> convs -> 1 canal (logit corredor)
|
||||
"""
|
||||
def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
in_ch = feat_ch + num_classes
|
||||
self.in_ch = in_ch
|
||||
self.net = nn.Sequential(
|
||||
nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(hidden),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout2d(dropout),
|
||||
nn.Conv2d(hidden, hidden, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(hidden),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout2d(dropout),
|
||||
nn.Conv2d(hidden, 1, kernel_size=1)
|
||||
)
|
||||
|
||||
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
|
||||
x = torch.cat([feat, logits_seg], dim=1)
|
||||
return self.net(x)
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# One epoch
|
||||
# -----------------------------
|
||||
def run_one_epoch_dual(base_model: nn.Module,
|
||||
corridor_head: nn.Module,
|
||||
loader: DataLoader,
|
||||
optimizer: Optional[torch.optim.Optimizer],
|
||||
device: torch.device,
|
||||
num_classes: int,
|
||||
ignore_index1: int,
|
||||
ignore_index2: int,
|
||||
criterion_seg: nn.Module,
|
||||
bce_corr: nn.Module,
|
||||
lambda_corr: float,
|
||||
amp: bool,
|
||||
scaler: Optional[GradScaler],
|
||||
train: bool,
|
||||
grad_accum: int = 1,
|
||||
corridor_thr: float = 0.5,
|
||||
corr_dice_mix: float = 0.0) -> Dict[str, Any]:
|
||||
|
||||
base_model.train(train)
|
||||
corridor_head.train(train)
|
||||
|
||||
total_loss = 0.0
|
||||
total_loss_seg = 0.0
|
||||
total_loss_corr = 0.0
|
||||
|
||||
cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)
|
||||
|
||||
corr_iou_accum = 0.0
|
||||
corr_acc_accum = 0.0
|
||||
|
||||
t0 = time.time()
|
||||
n_batches = 0
|
||||
|
||||
with torch.set_grad_enabled(train):
|
||||
if train and optimizer is not None:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
for step, (imgs, masks1, masks2) in enumerate(loader, start=1):
|
||||
imgs = imgs.to(device, non_blocking=True)
|
||||
masks1 = masks1.to(device, non_blocking=True)
|
||||
masks2 = masks2.to(device, non_blocking=True)
|
||||
|
||||
#print("masks2 unique:", torch.unique(masks2)[:10])
|
||||
#print("masks2 max:", masks2.max().item(), "min:", masks2.min().item())
|
||||
|
||||
imgs = normalize_img(imgs)
|
||||
|
||||
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
|
||||
out = base_model(pixel_values=imgs)
|
||||
logits1 = out.logits # [B,K,H,W]
|
||||
|
||||
if logits1.shape[-2:] != masks1.shape[-2:]:
|
||||
logits1 = F.interpolate(logits1, size=masks1.shape[-2:], mode="bilinear", align_corners=False)
|
||||
|
||||
loss_seg = criterion_seg(logits1, masks1)
|
||||
|
||||
# feat do encoder: pega hidden_states (se disponível), senão reusa logits como proxy
|
||||
feat = None
|
||||
if hasattr(out, "hidden_states") and out.hidden_states is not None:
|
||||
feat = out.hidden_states[-1] # [B, C, h, w]
|
||||
else:
|
||||
# fallback: usa logits como feature (não ideal, mas mantém vivo)
|
||||
feat = logits1
|
||||
|
||||
if feat.shape[-2:] != logits1.shape[-2:]:
|
||||
feat = F.interpolate(feat, size=logits1.shape[-2:], mode="bilinear", align_corners=False)
|
||||
|
||||
logits_corr = corridor_head(feat, logits1) # [B,1,H,W]
|
||||
|
||||
if logits_corr.shape[-2:] != masks2.shape[-2:]:
|
||||
logits_corr = F.interpolate(logits_corr, size=masks2.shape[-2:], mode="bilinear", align_corners=False)
|
||||
|
||||
#with torch.no_grad():
|
||||
# p = torch.sigmoid(logits_corr)
|
||||
# print("corr prob mean:", p.mean().item(), "min:", p.min().item(), "max:", p.max().item())
|
||||
|
||||
# BCE (com ignore via masking manual)
|
||||
valid = (masks2 != ignore_index2)
|
||||
if valid.sum().item() > 0:
|
||||
tgt = masks2.float().unsqueeze(1)
|
||||
bce_val = bce_corr(logits_corr[valid.unsqueeze(1)], tgt[valid.unsqueeze(1)])
|
||||
if corr_dice_mix > 0:
|
||||
d_val = dice_loss_with_logits(logits_corr, masks2, ignore_index=ignore_index2)
|
||||
loss_corr = (1.0 - corr_dice_mix) * bce_val + corr_dice_mix * d_val
|
||||
else:
|
||||
loss_corr = bce_val
|
||||
else:
|
||||
loss_corr = torch.zeros([], device=device, dtype=loss_seg.dtype)
|
||||
|
||||
loss = loss_seg + lambda_corr * loss_corr
|
||||
|
||||
if train and grad_accum > 1:
|
||||
loss = loss / grad_accum
|
||||
|
||||
if train and optimizer is not None:
|
||||
if amp and scaler is not None and device.type == "cuda":
|
||||
scaler.scale(loss).backward()
|
||||
if (step % grad_accum) == 0:
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
else:
|
||||
loss.backward()
|
||||
if (step % grad_accum) == 0:
|
||||
optimizer.step()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
total_loss += float(loss.item()) * (grad_accum if (train and grad_accum > 1) else 1.0)
|
||||
total_loss_seg += float(loss_seg.item())
|
||||
total_loss_corr += float(loss_corr.item())
|
||||
n_batches += 1
|
||||
|
||||
# metrics seg
|
||||
preds1 = torch.argmax(logits1.detach(), dim=1)
|
||||
update_confusion_matrix(cm, preds1, masks1, num_classes, ignore_index=ignore_index1)
|
||||
|
||||
# metrics corr
|
||||
ciou, cacc = binary_iou_and_acc_from_logits(logits_corr.detach(), masks2, thr=corridor_thr, ignore_index=ignore_index2)
|
||||
corr_iou_accum += ciou
|
||||
corr_acc_accum += cacc
|
||||
|
||||
miou, iou_per_class = compute_iou_from_cm(cm)
|
||||
acc = compute_pixel_acc_from_cm(cm)
|
||||
|
||||
out = {
|
||||
"loss": total_loss / max(1, n_batches),
|
||||
"loss_seg": total_loss_seg / max(1, n_batches),
|
||||
"loss_corr": total_loss_corr / max(1, n_batches),
|
||||
"acc": acc,
|
||||
"miou": miou,
|
||||
"iou_per_class": iou_per_class,
|
||||
"corr_iou": corr_iou_accum / max(1, n_batches),
|
||||
"corr_acc": corr_acc_accum / max(1, n_batches),
|
||||
"time_s": time.time() - t0
|
||||
}
|
||||
return out
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# LR / EarlyStop helpers
|
||||
# -----------------------------
|
||||
@dataclass
|
||||
class EarlyStopState:
|
||||
best: float = -1e9
|
||||
bad_epochs: int = 0
|
||||
stopped: bool = False
|
||||
|
||||
|
||||
def agg_metric(es_metric: str, miou: float, main_iou: float, corr_iou: float) -> float:
|
||||
if es_metric == "miou":
|
||||
return float(miou)
|
||||
if es_metric == "main":
|
||||
return float(main_iou)
|
||||
if es_metric == "corr":
|
||||
return float(corr_iou)
|
||||
if es_metric == "harmonic":
|
||||
eps = 1e-6
|
||||
a = max(eps, float(main_iou))
|
||||
b = max(eps, float(corr_iou))
|
||||
return float(2.0 * a * b / (a + b + eps))
|
||||
# default
|
||||
return float(miou)
|
||||
|
||||
|
||||
def apply_warmup(optimizer, base_lr: float, epoch: int, warmup_epochs: int):
|
||||
if warmup_epochs <= 0:
|
||||
return
|
||||
if epoch <= warmup_epochs:
|
||||
wf = epoch / float(warmup_epochs)
|
||||
for g in optimizer.param_groups:
|
||||
g["lr"] = base_lr * wf
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Main
|
||||
# -----------------------------
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
|
||||
ap.add_argument("--config", default="config_oak.json")
|
||||
ap.add_argument("--epochs", type=int, default=120)
|
||||
ap.add_argument("--batch", type=int, default=4)
|
||||
ap.add_argument("--num_workers", type=int, default=4)
|
||||
ap.add_argument("--seed", type=int, default=42)
|
||||
|
||||
ap.add_argument("--lr", type=float, default=6e-5)
|
||||
ap.add_argument("--wd", type=float, default=1e-2)
|
||||
ap.add_argument("--min_lr", type=float, default=6e-6)
|
||||
|
||||
ap.add_argument("--amp", action="store_true")
|
||||
ap.add_argument("--amp_val", action="store_true")
|
||||
ap.add_argument("--grad_accum", type=int, default=1)
|
||||
|
||||
ap.add_argument("--ignore_index", type=int, default=255)
|
||||
ap.add_argument("--ignore_index2", type=int, default=255)
|
||||
|
||||
ap.add_argument("--lambda_corr", type=float, default=0.5)
|
||||
ap.add_argument("--corr_thr", type=float, default=0.5)
|
||||
|
||||
# class weights
|
||||
ap.add_argument("--use_weights", action="store_true")
|
||||
ap.add_argument("--cw_max_samples", type=int, default=800)
|
||||
|
||||
# pos_weight corredor
|
||||
ap.add_argument("--auto_pos_weight", action="store_true")
|
||||
ap.add_argument("--pw_max_samples", type=int, default=800)
|
||||
|
||||
# mistura dice no corredor (rampa)
|
||||
ap.add_argument("--corr_dice_max", type=float, default=0.15)
|
||||
ap.add_argument("--corr_dice_ramp_epochs", type=int, default=10)
|
||||
|
||||
# schedulers
|
||||
ap.add_argument("--warmup_epochs", type=int, default=0)
|
||||
ap.add_argument("--cosine_epochs", type=int, default=0, help="se >0: usa cosine por N épocas e depois troca p/ plateau")
|
||||
ap.add_argument("--plateau_patience", type=int, default=3)
|
||||
ap.add_argument("--plateau_factor", type=float, default=0.6)
|
||||
ap.add_argument("--plateau_cooldown", type=int, default=1)
|
||||
|
||||
# early stopping
|
||||
ap.add_argument("--es_metric", default="harmonic", choices=["miou", "main", "corr", "harmonic"])
|
||||
ap.add_argument("--es_patience", type=int, default=12)
|
||||
ap.add_argument("--es_min_delta", type=float, default=2e-4)
|
||||
|
||||
# logs/ckpt
|
||||
ap.add_argument("--save_every", type=int, default=0)
|
||||
ap.add_argument("--resume", action="store_true")
|
||||
|
||||
ap.add_argument("--main_class", type=str, default=None)
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
seed_everything(args.seed)
|
||||
|
||||
with open(args.config, "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
MODEL_NAME = config["model_name"]
|
||||
RESOLUCAO = config["resolucao"]
|
||||
ROI_INICIO = config["roi_inicio"]
|
||||
ROI_TAMANHO = config["roi_tamanho"]
|
||||
BACKBONE = config["backbone"]
|
||||
MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower()
|
||||
if args.main_class is not None:
|
||||
MAIN_CLASS_NAME = args.main_class.lower()
|
||||
|
||||
dataset_path = os.path.join(MODELO, "dataset")
|
||||
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
||||
|
||||
# save paths
|
||||
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME + "_dual")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
last_ckpt_path = os.path.join(save_path, "last.pt")
|
||||
best_miou_path = os.path.join(save_path, "best_miou.pt")
|
||||
best_main_path = os.path.join(save_path, "best_main.pt")
|
||||
best_corr_path = os.path.join(save_path, "best_corr.pt")
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print("Device:", device)
|
||||
|
||||
# seus splits estão em train/group/... então o ROISegDataset deve saber ler de lá
|
||||
split_train = os.path.join(dataset_path, "split", "train")
|
||||
split_val = os.path.join(dataset_path, "split", "val")
|
||||
|
||||
ds_train_base = ROISegDataset(
|
||||
split_train,
|
||||
save_path, ROI_INICIO, ROI_TAMANHO,
|
||||
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
|
||||
)
|
||||
|
||||
ds_val_base = ROISegDataset(
|
||||
split_val,
|
||||
save_path, ROI_INICIO, ROI_TAMANHO,
|
||||
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
|
||||
)
|
||||
|
||||
ds_train = DualMaskROISegDataset(ds_train_base, split_train, ignore_index2=args.ignore_index2)
|
||||
ds_val = DualMaskROISegDataset(ds_val_base, split_val, ignore_index2=args.ignore_index2)
|
||||
|
||||
CLASS_NAMES = getattr(ds_train_base, "classes", None)
|
||||
if CLASS_NAMES is None:
|
||||
raise RuntimeError("ROISegDataset precisa expor .classes (list ou dict).")
|
||||
|
||||
if isinstance(CLASS_NAMES, list):
|
||||
num_classes = len(CLASS_NAMES)
|
||||
class_id_by_name = {n.lower(): i for i, n in enumerate(CLASS_NAMES)}
|
||||
class_name_by_id = {i: n for i, n in enumerate(CLASS_NAMES)}
|
||||
elif isinstance(CLASS_NAMES, dict):
|
||||
if all(isinstance(k, int) for k in CLASS_NAMES.keys()):
|
||||
num_classes = len(CLASS_NAMES)
|
||||
class_name_by_id = {int(k): str(v) for k, v in CLASS_NAMES.items()}
|
||||
class_id_by_name = {str(v).lower(): int(k) for k, v in CLASS_NAMES.items()}
|
||||
else:
|
||||
class_id_by_name = {str(k).lower(): int(v) for k, v in CLASS_NAMES.items()}
|
||||
num_classes = len(class_id_by_name)
|
||||
class_name_by_id = {v: k for k, v in class_id_by_name.items()}
|
||||
else:
|
||||
raise RuntimeError("Formato de .classes não reconhecido.")
|
||||
|
||||
main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None)
|
||||
if main_class_id is None:
|
||||
print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.")
|
||||
else:
|
||||
print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}")
|
||||
|
||||
|
||||
dl_train = DataLoader(
|
||||
ds_train,
|
||||
batch_size=args.batch,
|
||||
shuffle=True,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=(device.type == "cuda"),
|
||||
collate_fn=default_collate_dual
|
||||
)
|
||||
dl_val = DataLoader(
|
||||
ds_val,
|
||||
batch_size=args.batch,
|
||||
shuffle=False,
|
||||
num_workers=max(0, args.num_workers // 2),
|
||||
pin_memory=(device.type == "cuda"),
|
||||
collate_fn=default_collate_dual
|
||||
)
|
||||
|
||||
# model base SegFormer
|
||||
base_model = SegformerForSemanticSegmentation.from_pretrained(
|
||||
BACKBONE,
|
||||
num_labels=num_classes,
|
||||
ignore_mismatched_sizes=True,
|
||||
use_safetensors=True, # <- evita torch.load do .bin
|
||||
)
|
||||
base_model.config.output_hidden_states = True
|
||||
base_model.to(device)
|
||||
|
||||
# descobre feat_ch via dummy forward (mais robusto)
|
||||
with torch.no_grad():
|
||||
dummy = torch.zeros((1, 3, 512, 512), device=device)
|
||||
out = base_model(pixel_values=dummy)
|
||||
feat_ch = None
|
||||
if hasattr(out, "hidden_states") and out.hidden_states is not None:
|
||||
feat_ch = int(out.hidden_states[-1].shape[1])
|
||||
else:
|
||||
feat_ch = int(out.logits.shape[1]) # fallback
|
||||
corridor_head = CorridorHead(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1).to(device)
|
||||
print(f"[corridor_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {corridor_head.in_ch}")
|
||||
|
||||
# losses
|
||||
class_weights = None
|
||||
if args.use_weights:
|
||||
w = estimate_class_weights_from_ds(ds_train_base, num_classes=num_classes, ignore_index=args.ignore_index, max_samples=args.cw_max_samples)
|
||||
class_weights = w.to(device)
|
||||
print("[INFO] class_weights:", class_weights.detach().cpu().numpy().round(3).tolist())
|
||||
criterion_seg = nn.CrossEntropyLoss(ignore_index=args.ignore_index, weight=class_weights)
|
||||
|
||||
pos_weight = None
|
||||
if args.auto_pos_weight:
|
||||
pw = estimate_pos_weight_from_masks2(ds_train, ignore_index2=args.ignore_index2, max_samples=args.pw_max_samples)
|
||||
if pw is not None:
|
||||
pos_weight = pw.to(device)
|
||||
print("[INFO] pos_weight corredor (neg/pos):", float(pos_weight.item()))
|
||||
else:
|
||||
print("[INFO] pos_weight não estimável (sem positivos suficientes), usando None.")
|
||||
bce_corr = nn.BCEWithLogitsLoss(pos_weight=pos_weight) if pos_weight is not None else nn.BCEWithLogitsLoss()
|
||||
|
||||
# optimizer (base + head)
|
||||
params = list(base_model.parameters()) + list(corridor_head.parameters())
|
||||
optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wd)
|
||||
|
||||
scaler = GradScaler(enabled=(args.amp and device.type == "cuda"))
|
||||
|
||||
# schedulers
|
||||
cosine = None
|
||||
if args.cosine_epochs and args.cosine_epochs > 0:
|
||||
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
|
||||
optimizer,
|
||||
T_max=max(1, args.cosine_epochs),
|
||||
eta_min=args.min_lr
|
||||
)
|
||||
plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
||||
optimizer,
|
||||
mode="min",
|
||||
factor=args.plateau_factor,
|
||||
patience=args.plateau_patience,
|
||||
cooldown=args.plateau_cooldown,
|
||||
min_lr=args.min_lr,
|
||||
verbose=True
|
||||
)
|
||||
active_sched = "cosine" if cosine is not None else "plateau"
|
||||
|
||||
# resume
|
||||
start_epoch = 1
|
||||
best_miou = -1.0
|
||||
best_main_iou = -1.0
|
||||
best_corr_iou = -1.0
|
||||
|
||||
if args.resume and os.path.exists(last_ckpt_path):
|
||||
ckpt = load_checkpoint(last_ckpt_path, base_model, corridor_head, optimizer=optimizer, scaler=scaler, map_location="cpu")
|
||||
start_epoch = int(ckpt.get("epoch", 0)) + 1
|
||||
best_miou = float(ckpt.get("best_miou", -1.0))
|
||||
best_main_iou = float(ckpt.get("best_main_iou", -1.0))
|
||||
best_corr_iou = float(ckpt.get("best_corr_iou", -1.0))
|
||||
print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f} best_corr_iou={best_corr_iou:.4f}")
|
||||
|
||||
# early stopping
|
||||
es = EarlyStopState(best=-1e9, bad_epochs=0, stopped=False)
|
||||
|
||||
for epoch in range(start_epoch, args.epochs + 1):
|
||||
lr_now = optimizer.param_groups[0]["lr"]
|
||||
|
||||
# warmup
|
||||
apply_warmup(optimizer, args.lr, epoch, args.warmup_epochs)
|
||||
lr_now = optimizer.param_groups[0]["lr"]
|
||||
|
||||
# corr dice ramp
|
||||
if args.corr_dice_max > 0 and args.corr_dice_ramp_epochs > 0:
|
||||
corr_dice_mix = min(args.corr_dice_max, (epoch - 1) / float(args.corr_dice_ramp_epochs) * args.corr_dice_max)
|
||||
else:
|
||||
corr_dice_mix = 0.0
|
||||
|
||||
print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} | lambda_corr={args.lambda_corr} | corr_dice={corr_dice_mix:.3f} | sched={active_sched} ====")
|
||||
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
tr = run_one_epoch_dual(
|
||||
base_model=base_model,
|
||||
corridor_head=corridor_head,
|
||||
loader=dl_train,
|
||||
optimizer=optimizer,
|
||||
device=device,
|
||||
num_classes=num_classes,
|
||||
ignore_index1=args.ignore_index,
|
||||
ignore_index2=args.ignore_index2,
|
||||
criterion_seg=criterion_seg,
|
||||
bce_corr=bce_corr,
|
||||
lambda_corr=args.lambda_corr,
|
||||
amp=args.amp,
|
||||
scaler=scaler,
|
||||
train=True,
|
||||
grad_accum=max(1, args.grad_accum),
|
||||
corridor_thr=args.corr_thr,
|
||||
corr_dice_mix=corr_dice_mix
|
||||
)
|
||||
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
va = run_one_epoch_dual(
|
||||
base_model=base_model,
|
||||
corridor_head=corridor_head,
|
||||
loader=dl_val,
|
||||
optimizer=None,
|
||||
device=device,
|
||||
num_classes=num_classes,
|
||||
ignore_index1=args.ignore_index,
|
||||
ignore_index2=args.ignore_index2,
|
||||
criterion_seg=criterion_seg,
|
||||
bce_corr=bce_corr,
|
||||
lambda_corr=args.lambda_corr,
|
||||
amp=args.amp_val,
|
||||
scaler=None,
|
||||
train=False,
|
||||
grad_accum=1,
|
||||
corridor_thr=args.corr_thr,
|
||||
corr_dice_mix=corr_dice_mix
|
||||
)
|
||||
|
||||
# scheduler step
|
||||
if cosine is not None and epoch <= args.cosine_epochs:
|
||||
cosine.step()
|
||||
else:
|
||||
active_sched = "plateau"
|
||||
plateau.step(va["loss"])
|
||||
|
||||
main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id]
|
||||
corr_iou = float(va["corr_iou"])
|
||||
|
||||
print(f"TRAIN: loss={tr['loss']:.4f} (seg={tr['loss_seg']:.4f} corr={tr['loss_corr']:.4f}) "
|
||||
f"acc={tr['acc']:.4f} miou={tr['miou']:.4f} corr_iou={tr['corr_iou']:.4f} (t={tr['time_s']:.1f}s)")
|
||||
print(f"VAL : loss={va['loss']:.4f} (seg={va['loss_seg']:.4f} corr={va['loss_corr']:.4f}) "
|
||||
f"acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={float(main_iou):.4f} "
|
||||
f"corr_iou={va['corr_iou']:.4f} corr_acc={va['corr_acc']:.4f} (t={va['time_s']:.1f}s)")
|
||||
print("IoU per class:", _pretty_iou(va["iou_per_class"], class_name_by_id))
|
||||
|
||||
# save last
|
||||
save_checkpoint(
|
||||
last_ckpt_path,
|
||||
base_model,
|
||||
corridor_head,
|
||||
optimizer,
|
||||
scaler=scaler,
|
||||
epoch=epoch,
|
||||
best_miou=best_miou,
|
||||
best_main_iou=best_main_iou,
|
||||
best_corr_iou=best_corr_iou,
|
||||
extra={
|
||||
"val_loss": va["loss"],
|
||||
"val_miou": va["miou"],
|
||||
"val_main_iou": float(main_iou),
|
||||
"val_corr_iou": float(corr_iou),
|
||||
"lr": optimizer.param_groups[0]["lr"],
|
||||
"corr_dice_mix": corr_dice_mix
|
||||
}
|
||||
)
|
||||
|
||||
# save every
|
||||
if args.save_every > 0 and (epoch % args.save_every == 0):
|
||||
save_checkpoint(
|
||||
os.path.join(save_path, f"epoch_{epoch:04d}.pt"),
|
||||
base_model, corridor_head, optimizer, scaler=scaler,
|
||||
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou
|
||||
)
|
||||
|
||||
# bests
|
||||
if va["miou"] > best_miou:
|
||||
best_miou = va["miou"]
|
||||
save_checkpoint(best_miou_path, base_model, corridor_head, optimizer, scaler=scaler,
|
||||
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou)
|
||||
print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}")
|
||||
|
||||
if float(main_iou) > best_main_iou:
|
||||
best_main_iou = float(main_iou)
|
||||
save_checkpoint(best_main_path, base_model, corridor_head, optimizer, scaler=scaler,
|
||||
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou)
|
||||
print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}")
|
||||
|
||||
if corr_iou > best_corr_iou:
|
||||
best_corr_iou = corr_iou
|
||||
save_checkpoint(best_corr_path, base_model, corridor_head, optimizer, scaler=scaler,
|
||||
epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou)
|
||||
print(f"[BEST CORR] {best_corr_iou:.4f} -> saved: {best_corr_path}")
|
||||
|
||||
# early stopping
|
||||
score = agg_metric(args.es_metric, va["miou"], float(main_iou), corr_iou)
|
||||
if score > es.best + args.es_min_delta:
|
||||
es.best = score
|
||||
es.bad_epochs = 0
|
||||
print(f"[ES] improved {args.es_metric} -> {score:.6f}")
|
||||
else:
|
||||
es.bad_epochs += 1
|
||||
print(f"[ES] no improve ({es.bad_epochs}/{args.es_patience}) best={es.best:.6f} now={score:.6f}")
|
||||
if es.bad_epochs >= args.es_patience:
|
||||
print(f"🛑 Early stopping acionado (metric={args.es_metric}).")
|
||||
break
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -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) ===
|
||||
|
|
@ -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/<qualquer>/images/*.png
|
||||
split/test/group/<qualquer>/masks/*.png
|
||||
split/test/group/<qualquer>/masks2/*.png (opcional, mas recomendado)
|
||||
"""
|
||||
group_root = os.path.join(split_root, "group")
|
||||
if not os.path.isdir(group_root):
|
||||
raise RuntimeError(f"Não achei pasta: {group_root}")
|
||||
|
||||
samples: List[Sample] = []
|
||||
|
||||
# pega qualquer images/ em qualquer subpasta de group
|
||||
img_dirs = glob.glob(os.path.join(group_root, "**", "images"), recursive=True)
|
||||
img_dirs = [d for d in img_dirs if os.path.isdir(d)]
|
||||
|
||||
for idir in img_dirs:
|
||||
base = os.path.dirname(idir) # .../group/<grupo_ou_subgrupo>
|
||||
group_name = os.path.relpath(base, group_root).replace("\\", "/")
|
||||
|
||||
mdir = os.path.join(base, "masks")
|
||||
m2dir = os.path.join(base, "masks2")
|
||||
|
||||
img_paths = []
|
||||
for ext in IMG_EXTS:
|
||||
img_paths.extend(glob.glob(os.path.join(idir, f"*{ext}")))
|
||||
img_paths = sorted(img_paths)
|
||||
|
||||
def find_mask(dir_path: str, filename: str):
|
||||
stem, _ = os.path.splitext(filename)
|
||||
for ext in [".png", ".jpg", ".jpeg", ".bmp", ".tif"]:
|
||||
p = os.path.join(dir_path, stem + ext)
|
||||
if os.path.exists(p):
|
||||
return p
|
||||
return None
|
||||
|
||||
for ip in img_paths:
|
||||
fn = os.path.basename(ip)
|
||||
|
||||
mask_path = find_mask(mdir, fn)
|
||||
mask2_path = find_mask(m2dir, fn)
|
||||
|
||||
samples.append(Sample(
|
||||
img_path=ip,
|
||||
mask_path=mask_path,
|
||||
mask2_path=mask2_path,
|
||||
group_name=group_name,
|
||||
filename=fn
|
||||
))
|
||||
|
||||
if len(samples) == 0:
|
||||
raise RuntimeError(f"Nenhuma imagem encontrada em: {group_root}/**/images")
|
||||
return samples
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Model loading
|
||||
# -----------------------------
|
||||
def build_models(num_classes: int, device: torch.device, backbone: str):
|
||||
base_model = SegformerForSemanticSegmentation.from_pretrained(
|
||||
backbone,
|
||||
num_labels=num_classes,
|
||||
ignore_mismatched_sizes=True,
|
||||
use_safetensors=True, # <- evita torch.load do .bin
|
||||
)
|
||||
base_model.to(device)
|
||||
base_model.config.output_hidden_states = True
|
||||
|
||||
# descobre feat_ch real a partir de um forward fake
|
||||
with torch.no_grad():
|
||||
dummy = torch.zeros((1, 3, 512, 512), device=device)
|
||||
out = base_model(pixel_values=dummy, output_hidden_states=True)
|
||||
feat_ch = int(out.hidden_states[-1].shape[1])
|
||||
|
||||
in_ch = feat_ch + num_classes
|
||||
corridor_head = CorridorHead(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1)
|
||||
print(f"[corridor_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {in_ch}")
|
||||
return base_model, corridor_head
|
||||
|
||||
|
||||
def make_corridor_head(in_ch: int, mid: int = 256):
|
||||
"""
|
||||
Cria CorridorHead independente de qual assinatura a classe tem no script.
|
||||
Suporta:
|
||||
- CorridorHead(in_ch=..., mid=...)
|
||||
- CorridorHead(in_channels=..., hidden=...)
|
||||
- CorridorHead(feat_ch=..., num_classes=...) (quando in_ch = feat_ch + num_classes)
|
||||
"""
|
||||
# 1) assinatura direta in_ch/mid
|
||||
try:
|
||||
return CorridorHead(in_ch=in_ch, mid=mid)
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
# 2) variações comuns
|
||||
try:
|
||||
return CorridorHead(in_channels=in_ch, mid=mid)
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
return CorridorHead(in_channels=in_ch, hidden=mid)
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
# 3) assinatura "feat_ch + num_classes"
|
||||
# aqui só funciona se for exatamente "feat_ch, num_classes"
|
||||
# e se in_ch for decomponível (ex: 6 = 3+3 ou 515 = 512+3)
|
||||
for feat_ch_guess, num_classes_guess in [(512, 3), (3, 3)]:
|
||||
if feat_ch_guess + num_classes_guess == in_ch:
|
||||
try:
|
||||
return CorridorHead(feat_ch=feat_ch_guess, num_classes=num_classes_guess, hidden=mid)
|
||||
except TypeError:
|
||||
try:
|
||||
return CorridorHead(feat_ch=feat_ch_guess, num_classes=num_classes_guess)
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
raise RuntimeError(
|
||||
f"Não consegui instanciar CorridorHead para in_ch={in_ch}. "
|
||||
f"Verifique a assinatura do __init__() da classe CorridorHead no script."
|
||||
)
|
||||
|
||||
def load_dual_checkpoint(ckpt_path, base_model, corridor_head, device):
|
||||
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
||||
|
||||
base_model.load_state_dict(ckpt["model"], strict=True)
|
||||
|
||||
# >>> pega quantos canais o head do checkpoint espera
|
||||
sd = ckpt["corridor_head"]
|
||||
in_ch_ckpt = int(sd["net.0.weight"].shape[1]) # ex: 6 ou 515
|
||||
|
||||
# >>> recria head se o in_ch não bater
|
||||
in_ch_model = int(corridor_head.net[0].in_channels)
|
||||
if in_ch_model != in_ch_ckpt:
|
||||
print(f"[corridor_head] rebuild: in_ch_model={in_ch_model} -> in_ch_ckpt={in_ch_ckpt}")
|
||||
corridor_head = make_corridor_head(in_ch=in_ch_ckpt, mid=256)
|
||||
|
||||
corridor_head.load_state_dict(sd, strict=True)
|
||||
|
||||
base_model.to(device).eval()
|
||||
corridor_head.to(device).eval()
|
||||
meta = ckpt.get("meta", {})
|
||||
return meta, corridor_head
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Inference
|
||||
# -----------------------------
|
||||
def run_corridor_head(corridor_head, x_rgb, feat, logits_seg):
|
||||
import inspect
|
||||
import torch.nn.functional as F
|
||||
|
||||
n_args = len(inspect.signature(corridor_head.forward).parameters)
|
||||
|
||||
# descobre in_ch esperado pela primeira conv do head
|
||||
in_ch = None
|
||||
if hasattr(corridor_head, "net"):
|
||||
try:
|
||||
in_ch = int(corridor_head.net[0].in_channels)
|
||||
except Exception:
|
||||
in_ch = None
|
||||
|
||||
# prepara entradas conforme o head
|
||||
if in_ch == 6:
|
||||
# RGB(3) + probs(3)
|
||||
probs = torch.softmax(logits_seg, dim=1)
|
||||
if probs.shape[-2:] != x_rgb.shape[-2:]:
|
||||
probs = F.interpolate(probs, size=x_rgb.shape[-2:], mode="bilinear", align_corners=False)
|
||||
feat_in = x_rgb
|
||||
logits_in = probs
|
||||
else:
|
||||
# feat(512) + logits(3)
|
||||
if logits_seg.shape[-2:] != feat.shape[-2:]:
|
||||
logits_in = F.interpolate(logits_seg, size=feat.shape[-2:], mode="bilinear", align_corners=False)
|
||||
else:
|
||||
logits_in = logits_seg
|
||||
feat_in = feat
|
||||
|
||||
# chama do jeito que o head espera
|
||||
if n_args == 2:
|
||||
return corridor_head(feat_in, logits_in)
|
||||
elif n_args == 1:
|
||||
x2 = torch.cat([feat_in, logits_in], dim=1)
|
||||
return corridor_head(x2)
|
||||
|
||||
raise RuntimeError(f"Assinatura inesperada: {inspect.signature(corridor_head.forward)}")
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_dual(base_model, corridor_head, img_bgr: np.ndarray, device: torch.device,
|
||||
input_size: int = 512, corridor_thr: float = 0.5):
|
||||
"""
|
||||
Retorna:
|
||||
pred_ids_full: [H,W] uint8
|
||||
prob_corr_full: [H,W] float32 (0..1)
|
||||
bin_corr_full: [H,W] uint8 (0/255)
|
||||
"""
|
||||
h0, w0 = img_bgr.shape[:2]
|
||||
|
||||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||||
img_res = cv2.resize(img_rgb, (input_size, input_size), interpolation=cv2.INTER_AREA)
|
||||
|
||||
x = torch.from_numpy(img_res).float() / 255.0
|
||||
x = x.permute(2, 0, 1).unsqueeze(0).to(device)
|
||||
|
||||
out = base_model(pixel_values=x, output_hidden_states=True)
|
||||
logits_seg = out.logits # [B, 3, H, W]
|
||||
feat = out.hidden_states[-1] # [B, 512, h, w] (não vamos usar no modo 6)
|
||||
|
||||
logits_corr = run_corridor_head(corridor_head, x, feat, logits_seg)
|
||||
|
||||
prob_corr = torch.sigmoid(logits_corr)[0, 0].detach().float().cpu().numpy()
|
||||
bin_corr = (prob_corr >= corridor_thr).astype(np.uint8) * 255
|
||||
|
||||
pred_ids = torch.argmax(logits_seg, dim=1)[0].detach().cpu().numpy().astype(np.uint8)
|
||||
|
||||
# volta pro tamanho original
|
||||
pred_ids_full = cv2.resize(pred_ids, (w0, h0), interpolation=cv2.INTER_NEAREST)
|
||||
prob_corr_full = cv2.resize(prob_corr, (w0, h0), interpolation=cv2.INTER_LINEAR)
|
||||
bin_corr_full = cv2.resize(bin_corr, (w0, h0), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
return pred_ids_full, prob_corr_full, bin_corr_full
|
||||
|
||||
|
||||
# -----------------------------
|
||||
# Viewer
|
||||
# -----------------------------
|
||||
def make_panel(sample: Sample , lut_bgr: dict, img_bgr: np.ndarray,
|
||||
gt1: Optional[np.ndarray],
|
||||
pred1: np.ndarray,
|
||||
gt2: Optional[np.ndarray],
|
||||
prob2: np.ndarray,
|
||||
bin2: np.ndarray,
|
||||
corridor_thr: float):
|
||||
h, w = img_bgr.shape[:2]
|
||||
|
||||
pred1_col = colorize_seg(pred1, lut_bgr)
|
||||
pred1_ov = overlay_mask(img_bgr, pred1_col, 0.45)
|
||||
|
||||
if gt1 is not None:
|
||||
gt1_col = colorize_seg(gt1, lut_bgr)
|
||||
gt1_ov = overlay_mask(img_bgr, gt1_col, 0.45)
|
||||
else:
|
||||
gt1_ov = img_bgr.copy()
|
||||
gt1_col = np.zeros_like(img_bgr)
|
||||
|
||||
# --- PRED CORR (prob) sólido: verde=navegável, vermelho=bloqueado
|
||||
prob = np.clip(prob2, 0.0, 1.0).astype(np.float32)
|
||||
g = (prob * 255.0).astype(np.uint8)
|
||||
r = ((1.0 - prob) * 255.0).astype(np.uint8)
|
||||
b = np.zeros_like(g, dtype=np.uint8)
|
||||
prob_solid = cv2.merge([b, g, r]) # BGR
|
||||
|
||||
# (opcional) desenha contorno do binário por cima do prob
|
||||
edges = cv2.Canny(bin2, 50, 150)
|
||||
prob_solid[edges > 0] = (255, 255, 255)
|
||||
|
||||
# --- PRED CORR (bin) sólido: branco=navegável, preto=bloqueado
|
||||
bin2_solid = cv2.cvtColor(bin2, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
if gt2 is not None:
|
||||
gt2_vis = (gt2.copy()).astype(np.uint8)
|
||||
gt2_vis[gt2_vis == 1] = 255
|
||||
gt2_bgr = cv2.cvtColor(gt2_vis, cv2.COLOR_GRAY2BGR)
|
||||
gt2_ov = overlay_mask(img_bgr, gt2_bgr, 0.35)
|
||||
else:
|
||||
gt2_ov = img_bgr.copy()
|
||||
|
||||
# monta grid
|
||||
tile_w = 520
|
||||
tile_h = int(tile_w * h / w)
|
||||
|
||||
def fit(im):
|
||||
return cv2.resize(im, (tile_w, tile_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
row1 = np.concatenate([fit(img_bgr), fit(gt1_ov), fit(pred1_ov)], axis=1)
|
||||
row2 = np.concatenate([fit(gt2_ov), fit(prob_solid), fit(bin2_solid)], axis=1)
|
||||
panel = np.concatenate([row1, row2], axis=0)
|
||||
|
||||
hud = [
|
||||
f"{sample.group_name}/{sample.filename}",
|
||||
f"Corr thr={corridor_thr:.2f} | A/D navega | Q sai",
|
||||
"Topo: IMG | GT SEG | PRED SEG",
|
||||
"Baixo: GT CORR | PRED CORR (prob) | PRED CORR (bin)"
|
||||
]
|
||||
panel = put_hud(panel, hud)
|
||||
return panel
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--ckpt", type=str, default=None, help="Caminho do .pt (se não informar, usa best_miou.pt)")
|
||||
ap.add_argument("--split", default="test", choices=["test", "val"], help="qual split usar")
|
||||
ap.add_argument("--input_size", type=int, default=512)
|
||||
ap.add_argument("--corr_thr", type=float, default=0.5)
|
||||
args = ap.parse_args()
|
||||
|
||||
device = find_device()
|
||||
print("Device:", device)
|
||||
|
||||
with open("config_oak.json", "r") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
MODEL_NAME = config["model_name"] # ex: "segformer_b3"
|
||||
BACKBONE = config["backbone"]
|
||||
modelo_folder = config["modelo"] # ex: "dif"
|
||||
dataset_path = os.path.join(MODELO, "dataset")
|
||||
|
||||
split_root = os.path.join(dataset_path, "split", args.split)
|
||||
if not os.path.isdir(split_root):
|
||||
print(f"[WARN] split/{args.split} não existe. Vou usar split/val.")
|
||||
split_root = os.path.join(dataset_path, "split", "val")
|
||||
|
||||
samples = discover_samples(split_root)
|
||||
print(f"[INFO] samples: {len(samples)} | split_root={split_root}")
|
||||
|
||||
labelmap_path = os.path.join(dataset_path, "labelmap.txt") # ajusta se seu arquivo estiver em outro lugar
|
||||
lut_bgr, id_to_name, ignore_id = load_labelmap_lut_bgr(labelmap_path)
|
||||
print("[labelmap]", id_to_name)
|
||||
|
||||
# ---- Descobre checkpoint ----
|
||||
if args.ckpt is not None:
|
||||
ckpt_path = args.ckpt
|
||||
else:
|
||||
# Caminho padrão igual ao treino:
|
||||
# save_path = MODELO/backup/modelo/model_name/raw4
|
||||
save_path = os.path.join(MODELO, "backup", modelo_folder, f"{MODEL_NAME}_dual")
|
||||
ckpt_path = os.path.join(save_path, "best_miou.pt")
|
||||
|
||||
if not os.path.isfile(ckpt_path):
|
||||
raise SystemExit(f"Checkpoint não encontrado em: {ckpt_path}")
|
||||
|
||||
print(f"[model] ckpt = {ckpt_path}")
|
||||
|
||||
num_classes = len(id_to_name)
|
||||
base_model, corridor_head = build_models(num_classes=num_classes, device=device, backbone=BACKBONE)
|
||||
meta, corridor_head = load_dual_checkpoint(ckpt_path, base_model, corridor_head, device=device)
|
||||
print(f"... epoch={meta.get('epoch')} best_miou={meta.get('best_miou')} best_main={meta.get('best_main_iou')} best_corr={meta.get('best_corr_iou')}")
|
||||
|
||||
idx = 0
|
||||
win = "SegFormer Dual Viewer"
|
||||
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
s = samples[idx]
|
||||
img = imread_bgr(s.img_path)
|
||||
|
||||
gt1 = imread_gray(s.mask_path) if s.mask_path else None
|
||||
gt2 = imread_gray(s.mask2_path) if s.mask2_path else None
|
||||
|
||||
pred1, prob2, bin2 = infer_dual(
|
||||
base_model=base_model,
|
||||
corridor_head=corridor_head,
|
||||
img_bgr=img,
|
||||
device=device,
|
||||
input_size=args.input_size,
|
||||
corridor_thr=args.corr_thr
|
||||
)
|
||||
|
||||
def dbg_mask(name, m):
|
||||
if m is None:
|
||||
print(f"[{name}] None")
|
||||
else:
|
||||
u = np.unique(m)
|
||||
print(f"[{name}] shape={m.shape} dtype={m.dtype} unique={u[:20]}{'...' if len(u)>20 else ''}")
|
||||
|
||||
#dbg_mask("gt1(seg)", gt1)
|
||||
#dbg_mask("gt2(corr)", gt2)
|
||||
|
||||
panel = make_panel(
|
||||
sample=s,
|
||||
lut_bgr=lut_bgr,
|
||||
img_bgr=img,
|
||||
gt1=gt1,
|
||||
pred1=pred1,
|
||||
gt2=gt2,
|
||||
prob2=prob2,
|
||||
bin2=bin2,
|
||||
corridor_thr=args.corr_thr
|
||||
)
|
||||
|
||||
cv2.imshow(win, panel)
|
||||
k = cv2.waitKey(0) & 0xFF
|
||||
|
||||
# Q / ESC
|
||||
if k in (ord('q'), 27):
|
||||
break
|
||||
|
||||
# A (volta)
|
||||
if k in (ord('a'), ord('A')):
|
||||
idx = (idx - 1) % len(samples)
|
||||
continue
|
||||
|
||||
# D (avança)
|
||||
if k in (ord('d'), ord('D')):
|
||||
idx = (idx + 1) % len(samples)
|
||||
continue
|
||||
|
||||
# W/S muda threshold
|
||||
if k in (ord('w'), ord('W')):
|
||||
args.corr_thr = min(0.95, args.corr_thr + 0.05)
|
||||
continue
|
||||
if k in (ord('s'), ord('S')):
|
||||
args.corr_thr = max(0.05, args.corr_thr - 0.05)
|
||||
continue
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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"
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
|
||||
|
||||
modelo = "nvidia/mit-b0"
|
||||
pasta = r"C:\AgroBaseModels\Backbones\nvidia_mit_b0"
|
||||
|
||||
model = SegformerForSemanticSegmentation.from_pretrained(modelo)
|
||||
processor = SegformerImageProcessor.from_pretrained(modelo)
|
||||
|
||||
model.save_pretrained(pasta)
|
||||
processor.save_pretrained(pasta)
|
||||
|
||||
print("Backbone salvo em:", pasta)
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -230,6 +230,7 @@ def main():
|
|||
|
||||
try:
|
||||
with MultiSpectralClient(
|
||||
#mx_id="194430108133AC2F00",
|
||||
width=raw_w,
|
||||
height=raw_h,
|
||||
bayer=args.bayer,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -18,10 +18,15 @@ RESOLUCAO = tuple(config.get("resolucao"))
|
|||
|
||||
TENSOR_EXT = ".npy"
|
||||
MASK_NPY_SUFFIX = ".npy"
|
||||
AUX_MASK_DIRS = [
|
||||
"masks_vegetation",
|
||||
"masks_cana",
|
||||
]
|
||||
|
||||
RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
|
||||
RE_AUGMENTED_FAMILY = re.compile(r"^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$", re.IGNORECASE)
|
||||
RE_AUG_SUFFIX = re.compile(r"_aug[a-zA-Z0-9]*_\d+$", re.IGNORECASE)
|
||||
MULTI_HEAD = bool(config.get("multi_head", False))
|
||||
|
||||
|
||||
def garantir(p):
|
||||
|
|
@ -237,6 +242,46 @@ def copiar_optional(src_dir, dst_dir, base, ext):
|
|||
return dst
|
||||
|
||||
|
||||
def copiar_mask_dir_optional(src_group_dir, dst_group_dir, mask_dir_name, base):
|
||||
"""
|
||||
Copia uma pasta auxiliar de máscara, como:
|
||||
masks_vegetation/
|
||||
masks_cana/
|
||||
|
||||
Copia:
|
||||
<base>.npy obrigatório se existir
|
||||
<base>.png opcional se existir
|
||||
|
||||
Retorna caminhos de destino ou None.
|
||||
"""
|
||||
src_dir = os.path.join(src_group_dir, mask_dir_name)
|
||||
dst_dir = os.path.join(dst_group_dir, mask_dir_name)
|
||||
|
||||
npy_src = os.path.join(src_dir, base + ".npy")
|
||||
png_src = os.path.join(src_dir, base + ".png")
|
||||
|
||||
result = {
|
||||
"npy": None,
|
||||
"png": None,
|
||||
}
|
||||
|
||||
if not os.path.exists(npy_src):
|
||||
return result
|
||||
|
||||
garantir(dst_dir)
|
||||
|
||||
npy_dst = os.path.join(dst_dir, base + ".npy")
|
||||
shutil.copy2(npy_src, npy_dst)
|
||||
result["npy"] = npy_dst
|
||||
|
||||
if os.path.exists(png_src):
|
||||
png_dst = os.path.join(dst_dir, base + ".png")
|
||||
shutil.copy2(png_src, png_dst)
|
||||
result["png"] = png_dst
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def copiar(
|
||||
nomes,
|
||||
src_group_dir,
|
||||
|
|
@ -304,11 +349,37 @@ def copiar(
|
|||
shutil.copy2(cand, preview_dst)
|
||||
break
|
||||
|
||||
aux_masks = {}
|
||||
for aux_dir in AUX_MASK_DIRS:
|
||||
aux_masks[aux_dir] = copiar_mask_dir_optional(
|
||||
src_group_dir=src_group_dir,
|
||||
dst_group_dir=dst_group_dir,
|
||||
mask_dir_name=aux_dir,
|
||||
base=base,
|
||||
)
|
||||
|
||||
if MULTI_HEAD:
|
||||
for aux_dir in AUX_MASK_DIRS:
|
||||
aux_npy = aux_masks.get(aux_dir, {}).get("npy")
|
||||
if aux_npy is None:
|
||||
raise RuntimeError(
|
||||
f"multi_head=true, mas máscara auxiliar ausente: "
|
||||
f"{aux_dir}/{base}.npy em {src_group_dir}"
|
||||
)
|
||||
|
||||
rows.append({
|
||||
"base": base,
|
||||
"tensor": tensor_dst,
|
||||
|
||||
"mask_npy": mask_npy_dst,
|
||||
"mask_png": mask_png_dst,
|
||||
|
||||
"mask_vegetation_npy": aux_masks.get("masks_vegetation", {}).get("npy"),
|
||||
"mask_vegetation_png": aux_masks.get("masks_vegetation", {}).get("png"),
|
||||
|
||||
"mask_cana_npy": aux_masks.get("masks_cana", {}).get("npy"),
|
||||
"mask_cana_png": aux_masks.get("masks_cana", {}).get("png"),
|
||||
|
||||
"meta": meta_dst,
|
||||
"preview": preview_dst,
|
||||
})
|
||||
|
|
@ -447,8 +518,16 @@ def write_manifest(path, rows):
|
|||
"group",
|
||||
"base",
|
||||
"tensor",
|
||||
|
||||
"mask_npy",
|
||||
"mask_png",
|
||||
|
||||
"mask_vegetation_npy",
|
||||
"mask_vegetation_png",
|
||||
|
||||
"mask_cana_npy",
|
||||
"mask_cana_png",
|
||||
|
||||
"meta",
|
||||
"preview",
|
||||
]
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -45,6 +45,7 @@ class OakFcc3Service:
|
|||
|
||||
def get_config(self):
|
||||
return {
|
||||
"mx_id": self.manager.mx_id,
|
||||
"ok": True,
|
||||
"fps": self.manager.fps,
|
||||
"width": self.manager.width,
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
|
|
@ -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
|
||||
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -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::
|
||||
|
|
@ -1,37 +1,68 @@
|
|||
from core.oak_fcc3_service import OakFcc3Service
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
svc = OakFcc3Service(
|
||||
fps=15,
|
||||
width=640,
|
||||
height=400,
|
||||
frame_type="MULTISPEC",
|
||||
capture_mode="TRIPLE",
|
||||
raw_policy="require_triple",
|
||||
sync_mode="best",
|
||||
sync_tolerance_ms=25.0,
|
||||
)
|
||||
import depthai as dai
|
||||
|
||||
svc.connect()
|
||||
svc.begin(frame_type="RAW_BRUTO", output_dtype="uint8", capture_mode="TRIPLE")
|
||||
|
||||
frame, meta = svc.capture_frame(timeout=2.0)
|
||||
def ler_attr_ou_metodo(obj, nomes, default=None):
|
||||
for nome in nomes:
|
||||
try:
|
||||
valor = getattr(obj, nome, None)
|
||||
|
||||
print(meta["camera_info"])
|
||||
for cam_id, arr in frame.items():
|
||||
print(cam_id, arr.shape, arr.dtype, arr.size)
|
||||
if valor is None:
|
||||
continue
|
||||
|
||||
from core.raw_processor_core import RawProcessorCore
|
||||
from core.raw_processor_preview import RawProcessorPreview
|
||||
import cv2
|
||||
if callable(valor):
|
||||
return valor()
|
||||
|
||||
core = RawProcessorCore(sensor_width=1280, sensor_height=800, bayer_pattern="GBRG")
|
||||
preview = RawProcessorPreview(sensor_width=1280, sensor_height=800, bayer_pattern="GBRG")
|
||||
return valor
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raw16 = core.unpack_raw10_packed(frame["cam0"], sensor_width=1280, sensor_height=800)
|
||||
img = preview.raw16_to_preview_bgr(raw16, bit_depth=10)
|
||||
return default
|
||||
|
||||
cv2.imwrite("calibration/cam0_raw_preview.png", img)
|
||||
print(raw16.shape, raw16.dtype, raw16.min(), raw16.max())
|
||||
|
||||
svc.stop()
|
||||
svc.disconnect()
|
||||
def main():
|
||||
print("==========================================")
|
||||
print("Teste simples - listar dispositivos DepthAI")
|
||||
print("==========================================")
|
||||
|
||||
try:
|
||||
devices = dai.Device.getAllAvailableDevices()
|
||||
except Exception as e:
|
||||
print("[ERRO] Falha ao chamar dai.Device.getAllAvailableDevices()")
|
||||
print("Erro:", e)
|
||||
return
|
||||
|
||||
if not devices:
|
||||
print("[INFO] Nenhum dispositivo DepthAI/OAK encontrado.")
|
||||
return
|
||||
|
||||
print(f"[OK] Dispositivos encontrados: {len(devices)}")
|
||||
print()
|
||||
|
||||
for i, dev_info in enumerate(devices, start=1):
|
||||
device_id = ler_attr_ou_metodo(
|
||||
dev_info,
|
||||
["getMxId", "mxid", "getDeviceId", "deviceId"],
|
||||
default=None
|
||||
)
|
||||
|
||||
name = ler_attr_ou_metodo(dev_info, ["name"], default=None)
|
||||
state = ler_attr_ou_metodo(dev_info, ["state"], default=None)
|
||||
protocol = ler_attr_ou_metodo(dev_info, ["protocol"], default=None)
|
||||
platform = ler_attr_ou_metodo(dev_info, ["platform"], default=None)
|
||||
status = ler_attr_ou_metodo(dev_info, ["status"], default=None)
|
||||
|
||||
print(f"Dispositivo #{i}")
|
||||
print(f" DeviceId : {device_id}")
|
||||
print(f" Name : {name}")
|
||||
print(f" State : {state}")
|
||||
print(f" Protocol : {protocol}")
|
||||
print(f" Platform : {platform}")
|
||||
print(f" Status : {status}")
|
||||
print("-" * 42)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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ê:
|
||||
<grupo>/bins/<base>_CAM_A.bin
|
||||
<grupo>/bins/<base>_CAM_B.bin
|
||||
<grupo>/bins/<base>_CAM_C.bin
|
||||
<grupo>/metas/<base>.json
|
||||
|
||||
E monta o tensor final usando RawProcessorCore + module_params.json,
|
||||
passando pelo mesmo fluxo real do normalize/runtime:
|
||||
RAW10 -> linear_demosaic/bayer_planes -> rgb_calibration
|
||||
-> radiometric_normalization -> homography/crop/resize
|
||||
-> flatfield final_tensor_space -> tensor [R,G,B,RE,NIR]
|
||||
|
||||
Mostra:
|
||||
esquerda : RGB do tensor base, com module_params original
|
||||
centro : RGB do tensor ajustado pelos sliders
|
||||
direita : diff x4 entre base e ajustado
|
||||
|
||||
Também mostra p50/mean/p01/p99/std de cada canal do tensor ajustado.
|
||||
|
||||
Uso:
|
||||
python raw_tensor_calibration_tuner.py ^
|
||||
--src_dir dataset/original/group/chao ^
|
||||
--module_params calibration/module_params.json ^
|
||||
--res 1024x640 ^
|
||||
--raw_size 1280x800
|
||||
|
||||
Controles:
|
||||
A / seta esquerda -> amostra anterior
|
||||
D / seta direita -> próxima amostra
|
||||
S -> salva patch JSON
|
||||
P -> imprime patch JSON
|
||||
R -> reseta sliders para valores iniciais do module_params
|
||||
Q / ESC -> sair
|
||||
|
||||
Notas:
|
||||
- Sliders flat_* usam escala 1000 = strength 1.0.
|
||||
Para flat-field, neutro é 0.
|
||||
- Sliders rgb_gain_* usam escala 1000 = ganho 1.0.
|
||||
Para RGB gain, neutro é 1000.
|
||||
- auto_level/gamma afetam apenas a visualização, não o tensor.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple, List
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Import do core
|
||||
# ============================================================
|
||||
|
||||
def add_project_paths():
|
||||
here = Path.cwd()
|
||||
candidates = [
|
||||
here,
|
||||
here / "Python" / "Scripts" / "workers" / "camera_worker" / "oak_fcc3_core",
|
||||
here / "Scripts" / "workers" / "camera_worker" / "oak_fcc3_core",
|
||||
]
|
||||
|
||||
for c in candidates:
|
||||
if c.exists() and str(c) not in sys.path:
|
||||
sys.path.insert(0, str(c))
|
||||
|
||||
add_project_paths()
|
||||
|
||||
try:
|
||||
from core.raw_processor_core import RawProcessorCore
|
||||
except Exception as e:
|
||||
raise SystemExit(
|
||||
"[ERRO] Não consegui importar core.raw_processor_core.RawProcessorCore.\n"
|
||||
"Rode este script a partir da pasta do projeto onde existe Python/Scripts/workers/camera_worker/oak_fcc3_core,\n"
|
||||
"ou coloque este script dentro da pasta oak_fcc3_core.\n"
|
||||
f"Erro original: {type(e).__name__}: {e}"
|
||||
)
|
||||
|
||||
|
||||
CHANNEL_NAMES = ["R", "G", "B", "RE", "NIR"]
|
||||
|
||||
|
||||
# ============================================================
|
||||
# IO helpers
|
||||
# ============================================================
|
||||
|
||||
def load_json(path: Path) -> dict:
|
||||
with path.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def save_json(path: Path, data: dict):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def parse_size(txt: str) -> Tuple[int, int]:
|
||||
w, h = [int(x) for x in str(txt).lower().split("x")]
|
||||
return w, h
|
||||
|
||||
|
||||
@dataclass
|
||||
class SampleBundle:
|
||||
base: str
|
||||
meta_path: Path
|
||||
bin_paths: Dict[str, Path]
|
||||
|
||||
|
||||
def find_sample_bins(bins_dir: Path, base: str) -> Dict[str, Path]:
|
||||
out = {}
|
||||
for cam_id in ("CAM_A", "CAM_B", "CAM_C"):
|
||||
p = bins_dir / f"{base}_{cam_id}.bin"
|
||||
if p.exists():
|
||||
out[cam_id] = p
|
||||
return out
|
||||
|
||||
|
||||
def collect_samples(src_dir: Path) -> List[SampleBundle]:
|
||||
metas_dir = src_dir / "metas"
|
||||
bins_dir = src_dir / "bins"
|
||||
|
||||
if not metas_dir.is_dir() or not bins_dir.is_dir():
|
||||
raise RuntimeError(
|
||||
f"src_dir precisa conter bins/ e metas/. Veio: {src_dir}"
|
||||
)
|
||||
|
||||
samples = []
|
||||
for meta_path in sorted(metas_dir.glob("*.json")):
|
||||
base = meta_path.stem
|
||||
bin_paths = find_sample_bins(bins_dir, base)
|
||||
|
||||
if not all(cam in bin_paths for cam in ("CAM_A", "CAM_B", "CAM_C")):
|
||||
print(f"[WARN] Pulando {base}: bins incompletos {list(bin_paths.keys())}")
|
||||
continue
|
||||
|
||||
samples.append(SampleBundle(base=base, meta_path=meta_path, bin_paths=bin_paths))
|
||||
|
||||
if not samples:
|
||||
raise RuntimeError(f"Nenhuma amostra válida encontrada em {src_dir}")
|
||||
|
||||
return samples
|
||||
|
||||
|
||||
def load_frame_from_saved_bins(sample: SampleBundle, meta: dict) -> Dict[str, np.ndarray]:
|
||||
saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
|
||||
saved_shapes = meta.get("saved_payload_shapes", {}) or {}
|
||||
|
||||
frame = {}
|
||||
|
||||
for cam_id, bin_path in sample.bin_paths.items():
|
||||
dtype = saved_dtypes.get(cam_id)
|
||||
shape = saved_shapes.get(cam_id)
|
||||
|
||||
if dtype is None or shape is None:
|
||||
raise RuntimeError(f"Faltam saved_payload_dtypes/shapes para {cam_id} em {sample.base}")
|
||||
|
||||
arr = np.fromfile(str(bin_path), dtype=np.dtype(dtype)).reshape(tuple(shape))
|
||||
frame[cam_id] = arr
|
||||
|
||||
return frame
|
||||
|
||||
|
||||
def build_processing_meta(meta: dict) -> dict:
|
||||
stream_meta = dict(meta.get("stream_meta", {}) or {})
|
||||
stream_meta["frame_type"] = "RAW_BRUTO"
|
||||
|
||||
if "camera_info" not in stream_meta and isinstance(meta.get("camera_info"), dict):
|
||||
stream_meta["camera_info"] = meta.get("camera_info")
|
||||
|
||||
# Mantém compatibilidade com metadados antigos e novos.
|
||||
for key in (
|
||||
"actual_camera_controls",
|
||||
"startup_camera_controls",
|
||||
"camera_controls",
|
||||
"frame_controls",
|
||||
"camera_frames",
|
||||
):
|
||||
if meta.get(key) is not None:
|
||||
stream_meta[key] = meta.get(key)
|
||||
|
||||
return stream_meta
|
||||
|
||||
|
||||
def tensor_stats(tensor: np.ndarray) -> Dict[str, dict]:
|
||||
out = {}
|
||||
for i, ch in enumerate(CHANNEL_NAMES):
|
||||
if i >= tensor.shape[0]:
|
||||
break
|
||||
|
||||
vals = tensor[i].astype(np.float32).reshape(-1)
|
||||
vals = vals[np.isfinite(vals)]
|
||||
if vals.size == 0:
|
||||
out[ch] = {}
|
||||
continue
|
||||
|
||||
out[ch] = {
|
||||
"min": float(np.min(vals)),
|
||||
"p01": float(np.percentile(vals, 1)),
|
||||
"p50": float(np.percentile(vals, 50)),
|
||||
"p99": float(np.percentile(vals, 99)),
|
||||
"max": float(np.max(vals)),
|
||||
"mean": float(np.mean(vals)),
|
||||
"std": float(np.std(vals)),
|
||||
}
|
||||
|
||||
return out
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Config adjustment
|
||||
# ============================================================
|
||||
|
||||
def deep_merge(dst: dict, src: dict) -> dict:
|
||||
out = copy.deepcopy(dst)
|
||||
|
||||
def rec(a, b):
|
||||
for k, v in b.items():
|
||||
if isinstance(v, dict) and isinstance(a.get(k), dict):
|
||||
rec(a[k], v)
|
||||
else:
|
||||
a[k] = copy.deepcopy(v)
|
||||
|
||||
rec(out, src)
|
||||
return out
|
||||
|
||||
|
||||
def load_initial_values(module_params: dict):
|
||||
ff = module_params.get("flatfield_config", {}) or {}
|
||||
rgbc = module_params.get("rgb_calibration", {}) or {}
|
||||
|
||||
strengths = {ch: float(ff.get("strength", 0.0)) for ch in CHANNEL_NAMES}
|
||||
for k, v in (ff.get("strength_by_channel", {}) or {}).items():
|
||||
ku = str(k).upper()
|
||||
if ku in strengths:
|
||||
strengths[ku] = float(v)
|
||||
|
||||
# Contrato correto do RawProcessorCore: rgb_calibration.gains
|
||||
gains = {"R": 1.0, "G": 1.0, "B": 1.0}
|
||||
src = rgbc.get("gains", {}) if isinstance(rgbc, dict) else {}
|
||||
if isinstance(src, dict):
|
||||
for ch in ("R", "G", "B"):
|
||||
if ch in src:
|
||||
gains[ch] = float(src[ch])
|
||||
|
||||
# Compatibilidade com patch antigo do tuner.
|
||||
src2 = rgbc.get("gain_by_channel", {}) if isinstance(rgbc, dict) else {}
|
||||
if isinstance(src2, dict):
|
||||
for ch in ("R", "G", "B"):
|
||||
if ch in src2:
|
||||
gains[ch] = float(src2[ch])
|
||||
|
||||
return strengths, gains
|
||||
|
||||
|
||||
def make_patch(strengths: Dict[str, float], rgb_gains: Dict[str, float]) -> dict:
|
||||
return {
|
||||
"flatfield_config": {
|
||||
"strength_by_channel": {
|
||||
"R": float(strengths["R"]),
|
||||
"G": float(strengths["G"]),
|
||||
"B": float(strengths["B"]),
|
||||
"RE": float(strengths["RE"]),
|
||||
"NIR": float(strengths["NIR"]),
|
||||
}
|
||||
},
|
||||
"rgb_calibration": {
|
||||
"enabled": True,
|
||||
"gains": {
|
||||
"R": float(rgb_gains["R"]),
|
||||
"G": float(rgb_gains["G"]),
|
||||
"B": float(rgb_gains["B"]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def apply_patch_to_core(core: RawProcessorCore, patch: dict):
|
||||
# Flatfield.
|
||||
ff_patch = patch.get("flatfield_config", {}) or {}
|
||||
if isinstance(ff_patch, dict):
|
||||
core.flatfield_config = deep_merge(core.flatfield_config, ff_patch)
|
||||
|
||||
# RGB calibration.
|
||||
rgb_patch = patch.get("rgb_calibration", {}) or {}
|
||||
if isinstance(rgb_patch, dict):
|
||||
core.rgb_calibration = deep_merge(core.rgb_calibration, rgb_patch)
|
||||
|
||||
# Limpa caches cujo valor depende do flat gain/strength.
|
||||
if hasattr(core, "_flatfield_runtime_cache"):
|
||||
core._flatfield_runtime_cache.clear()
|
||||
|
||||
# Geometria/remap não depende do strength nem dos rgb gains,
|
||||
# então não precisa limpar.
|
||||
|
||||
|
||||
def create_core(raw_size: Tuple[int, int], bayer_pattern: str, module_params: Path) -> RawProcessorCore:
|
||||
raw_w, raw_h = raw_size
|
||||
core = RawProcessorCore(
|
||||
sensor_width=int(raw_w),
|
||||
sensor_height=int(raw_h),
|
||||
bayer_pattern=str(bayer_pattern or "RGGB").upper(),
|
||||
calibration_json_path=str(module_params),
|
||||
)
|
||||
return core
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Tensor generation
|
||||
# ============================================================
|
||||
|
||||
class TensorBuilder:
|
||||
def __init__(self, module_params: Path, raw_size: Tuple[int, int], res: Tuple[int, int]):
|
||||
self.module_params = module_params
|
||||
self.raw_size = raw_size
|
||||
self.res = res
|
||||
self.module_params_data = load_json(module_params)
|
||||
|
||||
self.base_core = None
|
||||
self.tune_core = None
|
||||
self.current_bayer = None
|
||||
|
||||
self._sample_cache = {}
|
||||
self._base_tensor_cache = {}
|
||||
|
||||
def ensure_cores(self, bayer_pattern: str):
|
||||
bayer = str(bayer_pattern or "RGGB").upper()
|
||||
|
||||
if self.base_core is not None and self.current_bayer == bayer:
|
||||
return
|
||||
|
||||
print(f"[CORE] criando cores bayer={bayer} raw={self.raw_size} res={self.res}")
|
||||
self.base_core = create_core(self.raw_size, bayer, self.module_params)
|
||||
self.tune_core = create_core(self.raw_size, bayer, self.module_params)
|
||||
self.current_bayer = bayer
|
||||
|
||||
def get_frame_meta(self, sample: SampleBundle):
|
||||
key = sample.base
|
||||
cached = self._sample_cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
meta = load_json(sample.meta_path)
|
||||
frame = load_frame_from_saved_bins(sample, meta)
|
||||
processing_meta = build_processing_meta(meta)
|
||||
bayer = str(meta.get("bayer_pattern") or meta.get("bayer") or "RGGB").upper()
|
||||
|
||||
cached = (frame, processing_meta, bayer, meta)
|
||||
self._sample_cache[key] = cached
|
||||
return cached
|
||||
|
||||
def build_base_tensor(self, sample: SampleBundle):
|
||||
if sample.base in self._base_tensor_cache:
|
||||
return self._base_tensor_cache[sample.base]
|
||||
|
||||
frame, processing_meta, bayer, _meta = self.get_frame_meta(sample)
|
||||
self.ensure_cores(bayer)
|
||||
|
||||
tensor = self.base_core.build_infer_tensor_from_stream(
|
||||
frame=frame,
|
||||
meta=processing_meta,
|
||||
channels_expected=5,
|
||||
target_size=self.res,
|
||||
)
|
||||
|
||||
tensor = np.ascontiguousarray(tensor.astype(np.float32, copy=False))
|
||||
self._base_tensor_cache[sample.base] = tensor
|
||||
return tensor
|
||||
|
||||
def build_tuned_tensor(self, sample: SampleBundle, strengths: Dict[str, float], rgb_gains: Dict[str, float]):
|
||||
frame, processing_meta, bayer, _meta = self.get_frame_meta(sample)
|
||||
self.ensure_cores(bayer)
|
||||
|
||||
patch = make_patch(strengths, rgb_gains)
|
||||
apply_patch_to_core(self.tune_core, patch)
|
||||
|
||||
tensor = self.tune_core.build_infer_tensor_from_stream(
|
||||
frame=frame,
|
||||
meta=processing_meta,
|
||||
channels_expected=5,
|
||||
target_size=self.res,
|
||||
)
|
||||
|
||||
return np.ascontiguousarray(tensor.astype(np.float32, copy=False))
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Display
|
||||
# ============================================================
|
||||
|
||||
def val_to_gain(v: int) -> float:
|
||||
return float(v) / 1000.0
|
||||
|
||||
|
||||
def gain_to_val(g: float) -> int:
|
||||
return int(round(max(0.0, min(3.0, float(g))) * 1000.0))
|
||||
|
||||
|
||||
def val_to_strength(v: int) -> float:
|
||||
return float(v) / 1000.0
|
||||
|
||||
|
||||
def strength_to_val(s: float) -> int:
|
||||
return int(round(max(0.0, min(2.0, float(s))) * 1000.0))
|
||||
|
||||
|
||||
def read_ui(win: str):
|
||||
strengths = {
|
||||
"R": val_to_strength(cv2.getTrackbarPos("flat_R", win)),
|
||||
"G": val_to_strength(cv2.getTrackbarPos("flat_G", win)),
|
||||
"B": val_to_strength(cv2.getTrackbarPos("flat_B", win)),
|
||||
"RE": val_to_strength(cv2.getTrackbarPos("flat_RE", win)),
|
||||
"NIR": val_to_strength(cv2.getTrackbarPos("flat_NIR", win)),
|
||||
}
|
||||
|
||||
gains = {
|
||||
"R": val_to_gain(cv2.getTrackbarPos("rgb_gain_R", win)),
|
||||
"G": val_to_gain(cv2.getTrackbarPos("rgb_gain_G", win)),
|
||||
"B": val_to_gain(cv2.getTrackbarPos("rgb_gain_B", win)),
|
||||
}
|
||||
|
||||
auto_level = bool(cv2.getTrackbarPos("auto_level", win))
|
||||
gamma = max(0.05, cv2.getTrackbarPos("gamma_x100", win) / 100.0)
|
||||
|
||||
return strengths, gains, auto_level, gamma
|
||||
|
||||
|
||||
def set_ui(win: str, strengths: Dict[str, float], gains: Dict[str, float], gamma: float):
|
||||
for ch in CHANNEL_NAMES:
|
||||
cv2.setTrackbarPos(f"flat_{ch}", win, strength_to_val(strengths.get(ch, 0.0)))
|
||||
for ch in ("R", "G", "B"):
|
||||
cv2.setTrackbarPos(f"rgb_gain_{ch}", win, gain_to_val(gains.get(ch, 1.0)))
|
||||
cv2.setTrackbarPos("gamma_x100", win, int(round(gamma * 100)))
|
||||
|
||||
|
||||
def preview_rgb_from_tensor(tensor: np.ndarray, auto_level: bool, gamma: float) -> np.ndarray:
|
||||
rgb = np.transpose(tensor[:3].astype(np.float32), (1, 2, 0))
|
||||
rgb = np.nan_to_num(rgb, nan=0.0, posinf=1.0, neginf=0.0)
|
||||
|
||||
if auto_level:
|
||||
# Auto-level visual separado por canal para facilitar inspeção.
|
||||
out = np.empty_like(rgb)
|
||||
for c in range(3):
|
||||
ch = rgb[..., c]
|
||||
lo = float(np.percentile(ch, 1.0))
|
||||
hi = float(np.percentile(ch, 99.0))
|
||||
if hi > lo:
|
||||
out[..., c] = (ch - lo) / (hi - lo)
|
||||
else:
|
||||
out[..., c] = ch
|
||||
rgb = out
|
||||
|
||||
rgb = np.clip(rgb, 0.0, 1.0)
|
||||
|
||||
if gamma > 0 and abs(gamma - 1.0) > 1e-6:
|
||||
rgb = np.power(rgb, gamma)
|
||||
|
||||
return (rgb * 255.0).clip(0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def put_text(img: np.ndarray, lines, x=12, y=25, scale=0.55):
|
||||
out = img.copy()
|
||||
yy = y
|
||||
for line in lines:
|
||||
cv2.putText(out, str(line), (x, yy), cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), 4, cv2.LINE_AA)
|
||||
cv2.putText(out, str(line), (x, yy), cv2.FONT_HERSHEY_SIMPLEX, scale, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
yy += int(22 * scale / 0.55)
|
||||
return out
|
||||
|
||||
|
||||
def resize_keep_width(img: np.ndarray, target_w: int) -> np.ndarray:
|
||||
h, w = img.shape[:2]
|
||||
if w <= target_w:
|
||||
return img
|
||||
scale = target_w / float(w)
|
||||
return cv2.resize(img, (target_w, int(h * scale)), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def stats_lines(stats: Dict[str, dict], prefix=""):
|
||||
lines = []
|
||||
for ch in CHANNEL_NAMES:
|
||||
s = stats.get(ch, {}) or {}
|
||||
if not s:
|
||||
continue
|
||||
lines.append(
|
||||
f"{prefix}{ch}: p50={s['p50']:.4f} mean={s['mean']:.4f} "
|
||||
f"p01={s['p01']:.4f} p99={s['p99']:.4f} std={s['std']:.4f}"
|
||||
)
|
||||
return lines
|
||||
|
||||
|
||||
def build_canvas(
|
||||
sample: SampleBundle,
|
||||
idx: int,
|
||||
total: int,
|
||||
base_tensor: np.ndarray,
|
||||
tuned_tensor: np.ndarray,
|
||||
strengths: Dict[str, float],
|
||||
gains: Dict[str, float],
|
||||
auto_level: bool,
|
||||
gamma: float,
|
||||
max_width: int,
|
||||
):
|
||||
base_rgb = preview_rgb_from_tensor(base_tensor, auto_level=auto_level, gamma=gamma)
|
||||
tuned_rgb = preview_rgb_from_tensor(tuned_tensor, auto_level=auto_level, gamma=gamma)
|
||||
|
||||
diff = np.abs(tuned_rgb.astype(np.int16) - base_rgb.astype(np.int16)).clip(0, 255).astype(np.uint8)
|
||||
diff = np.clip(diff.astype(np.uint16) * 4, 0, 255).astype(np.uint8)
|
||||
|
||||
h = base_rgb.shape[0]
|
||||
sep = np.zeros((h, 4, 3), dtype=np.uint8)
|
||||
canvas = np.hstack([base_rgb, sep, tuned_rgb, sep, diff])
|
||||
|
||||
st_base = tensor_stats(base_tensor)
|
||||
st_tuned = tensor_stats(tuned_tensor)
|
||||
|
||||
top_lines = [
|
||||
f"{idx + 1}/{total} | {sample.base}",
|
||||
"LEFT=base tensor | MID=tuned tensor | RIGHT=diff x4",
|
||||
f"flat R={strengths['R']:.3f} G={strengths['G']:.3f} B={strengths['B']:.3f} RE={strengths['RE']:.3f} NIR={strengths['NIR']:.3f}",
|
||||
f"rgb_gain R={gains['R']:.3f} G={gains['G']:.3f} B={gains['B']:.3f} | auto={int(auto_level)} gamma={gamma:.2f}",
|
||||
"A/D navega | S salva patch | P print | R reset | Q sai",
|
||||
]
|
||||
|
||||
canvas = put_text(canvas, top_lines, x=12, y=24, scale=0.56)
|
||||
|
||||
# Painel de stats na parte inferior, com fundo escuro.
|
||||
panel_h = 190
|
||||
panel = np.zeros((panel_h, canvas.shape[1], 3), dtype=np.uint8)
|
||||
|
||||
stat_lines = ["TUNED tensor stats:"] + stats_lines(st_tuned, "")
|
||||
# Também mostra relação G/R e G/B, útil para diagnosticar esverdeado.
|
||||
try:
|
||||
gr = st_tuned["G"]["mean"] / max(st_tuned["R"]["mean"], 1e-6)
|
||||
gb = st_tuned["G"]["mean"] / max(st_tuned["B"]["mean"], 1e-6)
|
||||
stat_lines.append(f"RGB balance mean: G/R={gr:.3f} | G/B={gb:.3f}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
panel = put_text(panel, stat_lines[:8], x=12, y=24, scale=0.58)
|
||||
|
||||
full = np.vstack([canvas, panel])
|
||||
return resize_keep_width(full, max_width)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--src_dir", required=True, help="Grupo original contendo bins/ e metas/. Ex: dataset/original/group/chao")
|
||||
ap.add_argument("--module_params", required=True, help="calibration/module_params.json")
|
||||
ap.add_argument("--res", default="1024x640", help="Resolução final WxH do tensor")
|
||||
ap.add_argument("--raw_size", default="1280x800", help="RAW size WxH")
|
||||
ap.add_argument("--out_patch", default="raw_tensor_tuned_patch.json")
|
||||
ap.add_argument("--max_width", type=int, default=1900)
|
||||
ap.add_argument("--gamma", type=float, default=1.0)
|
||||
args = ap.parse_args()
|
||||
|
||||
src_dir = Path(args.src_dir)
|
||||
module_params = Path(args.module_params)
|
||||
|
||||
if not src_dir.is_dir():
|
||||
raise SystemExit(f"[ERRO] src_dir não existe: {src_dir}")
|
||||
|
||||
if not module_params.is_file():
|
||||
raise SystemExit(f"[ERRO] module_params não existe: {module_params}")
|
||||
|
||||
res = parse_size(args.res)
|
||||
raw_size = parse_size(args.raw_size)
|
||||
|
||||
module_data = load_json(module_params)
|
||||
initial_strengths, initial_gains = load_initial_values(module_data)
|
||||
|
||||
samples = collect_samples(src_dir)
|
||||
|
||||
builder = TensorBuilder(
|
||||
module_params=module_params,
|
||||
raw_size=raw_size,
|
||||
res=res,
|
||||
)
|
||||
|
||||
win = "RAW Tensor Calibration Tuner"
|
||||
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
|
||||
|
||||
for ch in CHANNEL_NAMES:
|
||||
cv2.createTrackbar(f"flat_{ch}", win, 0, 2000, lambda v: None)
|
||||
|
||||
for ch in ("R", "G", "B"):
|
||||
cv2.createTrackbar(f"rgb_gain_{ch}", win, 1000, 3000, lambda v: None)
|
||||
|
||||
cv2.createTrackbar("auto_level", win, 0, 1, lambda v: None)
|
||||
cv2.createTrackbar("gamma_x100", win, 100, 300, lambda v: None)
|
||||
|
||||
set_ui(win, initial_strengths, initial_gains, args.gamma)
|
||||
|
||||
idx = 0
|
||||
last_key = None
|
||||
last_tuned_key = None
|
||||
base_tensor = None
|
||||
tuned_tensor = None
|
||||
|
||||
print("[INFO] RAW Tensor Calibration Tuner iniciado.")
|
||||
print("[INFO] O tensor base usa module_params original. O tensor tuned usa sliders.")
|
||||
print("[INFO] S salva patch JSON. P imprime valores. R reseta.")
|
||||
print("[INFO] Se ficar lento ao arrastar sliders, solte o mouse: ele recalcula o tensor real a cada mudança.")
|
||||
|
||||
while True:
|
||||
sample = samples[idx]
|
||||
|
||||
strengths, gains, auto_level, gamma = read_ui(win)
|
||||
|
||||
tuned_key = (
|
||||
sample.base,
|
||||
tuple(round(strengths[ch], 4) for ch in CHANNEL_NAMES),
|
||||
tuple(round(gains[ch], 4) for ch in ("R", "G", "B")),
|
||||
)
|
||||
|
||||
try:
|
||||
if last_key != sample.base:
|
||||
t0 = time.perf_counter()
|
||||
base_tensor = builder.build_base_tensor(sample)
|
||||
tuned_tensor = None
|
||||
last_tuned_key = None
|
||||
last_key = sample.base
|
||||
print(f"[LOAD] {sample.base} base_tensor={base_tensor.shape} {((time.perf_counter()-t0)*1000):.1f}ms")
|
||||
|
||||
if last_tuned_key != tuned_key or tuned_tensor is None:
|
||||
t0 = time.perf_counter()
|
||||
tuned_tensor = builder.build_tuned_tensor(sample, strengths, gains)
|
||||
last_tuned_key = tuned_key
|
||||
print(f"[TUNE] {sample.base} tuned_tensor={tuned_tensor.shape} {((time.perf_counter()-t0)*1000):.1f}ms")
|
||||
|
||||
canvas = build_canvas(
|
||||
sample=sample,
|
||||
idx=idx,
|
||||
total=len(samples),
|
||||
base_tensor=base_tensor,
|
||||
tuned_tensor=tuned_tensor,
|
||||
strengths=strengths,
|
||||
gains=gains,
|
||||
auto_level=auto_level,
|
||||
gamma=gamma,
|
||||
max_width=args.max_width,
|
||||
)
|
||||
|
||||
cv2.imshow(win, cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR))
|
||||
|
||||
except Exception as e:
|
||||
err = np.zeros((600, 1200, 3), dtype=np.uint8)
|
||||
err = put_text(err, [
|
||||
f"ERRO em {sample.base}",
|
||||
f"{type(e).__name__}: {e}",
|
||||
"A/D navega | Q sai",
|
||||
], x=20, y=40, scale=0.7)
|
||||
cv2.imshow(win, cv2.cvtColor(err, cv2.COLOR_RGB2BGR))
|
||||
|
||||
k = cv2.waitKey(30) & 0xFF
|
||||
|
||||
if k in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
|
||||
elif k in (ord("d"), ord("D"), 83):
|
||||
idx = (idx + 1) % len(samples)
|
||||
last_key = None
|
||||
last_tuned_key = None
|
||||
|
||||
elif k in (ord("a"), ord("A"), 81):
|
||||
idx = (idx - 1 + len(samples)) % len(samples)
|
||||
last_key = None
|
||||
last_tuned_key = None
|
||||
|
||||
elif k in (ord("r"), ord("R")):
|
||||
set_ui(win, initial_strengths, initial_gains, args.gamma)
|
||||
last_tuned_key = None
|
||||
|
||||
elif k in (ord("p"), ord("P")):
|
||||
patch = make_patch(strengths, gains)
|
||||
print(json.dumps(patch, ensure_ascii=False, indent=2))
|
||||
|
||||
elif k in (ord("s"), ord("S")):
|
||||
patch = make_patch(strengths, gains)
|
||||
out = Path(args.out_patch)
|
||||
save_json(out, patch)
|
||||
print(f"[SAVE] {out.resolve()}")
|
||||
print(json.dumps(patch, ensure_ascii=False, indent=2))
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue