Limpeza e refinamento na firmware do modulo multiespectral

This commit is contained in:
Diego Freitas 2026-04-22 17:08:49 -03:00
parent fc1289f177
commit 5c26d66378
34 changed files with 5506 additions and 3042 deletions

2
.gitignore vendored
View File

@ -65,6 +65,8 @@ Python/OAK/datasets/oak-d/dataset/
Python/OAK/datasets/oak-d/backup/ Python/OAK/datasets/oak-d/backup/
Python/OAK/datasets/gal5000/dataset/ Python/OAK/datasets/gal5000/dataset/
Python/OAK/datasets/gal5000/backup/ Python/OAK/datasets/gal5000/backup/
Python/OAK/datasets/multiespec_module/dataset/
Python/OAK/datasets/multiespec_module/backup/
Python/yolov8-seg/venv/ Python/yolov8-seg/venv/
Python/yolov8-seg/images/ Python/yolov8-seg/images/

View File

@ -1,328 +0,0 @@
import os
import time
import json
import argparse
from datetime import datetime
import numpy as np
import cv2
from raw_segformer_service import make_bgr_preview_from_raw
from gal5000.gal_service import Gal5000Camera # ajuste o nome do módulo se estiver diferente
# =========================
# Helpers gerais
# =========================
def clamp(v, lo, hi):
return lo if v < lo else hi if v > hi else v
def ts_name() -> str:
"""Timestamp legível e único para nome de arquivo."""
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
"""
Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
Pensado pra deixar o preview bonitinho sem estourar tudo.
"""
x = np.asarray(x)
if x.dtype != np.float32 and x.dtype != np.float64:
x = x.astype(np.float32)
# Se o canal já está em 0..1, escala pra 0..255 antes de cortar
if x.max() <= 1.5:
x = x * 255.0
lo = np.percentile(x, p_lo)
hi = np.percentile(x, p_hi)
if hi <= lo + 1e-3:
y = x
else:
y = (x - lo) * (255.0 / (hi - lo))
return np.clip(y, 0, 255).astype(np.uint8)
def overlay_hud(
img_bgr: np.ndarray,
lines: list[str],
base_h: int = 720,
base_font_scale: float = 0.75,
base_line_step: int = 28,
):
"""
Escreve textos empilhados no canto superior esquerdo,
ajustando o tamanho do texto de acordo com a altura da imagem.
base_h: altura de referência (ex: 720 ou a RAW_H original).
"""
h, w = img_bgr.shape[:2]
# Fator de escala com base na altura atual
scale = h / float(base_h)
# Evita ficar microscópico em resoluções muito baixas
scale = max(scale, 0.4)
font_scale = base_font_scale * scale
line_step = int(base_line_step * scale)
# Espessuras proporcionais
thick_outline = max(1, int(3 * scale))
thick_text = max(1, int(2 * scale))
# Margem superior / esquerda também escaladas
y = int(24 * scale)
x = int(12 * scale)
for s in lines:
# contorno preto
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
# texto branco
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
y += line_step
def save_sample_raw4(
base_dir: str,
raw4: np.ndarray,
preview_bgr: np.ndarray,
meta: dict,
):
"""
Salva:
- RAW4 como .raw float32 (4,H,W)
- preview RGB como .png
- metadados como .json
dentro de base_dir.
"""
os.makedirs(base_dir, exist_ok=True)
name = ts_name()
raw_path = os.path.join(base_dir, f"{name}.raw")
png_path = os.path.join(base_dir, f"{name}.png")
json_path = os.path.join(base_dir, f"{name}.json")
# RAW4
#np.save(raw_path, raw4.astype(np.float32))
raw4.astype(np.float32).tofile(raw_path)
# Preview
cv2.imwrite(png_path, preview_bgr)
# Metadados
with open(json_path, "w", encoding="utf-8") as f:
json.dump(meta, f, ensure_ascii=False, indent=2)
return raw_path, png_path, json_path
# =========================
# MAIN
# =========================
def main():
parser = argparse.ArgumentParser(
description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
parser.add_argument("--dll_dir", default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls", help="Pasta onde está a VT_SDK64.dll (usada pelo Gal5000Camera).")
parser.add_argument("--dll_name", default="VT_SDK64.dll", help="Nome da DLL da câmera.")
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
parser.add_argument("--no_ae", action="store_true", help="Desliga o AutoExposure do service (por padrão ele vem ligado).")
parser.add_argument("--upscale", type=int, default=2, help="Fator de upscale visual do preview.")
args = parser.parse_args()
# ===== config =====
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
RAW_W = config["raw_size"][0]
RAW_H = config["raw_size"][1]
# Define diretório de sessão:
# dataset/cana_<estado>/<horario>/<YYYYMMDD>/
session_dir = os.path.join(MODELO, args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
os.makedirs(session_dir, exist_ok=True)
print("============================================")
print("Coleta de dataset RAW4 - SegFormer B0")
print(f"Cana : {args.cana}")
print(f"Horário : {args.horario}")
print(f"Saída : {session_dir}")
print("============================================")
window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
auto_save = False
last_auto_t = 0.0
upscale = args.upscale
# Estatísticas simples de FPS
t_fps = time.time()
frames = 0
fps = 0.0
last_msg = ""
last_msg_t = 0.0
try:
cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
with cam:
print("[CAM] Status inicial:", cam.get_status())
# opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
cam.configure_fps(20)
cam.start_streaming()
apply_ir_comp = True
ir_k_r = 0.8
ir_k_g = 0.4
ir_k_b = 0.9
while True:
t0 = time.time()
raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
t1 = time.time()
ae_dbg = dbg.get("ae", {}) or {}
exp_raw = dbg.get("exp_raw", None)
gain_a = dbg.get("gain_a", None)
gain_d = dbg.get("gain_d", None)
bgr = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=upscale > 0, preview_scale=upscale, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
# FPS
frames += 1
dt_fps = time.time() - t_fps
if dt_fps >= 1.0:
fps = frames / dt_fps
frames = 0
t_fps = time.time()
# AE info
p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
sat_disp = ae_dbg.get("sat", 0.0)
hold = ae_dbg.get("hold", False)
ae_on = cam.is_auto_exposure_enabled()
# HUD principal
lines = [
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
f"AEdbg: p95={p95_disp:.1f} sat={sat_disp:.3f} hold={hold}",
"Keys: C/SPACE=save | A=auto-save | E=AE toggle | M=preview scale | Q/Esc=quit",
]
overlay_hud(bgr, lines, base_h=RAW_H)
# Mensagem rápida (ex: arquivo salvo)
if last_msg and (time.time() - last_msg_t) < 2.0:
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
cv2.imshow(window_name, bgr)
# Auto-save
now = time.time()
if auto_save and (now - last_auto_t) >= args.interval:
meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"cana": args.cana,
"horario": args.horario,
"raw4_shape": list(raw4_base.shape),
"out_h": RAW_H,
"out_w": RAW_W,
"ae_enabled": bool(ae_on),
"exp_raw": int(exp_raw) if exp_raw is not None else None,
"gain_a": int(gain_a) if gain_a is not None else None,
"gain_d": int(gain_d) if gain_d is not None else None,
"apply_ir_comp": apply_ir_comp,
"ir_k_r": ir_k_r,
"ir_k_g": ir_k_g,
"ir_k_b": ir_k_b,
"ae_dbg": {
k: (float(v) if isinstance(v, (int, float, np.floating)) else v)
for k, v in ae_dbg.items()
},
"note": "autosave",
}
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
last_msg = f"SALVO (auto): {os.path.basename(raw_path)}"
last_msg_t = now
last_auto_t = now
# Teclado
k = cv2.waitKey(1) & 0xFF
if k in (ord("q"), ord("Q"), 27): # Q ou ESC
break
elif k in (ord("a"), ord("A")):
auto_save = not auto_save
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
last_msg_t = time.time()
elif k in (ord("e"), ord("E")):
cam.enable_auto_exposure(not ae_on)
last_msg = f"AE -> {'ON' if cam.is_auto_exposure_enabled() else 'OFF'}"
last_msg_t = time.time()
elif k in (ord("m"), ord("M")):
upscale = 0 if upscale else args.upscale
last_msg = f"Preview UPSCALE -> {upscale}"
last_msg_t = time.time()
elif k in (ord("c"), ord("C"), 32): # C ou SPACE
meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"cana": args.cana,
"horario": args.horario,
"raw4_shape": list(raw4_base.shape),
"out_h": RAW_H,
"out_w": RAW_W,
"ae_enabled": bool(ae_on),
"exp_raw": int(exp_raw) if exp_raw is not None else None,
"gain_a": int(gain_a) if gain_a is not None else None,
"gain_d": int(gain_d) if gain_d is not None else None,
"apply_ir_comp": apply_ir_comp,
"ir_k_r": ir_k_r,
"ir_k_g": ir_k_g,
"ir_k_b": ir_k_b,
"ae_dbg": {
k2: (float(v2) if isinstance(v2, (int, float, np.floating)) else v2)
for k2, v2 in ae_dbg.items()
},
"note": "manual",
}
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
last_msg_t = time.time()
# você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)
# Só pra não ficar rodando a 1000 FPS na UI
# mas sem travar muito a captura
dt_loop = time.time() - t0
if dt_loop < 0.001:
time.sleep(0.001)
finally:
cv2.destroyAllWindows()
print("Fim da captura.")
if __name__ == "__main__":
main()

View File

@ -1,302 +1,176 @@
import os import os
import time import time
import json import json
import math import argparse
import ctypes as C
from ctypes import wintypes as W
from datetime import datetime from datetime import datetime
import numpy as np import numpy as np
import cv2 import cv2
# ========================= from raw_segformer_service import make_bgr_preview_from_raw
# CONFIG from gal_service import Gal5000Camera # ajuste o nome do módulo se estiver diferente
# =========================
SDK_DIR = os.path.join(os.path.dirname(__file__), "dlls")
DLL_NAME = "VT_SDK64.dll"
# Onde salvar o dataset
OUT_ROOT = os.path.join(os.path.dirname(__file__), "dataset")
SESSION_DIR = os.path.join(OUT_ROOT, datetime.now().strftime("%Y%m%d"))
os.makedirs(SESSION_DIR, exist_ok=True)
# Camera scan/open
DEVICE_UDEF = 0
DEVICE_INDEX = 0
DATA_RAW = 0
# RAW geometry (se mudar no futuro, ajuste)
RAW_W = 2592
RAW_H = 2056
TIMEOUT_MS = 2000
WINDOW_NAME = "GAL5000 Dataset Capture (C/SPACE=save | A=auto-save | E=AE toggle | Q=quit)"
# Preview
UPSCALE = 2
# Auto-save
CAPTURE_INTERVAL_S = 1.0
# Param IDs (VT_Param.h)
BUF_SIZE = 256
PARAM_ID_SENSOR_EXPOSURETIMERAW = 0x00003010
PARAM_ID_SENSOR_GAINANALOGRAW = 0x00003020
PARAM_ID_SENSOR_GAINDIGITRAW = 0x0000302A
# PARAM_VALUETYPE
VALUE_INT = 0
VALUE_FLOAT = 1
VALUE_STR = 2
# Exposure/Gain limits (ajuste depois conforme o sensor aceitar)
EXP_MIN = 1
EXP_MAX = 20000
GAIN_A_MIN, GAIN_A_MAX = 0, 255
GAIN_D_MIN, GAIN_D_MAX = 0, 255
# ========================= # =========================
# Helpers # Helpers gerais
# ========================= # =========================
def ck(ret: int, name: str):
if ret != 0:
raise RuntimeError(f"{name} falhou, ret={ret}")
def ts_name() -> str:
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
def clamp(v, lo, hi): def clamp(v, lo, hi):
return lo if v < lo else hi if v > hi else v return lo if v < lo else hi if v > hi else v
def norm8(x, p_lo=2, p_hi=98):
def ts_name() -> str:
"""Timestamp legível e único para nome de arquivo."""
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
"""
Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
Pensado pra deixar o preview bonitinho sem estourar tudo.
"""
x = np.asarray(x)
if x.dtype != np.float32 and x.dtype != np.float64:
x = x.astype(np.float32)
# Se o canal já está em 0..1, escala pra 0..255 antes de cortar
if x.max() <= 1.5:
x = x * 255.0
lo = np.percentile(x, p_lo) lo = np.percentile(x, p_lo)
hi = np.percentile(x, p_hi) hi = np.percentile(x, p_hi)
if hi <= lo + 1:
return x.astype(np.uint8) if hi <= lo + 1e-3:
y = (x.astype(np.float32) - lo) * (255.0 / (hi - lo)) y = x
else:
y = (x - lo) * (255.0 / (hi - lo))
return np.clip(y, 0, 255).astype(np.uint8) return np.clip(y, 0, 255).astype(np.uint8)
def make_rgb_preview(raw: np.ndarray, upscale=2) -> np.ndarray:
# pattern:
# R G
# IR B
R = raw[0::2, 0::2]
G = raw[0::2, 1::2]
B = raw[1::2, 1::2]
Rn, Gn, Bn = norm8(R), norm8(G), norm8(B) def overlay_hud(
bgr = np.dstack([Bn, Gn, Rn]) # OpenCV usa BGR img_bgr: np.ndarray,
if upscale and upscale != 1: lines: list[str],
bgr = cv2.resize(bgr, (bgr.shape[1]*upscale, bgr.shape[0]*upscale), interpolation=cv2.INTER_NEAREST) base_h: int = 720,
return bgr base_font_scale: float = 0.75,
base_line_step: int = 28,
def measure_raw_g_metrics(raw: np.ndarray): ):
""" """
Mede brilho no canal G cru usando uma ROI na base (mais parecido com chão). Escreve textos empilhados no canto superior esquerdo,
Retorna p90/p95 e fração saturada. ajustando o tamanho do texto de acordo com a altura da imagem.
base_h: altura de referência (ex: 720 ou a RAW_H original).
""" """
G = raw[0::2, 1::2] # H/2 x W/2 h, w = img_bgr.shape[:2]
h2, w2 = G.shape
# ROI: base da imagem, cortando laterais # Fator de escala com base na altura atual
y0, y1 = int(h2 * 0.55), int(h2 * 0.95) scale = h / float(base_h)
x0, x1 = int(w2 * 0.15), int(w2 * 0.85)
roi = G[y0:y1, x0:x1]
p90 = float(np.percentile(roi, 90)) # Evita ficar microscópico em resoluções muito baixas
p95 = float(np.percentile(roi, 95)) scale = max(scale, 0.4)
sat = float(np.mean(roi >= 250))
return p90, p95, sat
class RobustAE: font_scale = base_font_scale * scale
""" line_step = int(base_line_step * scale)
Controle soft de exposure (sem depender do GET da camera):
- mede p95 do canal G cru em ROI
- usa EMA + deadband (pra não ficar "descendo até 16" como você viu)
- passo multiplicativo em log, com limite de passo
"""
def __init__(self,
exp_min=EXP_MIN, exp_max=EXP_MAX,
target_p95=140.0,
deadband=6.0,
k=0.12,
max_step=0.10,
ema_alpha=0.20,
sat_limit=0.01):
self.exp_min = exp_min
self.exp_max = exp_max
self.target = target_p95
self.deadband = deadband
self.k = k
self.max_step = max_step
self.ema_alpha = ema_alpha
self.sat_limit = sat_limit
self.p95_ema = None
def step(self, raw, exp_raw): # Espessuras proporcionais
p90, p95, sat = measure_raw_g_metrics(raw) thick_outline = max(1, int(3 * scale))
thick_text = max(1, int(2 * scale))
# EMA do p95 (estabiliza) # Margem superior / esquerda também escaladas
if self.p95_ema is None: y = int(24 * scale)
self.p95_ema = p95 x = int(12 * scale)
else:
self.p95_ema = (1 - self.ema_alpha) * self.p95_ema + self.ema_alpha * p95
e = self.target - self.p95_ema # erro em nível de pixel
# deadband: segura a mão perto do alvo
if abs(e) <= self.deadband and sat <= self.sat_limit:
return exp_raw, {"p90": p90, "p95": p95, "p95_ema": self.p95_ema, "sat": sat, "hold": True}
# saturou: garante redução
if sat > self.sat_limit:
step = -min(self.max_step, 0.12)
else:
ratio = (self.target + 1e-6) / (self.p95_ema + 1e-6)
step = self.k * math.log(ratio)
step = max(-self.max_step, min(self.max_step, step))
new_exp = int(round(exp_raw * math.exp(step)))
new_exp = max(self.exp_min, min(self.exp_max, new_exp))
return new_exp, {"p90": p90, "p95": p95, "p95_ema": self.p95_ema, "sat": sat, "step": step, "hold": False}
# =========================
# STRUCTS + Param API
# =========================
class VT_FRAMEINFO(C.Structure):
_fields_ = [
("lFrameID", W.DWORD),
("lBufSize", W.DWORD),
("lWidth", W.DWORD),
("lHeight", W.DWORD),
("lPixBits", C.c_ubyte),
("_pad0", C.c_ubyte * 3),
("pBufPtr", C.POINTER(C.c_ubyte)),
("lFrameStatus", W.DWORD),
("lPixType", W.DWORD),
("lTimeStamp", W.DWORD),
("_reserve", W.DWORD * 8),
]
class VT_DEVPARAM(C.Structure):
_fields_ = [
("bUseName", W.BOOL),
("lParamByID", W.DWORD),
("lParamByName", C.c_char * BUF_SIZE),
]
def devparam_by_id(pid: int) -> VT_DEVPARAM:
p = VT_DEVPARAM()
p.bUseName = False
p.lParamByID = pid
p.lParamByName = b""
return p
# =========================
# DLL LOAD + prototypes
# =========================
os.add_dll_directory(SDK_DIR)
dll = C.WinDLL(os.path.join(SDK_DIR, DLL_NAME))
print("DLL carregada OK:", dll)
dll.VT_DeviceScan.argtypes = [C.POINTER(C.c_ubyte), C.c_int]
dll.VT_DeviceScan.restype = C.c_int
dll.VT_DeviceOpen.argtypes = [C.c_void_p, C.POINTER(W.HANDLE), C.c_int, C.c_int]
dll.VT_DeviceOpen.restype = C.c_int
dll.VT_SingleFrameCapture.argtypes = [W.HANDLE, C.POINTER(VT_FRAMEINFO), C.c_int, C.c_int, W.BOOL]
dll.VT_SingleFrameCapture.restype = C.c_int
dll.VT_DeviceClose.argtypes = [C.POINTER(W.HANDLE)]
dll.VT_DeviceClose.restype = C.c_int
dll.VT_ParamGetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
dll.VT_ParamGetValue.restype = C.c_int
dll.VT_ParamSetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
dll.VT_ParamSetValue.restype = C.c_int
def param_set_int(h: W.HANDLE, pid: int, value: int):
p = devparam_by_id(pid)
v = C.c_int(int(value))
ret = dll.VT_ParamSetValue(h, p, C.byref(v), VALUE_INT)
ck(ret, f"VT_ParamSetValue({hex(pid)})")
def capture_raw8(h: W.HANDLE) -> np.ndarray:
fi = VT_FRAMEINFO()
ret = dll.VT_SingleFrameCapture(h, C.byref(fi), DATA_RAW, TIMEOUT_MS, True)
ck(ret, "VT_SingleFrameCapture")
w, hh = int(fi.lWidth), int(fi.lHeight)
buf = C.string_at(fi.pBufPtr, fi.lBufSize)
arr = np.frombuffer(buf, dtype=np.uint8)
needed = w * hh
if arr.size < needed:
arr = np.pad(arr, (0, needed - arr.size), mode="constant", constant_values=0)
arr = arr[:needed].reshape(hh, w)
return arr
def overlay_hud(img_bgr, lines):
y = 28
for s in lines: for s in lines:
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 3, cv2.LINE_AA) # contorno preto
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255,255,255), 2, cv2.LINE_AA) cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
y += 28 # texto branco
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
y += line_step
def save_sample(raw: np.ndarray, bgr_preview: np.ndarray, meta: dict):
def save_sample_raw4(
base_dir: str,
raw4: np.ndarray,
preview_bgr: np.ndarray,
meta: dict,
):
"""
Salva:
- RAW4 como .raw float32 (4,H,W)
- preview RGB como .png
- metadados como .json
dentro de base_dir.
"""
os.makedirs(base_dir, exist_ok=True)
name = ts_name() name = ts_name()
raw_path = os.path.join(SESSION_DIR, f"{name}.raw")
png_path = os.path.join(SESSION_DIR, f"{name}.png")
json_path = os.path.join(SESSION_DIR, f"{name}.json")
raw.tofile(raw_path) raw_path = os.path.join(base_dir, f"{name}.raw")
cv2.imwrite(png_path, bgr_preview) png_path = os.path.join(base_dir, f"{name}.png")
json_path = os.path.join(base_dir, f"{name}.json")
# RAW4
#np.save(raw_path, raw4.astype(np.float32))
raw4.astype(np.float32).tofile(raw_path)
# Preview
cv2.imwrite(png_path, preview_bgr)
# Metadados
with open(json_path, "w", encoding="utf-8") as f: with open(json_path, "w", encoding="utf-8") as f:
json.dump(meta, f, ensure_ascii=False, indent=2) json.dump(meta, f, ensure_ascii=False, indent=2)
return raw_path, png_path, json_path return raw_path, png_path, json_path
# =========================
# MAIN
# =========================
def main(): def main():
# scan parser = argparse.ArgumentParser(
n = C.c_ubyte(0) description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
ck(dll.VT_DeviceScan(C.byref(n), DEVICE_UDEF), "VT_DeviceScan") formatter_class=argparse.ArgumentDefaultsHelpFormatter,
if n.value == 0: )
raise RuntimeError("Nenhuma câmera encontrada.")
# open parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
idx = C.c_ubyte(0) parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
h = W.HANDLE() parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
ck(dll.VT_DeviceOpen(C.byref(idx), C.byref(h), DEVICE_INDEX, DEVICE_UDEF), "VT_DeviceOpen") parser.add_argument("--dll_dir", default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls", help="Pasta onde está a VT_SDK64.dll (usada pelo Gal5000Camera).")
print("DeviceOpen OK, handle=", h.value) parser.add_argument("--dll_name", default="VT_SDK64.dll", help="Nome da DLL da câmera.")
print("Saving to:", SESSION_DIR) parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
parser.add_argument("--no_ae", action="store_true", help="Desliga o AutoExposure do service (por padrão ele vem ligado).")
parser.add_argument("--upscale", type=int, default=2, help="Fator de upscale visual do preview.")
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL) args = parser.parse_args()
# Estado local (não dependemos de GET) # ===== config =====
exp_raw = 1500 with open("config.json", "r", encoding="utf-8") as f:
gain_a = 0 config = json.load(f)
gain_d = 0 MODELO = config["camera"]
RAW_W = config["raw_size"][0]
RAW_H = config["raw_size"][1]
# Aplica estado inicial # Define diretório de sessão:
try: # dataset/cana_<estado>/<horario>/<YYYYMMDD>/
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw) session_dir = os.path.join(args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
param_set_int(h, PARAM_ID_SENSOR_GAINANALOGRAW, gain_a) os.makedirs(session_dir, exist_ok=True)
param_set_int(h, PARAM_ID_SENSOR_GAINDIGITRAW, gain_d)
except Exception as e: print("============================================")
print("[WARN] Falhou set inicial:", e) print("Coleta de dataset RAW4 - SegFormer B0")
print(f"Cana : {args.cana}")
print(f"Horário : {args.horario}")
print(f"Saída : {session_dir}")
print("============================================")
window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
ae = RobustAE(target_p95=140.0, deadband=6.0, k=0.12, max_step=0.10, ema_alpha=0.20, sat_limit=0.01)
ae_on = True
auto_save = False auto_save = False
last_auto_t = 0.0 last_auto_t = 0.0
upscale = args.upscale
# FPS # Estatísticas simples de FPS
t0 = time.time() t_fps = time.time()
frames = 0 frames = 0
fps = 0.0 fps = 0.0
@ -304,119 +178,151 @@ def main():
last_msg_t = 0.0 last_msg_t = 0.0
try: try:
cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
with cam:
print("[CAM] Status inicial:", cam.get_status())
# opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
cam.configure_fps(20)
cam.start_streaming()
apply_ir_comp = True
ir_k_r = 0.8
ir_k_g = 0.4
ir_k_b = 0.9
while True: while True:
raw = capture_raw8(h) t0 = time.time()
raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
t1 = time.time()
# soft AE ae_dbg = dbg.get("ae", {}) or {}
ae_dbg = {} exp_raw = dbg.get("exp_raw", None)
if ae_on: gain_a = dbg.get("gain_a", None)
new_exp, ae_dbg = ae.step(raw, exp_raw) gain_d = dbg.get("gain_d", None)
if new_exp != exp_raw:
exp_raw = new_exp
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
except Exception as e:
# se set falhar, desliga AE pra não ficar insistindo
print("[ERR] set exposure:", e)
ae_on = False
# preview RGB bonitão bgr = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=upscale > 0, preview_scale=upscale, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
rgb_clean = make_rgb_preview(raw, upscale=UPSCALE)
bgr = rgb_clean.copy()
# FPS # FPS
frames += 1 frames += 1
dt = time.time() - t0 dt_fps = time.time() - t_fps
if dt >= 1.0: if dt_fps >= 1.0:
fps = frames / dt fps = frames / dt_fps
frames = 0 frames = 0
t0 = time.time() t_fps = time.time()
# HUD # AE info
p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
sat_disp = ae_dbg.get("sat", 0.0)
hold = ae_dbg.get("hold", False)
ae_on = cam.is_auto_exposure_enabled()
# HUD principal
lines = [ lines = [
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Interval: {CAPTURE_INTERVAL_S:.1f}s", f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}", f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
f"AEdbg: p95={ae_dbg.get('p95_ema', ae_dbg.get('p95', 0)):.1f} sat={ae_dbg.get('sat', 0):.3f} hold={ae_dbg.get('hold', False)}", f"AEdbg: p95={p95_disp:.1f} sat={sat_disp:.3f} hold={hold}",
"Keys: C/SPACE=save | A=toggle autosave | E=toggle AE | +/- exp | Q/ESC quit", "Keys: C/SPACE=save | A=auto-save | E=AE toggle | M=preview scale | Q/Esc=quit",
] ]
overlay_hud(bgr, lines) overlay_hud(bgr, lines, base_h=RAW_H)
# msg pós-save # Mensagem rápida (ex: arquivo salvo)
if last_msg and (time.time() - last_msg_t) < 2.0: if last_msg and (time.time() - last_msg_t) < 2.0:
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2, cv2.LINE_AA)
cv2.imshow(WINDOW_NAME, bgr) cv2.imshow(window_name, bgr)
# autosave # Auto-save
now = time.time() now = time.time()
if auto_save and (now - last_auto_t) >= CAPTURE_INTERVAL_S: if auto_save and (now - last_auto_t) >= args.interval:
meta = { meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"), "ts": datetime.now().isoformat(timespec="milliseconds"),
"raw_w": RAW_W, "raw_h": RAW_H, "cana": args.cana,
"exp_raw": int(exp_raw), "horario": args.horario,
"gain_a": int(gain_a), "raw4_shape": list(raw4_base.shape),
"gain_d": int(gain_d), "out_h": RAW_H,
"ae_on": bool(ae_on), "out_w": RAW_W,
"ae_enabled": bool(ae_on),
"exp_raw": int(exp_raw) if exp_raw is not None else None,
"gain_a": int(gain_a) if gain_a is not None else None,
"gain_d": int(gain_d) if gain_d is not None else None,
"apply_ir_comp": apply_ir_comp,
"ir_k_r": ir_k_r,
"ir_k_g": ir_k_g,
"ir_k_b": ir_k_b,
"ae_dbg": {
k: (float(v) if isinstance(v, (int, float, np.floating)) else v)
for k, v in ae_dbg.items()
},
"note": "autosave", "note": "autosave",
} }
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta) rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
last_msg = f"SAVED: {os.path.basename(raw_path)}" raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
last_msg = f"SALVO (auto): {os.path.basename(raw_path)}"
last_msg_t = now last_msg_t = now
last_auto_t = now last_auto_t = now
# Teclado
k = cv2.waitKey(1) & 0xFF k = cv2.waitKey(1) & 0xFF
if k in (ord('q'), ord('Q'), 27): if k in (ord("q"), ord("Q"), 27): # Q ou ESC
break break
elif k in (ord('a'), ord('A')): elif k in (ord("a"), ord("A")):
auto_save = not auto_save auto_save = not auto_save
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}" last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
last_msg_t = time.time() last_msg_t = time.time()
elif k in (ord('e'), ord('E')): elif k in (ord("e"), ord("E")):
ae_on = not ae_on cam.enable_auto_exposure(not ae_on)
last_msg = f"AE -> {'ON' if ae_on else 'OFF'}" last_msg = f"AE -> {'ON' if cam.is_auto_exposure_enabled() else 'OFF'}"
last_msg_t = time.time() last_msg_t = time.time()
elif k in (ord('c'), ord('C'), 32): # C ou SPACE elif k in (ord("m"), ord("M")):
upscale = 0 if upscale else args.upscale
last_msg = f"Preview UPSCALE -> {upscale}"
last_msg_t = time.time()
elif k in (ord("c"), ord("C"), 32): # C ou SPACE
meta = { meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"), "ts": datetime.now().isoformat(timespec="milliseconds"),
"raw_w": RAW_W, "raw_h": RAW_H, "cana": args.cana,
"exp_raw": int(exp_raw), "horario": args.horario,
"gain_a": int(gain_a), "raw4_shape": list(raw4_base.shape),
"gain_d": int(gain_d), "out_h": RAW_H,
"ae_on": bool(ae_on), "out_w": RAW_W,
"ae_enabled": bool(ae_on),
"exp_raw": int(exp_raw) if exp_raw is not None else None,
"gain_a": int(gain_a) if gain_a is not None else None,
"gain_d": int(gain_d) if gain_d is not None else None,
"apply_ir_comp": apply_ir_comp,
"ir_k_r": ir_k_r,
"ir_k_g": ir_k_g,
"ir_k_b": ir_k_b,
"ae_dbg": {
k2: (float(v2) if isinstance(v2, (int, float, np.floating)) else v2)
for k2, v2 in ae_dbg.items()
},
"note": "manual", "note": "manual",
} }
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta) rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
last_msg = f"SAVED: {os.path.basename(raw_path)}" raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
last_msg_t = time.time() last_msg_t = time.time()
elif k in (ord('+'), ord('=')): # você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)
exp_raw = clamp(exp_raw + 200, EXP_MIN, EXP_MAX)
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
except Exception as e:
print("[ERR] manual exp +:", e)
elif k in (ord('-'), ord('_')): # Só pra não ficar rodando a 1000 FPS na UI
exp_raw = clamp(exp_raw - 200, EXP_MIN, EXP_MAX) # mas sem travar muito a captura
try: dt_loop = time.time() - t0
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw) if dt_loop < 0.001:
except Exception as e: time.sleep(0.001)
print("[ERR] manual exp -:", e)
finally: finally:
try:
ret = dll.VT_DeviceClose(C.byref(h))
if ret != 0:
print("VT_DeviceClose retornou:", ret)
except Exception as e:
print("Erro ao fechar:", e)
cv2.destroyAllWindows() cv2.destroyAllWindows()
print("Fim da captura.")
print("Fim.")
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View File

@ -1,11 +1,13 @@
import json import json
import os import os
import sys
import cv2 import cv2
import csv import csv
import shutil import shutil
import argparse import argparse
import numpy as np import numpy as np
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo # importa da sua utils.py from utils import carregar_labelmap_completo # importa da sua utils.py
# =================================================== # ===================================================

View File

@ -21,6 +21,7 @@ Saída:
""" """
import os import os
import sys
import cv2 import cv2
import csv import csv
import json import json
@ -28,6 +29,7 @@ import shutil
import argparse import argparse
import numpy as np import numpy as np
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:

View File

@ -21,11 +21,14 @@ RAW:
import argparse import argparse
import os import os
import json import json
import sys
import cv2 import cv2
import numpy as np import numpy as np
from typing import Dict, List, Tuple from typing import Dict, List, Tuple
from gal5000.gal_service import mosaic_to_raw4_resized_buf from gal5000.gal_service import mosaic_to_raw4_resized_buf
from raw_segformer_service import _infer_ignore_id from raw_segformer_service import _infer_ignore_id
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ===== config ===== # ===== config =====

View File

@ -21,6 +21,7 @@ Modos:
import os import os
import json import json
import sys
import time import time
import argparse import argparse
from collections import deque from collections import deque
@ -29,11 +30,12 @@ import cv2
import torch import torch
import numpy as np import numpy as np
from utils import converter_mask_ids_para_bgr, desenhar_legenda_horizontal
from raw_segformer_service import RawSegformerService, RawSegDataset # :contentReference[oaicite:0]{index=0} from raw_segformer_service import RawSegformerService, RawSegDataset # :contentReference[oaicite:0]{index=0}
from gal5000.gal_service import Gal5000Camera from gal5000.gal_service import Gal5000Camera
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import converter_mask_ids_para_bgr, desenhar_legenda_horizontal
# ============================================================ # ============================================================
# Helpers para modo "sem máscara" # Helpers para modo "sem máscara"

View File

@ -11,6 +11,7 @@
import os import os
import json import json
from contextlib import nullcontext from contextlib import nullcontext
import sys
from typing import List, Optional, Tuple from typing import List, Optional, Tuple
import threading import threading
import time import time
@ -23,6 +24,8 @@ from torch.utils.data import Dataset
import torch.nn.functional as F import torch.nn.functional as F
from transformers import SegformerForSemanticSegmentation from transformers import SegformerForSemanticSegmentation
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import (carregar_labelmap_completo, converter_mask_ids_para_bgr) from utils import (carregar_labelmap_completo, converter_mask_ids_para_bgr)
def build_raw_segformer_model( def build_raw_segformer_model(

View File

@ -0,0 +1,822 @@
import os
import time
import json
import argparse
from datetime import datetime
import cv2
import numpy as np
from multispectral_service import MultiSpectralService
from stream_receiver import StreamReceiver
from pi.raw_processor_core import RawProcessorCore
from pi.raw_processor_preview import RawProcessorPreview
STREAM_PORT = 6001
PI_HOST = "192.168.105.6"
PC_HOST = "192.168.105.5"
# =========================
# Helpers gerais
# =========================
def ts_name() -> str:
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
def overlay_hud(
img_bgr: np.ndarray,
lines: list[str],
base_h: int = 720,
base_font_scale: float = 0.75,
base_line_step: int = 28,
):
h, w = img_bgr.shape[:2]
scale = h / float(base_h)
scale = max(scale, 0.4)
font_scale = base_font_scale * scale
line_step = int(base_line_step * scale)
thick_outline = max(1, int(3 * scale))
thick_text = max(1, int(2 * scale))
y = int(24 * scale)
x = int(12 * scale)
for s in lines:
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
y += line_step
def save_sample(
base_dir: str,
frame_type: str,
preview_bgr: np.ndarray,
meta: dict,
raw_payload: np.ndarray | None = None,
packed_raw: np.ndarray | None = None,
packed_raw_by_camera: dict | None = None,
):
os.makedirs(base_dir, exist_ok=True)
name = ts_name()
png_path = os.path.join(base_dir, f"{name}.png")
json_path = os.path.join(base_dir, f"{name}.json")
if frame_type in ("RGB", "MULTISPEC"):
if raw_payload is None:
raise ValueError(f"raw_payload não pode ser None quando frame_type='{frame_type}'")
payload_path = os.path.join(base_dir, f"{name}.raw")
raw_payload.astype(np.float32).tofile(payload_path)
meta["saved_payload_type"] = frame_type.lower()
meta["saved_payload_path"] = os.path.basename(payload_path)
meta["saved_payload_dtype"] = "float32"
meta["saved_payload_shape"] = list(raw_payload.shape)
elif frame_type == "RAW_BRUTO":
if packed_raw_by_camera is not None:
payload_files = {}
payload_shapes = {}
payload_dtypes = {}
for cam_id, arr in packed_raw_by_camera.items():
path = os.path.join(base_dir, f"{name}_{cam_id}.bin")
arr.tofile(path)
payload_files[cam_id] = os.path.basename(path)
payload_shapes[cam_id] = list(arr.shape)
payload_dtypes[cam_id] = str(arr.dtype)
meta["saved_payload_type"] = "raw_native_multi"
meta["saved_payload_paths"] = payload_files
meta["saved_payload_shapes"] = payload_shapes
meta["saved_payload_dtypes"] = payload_dtypes
else:
if packed_raw is None:
raise ValueError("packed_raw não pode ser None quando frame_type='RAW_BRUTO'")
payload_path = os.path.join(base_dir, f"{name}.bin")
packed_raw.tofile(payload_path)
meta["saved_payload_type"] = "raw_native_single"
meta["saved_payload_path"] = os.path.basename(payload_path)
meta["saved_payload_dtype"] = str(packed_raw.dtype)
meta["saved_payload_shape"] = list(packed_raw.shape)
else:
raise ValueError(f"frame_type não suportado para save: {frame_type}")
cv2.imwrite(png_path, preview_bgr)
with open(json_path, "w", encoding="utf-8") as f:
json.dump(meta, f, ensure_ascii=False, indent=2)
return png_path, json_path
def get_camera_map_from_status(status: dict) -> dict:
result = {}
for cam in status.get("cameras", []):
result[cam.get("id")] = cam
return result
def validate_module_ready(status: dict, frame_type: str, raw_policy: str, capture_mode: str):
if not status.get("ok", True):
raise RuntimeError(f"Status inválido retornado pelo módulo: {status}")
active_ids = list(status.get("active_camera_ids", []))
active_count = int(status.get("camera_count_active", 0))
if frame_type == "RGB":
if "cam2" not in active_ids:
raise RuntimeError(
"Modo RGB requer cam2 ativa (USB RGB), mas o módulo não reportou cam2 como ativa."
)
return
if frame_type == "MULTISPEC":
if capture_mode == "TRIPLE":
missing = [cid for cid in ("cam0", "cam1", "cam2") if cid not in active_ids]
if missing:
raise RuntimeError(
f"Modo MULTISPEC/TRIPLE requer cam0, cam1 e cam2 ativas. "
f"Faltando: {missing}. Ativas atuais: {active_ids}"
)
return
if capture_mode == "DOUBLE":
has_rgb = "cam2" in active_ids
has_spec = ("cam0" in active_ids) or ("cam1" in active_ids)
if not has_rgb or not has_spec:
raise RuntimeError(
f"Modo MULTISPEC/DOUBLE requer cam2 + (cam0 ou cam1). "
f"Ativas atuais: {active_ids}"
)
return
# AUTO ou outros casos
has_rgb = "cam2" in active_ids
has_spec = ("cam0" in active_ids) or ("cam1" in active_ids)
if not (has_rgb and has_spec):
raise RuntimeError(
f"Modo MULTISPEC requer pelo menos RGB + 1 canal espectral. "
f"Ativas atuais: {active_ids}"
)
return
if frame_type == "RAW_BRUTO":
if raw_policy == "require_triple":
missing = [cid for cid in ("cam0", "cam1", "cam2") if cid not in active_ids]
if missing:
raise RuntimeError(
f"RAW_BRUTO com política require_triple exige três câmeras ativas. "
f"Faltando: {missing}. Ativas atuais: {active_ids}"
)
else:
if active_count < 1:
raise RuntimeError("RAW_BRUTO requer ao menos uma câmera ativa, mas nenhuma foi detectada.")
return
raise RuntimeError(f"frame_type desconhecido para validação: {frame_type}")
def build_preview_from_raw_payload(
frame,
meta: dict,
processor_core: RawProcessorCore,
processor_preview: RawProcessorPreview,
):
"""
Gera preview priorizando a câmera RGB (cam2).
Se cam2 não estiver presente, cai para fallback usando a primeira câmera mono disponível.
Retorna:
preview_bgr
payload_float_preview
preview_source_id
"""
payload_sources = meta.get("payload_sources", []) or []
# Caso multi-payload: tenta usar cam2 primeiro
if isinstance(frame, dict):
if "cam2" in frame:
rgb_frame = frame["cam2"]
if rgb_frame.ndim != 3 or rgb_frame.shape[2] != 3:
raise RuntimeError(f"cam2 recebida mas inválida para preview RGB: shape={rgb_frame.shape}")
preview_bgr = rgb_frame.copy()
payload_float = rgb_frame[:, :, ::-1].astype(np.float32) / 255.0
payload_float = np.transpose(payload_float, (2, 0, 1))
return preview_bgr, payload_float, "cam2"
# fallback: usa a primeira câmera mono disponível
fallback_id = None
for cid in ("cam0", "cam1"):
if cid in frame:
fallback_id = cid
break
if fallback_id is None:
raise RuntimeError("Nenhuma câmera disponível no payload para gerar preview")
packed = frame[fallback_id]
if packed.ndim == 3 and packed.shape[2] == 1:
packed = packed[:, :, 0]
cam_frames = meta.get("camera_frames", {}) or {}
cam_meta = cam_frames.get(fallback_id, {})
bit_depth = int(cam_meta.get("bit_depth", 10))
raw16 = processor_core.unpack_raw10_packed(packed)
preview_bgr = processor_preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)
payload_float = processor_core.build_training_rgb(
raw16,
output_dtype="float32",
bit_depth=bit_depth,
)
return preview_bgr, payload_float, fallback_id
# Caso single-payload
if isinstance(frame, np.ndarray):
# Se vier HWC/3ch, tratamos como RGB USB
if frame.ndim == 3 and frame.shape[2] == 3:
preview_bgr = frame.copy()
payload_float = frame[:, :, ::-1].astype(np.float32) / 255.0
payload_float = np.transpose(payload_float, (2, 0, 1))
return preview_bgr, payload_float, "cam2"
# Se vier mono packed, fallback antigo
packed = frame
if packed.ndim == 3 and packed.shape[2] == 1:
packed = packed[:, :, 0]
source_camera = meta.get("source_camera") or {}
bit_depth = int(source_camera.get("bit_depth", meta.get("source_bit_depth", 10)))
raw16 = processor_core.unpack_raw10_packed(packed)
preview_bgr = processor_preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)
payload_float = processor_core.build_training_rgb(
raw16,
output_dtype="float32",
bit_depth=bit_depth,
)
return preview_bgr, payload_float, source_camera.get("id", "unknown")
raise RuntimeError(f"Tipo de frame não suportado para preview: {type(frame)}")
def resolve_effective_capture_mode(frame_type: str, raw_policy: str, requested_mode: str) -> str:
"""
Decide o modo real que será pedido ao módulo.
Nova regra:
- Se o usuário pediu explicitamente SINGLE/DOUBLE/TRIPLE, respeitamos.
- Se pediu AUTO, deixamos AUTO ir para o módulo.
- A única exceção opcional é MULTISPEC + require_triple implícito,
mas mesmo assim podemos deixar o módulo resolver se preferirmos.
"""
if requested_mode in ("SINGLE", "DOUBLE", "TRIPLE"):
return requested_mode
# requested_mode == AUTO
return "AUTO"
# =========================
# MAIN
# =========================
def main():
parser = argparse.ArgumentParser(
description="Captura de dataset usando módulo multispectral Pi + StreamReceiver.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
parser.add_argument("--pi_host", default=PI_HOST, help="IP do servidor no Raspberry Pi.")
parser.add_argument("--pc_host", default=PC_HOST, help="IP local do notebook/PC que receberá o stream.")
parser.add_argument("--stream_port", type=int, default=STREAM_PORT, help="Porta TCP do receiver de stream.")
parser.add_argument("--server_port", type=int, default=5000, help="Porta TCP do servidor de comandos no Pi.")
parser.add_argument("--fps", type=int, default=20, help="FPS desejado.")
parser.add_argument("--width", type=int, default=640, help="Largura óptica da câmera.")
parser.add_argument("--height", type=int, default=480, help="Altura óptica da câmera.")
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
parser.add_argument("--preview_upscale", type=int, default=2, help="Fator de upscale visual do preview.")
parser.add_argument("--bayer", default="GBRG", choices=["GBRG", "GRBG", "RGGB", "BGGR"], help="Padrão Bayer das câmeras.")
parser.add_argument("--output_dtype", default="float32", choices=["uint8", "uint16", "float32"], help="Dtype do payload processado no Pi.")
parser.add_argument("--frame_type", default="RAW_BRUTO", choices=["RAW_BRUTO", "RGB", "MULTISPEC"], help="Tipo de payload pedido ao Pi.")
parser.add_argument("--capture_mode", default="AUTO", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"], help="Modo de captura desejado no módulo.")
parser.add_argument("--raw_policy", default="allow_single", choices=["allow_single", "require_triple"], help="Quando frame_type=RAW_BRUTO, define se o script aceita 1 câmera ou exige 3.")
args = parser.parse_args()
effective_capture_mode = resolve_effective_capture_mode(
frame_type=args.frame_type,
raw_policy=args.raw_policy,
requested_mode=args.capture_mode,
)
raw_w = args.width
raw_h = args.height
session_dir = os.path.join(
args.out_root,
"brutas",
f"cana_{args.cana}",
args.horario,
datetime.now().strftime("%Y%m%d"),
)
os.makedirs(session_dir, exist_ok=True)
print("============================================")
print("Coleta de dataset - Módulo Multiespectral")
print(f"Cana : {args.cana}")
print(f"Horário : {args.horario}")
print(f"Saída : {session_dir}")
print(f"Sensor : {raw_w}x{raw_h} | Bayer={args.bayer}")
print(f"FrameType : {args.frame_type}")
print(f"CaptureMode : {args.capture_mode} -> efetivo={effective_capture_mode}")
print(f"RAW policy : {args.raw_policy}")
print("============================================")
auto_save = False
last_auto_t = 0.0
preview_upscale = args.preview_upscale
t_view_fps = time.time()
view_frames = 0
fps_view = 0.0
t_stream_fps = time.time()
last_stream_frame_id = None
stream_frames_accum = 0
fps_stream = 0.0
last_msg = ""
last_msg_t = 0.0
receiver = StreamReceiver(host="0.0.0.0", port=args.stream_port)
svc = MultiSpectralService(host=args.pi_host, port=args.server_port, timeout=10)
print(f"[INFO] Verificando conexão com o módulo em {args.pi_host}:{args.server_port}...")
svc.ensure_alive()
if not svc.is_alive():
raise RuntimeError("Módulo não encontrado ou não respondeu ao ping.")
print("[OK] Módulo conectado e respondendo.")
window_name = "Dataset Capture (C/SPACE=save | A=auto-save | M=preview scale | Q=quit)"
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
processor_core_cam0 = RawProcessorCore(
sensor_width=raw_w,
sensor_height=raw_h,
bayer_pattern=args.bayer,
)
processor_preview_cam0 = RawProcessorPreview(
sensor_width=raw_w,
sensor_height=raw_h,
bayer_pattern=args.bayer,
)
last_frame_id = -1
last_payload_float = None
last_packed_raw = None
last_packed_raw_by_camera = None
last_preview_bgr = None
last_meta_stream = None
try:
receiver.start()
time.sleep(0.5)
svc.connect()
if args.frame_type in ("RAW_BRUTO", "MULTISPEC"):
modes_resp = svc.get_sensor_modes()
if not modes_resp.get("ok"):
print(f"[WARN] Falha ao obter sensor_modes: {modes_resp}")
else:
for mode in modes_resp.get("sensor_modes", []):
print(
f"[cam={mode.get('camera_id')} mode={mode.get('mode_index')}] "
f"size={mode.get('size')} "
f"format={mode.get('format')} "
f"bit_depth={mode.get('bit_depth')} "
f"fps={mode.get('fps')}"
)
else:
print("[INFO] get_sensor_modes pulado para frame_type=RGB")
print("SET CAM0 RES:", svc.set_camera_resolution(0, raw_w, raw_h))
print("SET CAM1 RES:", svc.set_camera_resolution(1, raw_w, raw_h))
print("SET CAM2 RES:", svc.set_camera_resolution(2, raw_w, raw_h))
print("SET CAM0 BAYER:", svc.set_camera_bayer(0, args.bayer))
print("SET CAM1 BAYER:", svc.set_camera_bayer(1, args.bayer))
print("SET FPS:", svc.set_fps(args.fps))
print("SET CAPTURE MODE:", svc.set_capture_mode(effective_capture_mode))
print("SET FRAME TYPE:", svc.set_frame_type(args.frame_type))
print("SET OUTPUT DTYPE:", svc.set_output_dtype(args.output_dtype))
begin_resp = svc.begin(
frame_type=args.frame_type,
output_dtype=args.output_dtype,
capture_mode=effective_capture_mode,
)
print("BEGIN:", begin_resp)
status = svc.get_status()
print("STATUS:", json.dumps({
"status": status.get("status"),
"detected_mode": status.get("detected_mode"),
"camera_count_active": status.get("camera_count_active"),
"active_camera_ids": status.get("active_camera_ids"),
}, ensure_ascii=False))
validate_module_ready(status, args.frame_type, args.raw_policy, effective_capture_mode)
print("START STREAM:", svc.start_stream(args.pc_host, args.stream_port, fps=args.fps))
camera_ctrl = svc.get_camera_controls()
ae_enabled = bool(camera_ctrl.get("ae_enable", True))
awb_enabled = bool(camera_ctrl.get("awb_enable", True))
manual_exposure_us = camera_ctrl.get("exposure_time_us", None)
manual_gain = camera_ctrl.get("analogue_gain", None)
manual_colour_gains = camera_ctrl.get("colour_gains", None)
while True:
t0 = time.time()
meta = receiver.last_meta
frame = receiver.last_frame
if meta is not None and frame is not None and meta.get("frame_id") != last_frame_id:
last_frame_id = meta["frame_id"]
try:
frame_type = meta.get("frame_type", "RAW_BRUTO")
dtype_str = meta.get("dtype") or meta.get("output_dtype", "uint8")
preview_source_id = "cam2"
if frame_type == "RAW_BRUTO":
if isinstance(frame, dict):
packed_by_camera = frame
preview_bgr, raw3_preview, preview_source_id = build_preview_from_raw_payload(
frame=frame,
meta=meta,
processor_core=processor_core_cam0,
processor_preview=processor_preview_cam0,
)
last_packed_raw = None
last_packed_raw_by_camera = {cam_id: arr.copy() for cam_id, arr in packed_by_camera.items()}
last_payload_float = raw3_preview.copy()
else:
preview_bgr, raw3_preview, preview_source_id = build_preview_from_raw_payload(
frame=frame,
meta=meta,
processor_core=processor_core_cam0,
processor_preview=processor_preview_cam0,
)
last_packed_raw = frame.copy()
last_packed_raw_by_camera = None
last_payload_float = raw3_preview.copy()
elif frame_type == "RGB":
rgb_chw = frame
if not isinstance(rgb_chw, np.ndarray) or rgb_chw.ndim != 3:
raise RuntimeError(f"Frame RGB inválido: type={type(rgb_chw)}")
if dtype_str == "uint8":
payload_float = rgb_chw.astype(np.float32) / 255.0
elif dtype_str == "float32":
payload_float = rgb_chw.astype(np.float32)
elif dtype_str == "uint16":
payload_float = rgb_chw.astype(np.float32) / 65535.0
else:
raise RuntimeError(f"dtype RGB não suportado: {dtype_str}")
preview_rgb = np.transpose(payload_float, (1, 2, 0))
preview_bgr = cv2.cvtColor(
np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8),
cv2.COLOR_RGB2BGR
)
last_payload_float = payload_float.copy()
last_packed_raw = None
last_packed_raw_by_camera = None
elif frame_type == "MULTISPEC":
multispec_chw = frame
if not isinstance(multispec_chw, np.ndarray) or multispec_chw.ndim != 3 or multispec_chw.shape[0] not in (4, 5):
raise RuntimeError(f"Frame MULTISPEC inválido: shape={getattr(multispec_chw, 'shape', None)}")
if dtype_str == "uint8":
payload_float = multispec_chw.astype(np.float32) / 255.0
elif dtype_str == "float32":
payload_float = multispec_chw.astype(np.float32)
elif dtype_str == "uint16":
payload_float = multispec_chw.astype(np.float32) / 65535.0
else:
raise RuntimeError(f"dtype MULTISPEC não suportado: {dtype_str}")
preview_rgb = np.transpose(payload_float[:3], (1, 2, 0))
preview_bgr = cv2.cvtColor(
np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8),
cv2.COLOR_RGB2BGR
)
last_payload_float = payload_float.copy()
last_packed_raw = None
last_packed_raw_by_camera = None
else:
raise RuntimeError(f"frame_type não suportado neste script: {frame_type}")
if preview_upscale and preview_upscale > 1:
preview_show = cv2.resize(
preview_bgr,
(preview_bgr.shape[1] * preview_upscale, preview_bgr.shape[0] * preview_upscale),
interpolation=cv2.INTER_NEAREST,
)
else:
preview_show = preview_bgr.copy()
curr_frame_id = meta.get("frame_id")
if curr_frame_id is not None:
if last_stream_frame_id != curr_frame_id:
stream_frames_accum += 1
last_stream_frame_id = curr_frame_id
dt_stream = time.time() - t_stream_fps
if dt_stream >= 1.0:
fps_stream = stream_frames_accum / dt_stream
stream_frames_accum = 0
t_stream_fps = time.time()
view_frames += 1
dt_view = time.time() - t_view_fps
if dt_view >= 1.0:
fps_view = view_frames / dt_view
view_frames = 0
t_view_fps = time.time()
active_sources = meta.get("payload_sources")
lines = [
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
f"Type={meta.get('frame_type')} | CaptureMode={effective_capture_mode} | RAW policy={args.raw_policy}",
f"Sources={active_sources} | FPS_STREAM={fps_stream:.1f} | FPS_VIEW={fps_view:.1f}",
f"frame_id={meta.get('frame_id')} | layout={meta.get('output_layout')} | dtype={meta.get('dtype') or meta.get('output_dtype')}",
f"codec={meta.get('codec_name', meta.get('codec_family', '-'))} | comp={meta.get('dt_comp', 0):.4f}s | send={meta.get('dt_send_payload_prev', 0):.4f}s",
f"AE={'ON' if ae_enabled else 'OFF'} | AWB={'ON' if awb_enabled else 'OFF'} | EXP={manual_exposure_us} | GAIN={manual_gain}",
"Keys: C/SPACE=save | A=auto-save | E=AE | W=AWB | I/K=exp | O/L=gain | R=reset | M=preview | Q/Esc=quit",
]
overlay_hud(preview_show, lines, base_h=raw_h)
if last_msg and (time.time() - last_msg_t) < 2.0:
cv2.putText(preview_show, last_msg, (12, preview_show.shape[0] - 18),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
cv2.imshow(window_name, preview_show)
last_preview_bgr = preview_bgr.copy()
last_meta_stream = dict(meta)
except Exception as e:
err = np.zeros((500, 1200, 3), dtype=np.uint8)
cv2.putText(err, f"Erro ao processar frame: {e}", (20, 60),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2, cv2.LINE_AA)
cv2.imshow(window_name, err)
print(f"[ERRO FRAME] {e}")
now = time.time()
can_save = (
last_meta_stream is not None and
last_preview_bgr is not None and
(
(last_meta_stream.get("frame_type") in ("RGB", "MULTISPEC") and last_payload_float is not None) or
(last_meta_stream.get("frame_type") == "RAW_BRUTO" and (last_packed_raw is not None or last_packed_raw_by_camera is not None))
)
)
if auto_save and can_save and (now - last_auto_t) >= args.interval:
frame_type_save = last_meta_stream.get("frame_type")
meta_save = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"cana": args.cana,
"horario": args.horario,
"sensor_width": raw_w,
"sensor_height": raw_h,
"bayer_pattern": args.bayer,
"fps_target": args.fps,
"frame_type": frame_type_save,
"capture_mode_requested": args.capture_mode,
"capture_mode_effective": effective_capture_mode,
"raw_policy": args.raw_policy,
"stream_meta": last_meta_stream,
"note": "autosave",
"raw_preview_reference_camera": preview_source_id,
}
save_sample(
session_dir,
frame_type=frame_type_save,
preview_bgr=last_preview_bgr,
meta=meta_save,
raw_payload=last_payload_float,
packed_raw=last_packed_raw,
packed_raw_by_camera=last_packed_raw_by_camera,
)
last_msg = "SALVO (auto)"
last_msg_t = now
last_auto_t = now
k = cv2.waitKey(1) & 0xFF
if k in (ord("q"), ord("Q"), 27):
break
elif k in (ord("a"), ord("A")):
auto_save = not auto_save
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
last_msg_t = time.time()
elif k in (ord("m"), ord("M")):
preview_upscale = 0 if preview_upscale else args.preview_upscale
last_msg = f"Preview UPSCALE -> {preview_upscale}"
last_msg_t = time.time()
elif k in (ord("e"), ord("E")):
ae_enabled = not ae_enabled
resp = svc.set_ae_enable(ae_enabled)
ae_enabled = bool(resp.get("ae_enable", ae_enabled))
last_msg = f"AE -> {'ON' if ae_enabled else 'OFF'}"
last_msg_t = time.time()
elif k in (ord("w"), ord("W")):
awb_enabled = not awb_enabled
resp = svc.set_awb_enable(awb_enabled)
awb_enabled = bool(resp.get("awb_enable", awb_enabled))
last_msg = f"AWB -> {'ON' if awb_enabled else 'OFF'}"
last_msg_t = time.time()
elif k in (ord("i"), ord("I")):
if manual_exposure_us is None:
manual_exposure_us = 15000
else:
manual_exposure_us = min(int(manual_exposure_us * 1.15), 200000)
if ae_enabled:
ae_enabled = False
svc.set_ae_enable(False)
resp = svc.set_exposure_time(int(manual_exposure_us))
manual_exposure_us = resp.get("exposure_time_us", manual_exposure_us)
last_msg = f"ExposureTime -> {manual_exposure_us} us"
last_msg_t = time.time()
elif k in (ord("k"), ord("K")):
if manual_exposure_us is None:
manual_exposure_us = 15000
else:
manual_exposure_us = max(int(manual_exposure_us / 1.15), 100)
if ae_enabled:
ae_enabled = False
svc.set_ae_enable(False)
resp = svc.set_exposure_time(int(manual_exposure_us))
manual_exposure_us = resp.get("exposure_time_us", manual_exposure_us)
last_msg = f"ExposureTime -> {manual_exposure_us} us"
last_msg_t = time.time()
elif k in (ord("o"), ord("O")):
if manual_gain is None:
manual_gain = 1.0
else:
manual_gain = min(float(manual_gain) * 1.10, 32.0)
if ae_enabled:
ae_enabled = False
svc.set_ae_enable(False)
resp = svc.set_analogue_gain(float(manual_gain))
manual_gain = resp.get("analogue_gain", manual_gain)
last_msg = f"AnalogueGain -> {manual_gain:.2f}"
last_msg_t = time.time()
elif k in (ord("l"), ord("L")):
if manual_gain is None:
manual_gain = 1.0
else:
manual_gain = max(float(manual_gain) / 1.10, 1.0)
if ae_enabled:
ae_enabled = False
svc.set_ae_enable(False)
resp = svc.set_analogue_gain(float(manual_gain))
manual_gain = resp.get("analogue_gain", manual_gain)
last_msg = f"AnalogueGain -> {manual_gain:.2f}"
last_msg_t = time.time()
elif k in (ord("r"), ord("R")):
svc.clear_exposure_time()
svc.clear_analogue_gain()
svc.clear_colour_gains()
manual_exposure_us = None
manual_gain = None
manual_colour_gains = None
last_msg = "Manual controls resetados"
last_msg_t = time.time()
elif k in (ord("c"), ord("C"), 32):
if can_save:
frame_type_save = last_meta_stream.get("frame_type")
meta_save = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"cana": args.cana,
"horario": args.horario,
"sensor_width": raw_w,
"sensor_height": raw_h,
"bayer_pattern": args.bayer,
"fps_target": args.fps,
"frame_type": frame_type_save,
"capture_mode_requested": args.capture_mode,
"capture_mode_effective": effective_capture_mode,
"raw_policy": args.raw_policy,
"stream_meta": last_meta_stream,
"camera_controls": {
"ae_enable": ae_enabled,
"awb_enable": awb_enabled,
"exposure_time_us": manual_exposure_us,
"analogue_gain": manual_gain,
"colour_gains": manual_colour_gains,
},
"note": "manual",
"raw_preview_reference_camera": preview_source_id,
}
save_sample(
session_dir,
frame_type=frame_type_save,
preview_bgr=last_preview_bgr,
meta=meta_save,
raw_payload=last_payload_float,
packed_raw=last_packed_raw,
packed_raw_by_camera=last_packed_raw_by_camera,
)
last_msg = "SALVO (manual)"
last_msg_t = time.time()
dt_loop = time.time() - t0
if dt_loop < 0.001:
time.sleep(0.001)
finally:
try:
print("STOP STREAM:", svc.stop_stream())
except Exception:
pass
try:
print("STOP:", svc.stop())
except Exception:
pass
svc.disconnect()
receiver.stop()
cv2.destroyAllWindows()
print("Fim da captura.")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,226 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Conta a porcentagem de pixels por classe dentro da ROI, por grupo e por split,
na estrutura nova baseada em tensors/masks .npy.
Estrutura esperada:
dataset/split/<split>/group/<grupo>/masks/
original_xxx.npy
original_xxx.png # opcional, debug
augmented_xxx.npy
augmented_xxx.png # opcional, debug
Regras:
- Prioriza máscaras .npy
- Se não encontrar .npy, pode usar .png como fallback
- Ignora pixels com ignore_id
- Normaliza a porcentagem somente sobre classes válidas
"""
import argparse
import os
import json
import sys
import numpy as np
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo, _infer_ignore_id
# ================= CONFIG =================
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
W, H = config["resolucao"][0], config["resolucao"][1]
ROI_INICIO = float(config["roi_inicio"])
ROI_TAMANHO = float(config["roi_tamanho"])
pasta_base = "dataset"
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
MAX_SAMPLES_PER_GROUP = 1000
# ================= HELPERS =================
def roi_slice(h: int):
"""
Replica a lógica antiga de ROI vertical.
"""
y_fim = int((1.0 - ROI_TAMANHO) * h)
y_ini = int(ROI_INICIO * h)
if y_ini <= y_fim:
y_fim = max(0, h - int(ROI_TAMANHO * h))
y_ini = h
return slice(y_fim, y_ini)
def load_mask(path: str) -> np.ndarray:
"""
Carrega máscara .npy ou .png.
A máscara oficial agora é .npy (IDs inteiros).
"""
ext = os.path.splitext(path)[1].lower()
if ext == ".npy":
arr = np.load(path)
if arr.ndim != 2:
raise RuntimeError(f"Máscara .npy inválida (esperado 2D): {path} shape={arr.shape}")
return arr.astype(np.int32)
if ext == ".png":
import cv2
m = cv2.imread(path, cv2.IMREAD_UNCHANGED)
if m is None:
raise RuntimeError(f"Falha ao ler PNG: {path}")
if m.ndim == 3:
# PNG de debug pode ter 3 canais; nesse caso ele não serve como máscara de IDs
raise RuntimeError(
f"PNG colorido encontrado em {path}. "
f"Para estatística, use a máscara .npy correspondente."
)
return m.astype(np.int32)
raise RuntimeError(f"Extensão de máscara não suportada: {path}")
def list_mask_files(mask_dir: str, prefer_npy: bool = True):
"""
Lista máscaras evitando contar duas vezes a mesma amostra
quando existem .npy e .png com o mesmo base.
"""
if not os.path.isdir(mask_dir):
return []
files = os.listdir(mask_dir)
by_base = {}
for fname in files:
low = fname.lower()
if not (low.endswith(".npy") or low.endswith(".png")):
continue
base, ext = os.path.splitext(fname)
ext = ext.lower()
full = os.path.join(mask_dir, fname)
if base not in by_base:
by_base[base] = full
else:
cur_ext = os.path.splitext(by_base[base])[1].lower()
if prefer_npy:
if cur_ext != ".npy" and ext == ".npy":
by_base[base] = full
else:
if cur_ext != ".png" and ext == ".png":
by_base[base] = full
return [by_base[k] for k in sorted(by_base.keys())]
def summarize_group(mask_dir, class_ids_sorted, classes, ignore_id, max_samples):
totals = {cid: 0 for cid in class_ids_sorted}
ignored_pixels = 0
valid_pixels = 0
n = 0
for path in list_mask_files(mask_dir, prefer_npy=True):
mask = load_mask(path)
rs = roi_slice(mask.shape[0])
roi = mask[rs, :]
ignored_pixels += int((roi == ignore_id).sum())
for cid in class_ids_sorted:
c = int((roi == cid).sum())
totals[cid] += c
valid_pixels += c
n += 1
if n >= max_samples:
break
return totals, valid_pixels, ignored_pixels, n
# ================= MAIN =================
def main(args):
_, _colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
class_ids_sorted = sorted(classes.keys())
root = os.path.join(pasta_base, "split", args.split, "group")
if not os.path.isdir(root):
raise SystemExit(f"Nenhum diretório encontrado em: {root}")
grupos = [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))]
global_totals = {cid: 0 for cid in class_ids_sorted}
global_valid_pixels = 0
global_ignored_pixels = 0
global_samples = 0
print(f"Split: {args.split}")
print(f"ROI: inicio={ROI_INICIO:.3f} tamanho={ROI_TAMANHO:.3f}")
print(f"Ignore ID: {ignore_id}")
print()
for g in sorted(grupos):
mdir = os.path.join(root, g, "masks")
if not os.path.isdir(mdir):
continue
totals, valid_pixels, ignored_pixels, n = summarize_group(
mask_dir=mdir,
class_ids_sorted=class_ids_sorted,
classes=classes,
ignore_id=ignore_id,
max_samples=args.max_samples,
)
denom = valid_pixels if valid_pixels > 0 else 1
parts = []
for cid in class_ids_sorted:
name = classes[cid]
perc = totals[cid] / denom
parts.append(f"{name}={perc:6.2%}")
ignore_ratio = ignored_pixels / max(1, (valid_pixels + ignored_pixels))
print(
f"{g:18s} "
+ " ".join(parts)
+ f" | ignore={ignore_ratio:6.2%} | amostras={n}"
)
for cid in class_ids_sorted:
global_totals[cid] += totals[cid]
global_valid_pixels += valid_pixels
global_ignored_pixels += ignored_pixels
global_samples += n
print("\nResumo global:")
denom = global_valid_pixels if global_valid_pixels > 0 else 1
parts = []
for cid in class_ids_sorted:
name = classes[cid]
perc = global_totals[cid] / denom
parts.append(f"{name}={perc:6.2%}")
ignore_ratio = global_ignored_pixels / max(1, (global_valid_pixels + global_ignored_pixels))
print(" " + " ".join(parts))
print(f" ignore={ignore_ratio:6.2%} | amostras={global_samples}")
if __name__ == "__main__":
ap = argparse.ArgumentParser(description="Verifica densidade de classes no dataset splitado.")
ap.add_argument("--split", type=str, default="train", help="Split para verificar: train, val ou test.")
ap.add_argument("--max-samples", type=int, default=MAX_SAMPLES_PER_GROUP, help="Máximo de amostras por grupo.")
args = ap.parse_args()
main(args)

View File

@ -1,134 +1,152 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
""" """
raw_triplet_sorter.py sample_bundle_sorter.py
------------------------------------------------- -------------------------------------------------
Ferramenta para classificar manualmente amostras RAW do dataset Classificador manual para o formato de dataset do módulo multiespectral.
(arquivos .raw + .png + .json) em pastas de labels,
usando atalhos de teclado. Cada amostra é formada por:
<base>.png
<base>.json
<base>_cam0.bin
<base>_cam1.bin
<base>_cam2.bin # opcional
Fluxo: Fluxo:
- Entrada: uma ou mais pastas de sessão (cada uma contendo N arquivos): - Entrada: uma ou mais pastas-raiz contendo amostras em subpastas.
- NOME.raw - O script mostra o preview (.png).
- NOME.png - Você usa teclas 1..9/0 para enviar o conjunto da amostra para uma label.
- NOME.json
- O script mostra o preview (PNG) e você usa teclas 1..9/0 para enviar
o TRIPLO (raw+preview+json) para uma pasta de saída organizada por label.
Estrutura de saída: Estrutura de saída:
out_root/ out_root/
<label>/ <label>/
raws/ bins/
previews/ previews/
metas/ metas/
masks/ masks/
Teclas: Regras:
- 1..9, 0 -> envia para a label correspondente - PNG vai para previews/
- Espaço / n / seta direita -> pular (skip) - JSON vai para metas/
- p / seta esquerda -> voltar - BINs vão para bins/
- b -> undo (desfaz última ação) - masks/ é criada vazia
- q / Esc -> sair
Depedências:
- Python 3.8+
- Pillow (PIL): pip install pillow
""" """
import argparse import argparse
import csv import csv
import json import json
import os
import shutil import shutil
import sys import sys
from dataclasses import dataclass
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import List, Tuple, Dict, Optional from typing import List, Optional
import tkinter as tk import tkinter as tk
from tkinter import messagebox, filedialog from tkinter import messagebox, filedialog
try: try:
from PIL import Image, ImageTk from PIL import Image, ImageTk
except Exception as e: except Exception:
print("ERROR: Pillow (PIL) is required. Install with: pip install pillow", file=sys.stderr) print("ERROR: Pillow (PIL) is required. Install with: pip install pillow", file=sys.stderr)
raise raise
# Sufixos/padrões esperados:
# - preview: NOME.png
# - raw: NOME.raw
# - meta: NOME.json
PREVIEW_EXT = ".png" PREVIEW_EXT = ".png"
RAW_EXT = ".raw"
META_EXT = ".json" META_EXT = ".json"
def find_triplets_in_folder(folder: Path) -> List[Tuple[Path, Path, Path, str]]: @dataclass
""" class SampleBundle:
Encontra tripletas (raw, preview, json) em uma pasta: sample_dir: Path
NOME.png preview_path: Path
NOME.raw meta_path: Path
NOME.json bin_paths: List[Path]
sample_id: str
Retorna lista de tuplas (preview_path, raw_path, json_path, base_name).
def find_sample_bundles_in_folder(folder: Path) -> List[SampleBundle]:
""" """
Procura amostras no formato:
<base>.png
<base>.json
<base>_cam0.bin
<base>_cam1.bin
<base>_cam2.bin
dentro da mesma pasta.
"""
import re
if not folder.is_dir(): if not folder.is_dir():
return [] return []
previews: Dict[str, Path] = {} bundles: List[SampleBundle] = []
raws: Dict[str, Path] = {} candidate_dirs = [folder]
metas: Dict[str, Path] = {} candidate_dirs.extend([p for p in folder.rglob("*") if p.is_dir()])
for p in folder.rglob("*"): bin_re = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
for d in candidate_dirs:
pngs = {}
jsons = {}
bins_by_base = {}
for p in d.iterdir():
if not p.is_file(): if not p.is_file():
continue continue
stem = p.stem
suffix = p.suffix.lower() suffix = p.suffix.lower()
# Preview
if suffix == PREVIEW_EXT: if suffix == PREVIEW_EXT:
base = stem pngs[p.stem] = p
previews[base] = p
continue continue
# RAW
if suffix == RAW_EXT:
base = stem
raws[base] = p
continue
# JSON
if suffix == META_EXT: if suffix == META_EXT:
base = stem jsons[p.stem] = p
metas[base] = p
continue continue
triplets = [] if suffix == ".bin":
for base, prev_path in previews.items(): m = bin_re.match(p.name)
raw_path = raws.get(base) if m:
meta_path = metas.get(base) base = m.group("base")
if raw_path is not None and meta_path is not None: cam = int(m.group("cam"))
triplets.append((prev_path, raw_path, meta_path, base)) bins_by_base.setdefault(base, []).append((cam, p))
# Ordena por base (se for número, ordena numérico) valid_bases = sorted(set(pngs.keys()) & set(jsons.keys()) & set(bins_by_base.keys()))
def sort_key(t):
b = t[3]
return (0, int(b)) if b.isdigit() else (1, b)
triplets.sort(key=sort_key) for base in valid_bases:
return triplets sorted_bins = [p for cam, p in sorted(bins_by_base[base], key=lambda x: x[0])]
if not sorted_bins:
continue
bundles.append(
SampleBundle(
sample_dir=d,
preview_path=pngs[base],
meta_path=jsons[base],
bin_paths=sorted_bins,
sample_id=base,
)
)
bundles.sort(key=lambda x: (str(x.sample_dir), x.sample_id))
return bundles
def collect_all_triplets(folders: List[Path]) -> List[Tuple[Path, Path, Path, str]]: def collect_all_bundles(folders: List[Path]) -> List[SampleBundle]:
all_tr = [] all_bundles: List[SampleBundle] = []
for f in folders: seen_dirs = set()
all_tr.extend(find_triplets_in_folder(f))
# já vem ordenado por pasta, mas garantimos ordenação global for folder in folders:
def sort_key(t): for bundle in find_sample_bundles_in_folder(folder):
b = t[3] key = str((bundle.sample_dir / bundle.sample_id).resolve())
return (0, int(b)) if b.isdigit() else (1, b) if key in seen_dirs:
all_tr.sort(key=sort_key) continue
return all_tr seen_dirs.add(key)
all_bundles.append(bundle)
all_bundles.sort(key=lambda x: (str(x.sample_dir), x.sample_id))
return all_bundles
class ActionLogger: class ActionLogger:
@ -136,7 +154,8 @@ class ActionLogger:
Loga ações em: Loga ações em:
- CSV: sorting_log.csv - CSV: sorting_log.csv
- JSONL: sorting_log.jsonl - JSONL: sorting_log.jsonl
Serve também para suportar "resume" (não reprocessar o que foi classificado/pulado).
O resume usa sample_dir como chave.
""" """
def __init__(self, out_root: Path): def __init__(self, out_root: Path):
@ -150,29 +169,29 @@ class ActionLogger:
with self.log_csv.open("r", newline="", encoding="utf-8") as f: with self.log_csv.open("r", newline="", encoding="utf-8") as f:
reader = csv.DictReader(f) reader = csv.DictReader(f)
for row in reader: for row in reader:
key = row.get("preview_path", "") key = row.get("meta_path", "")
if key: if key:
self.seen.add(key) self.seen.add(key)
except Exception: except Exception:
pass pass
def already_logged(self, preview_path: Path) -> bool: def already_logged(self, sample_key: Path) -> bool:
return str(preview_path) in self.seen return str(sample_key.resolve()) in self.seen
def log( def log(
self, self,
action: str, action: str,
sample_dir: Path,
preview_path: Path, preview_path: Path,
raw_path: Path,
meta_path: Path, meta_path: Path,
bin_paths: List[Path],
label: Optional[str] = None, label: Optional[str] = None,
dest_preview: Optional[Path] = None, dest_sample_dir: Optional[Path] = None,
dest_raw: Optional[Path] = None,
dest_meta: Optional[Path] = None,
): ):
self.out_root.mkdir(parents=True, exist_ok=True) sample_dir_str = str(sample_dir.resolve())
meta_key_str = str(meta_path.resolve())
bin_paths_str = [str(p.resolve()) for p in bin_paths]
# CSV
new_file = not self.log_csv.exists() new_file = not self.log_csv.exists()
with self.log_csv.open("a", newline="", encoding="utf-8") as f: with self.log_csv.open("a", newline="", encoding="utf-8") as f:
writer = csv.writer(f) writer = csv.writer(f)
@ -181,54 +200,50 @@ class ActionLogger:
"timestamp", "timestamp",
"action", "action",
"label", "label",
"sample_dir",
"preview_path", "preview_path",
"raw_path",
"meta_path", "meta_path",
"dest_preview", "bin_paths_json",
"dest_raw", "dest_sample_dir",
"dest_meta",
]) ])
writer.writerow([ writer.writerow([
datetime.now().isoformat(timespec="seconds"), datetime.now().isoformat(timespec="seconds"),
action, action,
label or "", label or "",
str(preview_path), sample_dir_str,
str(raw_path), str(preview_path.resolve()),
str(meta_path), str(meta_path.resolve()),
str(dest_preview or ""), json.dumps(bin_paths_str, ensure_ascii=False),
str(dest_raw or ""), str(dest_sample_dir.resolve()) if dest_sample_dir else "",
str(dest_meta or ""),
]) ])
# JSONL
with self.log_json.open("a", encoding="utf-8") as f: with self.log_json.open("a", encoding="utf-8") as f:
rec = { rec = {
"ts": datetime.now().isoformat(timespec="seconds"), "ts": datetime.now().isoformat(timespec="seconds"),
"action": action, "action": action,
"label": label, "label": label,
"preview_path": str(preview_path), "sample_dir": sample_dir_str,
"raw_path": str(raw_path), "preview_path": str(preview_path.resolve()),
"meta_path": str(meta_path), "meta_path": str(meta_path.resolve()),
"dest_preview": str(dest_preview) if dest_preview else None, "bin_paths": bin_paths_str,
"dest_raw": str(dest_raw) if dest_raw else None, "dest_sample_dir": str(dest_sample_dir.resolve()) if dest_sample_dir else None,
"dest_meta": str(dest_meta) if dest_meta else None,
} }
f.write(json.dumps(rec, ensure_ascii=False) + "\n") f.write(json.dumps(rec, ensure_ascii=False) + "\n")
self.seen.add(str(preview_path)) self.seen.add(str(meta_path.resolve()))
class RawTripletSorterApp: class SampleBundleSorterApp:
def __init__( def __init__(
self, self,
triplets: List[Tuple[Path, Path, Path, str]], bundles: List[SampleBundle],
labels: List[str], labels: List[str],
out_root: Path, out_root: Path,
move: bool, move: bool,
resume: bool, resume: bool,
display_height: int = 512, display_height: int = 512,
): ):
self.all_triplets = triplets self.all_bundles = bundles
self.labels = labels self.labels = labels
self.out_root = out_root self.out_root = out_root
self.move = move self.move = move
@ -236,28 +251,25 @@ class RawTripletSorterApp:
self.logger = ActionLogger(out_root) self.logger = ActionLogger(out_root)
if resume: if resume:
self.all_triplets = [ self.all_bundles = [
t for t in self.all_triplets if not self.logger.already_logged(t[0]) b for b in self.all_bundles if not self.logger.already_logged(b.meta_path)
] ]
# Estado de sessão
self.idx = 0 self.idx = 0
self.history: List[dict] = [] # para undo self.history = []
# Prepara pastas destino
for label in self.labels: for label in self.labels:
(out_root / label / "raws").mkdir(parents=True, exist_ok=True) label_root = out_root / label
(out_root / label / "previews").mkdir(parents=True, exist_ok=True) (label_root / "bins").mkdir(parents=True, exist_ok=True)
(out_root / label / "metas").mkdir(parents=True, exist_ok=True) (label_root / "previews").mkdir(parents=True, exist_ok=True)
(out_root / label / "masks").mkdir(parents=True, exist_ok=True) (label_root / "metas").mkdir(parents=True, exist_ok=True)
(label_root / "masks").mkdir(parents=True, exist_ok=True)
# UI
self.root = tk.Tk() self.root = tk.Tk()
self.root.title("Agrobotics RAW Triplet Sorter (preview | raw | json)") self.root.title("Agrobotics Sample Bundle Sorter")
self.root.geometry("1100x750") self.root.geometry("1100x750")
self.root.bind("<Key>", self.on_key) self.root.bind("<Key>", self.on_key)
# Topo: info + legenda
self.top_frame = tk.Frame(self.root) self.top_frame = tk.Frame(self.root)
self.top_frame.pack(side=tk.TOP, fill=tk.X) self.top_frame.pack(side=tk.TOP, fill=tk.X)
@ -267,19 +279,16 @@ class RawTripletSorterApp:
self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10)) self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10))
self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6) self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6)
# Área da imagem
self.img_frame = tk.Frame(self.root) self.img_frame = tk.Frame(self.root)
self.img_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) self.img_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
self.preview_label = tk.Label(self.img_frame) self.preview_label = tk.Label(self.img_frame)
self.preview_label.pack(side=tk.LEFT, expand=True, padx=6, pady=6) self.preview_label.pack(side=tk.LEFT, expand=True, padx=6, pady=6)
# Status (última ação)
self.status_var = tk.StringVar(value="Pronto.") 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 = 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.status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=4)
# Rodapé: ajuda
self.footer = tk.Label( self.footer = tk.Label(
self.root, self.root,
text="1..9/0 = labels | espaço/n/→ = próxima (skip) | p/← = anterior | b = undo | q/Esc = sair", text="1..9/0 = labels | espaço/n/→ = próxima (skip) | p/← = anterior | b = undo | q/Esc = sair",
@ -287,19 +296,31 @@ class RawTripletSorterApp:
) )
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2) self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
# Primeira renderização
self.render() self.render()
def unique_file(self, p: Path) -> Path:
if not p.exists():
return p
stem = p.stem
suffix = p.suffix
k = 1
while True:
cand = p.with_name(f"{stem}__{k}{suffix}")
if not cand.exists():
return cand
k += 1
def build_legend_text(self) -> str: def build_legend_text(self) -> str:
parts = [] parts = []
for i, label in enumerate(self.labels, start=1): for i, label in enumerate(self.labels, start=1):
key = i if i <= 9 else 0 # 0 = 10ª label key = i if i <= 9 else 0
parts.append(f"[{key}] {label}") parts.append(f"[{key}] {label}")
return " | ".join(parts) return " | ".join(parts)
def pil_load_preview(self, preview_path: Path) -> ImageTk.PhotoImage: def pil_load_preview(self, preview_path: Path) -> ImageTk.PhotoImage:
img = Image.open(preview_path).convert("RGB") img = Image.open(preview_path).convert("RGB")
# escala pra altura desejada mantendo proporção
h_target = self.display_height h_target = self.display_height
w, h = img.size w, h = img.size
new_w = int(w * (h_target / h)) new_w = int(w * (h_target / h))
@ -307,135 +328,140 @@ class RawTripletSorterApp:
return ImageTk.PhotoImage(img) return ImageTk.PhotoImage(img)
def render(self): def render(self):
if not self.all_triplets: if not self.all_bundles:
messagebox.showinfo("Fim", "Não há amostras para exibir (talvez tudo já foi classificado?).") messagebox.showinfo("Fim", "Não há amostras para exibir.")
self.root.destroy() self.root.destroy()
return return
self.idx = max(0, min(self.idx, len(self.all_triplets) - 1)) self.idx = max(0, min(self.idx, len(self.all_bundles) - 1))
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx] bundle = self.all_bundles[self.idx]
try: try:
tk_img = self.pil_load_preview(preview_path) tk_img = self.pil_load_preview(bundle.preview_path)
self.preview_tk = tk_img self.preview_tk = tk_img
self.preview_label.configure(image=self.preview_tk) self.preview_label.configure(image=self.preview_tk)
except Exception as e: except Exception as e:
self.info_label.configure(text=f"Erro ao abrir preview: {e}") self.info_label.configure(text=f"Erro ao abrir preview: {e}")
return return
bins_text = ", ".join(p.name for p in bundle.bin_paths)
self.info_label.configure( self.info_label.configure(
text=f"{self.idx+1}/{len(self.all_triplets)} | base='{base}' | PREVIEW: {preview_path.name}" text=(
f"{self.idx+1}/{len(self.all_bundles)} | "
f"sample='{bundle.sample_id}' | "
f"bins={len(bundle.bin_paths)} [{bins_text}]"
)
) )
def do_copy_or_move(self, src: Path, dst: Path):
if self.move:
shutil.move(str(src), str(dst))
else:
shutil.copy2(str(src), str(dst))
def unique_path(self, p: Path) -> Path:
"""Se o caminho já existe, gera um novo com sufixo __k."""
if not p.exists():
return p
stem, ext = p.stem, p.suffix
k = 1
while True:
cand = p.with_name(f"{stem}__{k}{ext}")
if not cand.exists():
return cand
k += 1
def send_to_label(self, label_index: int): def send_to_label(self, label_index: int):
if label_index < 0 or label_index >= len(self.labels): if label_index < 0 or label_index >= len(self.labels):
return return
label = self.labels[label_index] label = self.labels[label_index]
bundle = self.all_bundles[self.idx]
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx] label_root = self.out_root / label
dst_preview = self.unique_file(label_root / "previews" / bundle.preview_path.name)
dst_preview = self.out_root / label / "previews" / preview_path.name dst_meta = self.unique_file(label_root / "metas" / bundle.meta_path.name)
dst_raw = self.out_root / label / "raws" / raw_path.name dst_bins = [self.unique_file(label_root / "bins" / p.name) for p in bundle.bin_paths]
dst_meta = self.out_root / label / "metas" / meta_path.name
dst_preview = self.unique_path(dst_preview)
dst_raw = self.unique_path(dst_raw)
dst_meta = self.unique_path(dst_meta)
try: try:
self.do_copy_or_move(preview_path, dst_preview) if self.move:
self.do_copy_or_move(raw_path, dst_raw) shutil.move(str(bundle.preview_path), str(dst_preview))
self.do_copy_or_move(meta_path, dst_meta) shutil.move(str(bundle.meta_path), str(dst_meta))
for src_bin, dst_bin in zip(bundle.bin_paths, dst_bins):
shutil.move(str(src_bin), str(dst_bin))
else:
shutil.copy2(bundle.preview_path, dst_preview)
shutil.copy2(bundle.meta_path, dst_meta)
for src_bin, dst_bin in zip(bundle.bin_paths, dst_bins):
shutil.copy2(src_bin, dst_bin)
self.logger.log( self.logger.log(
"assign", action="assign",
preview_path=preview_path, sample_dir=bundle.sample_dir,
raw_path=raw_path, preview_path=bundle.preview_path,
meta_path=meta_path, meta_path=bundle.meta_path,
bin_paths=bundle.bin_paths,
label=label, label=label,
dest_preview=dst_preview, dest_sample_dir=label_root,
dest_raw=dst_raw,
dest_meta=dst_meta,
) )
self.history.append({ self.history.append({
"action": "assign", "action": "assign",
"label": label, "label": label,
"preview_src": preview_path, "sample_src": bundle.sample_dir,
"raw_src": raw_path, "preview_src": bundle.preview_path,
"meta_src": meta_path, "meta_src": bundle.meta_path,
"bin_srcs": list(bundle.bin_paths),
"preview_dst": dst_preview, "preview_dst": dst_preview,
"raw_dst": dst_raw,
"meta_dst": dst_meta, "meta_dst": dst_meta,
"bin_dsts": dst_bins,
"moved": self.move, "moved": self.move,
"index": self.idx, "index": self.idx,
}) })
self.status_var.set( self.status_var.set(
f"{'Movido' if self.move else 'Copiado'}'{label}': {preview_path.name} (+raw+json)" f"{'Movido' if self.move else 'Copiado'}'{label}': {bundle.sample_id}"
) )
self.idx += 1 self.idx += 1
if self.idx >= len(self.all_triplets): if self.idx >= len(self.all_bundles):
messagebox.showinfo("Concluído", "Você chegou ao final da fila!") messagebox.showinfo("Concluído", "Você chegou ao final da fila!")
self.root.destroy() self.root.destroy()
return return
self.render() self.render()
except Exception as e: except Exception as e:
messagebox.showerror("Erro", f"Falha ao copiar/mover: {e}") messagebox.showerror("Erro", f"Falha ao copiar/mover amostra: {e}")
self.status_var.set(f"ERRO: {e}") self.status_var.set(f"ERRO: {e}")
def undo(self): def undo(self):
if not self.history: if not self.history:
return return
last = self.history.pop() last = self.history.pop()
if last["action"] != "assign": if last["action"] != "assign":
return return
try: try:
preview_src = Path(last["preview_src"])
meta_src = Path(last["meta_src"])
bin_srcs = [Path(p) for p in last["bin_srcs"]]
preview_dst = Path(last["preview_dst"])
meta_dst = Path(last["meta_dst"])
bin_dsts = [Path(p) for p in last["bin_dsts"]]
if last["moved"]: if last["moved"]:
# mover de volta pro source shutil.move(str(preview_dst), str(preview_src))
shutil.move(str(last["preview_dst"]), str(last["preview_src"])) shutil.move(str(meta_dst), str(meta_src))
shutil.move(str(last["raw_dst"]), str(last["raw_src"])) for dst_bin, src_bin in zip(bin_dsts, bin_srcs):
shutil.move(str(last["meta_dst"]), str(last["meta_src"])) shutil.move(str(dst_bin), str(src_bin))
else: else:
# apagar os arquivos de destino if preview_dst.exists():
for p in [last["preview_dst"], last["raw_dst"], last["meta_dst"]]: preview_dst.unlink()
if p.exists(): if meta_dst.exists():
p.unlink() meta_dst.unlink()
for dst_bin in bin_dsts:
if dst_bin.exists():
dst_bin.unlink()
self.logger.log( self.logger.log(
"undo", action="undo",
preview_path=last["preview_src"], sample_dir=Path(last["sample_src"]),
raw_path=last["raw_src"], preview_path=preview_src,
meta_path=last["meta_src"], meta_path=meta_src,
bin_paths=bin_srcs,
label=last["label"], label=last["label"],
dest_preview=last["preview_dst"], dest_sample_dir=self.out_root / last["label"],
dest_raw=last["raw_dst"],
dest_meta=last["meta_dst"],
) )
self.idx = max(0, min(last.get("index", self.idx), len(self.all_triplets) - 1)) self.idx = max(0, min(last.get("index", self.idx), len(self.all_bundles) - 1))
self.status_var.set(f"Desfeito ← '{last['label']}': {last['preview_src'].name}") self.status_var.set(f"Desfeito ← '{last['label']}': {preview_src.stem}")
self.render() self.render()
except Exception as e: except Exception as e:
messagebox.showerror("Erro", f"Falha no undo: {e}") messagebox.showerror("Erro", f"Falha no undo: {e}")
self.status_var.set(f"ERRO: {e}") self.status_var.set(f"ERRO: {e}")
@ -443,38 +469,41 @@ class RawTripletSorterApp:
def on_key(self, event): def on_key(self, event):
ch = event.keysym.lower() ch = event.keysym.lower()
# 1..9 / 0 -> labels
if ch in [str(i) for i in range(1, 10)] or ch == "0": if ch in [str(i) for i in range(1, 10)] or ch == "0":
label_idx = 9 if ch == "0" else int(ch) - 1 label_idx = 9 if ch == "0" else int(ch) - 1
self.send_to_label(label_idx) self.send_to_label(label_idx)
return return
# espaço / n / seta direita = skip / próxima
if ch in ("space", "s", "n", "right"): if ch in ("space", "s", "n", "right"):
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx] bundle = self.all_bundles[self.idx]
self.logger.log("skip", preview_path, raw_path, meta_path, label=None) self.logger.log(
self.status_var.set(f"Pulada → {preview_path.name}") action="skip",
sample_dir=bundle.sample_dir,
preview_path=bundle.preview_path,
meta_path=bundle.meta_path,
bin_paths=bundle.bin_paths,
label=None,
dest_sample_dir=None,
)
self.status_var.set(f"Pulada → {bundle.sample_id}")
self.idx += 1 self.idx += 1
if self.idx >= len(self.all_triplets): if self.idx >= len(self.all_bundles):
messagebox.showinfo("Concluído", "Você chegou ao final da fila!") messagebox.showinfo("Concluído", "Você chegou ao final da fila!")
self.root.destroy() self.root.destroy()
return return
self.render() self.render()
return return
# p / seta esquerda = voltar
if ch in ("p", "left"): if ch in ("p", "left"):
self.idx = max(0, self.idx - 1) self.idx = max(0, self.idx - 1)
self.status_var.set("Voltou uma imagem.") self.status_var.set("Voltou uma amostra.")
self.render() self.render()
return return
# b = undo if ch == "b":
if ch in ("b",):
self.undo() self.undo()
return return
# q / Esc = sair
if ch in ("q", "escape"): if ch in ("q", "escape"):
self.root.destroy() self.root.destroy()
return return
@ -484,17 +513,12 @@ class RawTripletSorterApp:
class SetupWindow: class SetupWindow:
"""
Janela de setup (igual ao sorter anterior, mas voltada para RAW triplets).
"""
def __init__(self): def __init__(self):
self.root = tk.Tk() self.root = tk.Tk()
self.root.title("Configurar - RAW Triplet Sorter") self.root.title("Configurar - Sample Bundle Sorter")
self.root.geometry("720x520") self.root.geometry("720x520")
# Pastas de entrada frm_in = tk.LabelFrame(self.root, text="Pastas de entrada (raízes com subpastas de amostras)")
frm_in = tk.LabelFrame(self.root, text="Pastas de entrada (sessões de captura RAW)")
frm_in.pack(fill=tk.BOTH, expand=False, padx=10, pady=8) frm_in.pack(fill=tk.BOTH, expand=False, padx=10, pady=8)
self.inputs_listbox = tk.Listbox(frm_in, height=6) self.inputs_listbox = tk.Listbox(frm_in, height=6)
@ -506,20 +530,17 @@ class SetupWindow:
tk.Button(btns_in, text="Remover selecionada", command=self.remove_selected).pack(fill=tk.X, pady=2) tk.Button(btns_in, text="Remover selecionada", command=self.remove_selected).pack(fill=tk.X, pady=2)
tk.Button(btns_in, text="Limpar lista", command=self.clear_inputs).pack(fill=tk.X, pady=2) tk.Button(btns_in, text="Limpar lista", command=self.clear_inputs).pack(fill=tk.X, pady=2)
# Pasta de saída
frm_out = tk.LabelFrame(self.root, text="Pasta de saída (raiz do dataset rotulado)") frm_out = tk.LabelFrame(self.root, text="Pasta de saída (raiz do dataset rotulado)")
frm_out.pack(fill=tk.X, expand=False, padx=10, pady=8) frm_out.pack(fill=tk.X, expand=False, padx=10, pady=8)
self.out_root_var = tk.StringVar(value="") self.out_root_var = tk.StringVar(value="")
tk.Entry(frm_out, textvariable=self.out_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6) tk.Entry(frm_out, textvariable=self.out_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6)
tk.Button(frm_out, text="Escolher...", command=self.choose_out_root).pack(side=tk.RIGHT, padx=6, pady=6) tk.Button(frm_out, text="Escolher...", command=self.choose_out_root).pack(side=tk.RIGHT, padx=6, pady=6)
# Labels
frm_labels = tk.LabelFrame(self.root, text="Labels (classes) separadas por vírgula") frm_labels = tk.LabelFrame(self.root, text="Labels (classes) separadas por vírgula")
frm_labels.pack(fill=tk.X, expand=False, padx=10, pady=8) frm_labels.pack(fill=tk.X, expand=False, padx=10, pady=8)
self.labels_var = tk.StringVar(value="chao,chao_cana,chao_erva,chao_cana_erva,cana,cana_erva,erva") self.labels_var = tk.StringVar(value="chao,chao_cana,chao_erva,chao_cana_erva,cana,cana_erva,erva")
tk.Entry(frm_labels, textvariable=self.labels_var).pack(fill=tk.X, padx=6, pady=6) tk.Entry(frm_labels, textvariable=self.labels_var).pack(fill=tk.X, padx=6, pady=6)
# Opções
frm_opts = tk.LabelFrame(self.root, text="Opções") frm_opts = tk.LabelFrame(self.root, text="Opções")
frm_opts.pack(fill=tk.X, expand=False, padx=10, pady=8) frm_opts.pack(fill=tk.X, expand=False, padx=10, pady=8)
self.move_var = tk.BooleanVar(value=False) self.move_var = tk.BooleanVar(value=False)
@ -534,10 +555,8 @@ class SetupWindow:
tk.Label(frm_height, text="Altura de exibição (px):").pack(side=tk.LEFT) tk.Label(frm_height, text="Altura de exibição (px):").pack(side=tk.LEFT)
tk.Entry(frm_height, textvariable=self.height_var, width=6).pack(side=tk.LEFT, padx=6) tk.Entry(frm_height, textvariable=self.height_var, width=6).pack(side=tk.LEFT, padx=6)
# Botão iniciar
tk.Button(self.root, text="Iniciar classificação", command=self.start).pack(pady=10) tk.Button(self.root, text="Iniciar classificação", command=self.start).pack(pady=10)
# Hint
tk.Label( tk.Label(
self.root, self.root,
text="Teclas: 1..9 (0=10ª), espaço/n/→=pular, p/←=anterior, b=undo, q/Esc=sair", text="Teclas: 1..9 (0=10ª), espaço/n/→=pular, p/←=anterior, b=undo, q/Esc=sair",
@ -547,7 +566,7 @@ class SetupWindow:
self.result = None self.result = None
def add_input(self): def add_input(self):
p = filedialog.askdirectory(title="Selecione a pasta de sessão (contendo .raw, .png, .json)") p = filedialog.askdirectory(title="Selecione a pasta raiz contendo as amostras")
if p: if p:
self.inputs_listbox.insert(tk.END, p) self.inputs_listbox.insert(tk.END, p)
@ -600,18 +619,20 @@ def run_with_gui_setup():
res = setup.run() res = setup.run()
if not res: if not res:
return return
input_folders, labels, out_root, move, resume, height = res input_folders, labels, out_root, move, resume, height = res
triplets = collect_all_triplets(input_folders) bundles = collect_all_bundles(input_folders)
if not triplets: if not bundles:
messagebox.showinfo( messagebox.showinfo(
"Sem amostras", "Sem amostras",
"Nenhuma tripleta encontrada.\nCertifique-se de que existam arquivos .raw, .png, .json.", "Nenhuma amostra válida encontrada.\n"
"Cada amostra precisa ter <base>.png, <base>.json e um ou mais <base>_camX.bin.",
) )
return return
app = RawTripletSorterApp( app = SampleBundleSorterApp(
triplets=triplets, bundles=bundles,
labels=labels, labels=labels,
out_root=out_root, out_root=out_root,
move=move, move=move,
@ -623,39 +644,37 @@ def run_with_gui_setup():
def main(): def main():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Sortear manualmente amostras (raw + preview + json) em pastas de labels, usando hotkeys." description="Classificar manualmente amostras no novo formato do módulo multiespectral."
) )
parser.add_argument("--inputs", nargs="+", help="Pastas de sessão contendo .raw, .png e .json") parser.add_argument("--inputs", nargs="+", help="Pastas-raiz contendo subpastas de amostras")
parser.add_argument("--labels", nargs="+", help="Labels (classes) mapeadas para teclas 1..9/0") parser.add_argument("--labels", nargs="+", help="Labels (classes) mapeadas para teclas 1..9/0")
parser.add_argument("--out-root", help="Pasta raiz de saída") parser.add_argument("--out-root", help="Pasta raiz de saída")
parser.add_argument("--move", action="store_true", help="Mover em vez de copiar") parser.add_argument("--move", action="store_true", help="Mover em vez de copiar")
parser.add_argument("--resume", action="store_true", help="Pular itens já presentes no sorting_log.csv") parser.add_argument("--resume", action="store_true", help="Pular itens já presentes no sorting_log.csv")
parser.add_argument("--display-height", type=int, default=512, help="Altura de exibição do preview (px)") parser.add_argument("--display-height", type=int, default=512, help="Altura de exibição do preview (px)")
parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir a GUI de setup (usar apenas args CLI)") parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir a GUI de setup")
args = parser.parse_args() args = parser.parse_args()
# Se faltou algo e GUI está permitida, abre o setup.
if (not args.inputs or not args.labels or not args.out_root) and not args.no_gui_setup: if (not args.inputs or not args.labels or not args.out_root) and not args.no_gui_setup:
run_with_gui_setup() run_with_gui_setup()
return return
# Modo apenas CLI
if not args.inputs or not args.labels or not args.out_root: if not args.inputs or not args.labels or not args.out_root:
print("ERRO: É preciso informar --inputs, --labels e --out-root ou usar a GUI (sem --no-gui-setup).") print("ERRO: informe --inputs, --labels e --out-root ou use a GUI.")
sys.exit(1) sys.exit(1)
input_folders = [Path(p) for p in args.inputs] input_folders = [Path(p) for p in args.inputs]
labels = args.labels labels = args.labels
out_root = Path(args.out_root) out_root = Path(args.out_root)
triplets = collect_all_triplets(input_folders) bundles = collect_all_bundles(input_folders)
if not triplets: if not bundles:
print("Nenhuma tripleta encontrada nas pastas informadas.") print("Nenhuma amostra válida encontrada nas pastas informadas.")
sys.exit(1) sys.exit(1)
app = RawTripletSorterApp( app = SampleBundleSorterApp(
triplets=triplets, bundles=bundles,
labels=labels, labels=labels,
out_root=out_root, out_root=out_root,
move=args.move, move=args.move,

View File

@ -3,10 +3,11 @@ import os
import shutil import shutil
import argparse import argparse
import csv import csv
import re
# ================= CONFIG ================= # ================= CONFIG =================
with open("config.json", "r") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config["camera"] MODELO = config["camera"]
@ -14,18 +15,17 @@ MODELO = config["camera"]
# Raiz das brutas (todas as canas/horários/grupos) # Raiz das brutas (todas as canas/horários/grupos)
PASTA_BRUTAS_ROOT = os.path.join("dataset", "brutas") PASTA_BRUTAS_ROOT = os.path.join("dataset", "brutas")
# Onde você coloca os previews selecionados (tudo misturado) # Onde você coloca os previews selecionados manualmente
PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews") PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews")
# Destino final, organizado por grupo: # Destino final, organizado por grupo:
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks} # dataset/original/group/{GRUPO}/{previews,metas,bins,masks}
PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group") PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group")
EXT_PREVIEWS = (".png", ".jpg", ".jpeg") EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
EXT_RAWS = (".raw",)
EXT_MASKS = (".png",) EXT_MASKS = (".png",)
# Ajusta se suas metas tiverem outra extensão EXT_METAS = (".json",)
EXT_METAS = (".json", ".yml", ".yaml", ".txt", ".csv") BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
# ================= HELPERS ================= # ================= HELPERS =================
@ -51,7 +51,7 @@ def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
i += 1 i += 1
def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None: def buscar_por_base(pasta: str, base: str, exts):
""" """
Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'. Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'.
Retorna o caminho completo ou None. Retorna o caminho completo ou None.
@ -60,28 +60,54 @@ def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None:
return None return None
for nome in os.listdir(pasta): for nome in os.listdir(pasta):
nome_lower = nome.lower() root, ext = os.path.splitext(nome)
root, ext = os.path.splitext(nome_lower) if root.lower() == base.lower() and ext.lower() in tuple(e.lower() for e in exts):
if root == base.lower() and ext in exts:
return os.path.join(pasta, nome) return os.path.join(pasta, nome)
return None return None
def indexar_brutas_por_raw(root_brutas: str): def buscar_bins_por_base(pasta_bins: str, base: str):
""" """
Varre dataset/brutas recursivamente, olhando pastas 'raws' e montando índice: Procura todos os bins do tipo:
<base>_cam0.bin
<base>_cam1.bin
<base>_cam2.bin
Retorna lista ordenada por cam.
"""
encontrados = []
if not os.path.isdir(pasta_bins):
return encontrados
for nome in os.listdir(pasta_bins):
m = BIN_RE.match(nome)
if not m:
continue
if m.group("base").lower() != base.lower():
continue
cam = int(m.group("cam"))
encontrados.append((cam, os.path.join(pasta_bins, nome)))
encontrados.sort(key=lambda x: x[0])
return [p for _, p in encontrados]
def indexar_brutas_por_meta(root_brutas: str):
"""
Varre dataset/brutas recursivamente, olhando pastas 'metas' e montando índice:
base -> { base -> {
"raw": caminho_raw, "meta": caminho_meta,
"cana": "cana_alta" | "cana_baixa" | ... (se conseguir inferir), "bins": [cam0, cam1, ...],
"horario": "meio_dia" | "cedo" | ... (se conseguir inferir), "cana": "cana_alta" | "cana_baixa" | ...,
"horario": "meio_dia" | "cedo" | ...,
"grupo": "<GRUPO>", "grupo": "<GRUPO>",
"group_dir": caminho_da_pasta_do_grupo "group_dir": caminho_da_pasta_do_grupo
} }
Estrutura esperada (relativa a root_brutas): Estrutura esperada:
cana_alta/meio_dia/group/<GRUPO>/raws/*.raw cana_alta/meio_dia/group/<GRUPO>/{previews,metas,bins,masks}
""" """
index = {} index = {}
@ -91,59 +117,56 @@ def indexar_brutas_por_raw(root_brutas: str):
for dirpath, dirnames, filenames in os.walk(root_brutas): for dirpath, dirnames, filenames in os.walk(root_brutas):
base_dir = os.path.basename(dirpath).lower() base_dir = os.path.basename(dirpath).lower()
if base_dir != "metas":
if base_dir != "raws":
continue continue
# dirpath = .../cana_x/horario/group/<GRUPO>/raws # dirpath = .../cana_x/horario/group/<GRUPO>/metas
group_dir = os.path.dirname(dirpath) # .../cana_x/horario/group/<GRUPO> group_dir = os.path.dirname(dirpath)
bins_dir = os.path.join(group_dir, "bins")
rel = os.path.relpath(dirpath, root_brutas) rel = os.path.relpath(dirpath, root_brutas)
parts = rel.split(os.sep) parts = rel.split(os.sep)
# Defaults
cana = None cana = None
horario = None horario = None
grupo = None grupo = None
if len(parts) >= 5: if len(parts) >= 5:
# [0] = cana_<...>
# [1] = horario
# [2] = "group"
# [3] = <GRUPO>
cana = parts[0] cana = parts[0]
horario = parts[1] horario = parts[1]
# parts[2] deve ser "group"
grupo = parts[3] grupo = parts[3]
else: else:
# fallback bem genérico
if "group" in parts: if "group" in parts:
i = parts.index("group") i = parts.index("group")
if i + 1 < len(parts): if i + 1 < len(parts):
grupo = parts[i + 1] grupo = parts[i + 1]
for nome in filenames: for nome in filenames:
if not nome.lower().endswith(EXT_RAWS): if not nome.lower().endswith(EXT_METAS):
continue continue
base = os.path.splitext(nome)[0] base = os.path.splitext(nome)[0]
meta_path = os.path.join(dirpath, nome)
bin_paths = buscar_bins_por_base(bins_dir, base)
if base in index: if not bin_paths:
# Conflito (mesmo base em dois lugares) -> loga e mantém o primeiro print(f"[SKIP] Meta sem bins correspondentes: {meta_path}")
print(f"[CONFLITO] base repetida em raws: {base}")
continue continue
raw_path = os.path.join(dirpath, nome) if base in index:
print(f"[CONFLITO] base repetida em metas: {base}")
continue
index[base] = { index[base] = {
"raw": raw_path, "meta": meta_path,
"bins": bin_paths,
"cana": cana, "cana": cana,
"horario": horario, "horario": horario,
"grupo": grupo, "grupo": grupo,
"group_dir": group_dir, "group_dir": group_dir,
} }
print(f"[INDEX] Entradas indexadas por RAW: {len(index)}") print(f"[INDEX] Entradas indexadas por META: {len(index)}")
return index return index
@ -151,8 +174,8 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
if not os.path.isdir(PASTA_SELECTED_PREVIEWS): if not os.path.isdir(PASTA_SELECTED_PREVIEWS):
raise SystemExit(f"[ERRO] Pasta selected_previews não existe: {PASTA_SELECTED_PREVIEWS}") raise SystemExit(f"[ERRO] Pasta selected_previews não existe: {PASTA_SELECTED_PREVIEWS}")
# Indexa todas as brutas a partir dos RAWs # Indexa todas as brutas a partir das METAS + BINS
index_raws = indexar_brutas_por_raw(PASTA_BRUTAS_ROOT) index_samples = indexar_brutas_por_meta(PASTA_BRUTAS_ROOT)
registros = [] registros = []
total, movidos, ignorados = 0, 0, 0 total, movidos, ignorados = 0, 0, 0
@ -167,59 +190,54 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
continue continue
total += 1 total += 1
base = os.path.splitext(nome)[0] base = os.path.splitext(nome)[0]
info = index_raws.get(base) info = index_samples.get(base)
if info is None: if info is None:
ignorados += 1 ignorados += 1
print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via RAW)") print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via META+BINS)")
continue continue
grupo = info.get("grupo") or "unknown" grupo = info.get("grupo") or "unknown"
cana = info.get("cana") or "unknown" cana = info.get("cana") or "unknown"
horario = info.get("horario") or "unknown" horario = info.get("horario") or "unknown"
group_dir = info["group_dir"] group_dir = info["group_dir"]
raw_src = info["raw"] meta_src = info["meta"]
bins_src = info["bins"]
# Pastas irmãs em brutas # Pastas irmãs em brutas
metas_src_dir = os.path.join(group_dir, "metas")
masks_src_dir = os.path.join(group_dir, "masks") masks_src_dir = os.path.join(group_dir, "masks")
meta_src = buscar_por_base(metas_src_dir, base, EXT_METAS)
mask_src = buscar_por_base(masks_src_dir, base, EXT_MASKS) mask_src = buscar_por_base(masks_src_dir, base, EXT_MASKS)
# Destino: dataset/original/group/{GRUPO}/{previews,raws,metas,masks} # Destino: dataset/original/group/{GRUPO}/{previews,metas,bins,masks}
dest_group_root = os.path.join(PASTA_ORIGINAL_GROUP_ROOT, grupo) dest_group_root = os.path.join(PASTA_ORIGINAL_GROUP_ROOT, grupo)
dest_prev_dir = os.path.join(dest_group_root, "previews") dest_prev_dir = os.path.join(dest_group_root, "previews")
dest_raw_dir = os.path.join(dest_group_root, "raws")
dest_meta_dir = os.path.join(dest_group_root, "metas") dest_meta_dir = os.path.join(dest_group_root, "metas")
dest_bins_dir = os.path.join(dest_group_root, "bins")
dest_mask_dir = os.path.join(dest_group_root, "masks") dest_mask_dir = os.path.join(dest_group_root, "masks")
garantir_pasta(dest_prev_dir) garantir_pasta(dest_prev_dir)
garantir_pasta(dest_raw_dir)
garantir_pasta(dest_meta_dir) garantir_pasta(dest_meta_dir)
garantir_pasta(dest_bins_dir)
garantir_pasta(dest_mask_dir) garantir_pasta(dest_mask_dir)
# Define extensões
prev_ext_sel = os.path.splitext(nome)[1].lower() prev_ext_sel = os.path.splitext(nome)[1].lower()
raw_ext = os.path.splitext(raw_src)[1].lower() meta_ext = os.path.splitext(meta_src)[1].lower()
# Gera nome final único com base no preview # Nome-base final único definido a partir do preview
dst_preview = nome_disponivel(dest_prev_dir, base, prev_ext_sel) dst_preview = nome_disponivel(dest_prev_dir, base, prev_ext_sel)
new_base = os.path.splitext(os.path.basename(dst_preview))[0] new_base = os.path.splitext(os.path.basename(dst_preview))[0]
dst_raw = os.path.join(dest_raw_dir, new_base + raw_ext)
dst_meta = None
dst_mask = None
if meta_src is not None:
meta_ext = os.path.splitext(meta_src)[1].lower()
dst_meta = os.path.join(dest_meta_dir, new_base + meta_ext) dst_meta = os.path.join(dest_meta_dir, new_base + meta_ext)
dst_mask = os.path.join(dest_mask_dir, new_base + ".png") if mask_src else None
if mask_src is not None: dst_bins = []
mask_ext = os.path.splitext(mask_src)[1].lower() for src_bin in bins_src:
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext) m = BIN_RE.match(os.path.basename(src_bin))
if not m:
continue
cam = m.group("cam")
dst_bins.append(os.path.join(dest_bins_dir, f"{new_base}_cam{cam}.bin"))
# Copia/move preview selecionado # Copia/move preview selecionado
if copy_only: if copy_only:
@ -227,18 +245,20 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
else: else:
shutil.move(caminho_preview_sel, dst_preview) shutil.move(caminho_preview_sel, dst_preview)
# Copia RAW e metas/masks (se existirem) # Copia/move meta
if copy_only:
shutil.copy2(raw_src, dst_raw)
else:
shutil.move(raw_src, dst_raw)
if meta_src is not None:
if copy_only: if copy_only:
shutil.copy2(meta_src, dst_meta) shutil.copy2(meta_src, dst_meta)
else: else:
shutil.move(meta_src, dst_meta) shutil.move(meta_src, dst_meta)
# Copia/move bins
for src_bin, dst_bin in zip(bins_src, dst_bins):
if copy_only:
shutil.copy2(src_bin, dst_bin)
else:
shutil.move(src_bin, dst_bin)
# Copia/move mask se existir
if mask_src is not None: if mask_src is not None:
if copy_only: if copy_only:
shutil.copy2(mask_src, dst_mask) shutil.copy2(mask_src, dst_mask)
@ -253,16 +273,16 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
cana, cana,
horario, horario,
caminho_preview_sel, caminho_preview_sel,
raw_src, meta_src,
meta_src or "", json.dumps(bins_src, ensure_ascii=False),
mask_src or "", mask_src or "",
dst_preview, dst_preview,
dst_raw, dst_meta,
dst_meta or "", json.dumps(dst_bins, ensure_ascii=False),
dst_mask or "", dst_mask or "",
]) ])
print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario}") print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario} | bins={len(dst_bins)}")
if manifesto_csv and registros: if manifesto_csv and registros:
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f: with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
@ -273,12 +293,12 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
"cana", "cana",
"horario", "horario",
"src_preview_selected", "src_preview_selected",
"src_raw",
"src_meta", "src_meta",
"src_bins_json",
"src_mask", "src_mask",
"dst_preview", "dst_preview",
"dst_raw",
"dst_meta", "dst_meta",
"dst_bins_json",
"dst_mask", "dst_mask",
]) ])
w.writerows(registros) w.writerows(registros)
@ -294,7 +314,7 @@ def build_cli():
description=( description=(
"Organiza previews selecionadas (dataset/selected_previews) " "Organiza previews selecionadas (dataset/selected_previews) "
"descobrindo cana/horário/grupo em dataset/brutas e copiando/movendo " "descobrindo cana/horário/grupo em dataset/brutas e copiando/movendo "
"preview + raw + meta (+ mask se existir) para dataset/original/group/{GRUPO}." "preview + meta + bins (+ mask se existir) para dataset/original/group/{GRUPO}."
) )
) )
ap.add_argument( ap.add_argument(

View File

@ -1,76 +1,117 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
_2_create_full_mask_multispec.py
-------------------------------------------------
Gera máscaras sólidas para grupos de classe única no novo formato
multiespectral.
Estrutura esperada de entrada:
dataset/
brutas/
cana_<estado>/
<horario>/
group/
<grupo>/
previews/
metas/
bins/
masks/
Ou, usando --from-originals:
dataset/
original/
previews/
metas/
bins/
masks/
Cada amostra é formada por:
<base>.png
<base>.json
<base>_cam0.bin
<base>_cam1.bin
<base>_cam2.bin # opcional
O script:
1. percorre os previews do grupo
2. localiza o meta correspondente
3. localiza todos os bins correspondentes
4. cria uma máscara sólida com a cor da classe
5. opcionalmente copia preview + meta + bins + mask para dataset/original
com deduplicação de nome-base
"""
import json import json
import os import os
import sys
import cv2 import cv2
import csv import csv
import shutil import shutil
import argparse import argparse
import numpy as np import numpy as np
import re
from typing import List, Optional, Tuple
from utils import carregar_labelmap_completo # importa da sua utils.py sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo
# ===================================================
# ⚙️ Configurações base vindas do config.json
# ===================================================
with open("config.json", "r") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config["camera"] MODELO = config["camera"]
RESOLUCAO = config["raw_size"] RESOLUCAO = config["raw_size"]
RAW_WIDTH = RESOLUCAO[0] # largura real do raw (ex: 1296) RAW_WIDTH = RESOLUCAO[0]
RAW_HEIGHT = RESOLUCAO[1] # altura real do raw (ex: 1028) RAW_HEIGHT = RESOLUCAO[1]
RAW_EXTS = [".raw"] # se tiver outro, adiciona aqui
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews") PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks") PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws") PASTA_FINAL_METAS = os.path.join("dataset", "original", "metas")
PASTA_FINAL_BINS = os.path.join("dataset", "original", "bins")
# Cor da classe será obtida via labelmap
# COR_CLASSE_RGB = (128, 0, 0)
COPIAR_IMAGENS = True COPIAR_IMAGENS = True
EXT_IMAGENS = (".jpg", ".jpeg", ".png") EXT_IMAGENS = (".jpg", ".jpeg", ".png")
MASK_EXT_OUT = ".png" # saída das máscaras sempre PNG MASK_EXT_OUT = ".png"
FORCAR_SOBRESCRITA_NEW_MASK = False # sobrescrever máscara em new_masks se já existir FORCAR_SOBRESCRITA_NEW_MASK = False
VALIDAR_DIM_MASK_EXISTENTE = True # se existir, validar dimensões VALIDAR_DIM_MASK_EXISTENTE = True
MANIFESTO = "manifest.csv" MANIFESTO = "manifest.csv"
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
# ===================================================
# Utilidades gerais
# =================================================== # ===================================================
def garantir_pasta(p): def garantir_pasta(p: str):
os.makedirs(p, exist_ok=True) os.makedirs(p, exist_ok=True)
def criar_mask_solida(dim_h, dim_w, cor_rgb):
def criar_mask_solida(dim_h: int, dim_w: int, cor_rgb: Tuple[int, int, int]):
r, g, b = cor_rgb r, g, b = cor_rgb
mask_bgr = np.zeros((dim_h, dim_w, 3), dtype=np.uint8) mask_bgr = np.zeros((dim_h, dim_w, 3), dtype=np.uint8)
mask_bgr[:] = (b, g, r) # OpenCV usa BGR mask_bgr[:] = (b, g, r)
return mask_bgr return mask_bgr
def caminho_mask_new_para_img(pasta_new_masks, caminho_img):
base = os.path.splitext(os.path.basename(caminho_img))[0] def caminho_mask_new_para_base(pasta_new_masks: str, base: str):
return os.path.join(pasta_new_masks, base + MASK_EXT_OUT) return os.path.join(pasta_new_masks, base + MASK_EXT_OUT)
def ler_dim(caminho_img):
ext = os.path.splitext(caminho_img)[1].lower()
# Caso seja RAW: usamos dimensões fixas def ler_dim(caminho_img: str):
if ext in RAW_EXTS:
h, w = RAW_HEIGHT, RAW_WIDTH
return (h, w), None
# Caso normal: JPG/PNG/etc
img = cv2.imread(caminho_img, cv2.IMREAD_COLOR) img = cv2.imread(caminho_img, cv2.IMREAD_COLOR)
if img is None: if img is None:
raise RuntimeError(f"Erro ao abrir: {caminho_img}") raise RuntimeError(f"Erro ao abrir: {caminho_img}")
h, w = img.shape[:2] h, w = img.shape[:2]
return (h, w), img return (h, w), img
def salvar_mask_solidaria(pasta_new_masks, caminho_img, cor_rgb, forcar=False, validar_dim=True):
(h, w), _ = ler_dim(caminho_img) def salvar_mask_solidaria_por_base(pasta_new_masks: str, base: str, caminho_preview: str,
caminho_mask = caminho_mask_new_para_img(pasta_new_masks, caminho_img) cor_rgb: Tuple[int, int, int], forcar: bool = False,
validar_dim: bool = True):
(h, w), _ = ler_dim(caminho_preview)
caminho_mask = caminho_mask_new_para_base(pasta_new_masks, base)
if os.path.exists(caminho_mask) and not forcar: if os.path.exists(caminho_mask) and not forcar:
if validar_dim: if validar_dim:
@ -90,76 +131,130 @@ def salvar_mask_solidaria(pasta_new_masks, caminho_img, cor_rgb, forcar=False, v
print(f"[CRIADA] {os.path.basename(caminho_mask)} ({w}x{h})") print(f"[CRIADA] {os.path.basename(caminho_mask)} ({w}x{h})")
return caminho_mask return caminho_mask
# ========= NOVO: suporte a cópia tripla (img + mask + raw) =========
def _base_esta_ocupado(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw): # ===================================================
""" # Descoberta de arquivos da amostra
Verifica se algum dos arquivos (img/mask/raw) com esse base existe. # ===================================================
"""
if os.path.exists(os.path.join(dir_img, base + img_ext)): def localizar_meta_correspondente(pasta_metas: str, base: str) -> Optional[str]:
candidate = os.path.join(pasta_metas, base + ".json")
return candidate if os.path.exists(candidate) else None
def localizar_bins_correspondentes(pasta_bins: str, base: str) -> List[str]:
encontrados = []
if not os.path.isdir(pasta_bins):
return encontrados
for nome in os.listdir(pasta_bins):
m = BIN_RE.match(nome)
if not m:
continue
if m.group("base") != base:
continue
cam = int(m.group("cam"))
encontrados.append((cam, os.path.join(pasta_bins, nome)))
encontrados.sort(key=lambda x: x[0])
return [p for _, p in encontrados]
# ===================================================
# Deduplicação para cópia ao dataset final
# ===================================================
def _base_esta_ocupado_multispec(base: str, preview_ext: str, dir_preview: str,
dir_mask: str, dir_meta: str, dir_bins: str,
bin_src_paths: List[str]) -> bool:
if os.path.exists(os.path.join(dir_preview, base + preview_ext)):
return True return True
if os.path.exists(os.path.join(dir_mask, base + mask_ext)): if os.path.exists(os.path.join(dir_mask, base + ".png")):
return True return True
if raw_ext and dir_raw: if os.path.exists(os.path.join(dir_meta, base + ".json")):
if os.path.exists(os.path.join(dir_raw, base + raw_ext)):
return True return True
for src_bin in bin_src_paths:
nome = os.path.basename(src_bin)
m = BIN_RE.match(nome)
if not m:
continue
cam = m.group("cam")
bin_dest = os.path.join(dir_bins, f"{base}_cam{cam}.bin")
if os.path.exists(bin_dest):
return True
return False return False
def gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
""" def gerar_base_disponivel_multispec(base: str, preview_ext: str, dir_preview: str,
Gera um 'base' comum livre para img/mask/raw. dir_mask: str, dir_meta: str, dir_bins: str,
""" bin_src_paths: List[str]) -> str:
if not _base_esta_ocupado(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw): if not _base_esta_ocupado_multispec(base, preview_ext, dir_preview, dir_mask, dir_meta, dir_bins, bin_src_paths):
return base return base
i = 1 i = 1
while True: while True:
candidate = f"{base}_{i:03d}" candidate = f"{base}_{i:03d}"
if not _base_esta_ocupado(candidate, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw): if not _base_esta_ocupado_multispec(candidate, preview_ext, dir_preview, dir_mask, dir_meta, dir_bins, bin_src_paths):
return candidate return candidate
i += 1 i += 1
def copiar_triplo(caminho_img_src, caminho_mask_src, caminho_raw_src, dest_img_dir, dest_mask_dir, dest_raw_dir):
"""
Copia imagem, máscara e (se existir) o RAW correspondente
mantendo o mesmo base (com renome em caso de conflito).
Retorna (dst_img_path, dst_mask_path, dst_raw_path_ou_None). def copiar_conjunto_multispec(caminho_preview_src: str, caminho_mask_src: str,
""" caminho_meta_src: str, caminhos_bins_src: List[str],
garantir_pasta(dest_img_dir) dest_preview_dir: str, dest_mask_dir: str,
dest_meta_dir: str, dest_bins_dir: str):
garantir_pasta(dest_preview_dir)
garantir_pasta(dest_mask_dir) garantir_pasta(dest_mask_dir)
if dest_raw_dir: garantir_pasta(dest_meta_dir)
garantir_pasta(dest_raw_dir) garantir_pasta(dest_bins_dir)
base = os.path.splitext(os.path.basename(caminho_img_src))[0] base = os.path.splitext(os.path.basename(caminho_preview_src))[0]
preview_ext = os.path.splitext(caminho_preview_src)[1].lower()
img_ext = os.path.splitext(caminho_img_src)[1].lower() new_base = gerar_base_disponivel_multispec(
mask_ext = ".png" base,
raw_ext = os.path.splitext(caminho_raw_src)[1].lower() if caminho_raw_src else None preview_ext,
dest_preview_dir,
dest_mask_dir,
dest_meta_dir,
dest_bins_dir,
caminhos_bins_src,
)
new_base = gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dest_img_dir, dest_mask_dir, dest_raw_dir) dst_preview_path = os.path.join(dest_preview_dir, new_base + preview_ext)
dst_mask_path = os.path.join(dest_mask_dir, new_base + ".png")
dst_meta_path = os.path.join(dest_meta_dir, new_base + ".json")
dst_img_path = os.path.join(dest_img_dir, new_base + img_ext) shutil.copy2(caminho_preview_src, dst_preview_path)
dst_mask_path = os.path.join(dest_mask_dir, new_base + mask_ext)
dst_raw_path = os.path.join(dest_raw_dir, new_base + raw_ext) if (caminho_raw_src and dest_raw_dir and raw_ext) else None
shutil.copy2(caminho_img_src, dst_img_path)
shutil.copy2(caminho_mask_src, dst_mask_path) shutil.copy2(caminho_mask_src, dst_mask_path)
if caminho_raw_src and dst_raw_path: shutil.copy2(caminho_meta_src, dst_meta_path)
shutil.copy2(caminho_raw_src, dst_raw_path)
print(f"[COPIADO] {os.path.basename(dst_img_path)} | {os.path.basename(dst_mask_path)} | {os.path.basename(dst_raw_path)}")
else:
print(f"[COPIADO] {os.path.basename(dst_img_path)} | {os.path.basename(dst_mask_path)} (sem RAW)")
return dst_img_path, dst_mask_path, dst_raw_path dst_bins_paths = []
for src_bin in caminhos_bins_src:
nome = os.path.basename(src_bin)
m = BIN_RE.match(nome)
if not m:
continue
cam = m.group("cam")
dst_bin = os.path.join(dest_bins_dir, f"{new_base}_cam{cam}.bin")
shutil.copy2(src_bin, dst_bin)
dst_bins_paths.append(dst_bin)
print(
f"[COPIADO] {os.path.basename(dst_preview_path)} | "
f"{os.path.basename(dst_mask_path)} | "
f"{os.path.basename(dst_meta_path)} | bins={len(dst_bins_paths)}"
)
return dst_preview_path, dst_mask_path, dst_meta_path, dst_bins_paths
# ===================================================
# Labelmap
# =================================================== # ===================================================
def obter_cor_da_classe(nome_classe: str, caminho_labelmap: str): def obter_cor_da_classe(nome_classe: str, caminho_labelmap: str):
"""
o labelmap.txt e retorna o RGB correspondente à classe informada.
Usa carregar_labelmap_completo(utils.py).
"""
mapa_rgb, colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(caminho_labelmap) mapa_rgb, colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(caminho_labelmap)
alvo = nome_classe.strip().lower() alvo = nome_classe.strip().lower()
@ -186,31 +281,36 @@ def obter_cor_da_classe(nome_classe: str, caminho_labelmap: str):
print(f"[LABELMAP] Classe '{nome_classe}' -> ID={target_id} -> cor RGB={cor_rgb}") print(f"[LABELMAP] Classe '{nome_classe}' -> ID={target_id} -> cor RGB={cor_rgb}")
return cor_rgb return cor_rgb
def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None:
"""
Dado o nome do preview (ex: 20260122_091833_277.png),
tenta achar o RAW correspondente em PASTA_NEW_RAWS:
20260122_091833_277.raw (ou outras extensões de RAW_EXTS)
"""
base = os.path.splitext(nome_img)[0]
for ext in RAW_EXTS:
candidate = os.path.join(pasta_new_raws, base + ext)
if os.path.exists(candidate):
return candidate
return None
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False): # ===================================================
# Processamento principal
# ===================================================
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb,
fazer_copia_final=True, manifesto_csv=None,
orignais=False):
source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original") source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original")
if not orignais:
pasta_new_previews = os.path.join(source, "group", grupo, "previews") pasta_new_previews = os.path.join(source, "group", grupo, "previews")
pasta_new_masks = os.path.join(source, "group", grupo, "masks") pasta_new_masks = os.path.join(source, "group", grupo, "masks")
pasta_new_raws = os.path.join(source, "group", grupo, "raws") pasta_new_metas = os.path.join(source, "group", grupo, "metas")
pasta_new_bins = os.path.join(source, "group", grupo, "bins")
else:
pasta_new_previews = os.path.join(source, "previews")
pasta_new_masks = os.path.join(source, "masks")
pasta_new_metas = os.path.join(source, "metas")
pasta_new_bins = os.path.join(source, "bins")
garantir_pasta(pasta_new_previews) garantir_pasta(pasta_new_previews)
garantir_pasta(pasta_new_masks) garantir_pasta(pasta_new_masks)
garantir_pasta(pasta_new_raws) garantir_pasta(pasta_new_metas)
garantir_pasta(pasta_new_bins)
garantir_pasta(PASTA_FINAL_PREVIEWS) garantir_pasta(PASTA_FINAL_PREVIEWS)
garantir_pasta(PASTA_FINAL_MASKS) garantir_pasta(PASTA_FINAL_MASKS)
garantir_pasta(PASTA_FINAL_RAWS) garantir_pasta(PASTA_FINAL_METAS)
garantir_pasta(PASTA_FINAL_BINS)
registros = [] registros = []
total, criadas_mask, copiados, puladas, erros = 0, 0, 0, 0, 0 total, criadas_mask, copiados, puladas, erros = 0, 0, 0, 0, 0
@ -218,42 +318,53 @@ def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_fi
for nome in os.listdir(pasta_new_previews): for nome in os.listdir(pasta_new_previews):
if not nome.lower().endswith(EXT_IMAGENS): if not nome.lower().endswith(EXT_IMAGENS):
continue continue
total += 1 total += 1
caminho_img = os.path.join(pasta_new_previews, nome) caminho_preview = os.path.join(pasta_new_previews, nome)
base = os.path.splitext(nome)[0]
try: try:
# 1) criar máscara sólida em new_masks caminho_meta = localizar_meta_correspondente(pasta_new_metas, base)
antes = os.path.exists(caminho_mask_new_para_img(pasta_new_masks, caminho_img)) if caminho_meta is None:
caminho_mask_new = salvar_mask_solidaria( raise RuntimeError(f"Meta não encontrado para {nome} em {pasta_new_metas}")
caminhos_bins = localizar_bins_correspondentes(pasta_new_bins, base)
if not caminhos_bins:
raise RuntimeError(f"Nenhum bin encontrado para {nome} em {pasta_new_bins}")
antes = os.path.exists(caminho_mask_new_para_base(pasta_new_masks, base))
caminho_mask_new = salvar_mask_solidaria_por_base(
pasta_new_masks, pasta_new_masks,
caminho_img, base,
caminho_preview,
cor_classe_rgb, cor_classe_rgb,
forcar=FORCAR_SOBRESCRITA_NEW_MASK, forcar=FORCAR_SOBRESCRITA_NEW_MASK,
validar_dim=VALIDAR_DIM_MASK_EXISTENTE validar_dim=VALIDAR_DIM_MASK_EXISTENTE,
) )
if caminho_mask_new and not antes: if caminho_mask_new and not antes:
criadas_mask += 1 criadas_mask += 1
# 2) copiar imagem + máscara + RAW para as pastas finais (com renome se necessário)
if fazer_copia_final: if fazer_copia_final:
raw_src = localizar_raw_correspondente(pasta_new_raws, nome) dst_preview, dst_mask, dst_meta, dst_bins = copiar_conjunto_multispec(
if raw_src is None: caminho_preview,
print(f"[AVISO] RAW não encontrado para {nome} em {pasta_new_raws}, copiando só preview+mask.")
dst_img, dst_mask, dst_raw = copiar_triplo(
caminho_img,
caminho_mask_new, caminho_mask_new,
raw_src, caminho_meta,
caminhos_bins,
PASTA_FINAL_PREVIEWS, PASTA_FINAL_PREVIEWS,
PASTA_FINAL_MASKS, PASTA_FINAL_MASKS,
PASTA_FINAL_RAWS, PASTA_FINAL_METAS,
PASTA_FINAL_BINS,
) )
copiados += 1 copiados += 1
registros.append([ registros.append([
caminho_img, caminho_preview,
caminho_mask_new, caminho_mask_new,
raw_src or "", caminho_meta,
dst_img, json.dumps(caminhos_bins, ensure_ascii=False),
dst_preview,
dst_mask, dst_mask,
dst_raw or "", dst_meta,
json.dumps(dst_bins, ensure_ascii=False),
]) ])
else: else:
puladas += 1 puladas += 1
@ -266,44 +377,51 @@ def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_fi
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f: with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f) w = csv.writer(f)
w.writerow([ w.writerow([
"src_image", "src_preview",
"src_mask", "src_mask",
"src_raw", "src_meta",
"dst_image", "src_bins_json",
"dst_preview",
"dst_mask", "dst_mask",
"dst_raw", "dst_meta",
"dst_bins_json",
]) ])
w.writerows(registros) w.writerows(registros)
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).") print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
print(f"\nResumo: total_imgs={total} | masks_criadas={criadas_mask} | copiados={copiados} | puladas={puladas} | erros={erros}") print(
f"\nResumo: total_imgs={total} | masks_criadas={criadas_mask} | "
f"copiados={copiados} | puladas={puladas} | erros={erros}"
)
# ===================================================
# CLI
# ===================================================
def build_cli(): def build_cli():
ap = argparse.ArgumentParser( ap = argparse.ArgumentParser(
description="Gera máscaras sólidas para novas imagens de UMA classe (via labelmap) e copia para dataset final com dedup." description="Gera máscaras sólidas para grupos de classe única no novo formato multiespectral."
) )
ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.") ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.") ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
ap.add_argument("--classe", required=True, help="Nome da classe (como está no labelmap.txt). Ex: chao, cana, erva") ap.add_argument("--classe", required=True, help="Nome da classe como está no labelmap.txt. Ex: chao, cana, erva")
ap.add_argument("--grupo", required=True, help="Nome da pasta grupo. Ex: chao, cana, erva, chao_cana, chao_cana_erva") ap.add_argument("--grupo", required=True, help="Nome da pasta grupo. Ex: chao, cana, erva")
ap.add_argument("--no-copy", action="store_true", help="Não copia para as pastas finais (só cria masks em new_masks).") ap.add_argument("--no-copy", action="store_true", help="Não copia para as pastas finais, só cria masks no grupo.")
ap.add_argument("--from-originals", action="store_true", help="Faz o procedimento na pasta em originais") ap.add_argument("--from-originals", action="store_true", help="Executa na pasta dataset/original em vez de dataset/brutas/.../group/... .")
ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a gerar (ou vazio para não gerar).") ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a gerar.")
return ap return ap
if __name__ == "__main__": if __name__ == "__main__":
ap = build_cli() ap = build_cli()
args = ap.parse_args() args = ap.parse_args()
# Caminho padrão pro labelmap, caso não seja informado
labelmap_path = os.path.join("dataset", "labelmap.txt") labelmap_path = os.path.join("dataset", "labelmap.txt")
if not os.path.exists(labelmap_path): if not os.path.exists(labelmap_path):
raise SystemExit(f"Labelmap não encontrado em: {labelmap_path}") raise SystemExit(f"Labelmap não encontrado em: {labelmap_path}")
# Descobre a cor da classe no labelmap
cor_rgb = obter_cor_da_classe(args.classe, labelmap_path) cor_rgb = obter_cor_da_classe(args.classe, labelmap_path)
fazer_copia = not args.no_copy fazer_copia = not args.no_copy
manifesto_csv = args.manifest if args.manifest else None manifesto_csv = args.manifest if args.manifest else None
@ -314,5 +432,5 @@ if __name__ == "__main__":
cor_classe_rgb=cor_rgb, cor_classe_rgb=cor_rgb,
fazer_copia_final=fazer_copia, fazer_copia_final=fazer_copia,
manifesto_csv=manifesto_csv, manifesto_csv=manifesto_csv,
orignais=args.from_originals orignais=args.from_originals,
) )

View File

@ -3,10 +3,11 @@ import os
import shutil import shutil
import argparse import argparse
import csv import csv
import re
# ================= CONFIG ================= # ================= CONFIG =================
with open("config.json", "r") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config["camera"] MODELO = config["camera"]
@ -19,58 +20,20 @@ PASTA_NEW_MASKS = os.path.join("dataset", "new_masks")
# Destino final do dataset consolidado # Destino final do dataset consolidado
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews") PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws") PASTA_FINAL_METAS = os.path.join("dataset", "original", "metas")
PASTA_FINAL_BINS = os.path.join("dataset", "original", "bins")
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks") PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
EXT_PREVIEWS = (".png", ".jpg", ".jpeg") EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
EXT_MASKS = (".png",) EXT_MASKS = (".png",)
EXT_RAWS = (".raw",) EXT_METAS = (".json",)
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
# ========================================= # =========================================
def garantir_pasta(p: str): def garantir_pasta(p: str):
os.makedirs(p, exist_ok=True) os.makedirs(p, exist_ok=True)
def indexar_brutas_group(root_group: str):
"""
Varrre dataset/brutas/group recursivamente e monta dois dicionários:
base -> caminho_preview
base -> caminho_raw
Considera que a estrutura é:
brutas/group/<grupo>/previews/*.png|jpg
brutas/group/<grupo>/raws/*.raw
"""
previews_map = {}
raws_map = {}
if not os.path.isdir(root_group):
print(f"[AVISO] Pasta de grupos não existe: {root_group}")
return previews_map, raws_map
for dirpath, dirnames, filenames in os.walk(root_group):
base_dir = os.path.basename(dirpath).lower()
if base_dir == "previews":
for nome in filenames:
if not nome.lower().endswith(EXT_PREVIEWS):
continue
base = os.path.splitext(nome)[0]
# Se já existir, mantemos o primeiro encontrado (pode logar conflito se quiser)
if base not in previews_map:
previews_map[base] = os.path.join(dirpath, nome)
elif base_dir == "raws":
for nome in filenames:
if not nome.lower().endswith(EXT_RAWS):
continue
base = os.path.splitext(nome)[0]
if base not in raws_map:
raws_map[base] = os.path.join(dirpath, nome)
print(f"[INDEX] Previews indexados: {len(previews_map)}")
print(f"[INDEX] Raws indexados : {len(raws_map)}")
return previews_map, raws_map
def nome_disponivel(dest_dir: str, base: str, ext: str) -> str: def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
""" """
@ -88,16 +51,90 @@ def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
return p return p
i += 1 i += 1
def buscar_bins_por_base(pasta_bins: str, base: str):
encontrados = []
if not os.path.isdir(pasta_bins):
return encontrados
for nome in os.listdir(pasta_bins):
m = BIN_RE.match(nome)
if not m:
continue
if m.group("base").lower() != base.lower():
continue
cam = int(m.group("cam"))
encontrados.append((cam, os.path.join(pasta_bins, nome)))
encontrados.sort(key=lambda x: x[0])
return [p for _, p in encontrados]
def indexar_brutas_group(root_group: str):
"""
Varre dataset/brutas/<cana>/<horario>/group recursivamente e monta dois índices:
base -> caminho_preview
base -> {meta, bins}
Considera a estrutura:
brutas/cana_x/horario/group/<grupo>/previews/*.png|jpg
brutas/cana_x/horario/group/<grupo>/metas/*.json
brutas/cana_x/horario/group/<grupo>/bins/*_camX.bin
"""
previews_map = {}
amostras_map = {}
if not os.path.isdir(root_group):
print(f"[AVISO] Pasta de grupos não existe: {root_group}")
return previews_map, amostras_map
for dirpath, dirnames, filenames in os.walk(root_group):
base_dir = os.path.basename(dirpath).lower()
parent_group_dir = os.path.dirname(dirpath)
if base_dir == "previews":
for nome in filenames:
if not nome.lower().endswith(EXT_PREVIEWS):
continue
base = os.path.splitext(nome)[0]
if base not in previews_map:
previews_map[base] = os.path.join(dirpath, nome)
elif base_dir == "metas":
bins_dir = os.path.join(parent_group_dir, "bins")
for nome in filenames:
if not nome.lower().endswith(EXT_METAS):
continue
base = os.path.splitext(nome)[0]
meta_path = os.path.join(dirpath, nome)
bin_paths = buscar_bins_por_base(bins_dir, base)
if not bin_paths:
print(f"[SKIP] Meta sem bins correspondentes: {meta_path}")
continue
if base not in amostras_map:
amostras_map[base] = {
"meta": meta_path,
"bins": bin_paths,
}
print(f"[INDEX] Previews indexados : {len(previews_map)}")
print(f"[INDEX] Amostras indexadas : {len(amostras_map)}")
return previews_map, amostras_map
def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None): def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
garantir_pasta(PASTA_FINAL_PREVIEWS) garantir_pasta(PASTA_FINAL_PREVIEWS)
garantir_pasta(PASTA_FINAL_RAWS) garantir_pasta(PASTA_FINAL_METAS)
garantir_pasta(PASTA_FINAL_BINS)
garantir_pasta(PASTA_FINAL_MASKS) garantir_pasta(PASTA_FINAL_MASKS)
if not os.path.isdir(PASTA_NEW_MASKS): if not os.path.isdir(PASTA_NEW_MASKS):
raise SystemExit(f"[ERRO] Pasta new_masks não existe: {PASTA_NEW_MASKS}") raise SystemExit(f"[ERRO] Pasta new_masks não existe: {PASTA_NEW_MASKS}")
# 1) Indexa tudo que existe em brutas/group previews_map, amostras_map = indexar_brutas_group(PASTA_BRUTAS_GROUP_ROOT)
previews_map, raws_map = indexar_brutas_group(PASTA_BRUTAS_GROUP_ROOT)
registros = [] registros = []
total, copiados, ignorados = 0, 0, 0 total, copiados, ignorados = 0, 0, 0
@ -111,33 +148,45 @@ def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
caminho_mask_src = os.path.join(PASTA_NEW_MASKS, nome) caminho_mask_src = os.path.join(PASTA_NEW_MASKS, nome)
preview_src = previews_map.get(base) preview_src = previews_map.get(base)
raw_src = raws_map.get(base) amostra = amostras_map.get(base)
if preview_src is None or raw_src is None: if preview_src is None or amostra is None:
ignorados += 1 ignorados += 1
motivo = [] motivo = []
if preview_src is None: if preview_src is None:
motivo.append("preview") motivo.append("preview")
if raw_src is None: if amostra is None:
motivo.append("raw") motivo.append("meta/bins")
print(f"[SKIP] {base} -> faltando: {', '.join(motivo)}") print(f"[SKIP] {base} -> faltando: {', '.join(motivo)}")
# NÃO move a máscara, ela fica em new_masks pra você analisar depois
continue continue
meta_src = amostra["meta"]
bins_src = amostra["bins"]
prev_ext = os.path.splitext(preview_src)[1].lower() prev_ext = os.path.splitext(preview_src)[1].lower()
mask_ext = os.path.splitext(nome)[1].lower() mask_ext = os.path.splitext(nome)[1].lower()
raw_ext = os.path.splitext(raw_src)[1].lower() meta_ext = os.path.splitext(meta_src)[1].lower()
# Gera um base final único usando o preview como referência # Gera um base final único usando o preview como referência
dst_preview = nome_disponivel(PASTA_FINAL_PREVIEWS, base, prev_ext) dst_preview = nome_disponivel(PASTA_FINAL_PREVIEWS, base, prev_ext)
new_base = os.path.splitext(os.path.basename(dst_preview))[0] new_base = os.path.splitext(os.path.basename(dst_preview))[0]
dst_raw = os.path.join(PASTA_FINAL_RAWS, new_base + raw_ext) dst_meta = os.path.join(PASTA_FINAL_METAS, new_base + meta_ext)
dst_mask = os.path.join(PASTA_FINAL_MASKS, new_base + mask_ext) dst_mask = os.path.join(PASTA_FINAL_MASKS, new_base + mask_ext)
# Copia preview + raw, move ou copia a máscara dst_bins = []
for src_bin in bins_src:
m = BIN_RE.match(os.path.basename(src_bin))
if not m:
continue
cam = m.group("cam")
dst_bins.append(os.path.join(PASTA_FINAL_BINS, f"{new_base}_cam{cam}.bin"))
shutil.copy2(preview_src, dst_preview) shutil.copy2(preview_src, dst_preview)
shutil.copy2(raw_src, dst_raw) shutil.copy2(meta_src, dst_meta)
for src_bin, dst_bin in zip(bins_src, dst_bins):
shutil.copy2(src_bin, dst_bin)
if mover: if mover:
shutil.move(caminho_mask_src, dst_mask) shutil.move(caminho_mask_src, dst_mask)
else: else:
@ -146,24 +195,28 @@ def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
copiados += 1 copiados += 1
registros.append([ registros.append([
preview_src, preview_src,
raw_src, meta_src,
json.dumps(bins_src, ensure_ascii=False),
caminho_mask_src, caminho_mask_src,
dst_preview, dst_preview,
dst_raw, dst_meta,
json.dumps(dst_bins, ensure_ascii=False),
dst_mask, dst_mask,
]) ])
print(f"[OK] {new_base}") print(f"[OK] {new_base} | bins={len(dst_bins)}")
if manifesto_csv and registros: if manifesto_csv and registros:
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f: with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f) w = csv.writer(f)
w.writerow([ w.writerow([
"src_preview", "src_preview",
"src_raw", "src_meta",
"src_bins_json",
"src_mask", "src_mask",
"dst_preview", "dst_preview",
"dst_raw", "dst_meta",
"dst_bins_json",
"dst_mask", "dst_mask",
]) ])
w.writerows(registros) w.writerows(registros)
@ -171,13 +224,14 @@ def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
print(f"\nResumo: total_masks={total} | ingestas={copiados} | ignoradas={ignorados}") print(f"\nResumo: total_masks={total} | ingestas={copiados} | ignoradas={ignorados}")
# ================= CLI ================= # ================= CLI =================
def build_cli(): def build_cli():
ap = argparse.ArgumentParser( ap = argparse.ArgumentParser(
description=( description=(
"Ingere máscaras novas (dataset/new_masks) procurando previews e raws " "Ingere máscaras novas (dataset/new_masks) procurando previews, metas e bins "
"nas subpastas de dataset/brutas/group/** e copiando tudo para dataset/original." "nas subpastas de dataset/brutas/.../group/** e copiando tudo para dataset/original."
) )
) )
ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.") ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
@ -186,6 +240,7 @@ def build_cli():
ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).") ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).")
return ap return ap
if __name__ == "__main__": if __name__ == "__main__":
ap = build_cli() ap = build_cli()
args = ap.parse_args() args = ap.parse_args()

View File

@ -2,32 +2,41 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
""" """
Versão OTIMIZADA Agrupamento ultra-rápido por classes presentes na máscara. Agrupa automaticamente as amostras do dataset/original por classes presentes na máscara.
Substitui extrair_ids_presentes por uma versão 50× mais rápida.
Entrada: Entrada:
originals/ originals/
previews/ previews/
raws/ metas/
bins/
masks/ masks/
(opcional) masks2/ (opcional) masks2/
Saída: Saída:
originals/group/<grupo>/previews originals/group/<grupo>/previews
/raws /metas
/bins
/masks /masks
/masks2 /masks2
Cada amostra é:
<base>.png (preview)
<base>.json (meta)
<base>_cam0.bin, <base>_cam1.bin, ... (bins)
<base>.png (mask)
""" """
import os import os
import sys
import cv2 import cv2
import csv import csv
import json import json
import shutil import shutil
import argparse import argparse
import numpy as np import numpy as np
import re
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:
@ -38,13 +47,14 @@ USE_MASKS2 = config["dual_head"]
EXT_PREVIEWS = (".jpg", ".jpeg", ".png") EXT_PREVIEWS = (".jpg", ".jpeg", ".png")
EXT_MASKS = (".png", ".jpg", ".jpeg") EXT_MASKS = (".png", ".jpg", ".jpeg")
EXT_MASKS2 = (".png", ".jpg", ".jpeg") EXT_MASKS2 = (".png", ".jpg", ".jpeg")
EXT_RAW = (".raw",) EXT_METAS = (".json",)
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
MANIFESTO_DEFAULT = "manifest.csv" MANIFESTO_DEFAULT = "manifest.csv"
# ============================================================ # ============================================================
# 🔥 EXTRAÇÃO DE CLASSES — VERSÃO ULTRA OTIMIZADA # EXTRAÇÃO DE CLASSES DA MÁSCARA
# ============================================================ # ============================================================
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True): def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
@ -52,42 +62,30 @@ def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
Extração ultra-rápida de IDs presentes na máscara. Extração ultra-rápida de IDs presentes na máscara.
Estratégia: Estratégia:
1) Máscara 1 canal np.unique (instantâneo). 1) Máscara 1 canal -> np.unique
2) Máscara 3 canais: 2) Máscara 3 canais:
- se assume_rgb=True: labelmap está em RGB, - se assume_rgb=True: converte BGR -> RGB
convertemos a imagem BGR -> RGB para casar com cor_para_id. - senão usa BGR direto
- se assume_rgb=False: labelmap está em BGR, 3) Amostragem em grid + converter_mask_rgb_para_ids
mantemos a imagem em BGR. 4) Fallback full-scan
3) Amostragem em grid + converter_mask_rgb_para_ids.
4) Fallback full-scan se necessário.
""" """
m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if m is None: if m is None:
raise RuntimeError(f"Falha ao abrir máscara: {mask_path}") raise RuntimeError(f"Falha ao abrir máscara: {mask_path}")
# -------------------------------------------------------
# CASO 1: máscara indexada (1 canal) — instantâneo
# -------------------------------------------------------
if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1): if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1):
ids = np.unique(m) ids = np.unique(m)
return set(int(v) for v in ids) return set(int(v) for v in ids)
# -------------------------------------------------------
# CASO 2: máscara RGB
# -------------------------------------------------------
if assume_rgb: if assume_rgb:
# Labelmap em RGB, OpenCV em BGR -> converte
img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB) img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB)
else: else:
# Labelmap em BGR, OpenCV já em BGR -> usa direto
img = m img = m
# Agora cor_para_id e img estão no MESMO espaço de cor
mapa_rgb = cor_para_id mapa_rgb = cor_para_id
max_classes = len(cor_para_id) max_classes = len(cor_para_id)
# ---------- AMOSTRAGEM ---------- step = 8
step = 8 # pode virar 4 se quiser mais precisão, 16 se quiser mais velocidade
amostra = img[::step, ::step] 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, mapa_rgb, ignore_id=255)
ids = set(int(x) for x in np.unique(amostra_ids) if x != 255) ids = set(int(x) for x in np.unique(amostra_ids) if x != 255)
@ -95,19 +93,19 @@ def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
if len(ids) >= max_classes: if len(ids) >= max_classes:
return ids return ids
# ---------- FULL-SCAN OTIMIZADO (último caso) ----------
full_ids = converter_mask_rgb_para_ids(img, mapa_rgb, ignore_id=255) full_ids = converter_mask_rgb_para_ids(img, mapa_rgb, ignore_id=255)
ids = set(int(x) for x in np.unique(full_ids) if x != 255) ids = set(int(x) for x in np.unique(full_ids) if x != 255)
return ids return ids
# ============================================================ # ============================================================
# 🔧 Helpers # HELPERS
# ============================================================ # ============================================================
def garantir_pasta(p): def garantir_pasta(p):
os.makedirs(p, exist_ok=True) os.makedirs(p, exist_ok=True)
def nome_disponivel(dest_dir, base_name, ext): def nome_disponivel(dest_dir, base_name, ext):
cand = os.path.join(dest_dir, base_name + ext) cand = os.path.join(dest_dir, base_name + ext)
if not os.path.exists(cand): if not os.path.exists(cand):
@ -119,6 +117,7 @@ def nome_disponivel(dest_dir, base_name, ext):
return cand return cand
i += 1 i += 1
def mapear_por_base_priorizando_png(pasta, exts): def mapear_por_base_priorizando_png(pasta, exts):
if not pasta or not os.path.isdir(pasta): if not pasta or not os.path.isdir(pasta):
return {} return {}
@ -137,6 +136,42 @@ def mapear_por_base_priorizando_png(pasta, exts):
mapa[base] = cam mapa[base] = cam
return mapa return mapa
def mapear_metas_por_base(pasta, exts):
if not pasta or not os.path.isdir(pasta):
return {}
return {
os.path.splitext(nome)[0]: os.path.join(pasta, nome)
for nome in os.listdir(pasta)
if nome.lower().endswith(exts)
}
def mapear_bins_por_base(pasta_bins):
"""
Retorna:
{
base: [path_cam0, path_cam1, ...]
}
"""
mapa = {}
if not pasta_bins or not os.path.isdir(pasta_bins):
return mapa
for nome in os.listdir(pasta_bins):
m = BIN_RE.match(nome)
if not m:
continue
base = m.group("base")
cam = int(m.group("cam"))
mapa.setdefault(base, []).append((cam, os.path.join(pasta_bins, nome)))
for base in list(mapa.keys()):
mapa[base] = [p for _, p in sorted(mapa[base], key=lambda x: x[0])]
return mapa
def inferir_ignore_id(ignore_rgb, cor_para_id): def inferir_ignore_id(ignore_rgb, cor_para_id):
if ignore_rgb is None: if ignore_rgb is None:
return None return None
@ -146,11 +181,8 @@ def inferir_ignore_id(ignore_rgb, cor_para_id):
return ignore_rgb return ignore_rgb
return None return None
def montar_nome_grupo(ids_presentes, id_para_nome): 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
"""
if not ids_presentes: if not ids_presentes:
return "sem_classe" return "sem_classe"
nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)] nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)]
@ -158,54 +190,64 @@ def montar_nome_grupo(ids_presentes, id_para_nome):
# ============================================================ # ============================================================
# 📦 PIPELINE PRINCIPAL # CÓPIA / MOVIMENTAÇÃO DO CONJUNTO MULTISPEC
# ============================================================ # ============================================================
def copiar_ou_mover_tripla(preview_src, raw_src, mask_src, def copiar_ou_mover_conjunto(preview_src, meta_src, bins_src, mask_src,
dest_prev_dir, dest_raw_dir, dest_mask_dir, dest_prev_dir, dest_meta_dir, dest_bins_dir, dest_mask_dir,
mover=False, mover=False,
mask2_src=None, dest_mask2_dir=None): mask2_src=None, dest_mask2_dir=None):
garantir_pasta(dest_prev_dir) garantir_pasta(dest_prev_dir)
garantir_pasta(dest_raw_dir) garantir_pasta(dest_meta_dir)
garantir_pasta(dest_bins_dir)
garantir_pasta(dest_mask_dir) garantir_pasta(dest_mask_dir)
base_prev = os.path.splitext(os.path.basename(preview_src))[0] base_prev = os.path.splitext(os.path.basename(preview_src))[0]
prev_ext = os.path.splitext(preview_src)[1].lower() prev_ext = os.path.splitext(preview_src)[1].lower()
raw_ext = os.path.splitext(raw_src)[1].lower() meta_ext = os.path.splitext(meta_src)[1].lower()
mask_ext = os.path.splitext(mask_src)[1].lower() mask_ext = os.path.splitext(mask_src)[1].lower()
dst_prev = nome_disponivel(dest_prev_dir, base_prev, prev_ext) dst_prev = nome_disponivel(dest_prev_dir, base_prev, prev_ext)
new_base = os.path.splitext(os.path.basename(dst_prev))[0] new_base = os.path.splitext(os.path.basename(dst_prev))[0]
dst_raw = os.path.join(dest_raw_dir, new_base + raw_ext) dst_meta = os.path.join(dest_meta_dir, new_base + meta_ext)
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext) dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
if os.path.exists(dst_raw) or os.path.exists(dst_mask): dst_bins = []
dst_mask = nome_disponivel(dest_mask_dir, new_base, mask_ext) for src_bin in bins_src:
new_base = os.path.splitext(os.path.basename(dst_mask))[0] m = BIN_RE.match(os.path.basename(src_bin))
dst_prev = nome_disponivel(dest_prev_dir, new_base, prev_ext) if not m:
dst_raw = nome_disponivel(dest_raw_dir, new_base, raw_ext) continue
cam = m.group("cam")
dst_bins.append(os.path.join(dest_bins_dir, f"{new_base}_cam{cam}.bin"))
if mover: if mover:
shutil.move(preview_src, dst_prev) shutil.move(preview_src, dst_prev)
shutil.move(raw_src, dst_raw) shutil.move(meta_src, dst_meta)
shutil.move(mask_src, dst_mask) shutil.move(mask_src, dst_mask)
if mask2_src: for src_bin, dst_bin in zip(bins_src, dst_bins):
shutil.move(mask2_src, os.path.join(dest_mask2_dir, new_base + os.path.splitext(mask2_src)[1])) shutil.move(src_bin, dst_bin)
else:
shutil.copy2(preview_src, dst_prev)
shutil.copy2(raw_src, dst_raw)
shutil.copy2(mask_src, dst_mask)
if mask2_src: if mask2_src:
garantir_pasta(dest_mask2_dir) garantir_pasta(dest_mask2_dir)
shutil.copy2(mask2_src, os.path.join(dest_mask2_dir, new_base + os.path.splitext(mask2_src)[1])) mask2_ext = os.path.splitext(mask2_src)[1].lower()
shutil.move(mask2_src, os.path.join(dest_mask2_dir, new_base + mask2_ext))
else:
shutil.copy2(preview_src, dst_prev)
shutil.copy2(meta_src, dst_meta)
shutil.copy2(mask_src, dst_mask)
for src_bin, dst_bin in zip(bins_src, dst_bins):
shutil.copy2(src_bin, dst_bin)
if mask2_src:
garantir_pasta(dest_mask2_dir)
mask2_ext = os.path.splitext(mask2_src)[1].lower()
shutil.copy2(mask2_src, os.path.join(dest_mask2_dir, new_base + mask2_ext))
return dst_prev, dst_raw, dst_mask, None return dst_prev, dst_meta, dst_bins, dst_mask
# ============================================================ # ============================================================
# 🚀 PROCESSAMENTO # PIPELINE PRINCIPAL
# ============================================================ # ============================================================
def processar(originals_dir, labelmap_path, mover=False, def processar(originals_dir, labelmap_path, mover=False,
@ -214,12 +256,13 @@ def processar(originals_dir, labelmap_path, mover=False,
originals_dir = os.path.abspath(originals_dir) originals_dir = os.path.abspath(originals_dir)
previews_dir = os.path.join(originals_dir, "previews") previews_dir = os.path.join(originals_dir, "previews")
raws_dir = os.path.join(originals_dir, "raws") metas_dir = os.path.join(originals_dir, "metas")
bins_dir = os.path.join(originals_dir, "bins")
masks_dir = os.path.join(originals_dir, "masks") masks_dir = os.path.join(originals_dir, "masks")
masks2_dir = os.path.join(originals_dir, "masks2") masks2_dir = os.path.join(originals_dir, "masks2")
group_dir = os.path.join(originals_dir, "group") group_dir = os.path.join(originals_dir, "group")
if not os.path.isdir(previews_dir) or not os.path.isdir(raws_dir) or not os.path.isdir(masks_dir): if not os.path.isdir(previews_dir) or not os.path.isdir(metas_dir) or not os.path.isdir(masks_dir) or not os.path.isdir(bins_dir):
raise RuntimeError("Estrutura inválida em originals/") raise RuntimeError("Estrutura inválida em originals/")
usar_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir) usar_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir)
@ -231,8 +274,8 @@ def processar(originals_dir, labelmap_path, mover=False,
mapa_previews = mapear_por_base_priorizando_png(previews_dir, EXT_PREVIEWS) mapa_previews = mapear_por_base_priorizando_png(previews_dir, EXT_PREVIEWS)
mapa_masks = mapear_por_base_priorizando_png(masks_dir, EXT_MASKS) mapa_masks = mapear_por_base_priorizando_png(masks_dir, EXT_MASKS)
mapa_raws = {os.path.splitext(n)[0]: os.path.join(raws_dir, n) mapa_metas = mapear_metas_por_base(metas_dir, EXT_METAS)
for n in os.listdir(raws_dir) if n.lower().endswith(EXT_RAW)} mapa_bins = mapear_bins_por_base(bins_dir)
mapa_masks2 = mapear_por_base_priorizando_png(masks2_dir, EXT_MASKS2) if usar_masks2 else {} mapa_masks2 = mapear_por_base_priorizando_png(masks2_dir, EXT_MASKS2) if usar_masks2 else {}
registros = [] registros = []
@ -242,15 +285,13 @@ def processar(originals_dir, labelmap_path, mover=False,
totais["total_masks"] += 1 totais["total_masks"] += 1
prev_path = mapa_previews.get(base) prev_path = mapa_previews.get(base)
raw_path = mapa_raws.get(base) meta_path = mapa_metas.get(base)
bins_paths = mapa_bins.get(base)
if not prev_path or not raw_path: if not prev_path or not meta_path or not bins_paths:
totais["pulados"] += 1 totais["pulados"] += 1
continue continue
# ------------------------------
# EXTRAÇÃO ULTRA RÁPIDA
# ------------------------------
ids_presentes = extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=not labelmap_bgr) ids_presentes = extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=not labelmap_bgr)
if ignore_id in ids_presentes: if ignore_id in ids_presentes:
@ -258,32 +299,47 @@ def processar(originals_dir, labelmap_path, mover=False,
grupo = montar_nome_grupo(ids_presentes, id_para_nome) grupo = montar_nome_grupo(ids_presentes, id_para_nome)
# ------------------------------
# CRIA PASTAS
# ------------------------------
dest_prev_dir = os.path.join(group_dir, grupo, "previews") dest_prev_dir = os.path.join(group_dir, grupo, "previews")
dest_raw_dir = os.path.join(group_dir, grupo, "raws") dest_meta_dir = os.path.join(group_dir, grupo, "metas")
dest_bins_dir = os.path.join(group_dir, grupo, "bins")
dest_mask_dir = os.path.join(group_dir, grupo, "masks") dest_mask_dir = os.path.join(group_dir, grupo, "masks")
dest_mask2_dir = os.path.join(group_dir, grupo, "masks2") if usar_masks2 else None dest_mask2_dir = os.path.join(group_dir, grupo, "masks2") if usar_masks2 else None
dst_prev, dst_raw, dst_mask, _ = copiar_ou_mover_tripla( dst_prev, dst_meta, dst_bins, dst_mask = copiar_ou_mover_conjunto(
prev_path, raw_path, mask_path, prev_path, meta_path, bins_paths, mask_path,
dest_prev_dir, dest_raw_dir, dest_mask_dir, dest_prev_dir, dest_meta_dir, dest_bins_dir, dest_mask_dir,
mover=mover, mover=mover,
mask2_src=mapa_masks2.get(base), mask2_src=mapa_masks2.get(base),
dest_mask2_dir=dest_mask2_dir dest_mask2_dir=dest_mask2_dir,
) )
registros.append([grupo, prev_path, raw_path, mask_path, dst_prev, dst_raw, dst_mask]) registros.append([
grupo,
prev_path,
meta_path,
json.dumps(bins_paths, ensure_ascii=False),
mask_path,
dst_prev,
dst_meta,
json.dumps(dst_bins, ensure_ascii=False),
dst_mask,
])
totais["processados"] += 1 totais["processados"] += 1
# ------------------------------
# MANIFESTO
# ------------------------------
if manifesto: if manifesto:
with open(manifesto, "w", newline="", encoding="utf-8") as f: with open(manifesto, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f) w = csv.writer(f)
w.writerow(["grupo","src_preview","src_raw","src_mask","dst_preview","dst_raw","dst_mask"]) w.writerow([
"grupo",
"src_preview",
"src_meta",
"src_bins_json",
"src_mask",
"dst_preview",
"dst_meta",
"dst_bins_json",
"dst_mask",
])
for r in registros: for r in registros:
w.writerow(r) w.writerow(r)
@ -306,6 +362,7 @@ def build_cli():
ap.add_argument("--labels-bgr", action="store_true") ap.add_argument("--labels-bgr", action="store_true")
return ap return ap
if __name__ == "__main__": if __name__ == "__main__":
ap = build_cli() ap = build_cli()
args = ap.parse_args() args = ap.parse_args()
@ -318,5 +375,5 @@ if __name__ == "__main__":
manifesto=args.manifest, manifesto=args.manifest,
validar_dim=not args.no_validate, validar_dim=not args.no_validate,
estrito=args.strict, estrito=args.strict,
labelmap_bgr=args.labels_bgr labelmap_bgr=args.labels_bgr,
) )

File diff suppressed because it is too large Load Diff

View File

@ -18,308 +18,187 @@ RAW:
- Detecta dtype (uint8/uint16) pelo tamanho do arquivo - Detecta dtype (uint8/uint16) pelo tamanho do arquivo
- Carrega como (H,W), redimensiona, salva em .raw - Carrega como (H,W), redimensiona, salva em .raw
""" """
import argparse
import os import os
import json import json
import sys
import cv2 import cv2
import numpy as np import numpy as np
from typing import Dict, List, Tuple from pi.raw_processor_core import RawProcessorCore
from gal5000.gal_service import mosaic_to_raw4_resized_buf sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from raw_segformer_service import _infer_ignore_id from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids, _infer_ignore_id
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ===== config ===== # ===== CONFIG =====
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config["camera"] RES = tuple(config["resolucao"]) # (W,H)
USE_MASKS2 = config["dual_head"] pasta_base = "dataset"
RESOLUCAO = tuple(config["resolucao"]) # [W,H]
MODEL_NAME = config["model_name"] INPUTS = [
CHANNELS = int(config.get("channels", 4)) ("original", os.path.join(pasta_base, "original", "group")),
pasta_base = os.path.join("dataset") ("augmented", os.path.join(pasta_base, "augmented", "group")),
]
OUTPUT_ROOT = os.path.join(pasta_base, f"{RES[0]}x{RES[1]}", "group")
labelmap_path = os.path.join(pasta_base, "labelmap.txt") labelmap_path = os.path.join(pasta_base, "labelmap.txt")
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag)
RESOLUCOES = {f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1])} # ===== GLOBAL STATS =====
FONTES = ["original", "augmented"] GLOBAL_SUM = None
GLOBAL_SUMSQ = None
GLOBAL_PIXELS = 0
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg")
MSK2_EXTS = (".png", ".jpg", ".jpeg")
RAW_EXTS = (".raw",)
# === 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)
def garantir_dir(p): def garantir_dir(p):
os.makedirs(p, exist_ok=True) os.makedirs(p, exist_ok=True)
def list_groups_raw(root) -> List[str]:
"""
Lista grupos válidos no modo RAW:
tem masks e previews (raws opcional, mas esperado).
"""
if not os.path.isdir(root):
return []
grupos = []
for name in sorted(os.listdir(root)):
gdir = os.path.join(root, name)
if not os.path.isdir(gdir):
continue
if os.path.isdir(os.path.join(gdir, "masks")) and os.path.isdir(os.path.join(gdir, "previews")):
grupos.append(name)
return grupos
def map_by_base_priorizando_png(dir_path: str, exts: Tuple[str, ...]) -> Dict[str, str]: def list_groups(root):
by_base = {} return [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))]
if not os.path.isdir(dir_path):
return by_base
for fname in os.listdir(dir_path):
low = fname.lower()
if not low.endswith(exts):
continue
base, ext = os.path.splitext(fname)
cand = os.path.join(dir_path, 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 map_raws_by_base(raw_dir: str) -> Dict[str, str]:
by_base = {}
if not os.path.isdir(raw_dir):
return by_base
for fname in os.listdir(raw_dir):
if fname.lower().endswith(RAW_EXTS):
base, _ = os.path.splitext(fname)
by_base[base] = os.path.join(raw_dir, fname)
return by_base
def load_raw4_float(path: str, src_hw: Tuple[int,int]) -> np.ndarray: def extract_bins_for_base(bins_dir, base):
h, w = src_hw return [
npx = h * w os.path.join(bins_dir, f)
fsize = os.path.getsize(path) for f in os.listdir(bins_dir)
if f.startswith(base) and f.endswith(".bin")
]
num_floats = fsize // 4
if num_floats != 4 * npx:
raise RuntimeError(
f"RAW {path}: esperado 4 canais float32, "
f"mas num_floats={num_floats}, H*W={npx}"
)
data = np.fromfile(path, dtype=np.float32) def process():
return data.reshape(4, h, w) # (C,H,W)
def save_raw(path: str, arr: np.ndarray):
np.asarray(arr).tofile(path)
def normalize_mask_ids(mask_path: str, cor_para_id, ignore_id: int, dim: Tuple[int,int]) -> np.ndarray:
msk_bgr = cv2.imread(mask_path, cv2.IMREAD_COLOR)
if msk_bgr is None:
raise RuntimeError(f"Erro ao ler máscara: {mask_path}")
msk_rgb = cv2.cvtColor(msk_bgr, cv2.COLOR_BGR2RGB)
ids = converter_mask_rgb_para_ids(msk_rgb, cor_para_id, ignore_id)
ids_res = cv2.resize(ids, dim, interpolation=cv2.INTER_NEAREST)
return ids_res
def normalize_mask2(mask2_path: str, dim: Tuple[int,int]) -> np.ndarray:
m2 = cv2.imread(mask2_path, cv2.IMREAD_UNCHANGED)
if m2 is None:
raise RuntimeError(f"Erro ao ler máscara2: {mask2_path}")
if len(m2.shape) == 3:
m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY)
else:
m2g = m2
_, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY)
m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST)
return m2res
def normalize_group_raw(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: int, groups_except: str = "") -> int:
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
grupos = list_groups_raw(fonte_root)
if not grupos:
return 0
not_want = {g.strip() for g in groups_except.split(",") if g.strip()} cor_para_id, _, _, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = _infer_ignore_id(ignore_rgb, 255)
core = RawProcessorCore(640, 480)
total = 0 total = 0
for nome_res, dim in RESOLUCOES.items(): for fonte_nome, fonte_root in INPUTS:
out_root = os.path.join(pasta_base, nome_res, "group")
for grupo in grupos: if not os.path.isdir(fonte_root):
if grupo in not_want:
print(f"[WARN] Grupo desconsiderado: {grupo}")
continue continue
in_prev = os.path.join(fonte_root, grupo, "previews") for grupo in list_groups(fonte_root):
in_raw = os.path.join(fonte_root, grupo, "raws")
in_msk = os.path.join(fonte_root, grupo, "masks")
in_msk2 = os.path.join(fonte_root, grupo, "masks2")
if not (os.path.isdir(in_prev) and os.path.isdir(in_msk)): print(f"\n[{fonte_nome}] Grupo: {grupo}")
print(f"[WARN] Grupo inválido (sem previews/masks): {grupo}")
gpath = os.path.join(fonte_root, grupo)
bins_dir = os.path.join(gpath, "bins")
meta_dir = os.path.join(gpath, "metas")
mask_dir = os.path.join(gpath, "masks")
if not os.path.isdir(meta_dir):
continue continue
usar_raw = os.path.isdir(in_raw) out_dir = os.path.join(OUTPUT_ROOT, grupo)
usar_msk2 = USE_MASKS2 and os.path.isdir(in_msk2) out_tensor = os.path.join(out_dir, "tensors")
out_mask = os.path.join(out_dir, "masks")
out_prev = os.path.join(out_root, grupo, "previews") garantir_dir(out_tensor)
out_raw = os.path.join(out_root, grupo, "raws") if usar_raw else None garantir_dir(out_mask)
out_msk = os.path.join(out_root, grupo, "masks")
out_msk2 = os.path.join(out_root, grupo, "masks2") if usar_msk2 else None
garantir_dir(out_prev) for fname in os.listdir(meta_dir):
garantir_dir(out_msk)
if usar_raw and out_raw:
garantir_dir(out_raw)
if usar_msk2 and out_msk2:
garantir_dir(out_msk2)
prev_files = [f for f in os.listdir(in_prev) if os.path.splitext(f.lower())[1] in IMG_EXTS] base = fname.replace(".json", "")
msk_map = map_by_base_priorizando_png(in_msk, MSK_EXTS) meta_path = os.path.join(meta_dir, fname)
raw_map = map_raws_by_base(in_raw) if usar_raw else {}
msk2_map = map_by_base_priorizando_png(in_msk2, MSK2_EXTS) if usar_msk2 else {}
n = len(prev_files) with open(meta_path) as f:
for i, fname in enumerate(sorted(prev_files), 1): meta = json.load(f)
base, ext = os.path.splitext(fname)
prev_path = os.path.join(in_prev, fname)
msk_path = msk_map.get(base)
raw_path = raw_map.get(base) if usar_raw else None
msk2_path = msk2_map.get(base) if usar_msk2 else None
if not msk_path: bins_paths = extract_bins_for_base(bins_dir, base)
print(f"[WARN] [{fonte_nome} | {grupo}] Sem máscara p/ {fname}, pulando.")
if not bins_paths:
continue continue
# --- preview --- # ===== LOAD BINS =====
prev_bgr = cv2.imread(prev_path, cv2.IMREAD_COLOR) bins_data = []
if prev_bgr is None: bins_meta = []
print(f"[WARN] [{fonte_nome} | {grupo}] Falha ao ler preview: {prev_path}")
continue
prev_res = cv2.resize(prev_bgr, dim, interpolation=cv2.INTER_AREA)
# nomes saída com prefixo (igual o normalize atual) for path in bins_paths:
out_name_prev = f"{fonte_nome}_{fname}" cam_id = os.path.basename(path).split("_")[-1].replace(".bin", "")
out_name_base = os.path.splitext(out_name_prev)[0] # pra raw/masks
cv2.imwrite(os.path.join(out_prev, out_name_prev), prev_res) cam_meta = core.extract_camera_meta(meta, cam_id)
data = core.load_native_bin(path, cam_meta)
# --- mask ids --- bins_data.append(data)
ids_res = normalize_mask_ids(msk_path, cor_para_id, ignore_id, dim) bins_meta.append(cam_meta)
cv2.imwrite(os.path.join(out_msk, out_name_base + ".png"), ids_res)
# --- mask2 --- # ===== BUILD TENSOR =====
if usar_msk2 and out_msk2: tensor, channel_names = core.build_multispectral_tensor(bins_data, bins_meta)
if msk2_path:
m2res = normalize_mask2(msk2_path, dim)
cv2.imwrite(os.path.join(out_msk2, out_name_base + ".png"), m2res)
else:
print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname}")
# --- raw --- # ===== RESIZE =====
if usar_raw and out_raw:
if raw_path:
src_h, src_w = prev_bgr.shape[:2]
raw4 = load_raw4_float(raw_path, (src_h, src_w)) # (4, Hsrc, Wsrc)
out_w, out_h = dim
if (src_w, src_h) != (out_w, out_h):
# redimensiona cada canal
chans = [] chans = []
for k in range(raw4.shape[0]): for ch in tensor:
ch = raw4[k] ch_res = cv2.resize(ch, RES, interpolation=cv2.INTER_AREA)
ch_res = cv2.resize(ch, (out_w, out_h), interpolation=cv2.INTER_AREA)
chans.append(ch_res.astype(np.float32)) chans.append(ch_res.astype(np.float32))
raw4 = np.stack(chans, axis=0) # (4, out_h, out_w)
# Atualiza acumuladores de stats tensor = np.stack(chans, axis=0)
# raw4: (C,H,W) -> (C,N)
c, hh, ww = raw4.shape # ===== STATS =====
c, h, w = tensor.shape
if GLOBAL_SUM is None: if GLOBAL_SUM is None:
GLOBAL_SUM = np.zeros(c, dtype=np.float64) GLOBAL_SUM = np.zeros(c, dtype=np.float64)
GLOBAL_SUMSQ = np.zeros(c, dtype=np.float64) GLOBAL_SUMSQ = np.zeros(c, dtype=np.float64)
flat = raw4.reshape(c, -1).astype(np.float64) flat = tensor.reshape(c, -1).astype(np.float64)
GLOBAL_SUM += flat.sum(axis=1) GLOBAL_SUM += flat.sum(axis=1)
GLOBAL_SUMSQ += (flat ** 2).sum(axis=1) GLOBAL_SUMSQ += (flat ** 2).sum(axis=1)
GLOBAL_PIXELS += hh * ww # por canal é o mesmo H*W GLOBAL_PIXELS += h * w
# Salva como float32 "linearzão" (4 * H * W floats) # ===== SAVE TENSOR =====
save_raw(os.path.join(out_raw, out_name_base + ".raw"), raw4.astype(np.float32)) name = f"{fonte_nome}_{base}"
else:
print(f"[WARN] [{fonte_nome} | {grupo}] Sem RAW p/ {fname} (seguindo só preview+mask).") np.save(os.path.join(out_tensor, name + ".npy"), tensor)
# ===== MASK =====
mask_path = os.path.join(mask_dir, base + ".png")
if os.path.exists(mask_path):
msk = cv2.imread(mask_path)
msk = cv2.cvtColor(msk, cv2.COLOR_BGR2RGB)
ids = converter_mask_rgb_para_ids(msk, cor_para_id, ignore_id)
ids_res = cv2.resize(ids, RES, interpolation=cv2.INTER_NEAREST)
np.save(os.path.join(out_mask, name + ".npy"), ids_res)
# debug visual
cv2.imwrite(
os.path.join(out_mask, name + ".png"),
ids_res
)
total += 1 total += 1
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n}{fname}") print(f"OK: {name}")
return total # ===== SAVE STATS =====
if GLOBAL_SUM is not None:
def main(args): mean = GLOBAL_SUM / GLOBAL_PIXELS
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)
total_geral = 0
# ORIGINAL
orig_group = os.path.join(pasta_base, "original", "group")
if os.path.isdir(orig_group):
total_geral += normalize_group_raw(orig_group, "original", cor_para_id, ignore_id, groups_except=args.groups_except)
else:
print("[WARN] Não achei original/group (modo RAW).")
# AUGMENTED
aug_group = os.path.join(pasta_base, "augmented", "group")
if os.path.isdir(aug_group):
total_geral += normalize_group_raw(aug_group, "augmented", cor_para_id, ignore_id, groups_except=args.groups_except)
else:
print("[WARN] Não achei augmented/group (modo RAW).")
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 var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean**2
std = np.sqrt(np.maximum(var, 1e-6)) 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:
channel_names = ["R", "G", "IR", "B"]
stats = { stats = {
"channels": channel_names[:len(mean_list)], "channels": channel_names,
"mean": mean_list, "mean": mean.tolist(),
"std": std_list, "std": std.tolist(),
"pixels_per_channel": int(GLOBAL_PIXELS),
} }
save_path = os.path.join("backup", config["modelo"], config["model_name"], config["stats_source_tag"])
garantir_dir(save_path) garantir_dir(save_path)
stats_path = os.path.join(save_path, "norm_stats.json")
with open(stats_path, "w", encoding="utf-8") as f:
json.dump(stats, f, indent=2, ensure_ascii=False)
print(f"📁 Stats salvos em: {stats_path}") with open(os.path.join(save_path, "norm_stats.json"), "w") as f:
print(f" mean: {mean_list}") json.dump(stats, f, indent=2)
print(f" std : {std_list}")
else: print("\n📊 STATS:")
print("⚠️ Nenhum RAW processado, não há stats para salvar.") print(stats)
print(f"\n✅ FINALIZADO: {total} samples")
if __name__ == "__main__": if __name__ == "__main__":
ap = argparse.ArgumentParser(description="Normalize por grupos (RAW: previews/raws/masks)") process()
ap.add_argument("--groups-except", type=str, default="", help="Grupos para não usar, separados por vírgula.")
args = ap.parse_args()
main(args)

View File

@ -1,23 +1,6 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
"""
Split estratificado por GRUPO com **val/test do ORIGINAL** e
garantia de NÃO VAZAMENTO entre splits, agora para estrutura RAW:
de:
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
Escreve em:
dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}
Definições:
- "Família" = todas as variações da MESMA base original:
original_<base>.* e augmented_<base>_aug_XX.*
- Val/Test: **original_<base>** (sem augmented)
- Train: original_<base> **e** todos augmented_<base>_aug_XX
Baseado no _7_split.py original (versão images/masks). :contentReference[oaicite:1]{index=1}
"""
import os import os
import re import re
import json import json
@ -25,25 +8,20 @@ import shutil
import random import random
import argparse import argparse
# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
RESOLUCAO = tuple(config.get("resolucao"))
# Pastas (ajustadas para PREVIEWS/RAWS) RESOLUCAO = tuple(config.get("resolucao"))
pasta_origem = os.path.join("dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group") pasta_origem = os.path.join("dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
pasta_destino = os.path.join("dataset", "split") pasta_destino = os.path.join("dataset", "split")
IMG_EXTS = (".jpg", ".jpeg", ".png") TENSOR_EXT = ".npy"
MSK_EXT = ".png" # máscaras normalizadas em PNG MASK_NPY_SUFFIX = ".npy"
RAW_EXT = ".raw"
# Regex para identificar famílias RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
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_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)
RE_AUG_SUFFIX = re.compile(r'_aug[a-zA-Z0-9]*_\d+$', re.IGNORECASE)
def garantir(p): def garantir(p):
@ -51,7 +29,6 @@ def garantir(p):
def lista_grupos(root): def lista_grupos(root):
"""Lista grupos que têm previews + masks (raws opcional)."""
if not os.path.isdir(root): if not os.path.isdir(root):
return [] return []
out = [] out = []
@ -59,43 +36,27 @@ def lista_grupos(root):
gdir = os.path.join(root, g) gdir = os.path.join(root, g)
if not os.path.isdir(gdir): if not os.path.isdir(gdir):
continue continue
if os.path.isdir(os.path.join(gdir, "previews")) and os.path.isdir(os.path.join(gdir, "masks")): if os.path.isdir(os.path.join(gdir, "tensors")) and os.path.isdir(os.path.join(gdir, "masks")):
out.append(g) out.append(g)
return out return out
def listar_previews(prev_dir): def listar_tensors(tensor_dir):
if not os.path.isdir(prev_dir): if not os.path.isdir(tensor_dir):
return [] return []
fs = [] fs = []
for f in os.listdir(prev_dir): for f in os.listdir(tensor_dir):
ext = os.path.splitext(f.lower())[1] if f.lower().endswith(TENSOR_EXT):
if ext in IMG_EXTS:
fs.append(f) fs.append(f)
return sorted(fs) return sorted(fs)
def mask_from_image_name(img_name): def mask_npy_from_tensor_name(tensor_name):
base, _ = os.path.splitext(img_name) base, _ = os.path.splitext(tensor_name)
return base + MSK_EXT return base + MASK_NPY_SUFFIX
def mask2_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return base + MSK_EXT # masks2 também normalizadas em PNG
def raw_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return base + RAW_EXT
def classify_source_and_family(filename_no_ext): 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"
"""
m = RE_ORIGINAL_PREFIX.match(filename_no_ext) m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
if m: if m:
return "original", m.group(1) return "original", m.group(1)
@ -104,7 +65,6 @@ def classify_source_and_family(filename_no_ext):
if m: if m:
return "augmented", m.group(1) 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): if RE_AUG_SUFFIX.search(filename_no_ext):
fam = RE_AUG_SUFFIX.sub("", filename_no_ext) fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
return "augmented", fam return "augmented", fam
@ -112,40 +72,43 @@ def classify_source_and_family(filename_no_ext):
return "unknown", filename_no_ext return "unknown", filename_no_ext
def build_family_index(prev_dir, msk_dir): def build_family_index(tensor_dir, mask_dir):
""" """
Constroi índice de famílias a partir de previews/masks. family -> {
Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]} "original": tensor_name or None,
(strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe) "augmented": [tensor_name, ...],
"all": [...]
}
indexa se houver pelo menos mask .npy correspondente.
""" """
familias = {} familias = {}
imgs = listar_previews(prev_dir) tensors = listar_tensors(tensor_dir)
for img_name in imgs:
base_no_ext, ext = os.path.splitext(img_name) for tensor_name in tensors:
mask_name = mask_from_image_name(img_name) base_no_ext, _ = os.path.splitext(tensor_name)
if not os.path.exists(os.path.join(msk_dir, mask_name)): mask_npy_name = mask_npy_from_tensor_name(tensor_name)
continue # garante pareamento preview/mask
if not os.path.exists(os.path.join(mask_dir, mask_npy_name)):
continue
source, fam = classify_source_and_family(base_no_ext) source, fam = classify_source_and_family(base_no_ext)
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []}) d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
d["all"].append(img_name) d["all"].append(tensor_name)
if source == "original": if source == "original":
d["original"] = img_name d["original"] = tensor_name
elif source == "augmented": elif source == "augmented":
d["augmented"].append(img_name) d["augmented"].append(tensor_name)
else: else:
# trata como original desconhecido para não perder dado
if d["original"] is None: if d["original"] is None:
d["original"] = img_name d["original"] = tensor_name
else: else:
d["augmented"].append(img_name) d["augmented"].append(tensor_name)
return familias return familias
def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test): def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
"""
Mesmo esquema do script original: calcula quantas FAMÍLIAS vão para train/val/test.
"""
n_train = int(round(n * p_train)) n_train = int(round(n * p_train))
n_val = int(round(n * p_val)) n_val = int(round(n * p_val))
n_test = n - n_train - n_val n_test = n - n_train - n_val
@ -178,6 +141,7 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
else: else:
break break
total = n_train + n_val + n_test total = n_train + n_val + n_test
while total < n: while total < n:
if n_train - min_train <= n_val - min_val: if n_train - min_train <= n_val - min_val:
n_train += 1 n_train += 1
@ -190,11 +154,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_val = max(0, min(resto, min_val))
n_test = max(0, resto - n_val) n_test = max(0, resto - n_val)
# ajuste final
diff = n - (n_train + n_val + n_test) diff = n - (n_train + n_val + n_test)
if diff != 0: if diff != 0:
if diff > 0: if diff > 0:
# adiciona em train, depois val
take = min(diff, n - n_train) take = min(diff, n - n_train)
n_train += take n_train += take
diff -= take diff -= take
@ -202,7 +164,6 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
n_val += diff n_val += diff
else: else:
diff = -diff diff = -diff
# tira de test, depois val
take = min(diff, n_test) take = min(diff, n_test)
n_test -= take n_test -= take
diff -= take diff -= take
@ -214,71 +175,47 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
def copiar( def copiar(
nomes, nomes,
src_prev_dir, src_tensor_dir,
src_msk_dir, src_mask_dir,
dst_prev_dir, dst_tensor_dir,
dst_msk_dir, dst_mask_dir,
src_msk2_dir=None,
dst_msk2_dir=None,
src_raw_dir=None,
dst_raw_dir=None,
): ):
""" """
Copia preview + mask (+ mask2) (+ raw) para o split. Copia:
- tensor .npy
- mask .npy obrigatória
- mask .png opcional (debug)
""" """
garantir(dst_prev_dir) garantir(dst_tensor_dir)
garantir(dst_msk_dir) garantir(dst_mask_dir)
use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir))
use_raw = bool(src_raw_dir and dst_raw_dir and os.path.isdir(src_raw_dir))
if use_msk2:
garantir(dst_msk2_dir)
if use_raw:
garantir(dst_raw_dir)
moved = 0 moved = 0
for nome in nomes: for nome in nomes:
mask_name = mask_from_image_name(nome) tensor_src = os.path.join(src_tensor_dir, nome)
src_prev = os.path.join(src_prev_dir, nome) mask_npy_name = mask_npy_from_tensor_name(nome)
src_msk = os.path.join(src_msk_dir, mask_name)
if not (os.path.exists(src_prev) and os.path.exists(src_msk)): mask_npy_src = os.path.join(src_mask_dir, mask_npy_name)
if not (os.path.exists(tensor_src) and os.path.exists(mask_npy_src)):
continue continue
shutil.copy2(src_prev, os.path.join(dst_prev_dir, nome)) shutil.copy2(tensor_src, os.path.join(dst_tensor_dir, nome))
shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name)) shutil.copy2(mask_npy_src, os.path.join(dst_mask_dir, mask_npy_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_raw:
raw_name = raw_from_image_name(nome)
src_raw = os.path.join(src_raw_dir, raw_name)
if os.path.exists(src_raw):
shutil.copy2(src_raw, os.path.join(dst_raw_dir, raw_name))
moved += 1 moved += 1
return moved return moved
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None): def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
src_prev_dir = os.path.join(pasta_origem, group_name, "previews") src_tensor_dir = os.path.join(pasta_origem, group_name, "tensors")
src_msk_dir = os.path.join(pasta_origem, group_name, "masks") src_mask_dir = os.path.join(pasta_origem, group_name, "masks")
src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2")
src_raw_dir = os.path.join(pasta_origem, group_name, "raws")
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir) familias = build_family_index(src_tensor_dir, src_mask_dir)
use_raw = os.path.isdir(src_raw_dir)
familias = build_family_index(src_prev_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] familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
total_familias = len(familias_originais) total_familias = len(familias_originais)
if total_familias == 0: if total_familias == 0:
print(f"[{group_name}] 0 famílias com original, pulando.") print(f"[{group_name}] 0 famílias com original, pulando.")
return {"train": 0, "val": 0, "test": 0, "familias": 0} return {"train": 0, "val": 0, "test": 0, "familias": 0}
@ -295,7 +232,6 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
fam_val = set(familias_originais[n_tr:n_tr+n_va]) 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_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: if caps_map and group_name in caps_map:
cap = caps_map[group_name] cap = caps_map[group_name]
if len(fam_train) > cap: if len(fam_train) > cap:
@ -320,66 +256,40 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
elif fam in fam_test: elif fam in fam_test:
if d["original"]: if d["original"]:
nomes_test.append(d["original"]) nomes_test.append(d["original"])
else:
pass
# dest dirs dst_train_tensor = os.path.join(pasta_destino, "train", "group", group_name, "tensors")
dest_train_prev = os.path.join(pasta_destino, "train", "group", group_name, "previews") dst_train_mask = os.path.join(pasta_destino, "train", "group", group_name, "masks")
dest_train_msk = os.path.join(pasta_destino, "train", "group", group_name, "masks")
dest_val_prev = os.path.join(pasta_destino, "val", "group", group_name, "previews")
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
dest_test_prev = os.path.join(pasta_destino, "test", "group", group_name, "previews")
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 dst_val_tensor = os.path.join(pasta_destino, "val", "group", group_name, "tensors")
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None dst_val_mask = os.path.join(pasta_destino, "val", "group", group_name, "masks")
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
dest_train_raw = os.path.join(pasta_destino, "train", "group", group_name, "raws") if use_raw else None dst_test_tensor = os.path.join(pasta_destino, "test", "group", group_name, "tensors")
dest_val_raw = os.path.join(pasta_destino, "val", "group", group_name, "raws") if use_raw else None dst_test_mask = os.path.join(pasta_destino, "test", "group", group_name, "masks")
dest_test_raw = os.path.join(pasta_destino, "test", "group", group_name, "raws") if use_raw else None
m_train = copiar( m_train = copiar(nomes_train, src_tensor_dir, src_mask_dir, dst_train_tensor, dst_train_mask)
nomes_train, m_val = copiar(nomes_val, src_tensor_dir, src_mask_dir, dst_val_tensor, dst_val_mask)
src_prev_dir, src_msk_dir, m_test = copiar(nomes_test, src_tensor_dir, src_mask_dir, dst_test_tensor, dst_test_mask)
dest_train_prev, dest_train_msk,
src_msk2_dir, dest_train_msk2,
src_raw_dir, dest_train_raw
)
m_val = copiar(
nomes_val,
src_prev_dir, src_msk_dir,
dest_val_prev, dest_val_msk,
src_msk2_dir, dest_val_msk2,
src_raw_dir, dest_val_raw
)
m_test = copiar(
nomes_test,
src_prev_dir, src_msk_dir,
dest_test_prev, dest_test_msk,
src_msk2_dir, dest_test_msk2,
src_raw_dir, dest_test_raw
)
print(f"[{group_name}] famílias={total_familias} → train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}") print(f"[{group_name}] famílias={total_familias} → train(tensors)={m_train}, val(tensors)={m_val}, test(tensors)={m_test}")
return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias} return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}
def main(): def main():
ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (RAW: previews/raws/masks).") ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (tensors/masks).")
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).") ap.add_argument("--train", type=float, default=0.70)
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação (default=0.29).") ap.add_argument("--val", type=float, default=0.29)
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste (default=0.01).") ap.add_argument("--test", type=float, default=0.01)
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento (default=42).") ap.add_argument("--seed", type=int, 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-train", type=int, 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-val", type=int, default=1)
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).") ap.add_argument("--min-test", type=int, default=0)
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).") ap.add_argument("--resolucao", type=str, default=None,
help="Sobrescreve resolução no formato WxH (ex: 1024x800).")
ap.add_argument("--cap-train-families", type=str, default="", 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'") help="Mapa 'grupo:cap,...' para limitar famílias no TRAIN. Ex.: 'chao:350'")
args = ap.parse_args() args = ap.parse_args()
@ -410,6 +320,7 @@ def main():
soma = args.train + args.val + args.test soma = args.train + args.val + args.test
if soma <= 0: if soma <= 0:
raise ValueError("Soma de proporções deve ser > 0.") raise ValueError("Soma de proporções deve ser > 0.")
p_train = args.train / soma p_train = args.train / soma
p_val = args.val / soma p_val = args.val / soma
p_test = args.test / soma p_test = args.test / soma
@ -439,12 +350,12 @@ def main():
for k in total_global.keys(): for k in total_global.keys():
total_global[k] += res.get(k, 0) total_global[k] += res.get(k, 0)
print("\nResumo global (imagens copiadas):") print("\nResumo global (tensors copiados):")
print(f" train: {total_global['train']}") print(f" train: {total_global['train']}")
print(f" val: {total_global['val']}") print(f" val: {total_global['val']}")
print(f" test: {total_global['test']}") 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 RAW sem vazamento concluído!") print("\n✅ Split sem vazamento concluído!")
if __name__ == "__main__": if __name__ == "__main__":

View File

@ -16,7 +16,14 @@ from torch.utils.data import DataLoader, Dataset
from torch.amp import autocast, GradScaler from torch.amp import autocast, GradScaler
import torch.nn.functional as F import torch.nn.functional as F
from raw_segformer_service import (RawSegDataset, normalize_raw, patch_segformer_input_channels, build_raw_segformer_model, build_dual_branch_segformer_model) from multispec_segformer_service import (
MultispecSegDataset,
normalize_per_batch,
patch_segformer_input_channels,
build_multispec_segformer_model,
build_dual_branch_segformer_model,
build_fixed_normalizer,
)
# ---------------------------- # ----------------------------
@ -208,7 +215,7 @@ def run_one_epoch(
if normalizer is not None: if normalizer is not None:
imgs = normalizer(imgs) imgs = normalizer(imgs)
else: else:
imgs = normalize_raw(imgs) imgs = normalize_per_batch(imgs)
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
out = model(pixel_values=imgs) out = model(pixel_values=imgs)
@ -374,15 +381,11 @@ def main():
# ========================== # ==========================
norm_stats = None norm_stats = None
normalizer = None normalizer = 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")
norm_stats_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag, "norm_stats.json") norm_stats_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag, "norm_stats.json")
if args.norm_stats is not None: if args.norm_stats is not None:
norm_stats_path = args.norm_stats norm_stats_path = args.norm_stats
if norm_stats_path is not None and os.path.exists(norm_stats_path): if norm_stats_path and os.path.exists(norm_stats_path):
with open(norm_stats_path, "r", encoding="utf-8") as f: with open(norm_stats_path, "r", encoding="utf-8") as f:
norm_stats = json.load(f) norm_stats = json.load(f)
@ -390,89 +393,22 @@ def main():
stats_mean = norm_stats.get("mean", []) stats_mean = norm_stats.get("mean", [])
stats_std = norm_stats.get("std", []) stats_std = norm_stats.get("std", [])
print(f"[NORM] usando stats fixos de: {norm_stats_path}") if len(stats_mean) != CHANNELS or len(stats_std) != CHANNELS:
raise RuntimeError(
f"norm_stats incompatível com channels={CHANNELS}: "
f"mean={len(stats_mean)} std={len(stats_std)}"
)
if len(stats_channels) != CHANNELS:
raise RuntimeError(
f"norm_stats.channels incompatível com channels={CHANNELS}: {stats_channels}"
)
normalizer = build_fixed_normalizer(stats_mean, stats_std, device)
print(f"[NORM] usando stats fixos: {norm_stats_path}")
print(f"[NORM] channels={stats_channels}") print(f"[NORM] channels={stats_channels}")
print(f"[NORM] mean={stats_mean}")
print(f"[NORM] std ={stats_std}")
# JSON veio da normalização RAW4: ["R","G","IR","B"] em 0..1
# Nosso tensor de treino já vem em contrato interno:
# 4 canais: [R,G,B,IR]
# 5 canais: [R,G,B,IR,NDVI01]
# Então reordenamos os stats para [R,G,B,IR] e,
# se tiver NDVI, deixamos o 5º canal sem normalizar (mean=0, std=1).
# Garante que temos pelo menos R,G,IR,B
idx_by_name = {name: i for i, name in enumerate(stats_channels)}
required = ["R", "G", "IR", "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,IR,B. Mantendo normalize_raw dinâmico.")
normalizer = None
else: else:
# ordem interna bruta do JSON: R,G,IR,B print("[NORM] sem stats fixos, usando normalize_per_batch.")
mean_arr = np.array(stats_mean, dtype=np.float32)
std_arr = np.array(stats_std, dtype=np.float32)
# reordena para o contrato interno do modelo: [R,G,B,IR]
desired_order = ["R", "G", "B", "IR"]
mean_rgbi = []
std_rgbi = []
for ch_name in desired_order:
i = idx_by_name[ch_name]
mean_rgbi.append(mean_arr[i])
std_rgbi.append(std_arr[i])
mean_rgbi = np.array(mean_rgbi, dtype=np.float32)
std_rgbi = np.array(std_rgbi, dtype=np.float32)
if CHANNELS == 3 and not USE_NDVI:
mean3 = mean_rgbi[:3]
std3 = std_rgbi[:3]
mean_t = torch.tensor(mean3, dtype=torch.float32, device=device).view(1, 3, 1, 1)
std_t = torch.tensor(std3, dtype=torch.float32, device=device).view(1, 3, 1, 1)
def normalizer(x: torch.Tensor) -> torch.Tensor:
return (x - mean_t) / std_t
print("[NORM] Normalização fixa por canal ativada para [R,G,B].")
elif CHANNELS == 4 and not USE_NDVI:
# [R,G,B,IR]
mean_t = torch.tensor(mean_rgbi, dtype=torch.float32, device=device).view(1, 4, 1, 1)
std_t = torch.tensor(std_rgbi, dtype=torch.float32, device=device).view(1, 4, 1, 1)
def normalizer(x: torch.Tensor) -> torch.Tensor:
# x: (B,4,H,W) em 0..1
return (x - mean_t) / std_t
print("[NORM] Normalização fixa por canal ativada para [R,G,B,IR].")
elif CHANNELS == 5 and USE_NDVI:
# [R,G,B,IR,NDVI01]
# NDVI já está em 0..1 e não queremos mexer, então:
# mean_ndvi = 0, std_ndvi = 1 -> x_ndvi sai intacto.
mean5 = np.concatenate([mean_rgbi, np.array([0.0], dtype=np.float32)], axis=0)
std5 = np.concatenate([std_rgbi, np.array([1.0], dtype=np.float32)], axis=0)
mean_t = torch.tensor(mean5, dtype=torch.float32, device=device).view(1, 5, 1, 1)
std_t = torch.tensor(std5, dtype=torch.float32, device=device).view(1, 5, 1, 1)
def normalizer(x: torch.Tensor) -> torch.Tensor:
# x: (B,5,H,W) em 0..1, com NDVI no último canal
return (x - mean_t) / std_t
print("[NORM] Normalização fixa por canal ativada para [R,G,B,IR] (NDVI permanece inalterado).")
else:
print("[NORM] AVISO: norm_stats JSON é para 4 canais RAW (R,G,IR,B), "
f"mas config está channels={CHANNELS}, use_ndvi={USE_NDVI}. "
"Mantendo normalize_raw dinâmico.")
normalizer = None
else:
if norm_stats_path:
print(f"[NORM] Caminho de norm_stats não encontrado: {norm_stats_path}. Usando normalize_raw dinâmico.")
else:
print("[NORM] norm_stats não informado. Usando normalize_raw dinâmico.")
resize_hw = None resize_hw = None
if args.resize_h is not None and args.resize_w is not None: if args.resize_h is not None and args.resize_w is not None:
@ -483,23 +419,20 @@ def main():
pass pass
# Datasets (RAW com 4 ou 5 canais) # Datasets (RAW com 4 ou 5 canais)
ds_train = RawSegDataset( ds_train = MultispecSegDataset(
os.path.join(dataset_path, "split", "train"), os.path.join(dataset_path, "split", "train"),
labelmap_path=labelmap_path, labelmap_path=labelmap_path,
max_value=args.raw_max,
resize_hw=resize_hw, resize_hw=resize_hw,
raw_hw=(H, W),
use_ndvi=USE_NDVI,
channels=CHANNELS, channels=CHANNELS,
strict_channels=True,
) )
ds_val = RawSegDataset(
ds_val = MultispecSegDataset(
os.path.join(dataset_path, "split", "val"), os.path.join(dataset_path, "split", "val"),
labelmap_path=labelmap_path, labelmap_path=labelmap_path,
max_value=args.raw_max,
resize_hw=resize_hw, resize_hw=resize_hw,
raw_hw=(H, W),
use_ndvi=USE_NDVI,
channels=CHANNELS, channels=CHANNELS,
strict_channels=True,
) )
# classes # classes
@ -538,7 +471,7 @@ def main():
) )
if FUSION_MODE == "stacked": if FUSION_MODE == "stacked":
model = build_raw_segformer_model( model = build_multispec_segformer_model(
num_classes=num_classes, num_classes=num_classes,
channels=CHANNELS, channels=CHANNELS,
backbone=BACKBONE, backbone=BACKBONE,

View File

@ -0,0 +1,413 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Teste/visualização do SegFormer para o novo pipeline multiespectral.
Modos:
1) Dataset (com GT):
- Usa MultispecSegDataset a partir de split (train/val/test)
- Mostra: [preview tensor] | [mask GT] | [overlay predito]
2) Pasta sem máscara (--test_folder sem masks):
- Faz inferência com tensors .npy (sem GT)
- Mostra: [preview tensor] | [overlay predito]
3) Câmera ao vivo (--camera):
- Mantido como placeholder para integração futura com FrameService/MULTISPEC.
- Neste script, o foco principal é dataset + tensors finais.
"""
import os
import json
import sys
import time
import argparse
import cv2
import torch
import numpy as np
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import converter_mask_ids_para_bgr, desenhar_legenda_horizontal
from multispec_segformer_service import (
MultispecSegformerService,
MultispecSegDataset,
)
# ============================================================
# Helpers para modo sem máscara
# ============================================================
def _collect_tensor_paths(root: str) -> list[str]:
"""
Coleta caminhos de tensors .npy em 'root' usando a convenção nova:
- se root/group existe: procura group/*/tensors
- senão: procura root/tensors
"""
paths = []
dir_group = os.path.join(root, "group")
if os.path.isdir(dir_group):
for g in sorted(os.listdir(dir_group)):
gdir = os.path.join(dir_group, g)
if not os.path.isdir(gdir):
continue
g_tensors = os.path.join(gdir, "tensors")
if not os.path.isdir(g_tensors):
continue
for fn in sorted(os.listdir(g_tensors)):
if fn.lower().endswith(".npy"):
paths.append(os.path.join(g_tensors, fn))
else:
dir_tensors = os.path.join(root, "tensors")
if not os.path.isdir(dir_tensors):
raise RuntimeError(f"Modo sem máscara: não achei pasta 'tensors' em {root}")
for fn in sorted(os.listdir(dir_tensors)):
if fn.lower().endswith(".npy"):
paths.append(os.path.join(dir_tensors, fn))
return paths
# ============================================================
# Script principal
# ============================================================
def main():
parser = argparse.ArgumentParser()
# Config / treino
parser.add_argument("--config", default="config.json", help="Caminho do config.json (mesmo do treino)")
parser.add_argument("--split_folder", type=str, default="val", help="train, val ou test (dentro de dataset/split)")
parser.add_argument("--root_override", type=str, default=None, help="Aponta direto para uma pasta split (sobrepõe split_folder)")
parser.add_argument(
"--test_folder",
type=str,
default=None,
help=(
"Pasta de teste para visualização. "
"Se tiver masks, usa dataset normal (mostra GT). "
"Se tiver apenas tensors, entra em modo inferência-only."
)
)
parser.add_argument("--ckpt", type=str, default=None, help="Caminho do .pt (se não informar, usa best_miou.pt da pasta rawX)")
parser.add_argument("--camera", action="store_true", help="Placeholder para integração futura com câmera ao vivo")
parser.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal")
parser.add_argument("--resize_h", type=int, default=None, help="Altura para inferência (override)")
parser.add_argument("--resize_w", type=int, default=None, help="Largura para inferência (override)")
parser.add_argument("--alpha", type=float, default=0.45, help="Alpha do overlay da máscara")
parser.add_argument("--camera_frame_type", type=str, default="MULTISPEC", choices=["RAW_BRUTO", "RGB", "MULTISPEC"])
parser.add_argument("--camera_capture_mode", type=str, default="AUTO", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"])
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
# ---- Carrega config ----
with open(args.config, "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
modelo_folder = config["modelo"]
CHANNELS = int(config.get("channels", 5))
FUSION_MODE = config.get("fusion_mode", "stacked")
RESOLUCAO = config["resolucao"]
W, H = RESOLUCAO[0], RESOLUCAO[1]
if args.resize_h is not None:
H = args.resize_h
if args.resize_w is not None:
W = args.resize_w
print(f"[cfg] resolucao nominal tensor: {W}x{H}")
print(f"[cfg] channels={CHANNELS}")
dataset_path = os.path.join("dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
# ==========================
# Normalização fixa (igual treino)
# ==========================
norm_mean = None
norm_std = None
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
norm_stats_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag, "norm_stats.json")
if args.norm_stats is not None:
norm_stats_path = args.norm_stats
if os.path.isfile(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}")
if len(stats_mean) != CHANNELS or len(stats_std) != CHANNELS:
raise RuntimeError(
f"norm_stats incompatível com channels={CHANNELS}: mean={len(stats_mean)} std={len(stats_std)}"
)
if len(stats_channels) != CHANNELS:
raise RuntimeError(
f"norm_stats.channels incompatível com channels={CHANNELS}: {stats_channels}"
)
norm_mean = stats_mean
norm_std = stats_std
else:
print(f"[NORM] norm_stats.json não encontrado em {norm_stats_path}. Usando normalização dinâmica por frame.")
# ---- Descobre checkpoint ----
if args.ckpt is not None:
ckpt_path = args.ckpt
else:
save_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag)
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}")
# ---- Instancia service do modelo ----
model_svc = MultispecSegformerService(
config_path=args.config,
ckpt_path=ckpt_path,
device=device,
use_amp=True,
mean=norm_mean,
std=norm_std,
)
classes = model_svc.get_classes()
colormap_rgb = model_svc.get_colormap()
ignore_id = model_svc.get_ignore_id()
print(f"[svc] classes: {classes}")
print(f"[svc] ignore_id={ignore_id}")
# ======================================================================
# MODO CÂMERA AO VIVO (placeholder)
# ======================================================================
if args.camera:
print("[mode] Câmera MULTISPEC + SegFormer")
from multispectral_service import MultiSpectralService
from stream_receiver import StreamReceiver
from pi.raw_processor_core import RawProcessorCore
STREAM_PORT = 6001
PI_HOST = "192.168.105.6"
PC_HOST = "192.168.105.5"
receiver = StreamReceiver(host="0.0.0.0", port=STREAM_PORT)
svc = MultiSpectralService(host=PI_HOST, port=5000, timeout=10)
receiver.start()
time.sleep(0.5)
svc.connect()
core = RawProcessorCore(sensor_width=W, sensor_height=H)
camera_frame_type = args.camera_frame_type
camera_output_dtype = "uint8"
camera_capture_mode = args.camera_capture_mode
print("SET FRAME TYPE:", svc.set_frame_type(camera_frame_type))
print("SET OUTPUT DTYPE:", svc.set_output_dtype(camera_output_dtype))
print("SET FPS:", svc.set_fps(15))
print("BEGIN:", svc.begin(
frame_type=camera_frame_type,
output_dtype=camera_output_dtype,
capture_mode=camera_capture_mode,
))
print("START STREAM:", svc.start_stream(PC_HOST, STREAM_PORT, fps=15))
win = "CAMERA MULTISPEC + SEGFORMER (Q=quit)"
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
last_frame_id = -1
try:
while True:
meta = receiver.last_meta
frame = receiver.last_frame
if meta is None or frame is None:
continue
frame_id = meta.get("frame_id")
if frame_id == last_frame_id:
continue
last_frame_id = frame_id
try:
raw_np = core.build_infer_tensor_from_stream(frame, meta, channels_expected=CHANNELS)
# =========================================================
# INFERÊNCIA
# =========================================================
pred_ids, preview_bgr, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_np, alpha=args.alpha)
# =========================================================
# HUD
# =========================================================
txt = f"C={CHANNELS} | inf={t_inf:.1f}ms | pvw={t_pvw:.1f}ms"
cv2.putText(
overlay,
txt,
(10, 24),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
(0, 255, 0),
2,
)
cv2.imshow(win, overlay)
except Exception as e:
print("[ERRO FRAME]", e)
key = cv2.waitKey(1) & 0xFF
if key in (ord("q"), ord("Q"), 27):
break
finally:
try:
svc.stop_stream()
except:
pass
try:
svc.stop()
except:
pass
svc.disconnect()
receiver.stop()
cv2.destroyAllWindows()
return
# ======================================================================
# MODO DATASET / TEST_FOLDER
# ======================================================================
if args.test_folder is not None:
root = args.test_folder
elif args.root_override is not None:
root = args.root_override
else:
root = os.path.join(dataset_path, "split", args.split_folder)
print(f"[data] root = {root}")
has_gt = True
ds = None
tensor_paths = []
try:
ds = MultispecSegDataset(
root,
labelmap_path=labelmap_path,
resize_hw=None,
channels=CHANNELS,
strict_channels=True,
)
print("[mode] Dataset com GT (masks) detectado. Mostrando GT + overlay.")
except RuntimeError as e:
msg = str(e)
if "Nenhum par tensor/mask encontrado" in msg:
print("[mode] Nenhum par tensor/mask encontrado. Entrando em modo inferência-only (sem GT).")
has_gt = False
tensor_paths = _collect_tensor_paths(root)
if not tensor_paths:
raise RuntimeError(f"Modo inferência-only: não encontrei nenhum tensor em {root}")
print(f"[data] Tensors encontrados: {len(tensor_paths)}")
else:
raise
n = len(ds) if has_gt else len(tensor_paths)
idx = 0
print(f"[data] total de amostras: {n}")
print("Controles: D=próxima, A=anterior, Q=sair")
win_name = "Tensor preview | GT | Overlay (SegFormer)" if has_gt else "Tensor preview | Overlay (SegFormer)"
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
while True:
if has_gt:
sample = ds[idx]
img_tensor = sample["image"]
mask_gt = sample["mask"]
raw_np = img_tensor.cpu().numpy().astype(np.float32)
pred_ids, preview_bgr, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_np, alpha=args.alpha)
mask_gt_np = mask_gt.cpu().numpy().astype(np.uint8)
gt_rgb = converter_mask_ids_para_bgr(mask_gt_np, colormap_rgb, ignore_id)
h, w, _ = preview_bgr.shape
gt_resized = cv2.resize(gt_rgb, (w, h), interpolation=cv2.INTER_NEAREST)
overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_NEAREST)
resultado = np.concatenate([preview_bgr, gt_resized, overlay_resized], axis=1)
legenda = desenhar_legenda_horizontal(colormap_rgb, classes)
legenda_resized = cv2.resize(legenda, (resultado.shape[1], legenda.shape[0]), interpolation=cv2.INTER_NEAREST)
resultado_completo = np.concatenate([resultado, legenda_resized], axis=0)
header_txt = f"idx {idx + 1}/{n} | C={CHANNELS} | inf={t_inf:.1f}ms | pvw={t_pvw:.1f}ms"
else:
tensor_path = tensor_paths[idx]
raw_np = np.load(tensor_path).astype(np.float32)
pred_ids, preview_bgr, _, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_np, alpha=args.alpha)
h, w, _ = preview_bgr.shape
overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_NEAREST)
resultado = np.concatenate([preview_bgr, overlay_resized], axis=1)
resultado_completo = resultado
header_txt = f"idx {idx + 1}/{n} | C={CHANNELS} | {os.path.basename(tensor_path)} | inf={t_inf:.1f}ms | pvw={t_pvw:.1f}ms"
cv2.putText(resultado_completo, header_txt, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
try:
_x, _y, win_w, win_h = cv2.getWindowImageRect(win_name)
except Exception:
win_w, win_h = 0, 0
if win_w > 0 and win_h > 0:
display = cv2.resize(resultado_completo, (win_w, win_h), interpolation=cv2.INTER_NEAREST)
else:
display = resultado_completo
display_bgr = cv2.cvtColor(display, cv2.COLOR_RGB2BGR)
cv2.imshow(win_name, display_bgr)
key = cv2.waitKey(0) & 0xFF
if key in (ord("q"), ord("Q"), 27):
break
elif key in (ord("d"), ord("D")):
idx = (idx + 1) % n
elif key in (ord("a"), ord("A")):
idx = (idx - 1 + n) % n
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

@ -1,535 +0,0 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Teste/visualização do SegFormer RAW (4 ou 5 canais) treinado no _8_train_segformer_b3_raw.py
Modos:
1) Dataset (com GT):
- Usa RawSegDataset (raws + masks) a partir de split (train/val/test)
- Mostra: [preview RAW (RGB)] | [mask GT] | [overlay predito]
2) Pasta sem máscara (--test_folder sem masks):
- Faz inferência com os RAWs (sem GT)
- Mostra: [preview RAW (RGB)] | [overlay predito]
3) Câmera ao vivo (--camera):
- Usa Gal5000Camera para pegar RAW4 (R,G,IR,B)
- Monta os 4/5 canais conforme config
- Mostra overlay em tempo real
"""
import os
import json
import time
import argparse
from collections import deque
import cv2
import torch
import numpy as np
from utils import converter_mask_ids_para_bgr, desenhar_legenda_horizontal
from raw_segformer_service import RawSegformerService, RawSegDataset # :contentReference[oaicite:0]{index=0}
from gal5000.gal_service import Gal5000Camera
# ============================================================
# Helpers para modo "sem máscara"
# (copia condensada da lógica do RawSegDataset)
# ============================================================
_RAW_EXTS = (".npy", ".npz", ".raw")
def _load_raw_4ch(path: str, raw_hw):
"""
um RAW no formato usado pelo pipeline:
- mosaico uint8 antigo, ou
- RAW4 float32 novo, shape (4,H,W), salvo pelo normalize
Retorna: arr (H,W,4) na ordem [R,G,B,IR] em dtype original.
"""
import os
H, W = raw_hw
ext = os.path.splitext(path)[1].lower()
if ext == ".npz":
z = np.load(path)
key = list(z.keys())[0]
arr = z[key]
elif ext == ".npy":
arr = np.load(path)
elif ext == ".raw":
size_bytes = os.path.getsize(path)
mosa_bytes = H * W # mosaico uint8
raw4_bytes = 4 * H * W * 4 # 4 canais * H * W * 4 bytes (float32)
if size_bytes == mosa_bytes:
# Modo antigo: mosaico 2x2 uint8
arr_flat = np.fromfile(path, dtype=np.uint8)
raw2d = arr_flat.reshape(H, W)
if (H % 2) != 0 or (W % 2) != 0:
raise RuntimeError(f"raw_hw deve ser par em H e W (veio H={H}, W={W})")
H2, W2 = H // 2, W // 2
r_sub = raw2d[0::2, 0::2]
g_sub = raw2d[0::2, 1::2]
ir_sub = raw2d[1::2, 0::2]
b_sub = raw2d[1::2, 1::2]
from PIL import Image
def upsample(ch_2d: np.ndarray) -> np.ndarray:
im = Image.fromarray(ch_2d) # uint8
im = im.resize((W, H), resample=Image.BILINEAR)
return np.array(im)
r_full = upsample(r_sub)
g_full = upsample(g_sub)
ir_full = upsample(ir_sub)
b_full = upsample(b_sub)
# (H, W, 4) uint8 [R,G,B,IR]
arr = np.stack([r_full, g_full, b_full, ir_full], axis=-1).astype(np.uint8)
elif size_bytes == raw4_bytes:
# Novo modo: RAW4 float32 salvo pelo normalize_raw
arr_f32 = np.fromfile(path, dtype=np.float32)
raw4 = arr_f32.reshape(4, H, W) # (C,H,W) [R,G,IR,B]
arr = np.transpose(raw4, (1, 2, 0)) # (H,W,4) [R,G,IR,B]
# Reorganiza pra contrato interno: [R,G,B,IR]
arr = arr[..., [0, 1, 3, 2]]
else:
raise RuntimeError(
f"Tamanho inesperado em {path}: {size_bytes} bytes "
f"(esperado {mosa_bytes} ou {raw4_bytes})"
)
else:
raise RuntimeError(f"Extensão não suportada para RAW: {ext} ({path})")
if arr.ndim != 3:
raise RuntimeError(f"RAW precisa ser 3D, veio {arr.shape} em {path}")
if arr.shape[0] == 4 and arr.shape[-1] != 4:
arr = np.transpose(arr, (1, 2, 0)) # (H,W,4)
if arr.shape[-1] != 4:
raise RuntimeError(f"Esperava 4 canais base, veio shape={arr.shape} em {path}")
return arr # (H,W,4) [R,G,B,IR]
def _scale_to_float01(raw: np.ndarray, max_value: float | None = None) -> np.ndarray:
"""
Normaliza para float32 0..1 (igual RawSegDataset._scale_to_float01):
- uint8 / uint16: divide por max_value ou 255/65535
- float: usa max_value se fornecido; senão assume 0..1
"""
if raw.dtype == np.uint16:
mv = float(max_value) if max_value is not None else 65535.0
elif raw.dtype == np.uint8:
mv = float(max_value) if max_value is not None else 255.0
else:
mv = float(max_value) if max_value is not None else None
raw_f = raw.astype(np.float32)
if mv is not None and mv > 0:
raw_f /= mv
raw_f = np.clip(raw_f, 0.0, 1.0)
return raw_f
def _collect_raw_paths(root: str) -> list[str]:
"""
Coleta caminhos de RAWs em 'root' usando a mesma convenção do dataset:
- se root/group existe: procura group/*/raws
- senão: procura root/raws
"""
paths = []
dir_group = os.path.join(root, "group")
if os.path.isdir(dir_group):
# modo com grupos
for g in sorted(os.listdir(dir_group)):
gdir = os.path.join(dir_group, g)
if not os.path.isdir(gdir):
continue
g_raw = os.path.join(gdir, "raws")
if not os.path.isdir(g_raw):
continue
for fn in sorted(os.listdir(g_raw)):
if fn.lower().endswith(_RAW_EXTS):
paths.append(os.path.join(g_raw, fn))
else:
# modo simples
dir_raws = os.path.join(root, "raws")
if not os.path.isdir(dir_raws):
raise RuntimeError(f"Modo sem máscara: não achei pasta 'raws' em {root}")
for fn in sorted(os.listdir(dir_raws)):
if fn.lower().endswith(_RAW_EXTS):
paths.append(os.path.join(dir_raws, fn))
return paths
# ============================================================
# Script principal
# ============================================================
def main():
parser = argparse.ArgumentParser()
# Config / treino
parser.add_argument("--config", default="config.json", help="Caminho do config.json (mesmo do treino)")
parser.add_argument("--split_folder", type=str, default="val", help="train, val ou test (dentro de dataset/split)")
parser.add_argument("--root_override", type=str, default=None, help="Aponta direto para uma pasta split (sobrepõe split_folder)")
parser.add_argument(
"--test_folder",
type=str,
default=None,
help=(
"Pasta de teste para visualização. "
"Se tiver masks, usa dataset normal (mostra GT). "
"Se tiver apenas raws, entra em modo inferência-only."
)
)
parser.add_argument("--ckpt", type=str, default=None, help="Caminho do .pt (se não informar, usa best_miou.pt da pasta rawX)")
parser.add_argument("--raw_max", type=float, default=None, help="Mesmo raw_max do treino (se usou).")
# Modo câmera
parser.add_argument("--camera", action="store_true", help="Usa a GAL5000 ao vivo via DLL (realtime)")
parser.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal (ex: norm_stats.json).")
parser.add_argument(
"--dll_dir",
type=str,
default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls",
help="Pasta onde fica VT_SDK64.dll e dependências",
)
parser.add_argument("--dll_name", type=str, default="VT_SDK64.dll")
parser.add_argument("--resize_h", type=int, default=None, help="Altura para inferência (override)")
parser.add_argument("--resize_w", type=int, default=None, help="Largura para inferência (override)")
parser.add_argument("--alpha", type=float, default=0.45, help="Alpha do overlay da máscara")
parser.add_argument("--fps_win", type=int, default=30, help="Janela para FPS médio")
parser.add_argument("--timeout_ms", type=int, default=2000, help="Timeout da captura da câmera em ms")
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
# ---- Carrega config ----
with open(args.config, "r") as f:
config = json.load(f)
MODELO = config["camera"] # ex: "gal5000"
MODEL_NAME = config["model_name"] # ex: "segformer_b0"
#MODEL_NAME = "pulv_s3_1008x800"
modelo_folder = config["modelo"] # ex: "weed_1008x800"
USE_NDVI = bool(config.get("use_ndvi", False))
CHANNELS = int(config.get("channels", 4))
FUSION_MODE = config.get("fusion_mode", "stacked")
RESOLUCAO = config["resolucao"]
W, H = RESOLUCAO[0], RESOLUCAO[1]
if args.resize_h is not None:
H = args.resize_h
if args.resize_w is not None:
W = args.resize_w
print(f"[cfg] resolucao nominal RAW: {W}x{H}")
print(f"[cfg] channels={CHANNELS} use_ndvi={USE_NDVI}")
dataset_path = os.path.join("dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
# ==========================
# Normalização fixa (igual treino)
# ==========================
norm_mean = None
norm_std = None
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
norm_stats_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag, "norm_stats.json")
if args.norm_stats is not None:
norm_stats_path = args.norm_stats
if os.path.isfile(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}")
idx_by_name = {name: i for i, name in enumerate(stats_channels)}
m_R = stats_mean[idx_by_name["R"]]
m_G = stats_mean[idx_by_name["G"]]
m_IR = stats_mean[idx_by_name["IR"]]
m_B = stats_mean[idx_by_name["B"]]
s_R = stats_std[idx_by_name["R"]]
s_G = stats_std[idx_by_name["G"]]
s_IR = stats_std[idx_by_name["IR"]]
s_B = stats_std[idx_by_name["B"]]
if CHANNELS == 3 and not USE_NDVI:
norm_mean = [m_R, m_G, m_B]
norm_std = [s_R, s_G, s_B]
elif CHANNELS == 4 and not USE_NDVI:
norm_mean = [m_R, m_G, m_B, m_IR]
norm_std = [s_R, s_G, s_B, s_IR]
elif CHANNELS == 4 and USE_NDVI:
norm_mean = [m_R, m_G, m_B, 0.0]
norm_std = [s_R, s_G, s_B, 1.0]
elif CHANNELS == 5 and USE_NDVI:
norm_mean = [m_R, m_G, m_B, m_IR, 0.0]
norm_std = [s_R, s_G, s_B, s_IR, 1.0]
else:
print(f"[NORM] Config (channels={CHANNELS}, use_ndvi={USE_NDVI}) "
f"não compatível com norm_stats fixo. Mantendo normalização dinâmica.")
norm_mean = None
norm_std = None
else:
print(f"[NORM] norm_stats.json não encontrado em {norm_stats_path}. "
f"Usando normalização dinâmica por frame.")
# ---- Descobre checkpoint ----
if args.ckpt is not None:
ckpt_path = args.ckpt
else:
save_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag)
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}")
# ---- Instancia service do modelo RAW ----
model_svc = RawSegformerService(
config_path=args.config,
ckpt_path=ckpt_path,
device=device,
use_amp=True,
mean=norm_mean,
std=norm_std,
)
classes = model_svc.get_classes()
colormap_rgb = model_svc.get_colormap()
ignore_id = model_svc.get_ignore_id()
print(f"[svc] classes: {classes}")
print(f"[svc] ignore_id={ignore_id}")
# ======================================================================
# MODO CÂMERA AO VIVO
# ======================================================================
if args.camera:
print("[mode] Câmera ao vivo (GAL5000 + SegFormer RAW)")
cam = Gal5000Camera(
dll_dir=args.dll_dir,
dll_name=args.dll_name,
raw_w=W,
raw_h=H,
use_auto_exposure=True
)
win = "GAL5000 RAW + SegFormer (Q=quit)"
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
tq = deque(maxlen=max(5, int(args.fps_win)))
with cam:
cam.configure_fps(20)
cam.start_streaming()
while True:
t0 = time.time()
raw4_base, dbg = cam.grab_raw4(out_h=H, out_w=W, timeout_ms=args.timeout_ms)
t1 = time.time()
r = raw4_base[0]
g = raw4_base[1]
ir = raw4_base[2]
b = raw4_base[3]
raw_input = model_svc.build_raw_input(r, g, ir, b)
t2 = time.time()
pred_ids = model_svc.infer_raw(raw_input)
t3 = time.time()
_, _, overlay, _, t_pvw = model_svc.preview_infer_cached(raw_input, pred_ids, alpha=args.alpha)
t4 = time.time()
tq.append(t4 - t0)
fps = 1.0 / (sum(tq) / len(tq))
status = cam.get_status()
shape = dbg["raw_shape"]
tc = dbg.get("t_capture", 0) * 1000
tconv = dbg.get("t_convert", 0) * 1000
t_ae = dbg.get("t_ae", 0) * 1000
t_raw = (t2 - t1) * 1000
t_inf = (t3 - t2) * 1000
lines = [
f"{W}x{H} | FPS~{fps:.1f} | {shape}",
f"EX={status['exp_raw']} G={status['gain_a']}/{status['gain_d']}",
f"C={CHANNELS} NDVI={int(USE_NDVI)}",
f"cap={tc:.1f}ms ae={t_ae:.1f}ms conv={tconv:.1f}ms raw={t_raw:.1f}ms inf={t_inf:.1f}ms pvw={t_pvw:.1f}ms",
]
y = 24
for line in lines:
cv2.putText(overlay, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(255, 255, 255), 1, cv2.LINE_AA)
y += 22
cv2.imshow(win, overlay)
k = cv2.waitKey(1) & 0xFF
if k in (ord("q"), ord("Q"), 27):
break
cv2.destroyAllWindows()
return
# ======================================================================
# MODO DATASET / TEST_FOLDER
# ======================================================================
# ---- Define raiz do split/test ----
if args.test_folder is not None:
raw_root = args.test_folder
elif args.root_override is not None:
raw_root = args.root_override
else:
raw_root = os.path.join(dataset_path, "split", args.split_folder)
print(f"[data] raw_root = {raw_root}")
has_gt = True
ds = None
raw_paths = []
try:
ds = RawSegDataset(
raw_root,
labelmap_path=labelmap_path,
max_value=args.raw_max,
resize_hw=None,
raw_hw=(H, W),
use_ndvi=USE_NDVI,
channels=CHANNELS,
)
print("[mode] Dataset com GT (masks) detectado. Mostrando GT + overlay.")
except RuntimeError as e:
msg = str(e)
if "Nenhum par raw/mask encontrado" in msg:
print("[mode] Nenhum par raw/mask encontrado. Entrando em modo inferência-only (sem GT).")
has_gt = False
raw_paths = _collect_raw_paths(raw_root)
if not raw_paths:
raise RuntimeError(f"Modo inferência-only: não encontrei nenhum RAW em {raw_root}")
print(f"[data] RAWs encontrados: {len(raw_paths)}")
else:
raise
if has_gt:
n = len(ds)
else:
n = len(raw_paths)
idx = 0
print(f"[data] total de amostras: {n}")
print("Controles: D=próxima, A=anterior, Q=sair")
win_name = "RAW preview | GT | Overlay (SegFormer)" if has_gt else "RAW preview | Overlay (SegFormer)"
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
while True:
if has_gt:
# ------- Caminho com GT -------
sample = ds[idx]
img_raw = sample["image"] # tensor (C,H,W) em 0..1
mask_gt = sample["mask"] # tensor (H,W) long
raw_np = img_raw.cpu().numpy().astype(np.float32)
pred_ids, preview_rgb, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(
raw_np, alpha=args.alpha
)
mask_gt_np = mask_gt.cpu().numpy().astype(np.uint8)
gt_rgb = converter_mask_ids_para_bgr(mask_gt_np, colormap_rgb, ignore_id)
h, w, _ = preview_rgb.shape
gt_resized = cv2.resize(gt_rgb, (w, h), interpolation=cv2.INTER_NEAREST)
overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_NEAREST)
resultado = np.concatenate([preview_rgb, gt_resized, overlay_resized], axis=1)
legenda = desenhar_legenda_horizontal(colormap_rgb, classes)
legenda_resized = cv2.resize(legenda, (resultado.shape[1], legenda.shape[0]), interpolation=cv2.INTER_NEAREST)
resultado_completo = np.concatenate([resultado, legenda_resized], axis=0)
header_txt = f"idx {idx + 1}/{n} | C={CHANNELS} NDVI={int(USE_NDVI)}"
else:
# ------- Caminho sem GT (só RAW) -------
raw_path = raw_paths[idx]
raw4 = _load_raw_4ch(raw_path, raw_hw=(H, W)) # (H,W,4) [R,G,B,IR]
raw4_f = _scale_to_float01(raw4, max_value=args.raw_max) # float32 0..1
# Monta canais finais para o modelo (4 ou 5) via serviço
r = raw4_f[..., 0]
g = raw4_f[..., 1]
b = raw4_f[..., 2]
ir = raw4_f[..., 3]
raw_input = model_svc.build_raw_input(r, g, ir, b) # (C,H,W) float32
pred_ids = model_svc.infer_raw(raw_input)
_, preview_rgb, _, overlay, _, _ = model_svc.preview_infer_cached(raw_input, pred_ids, alpha=args.alpha)
h, w, _ = preview_rgb.shape
overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_NEAREST)
resultado = np.concatenate([preview_rgb, overlay_resized], axis=1)
header_txt = f"idx {idx + 1}/{n} | C={CHANNELS} NDVI={int(USE_NDVI)} | {os.path.basename(raw_path)}"
resultado_completo = resultado
# Escreve header
cv2.putText(resultado_completo, header_txt, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
# Ajuste dinâmico da janela
try:
_x, _y, win_w, win_h = cv2.getWindowImageRect(win_name)
except Exception:
win_w, win_h = 0, 0
if win_w > 0 and win_h > 0:
display = cv2.resize(resultado_completo, (win_w, win_h), interpolation=cv2.INTER_NEAREST)
else:
display = resultado_completo
cv2.imshow(win_name, display)
key = cv2.waitKey(0) & 0xFF
if key in (ord("q"), ord("Q"), 27):
break
elif key in (ord("d"), ord("D")):
idx = (idx + 1) % n
elif key in (ord("a"), ord("A")):
idx = (idx - 1 + n) % n
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

@ -6,14 +6,14 @@
"main_class_name": "cana", "main_class_name": "cana",
"es_classes": "", "es_classes": "",
"model_to_use": "geral", "model_to_use": "geral",
"raw_size": [1296, 1028], "raw_size": [640, 480],
"resolucao": [1024, 800], "resolucao": [1024, 800],
"roi_inicio": 0.0, "roi_inicio": 0.0,
"roi_tamanho": 1.0, "roi_tamanho": 1.0,
"shaves": 3, "shaves": 3,
"channels": 5, "channels": 4,
"use_ndvi": false, "use_ndvi": false,
"backbone": "nvidia/mit-b1", "backbone": "nvidia/mit-b1",
"fusion_mode": "stacked", "fusion_mode": "stacked",
"stats_source_tag": "stacked_raw5" "stats_source_tag": "stacked_raw4"
} }

View File

@ -0,0 +1,683 @@
# multispec_segformer_service.py
# Serviço e dataset para SegFormer no novo pipeline multiespectral.
#
# Contrato principal do tensor salvo:
# (C,H,W) float32 em 0..1
# onde C pode ser 3, 4 ou 5, conforme config/channels.
#
# Estruturas aceitas:
# 1) Simples:
# root/
# tensors/*.npy
# masks/*.npy
#
# 2) Com grupos:
# root/
# group/
# <grupo>/
# tensors/*.npy
# masks/*.npy
#
# A máscara oficial é .npy (IDs inteiros). O .png fica apenas para debug.
import os
import json
from contextlib import nullcontext
import sys
from typing import List, Optional, Tuple
import threading
import time
import numpy as np
import torch
import torch.nn as nn
from torch import amp
from torch.utils.data import Dataset
import torch.nn.functional as F
from transformers import SegformerForSemanticSegmentation
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo, converter_mask_ids_para_bgr
# ============================================================
# Model builders
# ============================================================
def build_multispec_segformer_model(
num_classes: int,
channels: int,
backbone: str = "nvidia/segformer-b3-finetuned-ade-512-512",
device: torch.device | None = None,
ckpt_path: str | None = None,
strict: bool = True,
) -> SegformerForSemanticSegmentation:
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SegformerForSemanticSegmentation.from_pretrained(
backbone,
num_labels=num_classes,
ignore_mismatched_sizes=True,
use_safetensors=True,
)
patch_segformer_input_channels(model, in_ch=channels)
if ckpt_path is not None:
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
if "model" not in ckpt:
raise RuntimeError(f"Checkpoint {ckpt_path} não contém chave 'model'.")
model.load_state_dict(ckpt["model"], strict=strict)
model.to(device)
return model
def build_dual_branch_segformer_model(
num_classes: int,
channels: int,
backbone: str = "nvidia/segformer-b1-finetuned-ade-512-512",
device: torch.device | None = None,
ckpt_path: str | None = None,
strict: bool = True,
) -> nn.Module:
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if channels < 4:
raise ValueError(f"Dual branch espera channels>=4, veio {channels}")
spec_channels = channels - 3
model = DualBranchSegformerV2(
num_classes=num_classes,
backbone=backbone,
spec_channels=spec_channels,
)
if ckpt_path is not None:
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
if "model" not in ckpt:
raise RuntimeError(f"Checkpoint {ckpt_path} não contém chave 'model'.")
model.load_state_dict(ckpt["model"], strict=strict)
model.to(device)
return model
# ============================================================
# Dual-branch blocks
# ============================================================
class ResidualBlock(nn.Module):
def __init__(self, ch: int):
super().__init__()
self.conv1 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(ch)
self.act = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(ch)
def forward(self, x):
identity = x
x = self.act(self.bn1(self.conv1(x)))
x = self.bn2(self.conv2(x))
x = self.act(x + identity)
return x
class SpectralEncoderV2(nn.Module):
def __init__(self, in_ch: int, out_ch: int = 256):
super().__init__()
self.stem = nn.Sequential(
nn.Conv2d(in_ch, 32, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
)
self.stage1 = nn.Sequential(
ResidualBlock(32),
nn.Conv2d(32, 64, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
self.stage2 = nn.Sequential(
ResidualBlock(64),
nn.Conv2d(64, 128, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
)
self.stage3 = nn.Sequential(
ResidualBlock(128),
nn.Conv2d(128, out_ch, 1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
x = self.stem(x)
x = self.stage1(x)
x = self.stage2(x)
x = self.stage3(x)
return x
class FeatureFusionBlock(nn.Module):
def __init__(self, ch_rgb: int, ch_spec: int, ch_fused: int):
super().__init__()
self.rgb_proj = nn.Conv2d(ch_rgb, ch_fused, 1, bias=False)
self.spec_proj = nn.Conv2d(ch_spec, ch_fused, 1, bias=False)
self.fuse = nn.Sequential(
nn.Conv2d(ch_fused * 2, ch_fused, 3, padding=1, bias=False),
nn.BatchNorm2d(ch_fused),
nn.ReLU(inplace=True),
nn.Conv2d(ch_fused, ch_fused, 3, padding=1, bias=False),
nn.BatchNorm2d(ch_fused),
nn.ReLU(inplace=True),
)
self.gate = nn.Sequential(
nn.Conv2d(ch_fused, ch_fused, 1),
nn.Sigmoid(),
)
def forward(self, rgb_feat, spec_feat):
rgb_feat = self.rgb_proj(rgb_feat)
spec_feat = self.spec_proj(spec_feat)
if spec_feat.shape[-2:] != rgb_feat.shape[-2:]:
spec_feat = F.interpolate(spec_feat, size=rgb_feat.shape[-2:], mode="bilinear", align_corners=False)
fused = torch.cat([rgb_feat, spec_feat], dim=1)
fused = self.fuse(fused)
gate = self.gate(fused)
fused = fused * gate
return fused
class DualBranchSegformerV2(nn.Module):
def __init__(self, num_classes: int, backbone: str, spec_channels: int, fused_channels: int = 256):
super().__init__()
self.rgb_model = SegformerForSemanticSegmentation.from_pretrained(
backbone,
num_labels=num_classes,
ignore_mismatched_sizes=True,
use_safetensors=True,
)
patch_segformer_input_channels(self.rgb_model, in_ch=3)
self.spec_branch = SpectralEncoderV2(in_ch=spec_channels, out_ch=fused_channels)
self.fusion = FeatureFusionBlock(
ch_rgb=self.rgb_model.config.hidden_sizes[-1],
ch_spec=fused_channels,
ch_fused=fused_channels,
)
self.classifier = nn.Sequential(
nn.Conv2d(fused_channels, fused_channels, 3, padding=1, bias=False),
nn.BatchNorm2d(fused_channels),
nn.ReLU(inplace=True),
nn.Conv2d(fused_channels, num_classes, 1, bias=True),
)
def forward(self, pixel_values: torch.Tensor):
rgb = pixel_values[:, :3, :, :]
spec = pixel_values[:, 3:, :, :]
rgb_out = self.rgb_model(pixel_values=rgb, output_hidden_states=True)
rgb_logits = rgb_out.logits
rgb_feat = rgb_out.hidden_states[-1]
spec_feat = self.spec_branch(spec)
fused_feat = self.fusion(rgb_feat, spec_feat)
logits = self.classifier(fused_feat)
if logits.shape[-2:] != rgb_logits.shape[-2:]:
logits = F.interpolate(logits, size=rgb_logits.shape[-2:], mode="bilinear", align_corners=False)
return type("DualBranchOutput", (), {"logits": logits})
# ============================================================
# Dataset
# ============================================================
class MultispecSegDataset(Dataset):
"""
Espera tensores finais prontos em .npy.
1) Simples:
root/
tensors/*.npy
masks/*.npy
2) Com grupos:
root/
group/
<grupo>/
tensors/*.npy
masks/*.npy
Saída:
- image: tensor (C,H,W) float32
- mask : (H,W) long
"""
def __init__(
self,
root: str,
labelmap_path: Optional[str] = None,
resize_hw: Optional[Tuple[int, int]] = None,
channels: int = 5,
strict_channels: bool = True,
):
super().__init__()
self.root = root
self.resize_hw = resize_hw
self.channels = int(channels)
self.strict_channels = strict_channels
self.items: List[Tuple[str, str]] = []
if self.channels not in (3, 4, 5, 6, 7):
raise ValueError(f"[dataset] channels inválido: {self.channels}")
dir_group = os.path.join(root, "group")
dir_tensor_simple = os.path.join(root, "tensors")
dir_mask_simple = os.path.join(root, "masks")
if os.path.isdir(dir_group):
for g in sorted(os.listdir(dir_group)):
gdir = os.path.join(dir_group, g)
if not os.path.isdir(gdir):
continue
g_tensors = os.path.join(gdir, "tensors")
g_masks = os.path.join(gdir, "masks")
if not (os.path.isdir(g_tensors) and os.path.isdir(g_masks)):
continue
self._append_pairs(g_tensors, g_masks)
else:
if not os.path.isdir(dir_tensor_simple):
raise RuntimeError(f"Não achei pasta tensors: {dir_tensor_simple}")
if not os.path.isdir(dir_mask_simple):
raise RuntimeError(f"Não achei pasta masks: {dir_mask_simple}")
self._append_pairs(dir_tensor_simple, dir_mask_simple)
if not self.items:
raise RuntimeError(f"Nenhum par tensor/mask encontrado em {root} (nomes base precisam bater).")
_, _, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
self.classes = classes
self.class_ids = sorted(classes.keys())
self.ignore_id = ignore_id
print(f"[dataset] root={root} items={len(self.items)} channels={self.channels}")
def _append_pairs(self, tensor_dir: str, mask_dir: str):
tensors = [fn for fn in sorted(os.listdir(tensor_dir)) if fn.lower().endswith(".npy")]
for fn in tensors:
stem = os.path.splitext(fn)[0]
m = os.path.join(mask_dir, stem + ".npy")
if os.path.exists(m):
self.items.append((os.path.join(tensor_dir, fn), m))
def __len__(self):
return len(self.items)
def _resize_tensor(self, x: np.ndarray, new_hw: Tuple[int, int]) -> np.ndarray:
h, w = new_hw
chans = []
for ch in x:
ch_res = F.interpolate(
torch.from_numpy(ch[None, None, :, :]).float(),
size=(h, w),
mode="bilinear",
align_corners=False,
)[0, 0].numpy()
chans.append(ch_res.astype(np.float32))
return np.stack(chans, axis=0)
def _resize_mask(self, y: np.ndarray, new_hw: Tuple[int, int]) -> np.ndarray:
h, w = new_hw
out = F.interpolate(
torch.from_numpy(y[None, None, :, :]).float(),
size=(h, w),
mode="nearest",
)[0, 0].numpy()
return out.astype(np.int64)
def __getitem__(self, idx: int):
tensor_path, mask_path = self.items[idx]
x = np.load(tensor_path).astype(np.float32) # (C,H,W)
y = np.load(mask_path).astype(np.int64) # (H,W)
if x.ndim != 3:
raise RuntimeError(f"Tensor precisa ser (C,H,W), veio {x.shape} em {tensor_path}")
if y.ndim != 2:
raise RuntimeError(f"Mask precisa ser (H,W), veio {y.shape} em {mask_path}")
if self.strict_channels and x.shape[0] != self.channels:
raise RuntimeError(
f"Tensor com número de canais inesperado: {x.shape[0]} em {tensor_path}. "
f"Esperado {self.channels}."
)
if self.resize_hw is not None:
x = self._resize_tensor(x, self.resize_hw)
y = self._resize_mask(y, self.resize_hw)
return {
"image": torch.from_numpy(x).float(),
"mask": torch.from_numpy(y).long(),
}
# ============================================================
# Normalização
# ============================================================
def normalize_per_batch(imgs: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
mean = imgs.mean(dim=(0, 2, 3), keepdim=True)
std = imgs.std(dim=(0, 2, 3), keepdim=True).clamp_min(eps)
return (imgs - mean) / std
def build_fixed_normalizer(mean: List[float], std: List[float], device: torch.device):
mean_t = torch.tensor(mean, dtype=torch.float32, device=device).view(1, -1, 1, 1)
std_t = torch.tensor(std, dtype=torch.float32, device=device).view(1, -1, 1, 1).clamp_min(1e-6)
def normalizer(x: torch.Tensor) -> torch.Tensor:
return (x - mean_t) / std_t
return normalizer
# ============================================================
# Preview helpers
# ============================================================
def make_bgr_preview_from_tensor(raw_np: np.ndarray, preview_fast: bool = False, preview_scale: int = 2) -> np.ndarray:
"""
Preview sempre usando os três primeiros canais como RGB.
Espera raw_np (C,H,W) float32 em 0..1.
"""
assert raw_np.ndim == 3 and raw_np.shape[0] >= 3
r = raw_np[0].astype(np.float32, copy=False)
g = raw_np[1].astype(np.float32, copy=False)
b = raw_np[2].astype(np.float32, copy=False)
if preview_fast and preview_scale > 1:
r = r[::preview_scale, ::preview_scale]
g = g[::preview_scale, ::preview_scale]
b = b[::preview_scale, ::preview_scale]
def stretch_channel(x: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray:
lo = np.percentile(x, p_low)
hi = np.percentile(x, p_high)
if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
return np.clip(x, 0.0, 1.0)
x = (x - lo) / (hi - lo)
return np.clip(x, 0.0, 1.0)
if not preview_fast:
r = stretch_channel(r)
g = stretch_channel(g)
b = stretch_channel(b)
bgr = np.stack([b, g, r], axis=0)
gamma = 1 / (1.8 if preview_fast else 2.0)
bgr = np.power(np.clip(bgr, 0.0, 1.0), gamma)
bgr8 = (bgr * 255.0).clip(0, 255).astype(np.uint8)
return np.transpose(bgr8, (1, 2, 0)).copy()
# ============================================================
# Patch input channels
# ============================================================
def patch_segformer_input_channels(model: SegformerForSemanticSegmentation, in_ch: int = 4):
enc = model.segformer.encoder
proj = enc.patch_embeddings[0].proj
if proj.in_channels == in_ch:
return
old_w = proj.weight.data
old_b = proj.bias.data 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,
).to(proj.weight.device)
with torch.no_grad():
if old_w.shape[1] == 3 and in_ch >= 3:
new_proj.weight[:, :3, :, :].copy_(old_w)
if in_ch > 3:
extra = old_w.mean(dim=1, keepdim=True)
for c in range(3, in_ch):
new_proj.weight[:, c:c+1, :, :].copy_(extra)
else:
nn.init.kaiming_normal_(new_proj.weight, mode="fan_out", nonlinearity="relu")
if old_b is not None:
new_proj.bias.copy_(old_b)
enc.patch_embeddings[0].proj = new_proj
# ============================================================
# Service
# ============================================================
class MultispecSegformerService:
"""
Serviço para inferência no novo pipeline multiespectral.
Espera tensores finais (C,H,W) prontos, em 0..1.
"""
def __init__(
self,
config_path: str,
ckpt_path: str | None = None,
labelmap_path: str | None = None,
device: torch.device | None = None,
use_amp: bool = True,
mean=None,
std=None,
load_model=True,
):
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.device = device
self.use_amp = use_amp and (self.device.type == "cuda")
torch.backends.cudnn.benchmark = True
if device.type == "cuda":
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
self._norm_mean = None
self._norm_std = None
if mean is not None and std is not None:
self.set_norm_stats(mean, std)
self._buf_shape = None
self._cpu_pinned = None
self._gpu_input = None
with open(config_path, "r", encoding="utf-8") as f:
cfg = json.load(f)
self.cfg = cfg
self.MODELO = cfg["camera"]
self.MODEL_NAME = cfg["model_name"]
self.modelo_folder = cfg["modelo"]
self.res_w, self.res_h = cfg["resolucao"]
self.backbone = cfg["backbone"]
self.fusion_mode = cfg.get("fusion_mode", "stacked")
self.channels = int(cfg.get("channels", 5))
self._preview_lock = threading.Lock()
self._preview_busy = False
self._last_preview = None
self._last_preview_ts = 0.0
self._last_preview_dt = 0.0
if labelmap_path is None:
labelmap_path = os.path.join("dataset", "labelmap.txt")
self.labelmap_path = labelmap_path
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
self.colormap_rgb = colormap_rgb
self.classes = classes
self.ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
self.num_classes = len(classes)
if load_model:
if ckpt_path is None:
save_path = os.path.join("backup", self.modelo_folder, self.MODEL_NAME, f"{self.fusion_mode}_raw{self.channels}")
ckpt_path = os.path.join(save_path, "best_miou.pt")
if not os.path.isfile(ckpt_path):
raise FileNotFoundError(f"Checkpoint não encontrado: {ckpt_path}")
self.ckpt_path = ckpt_path
self.model = self._load_model(backbone=self.backbone)
self.model.eval()
def set_norm_stats(self, mean, std):
mean = torch.tensor(mean, dtype=torch.float32, device=self.device).view(1, -1, 1, 1)
std = torch.tensor(std, dtype=torch.float32, device=self.device).view(1, -1, 1, 1).clamp_min(1e-6)
self._norm_mean = mean
self._norm_std = std
def prepare_infer_buffers(self, C, H, W):
shape = (1, C, H, W)
if self._buf_shape == shape:
return
self._buf_shape = shape
self._cpu_pinned = torch.empty(shape, dtype=torch.float32, pin_memory=True)
self._gpu_input = torch.empty(shape, dtype=torch.float32, device=self.device)
def _load_model(self, backbone: str) -> torch.nn.Module:
if self.fusion_mode == "dual_branch":
model = build_dual_branch_segformer_model(
num_classes=self.num_classes,
channels=self.channels,
backbone=backbone,
device=self.device,
ckpt_path=self.ckpt_path,
strict=True,
)
elif self.fusion_mode == "stacked":
model = build_multispec_segformer_model(
num_classes=self.num_classes,
channels=self.channels,
backbone=backbone,
device=self.device,
ckpt_path=self.ckpt_path,
strict=True,
)
else:
raise ValueError(f"fusion_mode inválido: {self.fusion_mode}")
return model
def _preprocess_tensor(self, raw_np: np.ndarray) -> torch.Tensor:
if raw_np.dtype != np.float32:
raw_np = raw_np.astype(np.float32)
x = torch.from_numpy(raw_np).unsqueeze(0).to(self.device)
if self._norm_mean is not None and self._norm_std is not None:
x = (x - self._norm_mean) / self._norm_std
else:
x = normalize_per_batch(x)
return x
def _build_preview(self, raw_np, pred_ids, alpha, fast):
t0 = time.time()
bgr = make_bgr_preview_from_tensor(raw_np, preview_fast=fast)
pred_bgr = converter_mask_ids_para_bgr(pred_ids, self.colormap_rgb, self.ignore_id)
a = float(np.clip(alpha, 0.0, 1.0))
overlay = (bgr.astype(np.float32) * (1 - a) + pred_bgr.astype(np.float32) * a)
overlay = np.clip(overlay, 0, 255).astype(np.uint8)
dt = (time.time() - t0) * 1000.0
return bgr, pred_bgr, overlay, dt
@torch.inference_mode()
def infer_ids(self, raw_np: np.ndarray, use_amp: bool | None = None) -> np.ndarray:
if use_amp is None:
use_amp = self.use_amp and (self.device.type == "cuda")
if raw_np.dtype != np.float32:
raw_np = raw_np.astype(np.float32, copy=False)
C, H, W = raw_np.shape
self.prepare_infer_buffers(C, H, W)
self._cpu_pinned[0].copy_(torch.from_numpy(raw_np), non_blocking=True)
x = self._gpu_input
x.copy_(self._cpu_pinned, non_blocking=True)
if self._norm_mean is not None and self._norm_std is not None:
x = (x - self._norm_mean) / self._norm_std
else:
mean = x.mean(dim=(2, 3), keepdim=True)
std = x.std(dim=(2, 3), keepdim=True).clamp_min(1e-6)
x = (x - mean) / std
ctx = amp.autocast("cuda") if use_amp else nullcontext()
with ctx:
out = self.model(pixel_values=x)
logits = out.logits
if logits.shape[-2:] != (H, W):
logits = F.interpolate(logits, size=(H, W), mode="bilinear", align_corners=False)
pred = torch.argmax(logits, dim=1)
pred_ids = pred[0].to(torch.uint8).cpu().numpy()
return pred_ids
@torch.no_grad()
def infer_raw(self, raw_np: np.ndarray):
return self.infer_ids(raw_np)
@torch.no_grad()
def infer_and_preview(self, raw_np: np.ndarray, alpha=0.5):
t0 = time.time()
pred_ids = self.infer_ids(raw_np)
t1 = time.time()
rgb, pred_rgb, overlay, dt = self._build_preview(raw_np, pred_ids, alpha, False)
t2 = time.time()
t_inf = (t1 - t0) * 1000
t_pvw = (t2 - t1) * 1000
return pred_ids, rgb, pred_rgb, overlay, t_inf, t_pvw
def get_classes(self):
return self.classes
def get_colormap(self):
return self.colormap_rgb
def get_ignore_id(self):
return self.ignore_id
# ============================================================
# Utils
# ============================================================
def _infer_ignore_id(ignore_rgb, default_id=255):
import numpy as _np
if isinstance(ignore_rgb, (list, tuple)):
if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, _np.integer)):
return int(ignore_rgb[0])
if len(ignore_rgb) == 3:
return default_id
if isinstance(ignore_rgb, (int, _np.integer)):
return int(ignore_rgb)
return default_id

View File

@ -0,0 +1,405 @@
import socket
import json
import numpy as np
import base64
import time
from typing import Optional, Any
class MultiSpectralService:
def __init__(self, host="192.168.105.6", port=5000, timeout=5):
self.host = host
self.port = port
self.timeout = timeout
self.sock = None
self.file = None
def __enter__(self):
self.connect()
return self
def __exit__(self, exc_type, exc, tb):
self.disconnect()
# =========================================================
# Conexão
# =========================================================
def connect(self):
if self.sock is not None:
return
self.sock = socket.create_connection((self.host, self.port), timeout=self.timeout)
self.sock.settimeout(self.timeout)
self.file = self.sock.makefile("r", encoding="utf-8")
def disconnect(self):
try:
if self.file:
self.file.close()
except Exception:
pass
try:
if self.sock:
self.sock.close()
except Exception:
pass
self.file = None
self.sock = None
def check_connection(self, timeout: float = None) -> bool:
try:
self.connect()
resp = self._send_command({"cmd": "ping"})
return resp.get("ok") and resp.get("reply") == "pong"
except Exception:
return False
def is_alive(self) -> bool:
try:
self.connect()
resp = self._send_command({"cmd": "ping"})
return resp.get("ok") and resp.get("reply") == "pong"
except Exception:
return False
def ensure_alive(self):
try:
self.connect()
except Exception as e:
raise RuntimeError(
f"Não foi possível conectar ao módulo em {self.host}:{self.port}. "
f"Verifique rede, IP e se o Pi está ligado. Erro: {e}"
) from e
try:
resp = self._send_command({"cmd": "ping"})
except Exception as e:
raise RuntimeError(
f"Conectou ao endereço {self.host}:{self.port}, mas o módulo não respondeu ao ping. "
f"Verifique se o serviço está rodando no Pi. Erro: {e}"
) from e
if not resp.get("ok") or resp.get("reply") != "pong":
raise RuntimeError(f"Resposta inválida do módulo ao ping: {resp}")
def _send_command(self, payload: dict) -> dict:
if self.sock is None:
self.connect()
data = (json.dumps(payload) + "\n").encode("utf-8")
self.sock.sendall(data)
line = self.file.readline()
if not line:
self.disconnect()
raise RuntimeError("Conexão encerrada pelo servidor")
return json.loads(line.strip())
# =========================================================
# Helpers numpy
# =========================================================
def _numpy_dtype_from_string(self, dtype_str: str):
mapping = {
"uint8": np.uint8,
"uint16": np.uint16,
"float32": np.float32,
}
if dtype_str not in mapping:
raise RuntimeError(f"dtype não suportado recebido do Pi: {dtype_str}")
return mapping[dtype_str]
def _reshape_array_from_shape(self, raw_bytes: bytes, dtype_str: str, shape: list | tuple):
np_dtype = self._numpy_dtype_from_string(dtype_str)
arr = np.frombuffer(raw_bytes, dtype=np_dtype)
return arr.reshape(tuple(shape))
def _reshape_array(self, raw_bytes: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
np_dtype = self._numpy_dtype_from_string(dtype_str)
arr = np.frombuffer(raw_bytes, dtype=np_dtype)
if layout == "HW":
return arr.reshape(height, width)
if layout == "CHW":
return arr.reshape(channels, height, width)
if layout == "HWC":
return arr.reshape(height, width, channels)
raise RuntimeError(f"Layout não suportado recebido do Pi: {layout}")
def _decode_single_array(self, raw: bytes, resp: dict):
dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8"
shape = resp.get("shape")
if shape:
return self._reshape_array_from_shape(raw, dtype_str, shape)
output_layout = resp.get("output_layout", "HW")
width = int(resp.get("output_width", resp.get("width", 0)))
height = int(resp.get("output_height", resp.get("height", 0)))
channels = int(resp.get("output_channels", resp.get("channels", 1)))
return self._reshape_array(
raw_bytes=raw,
dtype_str=dtype_str,
layout=output_layout,
width=width,
height=height,
channels=channels,
)
def _decode_multi_frames_base64(self, resp: dict):
frames_resp = resp.get("frames", {})
payload_parts = resp.get("payload_parts", []) or []
camera_frames = resp.get("camera_frames", {}) or {}
parts_by_cam = {}
for part in payload_parts:
cam_id = part.get("camera_id")
if cam_id:
parts_by_cam[cam_id] = part
frames = {}
for cam_id, item in frames_resp.items():
raw = base64.b64decode(item["data"])
part_meta = parts_by_cam.get(cam_id, {})
cam_meta = camera_frames.get(cam_id, {})
dtype_str = part_meta.get("dtype") or resp.get("dtype") or resp.get("output_dtype")
shape = part_meta.get("shape")
if dtype_str == "multi" or dtype_str is None:
# fallback conservador
dtype_str = "uint8" if int(cam_meta.get("channels", 1)) > 1 else "uint16"
if shape:
arr = self._reshape_array_from_shape(raw, dtype_str, shape)
else:
width = int(part_meta.get("width", cam_meta.get("width", 0)))
height = int(part_meta.get("height", cam_meta.get("height", 0)))
channels = int(part_meta.get("channels", cam_meta.get("channels", 1)))
layout = "HWC" if channels > 1 else "HW"
arr = self._reshape_array(
raw_bytes=raw,
dtype_str=dtype_str,
layout=layout,
width=width,
height=height,
channels=channels,
)
frames[cam_id] = arr
return frames
def _extract_meta(self, resp: dict, t0: float) -> dict:
dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8"
return {
"frame_id": resp.get("frame_id"),
"frame_type": resp.get("frame_type"),
"payload_format_version": resp.get("payload_format_version"),
"output_dtype": resp.get("output_dtype"),
"dtype": dtype_str,
"output_layout": resp.get("output_layout"),
"output_channels": resp.get("output_channels"),
"output_channel_names": resp.get("output_channel_names"),
"output_width": resp.get("output_width"),
"output_height": resp.get("output_height"),
"payload_sources": resp.get("payload_sources"),
"payload_complete": resp.get("payload_complete"),
"source_camera": resp.get("source_camera"),
"source_cameras": resp.get("source_cameras"),
"camera_frames": resp.get("camera_frames"),
"multi_payload": resp.get("multi_payload", False),
"payload_kind": resp.get("payload_kind"),
"payload_parts": resp.get("payload_parts"),
"packed_width": resp.get("packed_width"),
"packed_height": resp.get("packed_height"),
"source_width": resp.get("source_width"),
"source_height": resp.get("source_height"),
"source_bayer_pattern": resp.get("source_bayer_pattern"),
"source_bit_depth": resp.get("source_bit_depth"),
"size": resp.get("size"),
"ts_pi": resp.get("ts_pi"),
"ts_pi_monotonic": resp.get("ts_pi_monotonic"),
"dt_trigger": resp.get("dt_trigger"),
"dt_settle": resp.get("dt_settle"),
"dt_capture": resp.get("dt_capture"),
"dt_process": resp.get("dt_process"),
"dt_total_pi": resp.get("dt_total_pi"),
"dt_total_pc": time.perf_counter() - t0,
}
# =========================================================
# Comandos básicos
# =========================================================
def ping(self):
return self._send_command({"cmd": "ping"})
def get_status(self):
return self._send_command({"cmd": "get_status"})
def get_config(self):
return self._send_command({"cmd": "get_config"})
def begin(self, frame_type: str = "RAW_BRUTO", output_dtype: str = "uint8", capture_mode: str = "AUTO"):
return self._send_command({
"cmd": "begin",
"frame_type": frame_type,
"output_dtype": output_dtype,
"capture_mode": capture_mode,
})
def stop(self):
return self._send_command({"cmd": "stop"})
def set_fps(self, fps: int):
return self._send_command({"cmd": "set_fps", "value": fps})
def set_jpeg_quality(self, quality: int):
return self._send_command({"cmd": "set_jpeg_quality", "value": quality})
def set_frame_type(self, frame_type: str):
return self._send_command({"cmd": "set_frame_type", "value": frame_type})
def set_output_dtype(self, output_dtype: str):
return self._send_command({"cmd": "set_output_dtype", "value": output_dtype})
def set_capture_mode(self, capture_mode: str):
return self._send_command({"cmd": "set_capture_mode", "value": capture_mode})
def set_camera_enabled(self, index: int, enabled: bool):
return self._send_command({
"cmd": "set_camera_enabled",
"index": index,
"enabled": bool(enabled)
})
def set_camera_bayer(self, index: int, bayer_pattern: str):
return self._send_command({
"cmd": "set_camera_bayer",
"index": index,
"pattern": bayer_pattern
})
def set_camera_resolution(self, index: int, width: int, height: int):
return self._send_command({
"cmd": "set_camera_resolution",
"index": index,
"width": width,
"height": height
})
# =========================================================
# Captura
# =========================================================
def capture_frame(self):
t0 = time.perf_counter()
resp = self._send_command({"cmd": "capture_frame"})
if not resp.get("ok"):
raise RuntimeError(resp.get("error", "Falha ao capturar frame"))
encoding = resp.get("encoding", "base64")
meta = self._extract_meta(resp, t0)
if encoding == "base64":
raw = base64.b64decode(resp["data"])
arr = self._decode_single_array(raw, resp)
meta.update({
"shape": list(arr.shape),
"channels": int(arr.shape[0]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
else (int(arr.shape[2]) if arr.ndim == 3 else 1),
"width": int(arr.shape[2]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
else (int(arr.shape[1]) if arr.ndim == 3 else int(arr.shape[1])),
"height": int(arr.shape[1]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
else int(arr.shape[0]),
})
return arr, meta
if encoding == "base64-multi":
frames = self._decode_multi_frames_base64(resp)
meta.update({
"frames_meta": resp.get("camera_frames", {}),
"decoded_shapes": {cam_id: list(arr.shape) for cam_id, arr in frames.items()},
})
return frames, meta
raise RuntimeError(f"encoding não suportado recebido do Pi: {encoding}")
def capture_frame_array(self):
return self.capture_frame()
# =========================================================
# Stream
# =========================================================
def start_stream(self, host: str, port: int, fps: float):
return self._send_command({
"cmd": "start_stream",
"host": host,
"port": port,
"fps": fps
})
def stop_stream(self):
return self._send_command({"cmd": "stop_stream"})
# =========================================================
# Controles de câmera
# =========================================================
def get_camera_controls(self):
return self._send_command({"cmd": "get_camera_controls"})
def set_ae_enable(self, value: bool):
return self._send_command({"cmd": "set_ae_enable", "value": bool(value)})
def set_awb_enable(self, value: bool):
return self._send_command({"cmd": "set_awb_enable", "value": bool(value)})
def set_exposure_time(self, exposure_time_us: Optional[int] = None):
return self._send_command({"cmd": "set_exposure_time", "value": exposure_time_us})
def clear_exposure_time(self):
return self._send_command({"cmd": "clear_exposure_time"})
def set_analogue_gain(self, gain: float | None):
return self._send_command({"cmd": "set_analogue_gain", "value": gain})
def clear_analogue_gain(self):
return self._send_command({"cmd": "clear_analogue_gain"})
def set_colour_gains(self, r_gain: float, b_gain: float):
return self._send_command({
"cmd": "set_colour_gains",
"r_gain": r_gain,
"b_gain": b_gain
})
def clear_colour_gains(self):
return self._send_command({"cmd": "clear_colour_gains"})
def get_sensor_modes(self):
return self._send_command({"cmd": "get_sensor_modes"})

View File

@ -0,0 +1,453 @@
import os
import numpy as np
import math
from typing import Optional
class RawProcessorCore:
def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "GBRG"):
self.sensor_width = sensor_width
self.sensor_height = sensor_height
self.bayer_pattern = bayer_pattern.upper()
def unpack_raw10_packed(
self,
packed_frame: np.ndarray,
sensor_width: Optional[int] = None,
sensor_height: Optional[int] = None
):
if packed_frame.ndim == 3 and packed_frame.shape[2] == 1:
packed_frame = packed_frame[:, :, 0]
width = sensor_width if sensor_width is not None else self.sensor_width
height = sensor_height if sensor_height is not None else self.sensor_height
if width % 4 != 0:
raise ValueError(f"Largura {width} não é múltipla de 4 para RAW10 packed")
expected_packed_width = math.ceil(width * 10 / 8)
actual_h, actual_w = packed_frame.shape[:2]
padding = actual_w - expected_packed_width
if actual_h != height:
raise ValueError(
f"[ERRO FRAME] Altura packed inesperada: {packed_frame.shape}, "
f"esperado altura={height}"
)
if actual_w < expected_packed_width:
raise ValueError(
f"[ERRO FRAME] Largura packed menor que a útil esperada: {packed_frame.shape}, "
f"esperado pelo menos ({height}, {expected_packed_width})"
)
if padding > 64:
raise ValueError(
f"[ERRO FRAME] Padding excessivo no packed: {packed_frame.shape}, "
f"esperado útil ({height}, {expected_packed_width}), padding={padding}"
)
packed_frame = packed_frame[:, :expected_packed_width]
groups = packed_frame.reshape(height, width // 4, 5).astype(np.uint16)
b0 = groups[:, :, 0]
b1 = groups[:, :, 1]
b2 = groups[:, :, 2]
b3 = groups[:, :, 3]
b4 = groups[:, :, 4]
p0 = (b0 << 2) | ((b4 >> 0) & 0x03)
p1 = (b1 << 2) | ((b4 >> 2) & 0x03)
p2 = (b2 << 2) | ((b4 >> 4) & 0x03)
p3 = (b3 << 2) | ((b4 >> 6) & 0x03)
raw16 = np.empty((height, width), dtype=np.uint16)
raw16[:, 0::4] = p0
raw16[:, 1::4] = p1
raw16[:, 2::4] = p2
raw16[:, 3::4] = p3
return raw16
def extract_bayer_channels(self, raw16: np.ndarray) -> dict:
p = self.bayer_pattern
if p == "GBRG":
g1 = raw16[0::2, 0::2]
b = raw16[0::2, 1::2]
r = raw16[1::2, 0::2]
g2 = raw16[1::2, 1::2]
elif p == "GRBG":
g1 = raw16[0::2, 0::2]
r = raw16[0::2, 1::2]
b = raw16[1::2, 0::2]
g2 = raw16[1::2, 1::2]
elif p == "RGGB":
r = raw16[0::2, 0::2]
g1 = raw16[0::2, 1::2]
g2 = raw16[1::2, 0::2]
b = raw16[1::2, 1::2]
elif p == "BGGR":
b = raw16[0::2, 0::2]
g1 = raw16[0::2, 1::2]
g2 = raw16[1::2, 0::2]
r = raw16[1::2, 1::2]
else:
raise ValueError(f"Padrão Bayer não suportado: {p}")
return {"R": r, "G1": g1, "G2": g2, "B": b}
def build_training_rgb(
self,
raw16: np.ndarray,
output_dtype: str = "float32",
bit_depth: int = 10,
) -> np.ndarray:
ch = self.extract_bayer_channels(raw16)
max_val = float((1 << bit_depth) - 1)
r = ch["R"].astype(np.float32) / max_val
g = ((ch["G1"].astype(np.float32) + ch["G2"].astype(np.float32)) * 0.5) / max_val
b = ch["B"].astype(np.float32) / max_val
chw = np.stack([r, g, b], axis=0).astype(np.float32)
chw = np.clip(chw, 0.0, 1.0)
if output_dtype == "float32":
return chw
if output_dtype == "uint8":
return (chw * 255.0).clip(0, 255).astype(np.uint8)
if output_dtype == "uint16":
return (chw * 65535.0).clip(0, 65535).astype(np.uint16)
raise ValueError(f"output_dtype não suportado: {output_dtype}")
def build_multispectral_tensor(self, bins_data, bins_meta):
rgb = None
re = None
nir = None
for data, meta in zip(bins_data, bins_meta):
role = meta.get("role")
if role == "rgb":
rgb = data.astype(np.float32) / 255.0
rgb = rgb.transpose(2, 0, 1)
elif role == "re":
re = data.astype(np.float32) / 1023.0
re = re[None, :, :]
elif role == "nir":
nir = data.astype(np.float32) / 1023.0
nir = nir[None, :, :]
if rgb is None:
raise RuntimeError("RGB obrigatório")
arrays = [rgb]
channel_names = ["R", "G", "B"]
if re is not None:
arrays.append(re)
channel_names.append("RE")
if nir is not None:
arrays.append(nir)
channel_names.append("NIR")
min_h = min(a.shape[1] for a in arrays)
min_w = min(a.shape[2] for a in arrays)
arrays = [a[:, :min_h, :min_w] for a in arrays]
tensor = np.concatenate(arrays, axis=0)
return tensor, channel_names
def build_infer_tensor_from_stream(self, frame, meta, channels_expected):
"""
Converte o frame vindo do stream do Pi em tensor (C,H,W) float32 0..1
compatível com o modelo.
Suporta:
- RGB uint8/float32 pronto
- MULTISPEC uint8/float32 pronto
- RAW_BRUTO multi_payload (cam2 RGB + cam0/cam1 packed)
"""
frame_type = meta.get("frame_type")
dtype_str = meta.get("dtype") or meta.get("output_dtype", "uint8")
camera_frames = meta.get("camera_frames", {}) or {}
# -------------------------------------------------
# RAW_BRUTO multi_payload
# -------------------------------------------------
if frame_type == "RAW_BRUTO":
if not isinstance(frame, dict):
raise RuntimeError("RAW_BRUTO esperado como dict de câmeras no modo multi")
arrays = []
channel_names = []
# RGB USB
if "cam2" in frame:
rgb_bgr = frame["cam2"]
if rgb_bgr.ndim != 3 or rgb_bgr.shape[2] != 3:
raise RuntimeError(f"cam2 RGB inválida: shape={rgb_bgr.shape}")
rgb = rgb_bgr[:, :, ::-1].astype(np.float32) / 255.0
rgb_chw = np.transpose(rgb, (2, 0, 1))
arrays.append(rgb_chw)
channel_names.extend(["R", "G", "B"])
else:
raise RuntimeError("RAW_BRUTO para inferência precisa incluir cam2 (RGB)")
# RE / NIR
for cam_id, spec_name in (("cam0", "RE"), ("cam1", "NIR")):
if cam_id not in frame:
continue
packed = frame[cam_id]
if packed.ndim == 3 and packed.shape[2] == 1:
packed = packed[:, :, 0]
cam_meta = camera_frames.get(cam_id, {})
packed_width = int(cam_meta.get("width", packed.shape[1]))
height = int(cam_meta.get("height", packed.shape[0]))
bayer = cam_meta.get("bayer_pattern", self.bayer_pattern)
bit_depth = int(cam_meta.get("bit_depth", 10))
if bit_depth == 10:
real_width = int((packed_width * 8) / 10)
else:
real_width = packed_width
rp = RawProcessorCore(
sensor_width=real_width,
sensor_height=height,
bayer_pattern=bayer,
)
raw16 = rp.unpack_raw10_packed(packed)
max_val = float((1 << bit_depth) - 1)
single = np.clip(raw16.astype(np.float32) / max_val, 0.0, 1.0)[None, :, :]
arrays.append(single)
channel_names.append(spec_name)
if len(arrays) < 2:
raise RuntimeError("RAW_BRUTO requer RGB + pelo menos um canal espectral para inferência")
min_h = min(a.shape[1] for a in arrays)
min_w = min(a.shape[2] for a in arrays)
arrays = [a[:, :min_h, :min_w] for a in arrays]
raw_np = np.concatenate(arrays, axis=0)
if raw_np.shape[0] != channels_expected:
raise RuntimeError(
f"Tensor RAW_BRUTO montado com canais inesperados: {raw_np.shape[0]} | esperado={channels_expected} | got={channel_names}"
)
return raw_np
# -------------------------------------------------
# RGB ou MULTISPEC já pronto
# -------------------------------------------------
if frame_type == "RGB" or frame_type == "MULTISPEC":
if not isinstance(frame, np.ndarray):
raise RuntimeError(f"Frame {frame_type} esperado como ndarray")
if frame.ndim != 3:
raise RuntimeError(f"Frame {frame_type} inválido: shape={frame.shape}")
if dtype_str == "uint8":
raw_np = frame.astype(np.float32) / 255.0
elif dtype_str == "float32":
raw_np = frame.astype(np.float32)
elif dtype_str == "uint16":
raw_np = frame.astype(np.float32) / 65535.0
else:
raise RuntimeError(f"dtype {frame_type} não suportado: {dtype_str}")
if raw_np.shape[0] != channels_expected:
raise RuntimeError(
f"Frame {frame_type} com canais inesperados: {raw_np.shape[0]} | esperado={channels_expected}"
)
return raw_np
raise RuntimeError(f"frame_type não suportado para inferência: {frame_type}")
def extract_camera_meta(self, meta_json: dict, cam_id: str) -> dict:
cam_frames = meta_json.get("camera_frames", {}) or meta_json.get("stream_meta", {}).get("camera_frames", {})
cam = cam_frames.get(cam_id)
if not cam:
raise ValueError(f"Camera {cam_id} não encontrada no meta")
# Detecta RAW10 packed mono
if int(cam.get("channels", 1)) == 1 and int(cam.get("bit_depth", 10)) == 10:
packed_width = int(cam.get("width"))
height = int(cam.get("height"))
# 🔥 converte packed → real
real_width = int((packed_width * 8) / 10)
return {
"camera_id": cam_id,
"width": real_width,
"height": height,
"channels": 1,
"bit_depth": 10,
"shape": [height, packed_width], # packed shape
"role": cam.get("role")
}
# RGB
else:
width = int(cam.get("width"))
height = int(cam.get("height"))
channels = int(cam.get("channels", 3))
return {
"camera_id": cam_id,
"width": width,
"height": height,
"channels": channels,
"bit_depth": int(cam.get("bit_depth", 8)),
"shape": [height, width, channels], # 🔥 AQUI está a correção
"role": cam.get("role")
}
# ============================================================
# RAW10 PACKED
# ============================================================
def packed_width_for_raw10(self, sensor_width: int = None) -> int:
width = sensor_width if sensor_width is not None else self.sensor_width
return math.ceil(width * 10 / 8)
def pack_raw10_packed(self, raw16: np.ndarray) -> np.ndarray:
h, w = raw16.shape
if w % 4 != 0:
raise ValueError(f"Width precisa ser múltiplo de 4 para pack otimizado. Veio {w}")
raw16 = np.clip(raw16, 0, 1023).astype(np.uint16)
p0 = raw16[:, 0::4]
p1 = raw16[:, 1::4]
p2 = raw16[:, 2::4]
p3 = raw16[:, 3::4]
b0 = (p0 >> 2).astype(np.uint8)
b1 = (p1 >> 2).astype(np.uint8)
b2 = (p2 >> 2).astype(np.uint8)
b3 = (p3 >> 2).astype(np.uint8)
b4 = (
((p0 & 0x03) << 0) |
((p1 & 0x03) << 2) |
((p2 & 0x03) << 4) |
((p3 & 0x03) << 6)
).astype(np.uint8)
packed = np.empty((h, w // 4, 5), dtype=np.uint8)
packed[:, :, 0] = b0
packed[:, :, 1] = b1
packed[:, :, 2] = b2
packed[:, :, 3] = b3
packed[:, :, 4] = b4
return packed.reshape(h, w // 4 * 5)
def load_raw10_packed_file(self, path: str, width: int, height: int) -> np.ndarray:
packed_width = self.packed_width_for_raw10(width)
expected_size = height * packed_width
actual_size = os.path.getsize(path)
if actual_size != expected_size:
raise ValueError(
f"Tamanho inválido RAW10: {actual_size}, esperado {expected_size} em {path}"
)
packed = np.fromfile(path, dtype=np.uint8).reshape(height, packed_width)
return self.unpack_raw10_packed(packed, sensor_width=width, sensor_height=height)
def save_raw10_packed_file(self, path: str, raw16: np.ndarray):
packed = self.pack_raw10_packed(raw16)
packed.tofile(path)
# ============================================================
# RGB UINT8
# ============================================================
def load_rgb_u8_file(self, path: str, shape) -> np.ndarray:
arr = np.fromfile(path, dtype=np.uint8)
expected = np.prod(shape)
if arr.size != expected:
raise ValueError(
f"Tamanho inválido RGB: {arr.size}, esperado {expected} em {path}"
)
return arr.reshape(shape)
def save_rgb_u8_file(self, path: str, arr: np.ndarray):
arr.astype(np.uint8).tofile(path)
# ============================================================
# DISPATCHER (O MAIS IMPORTANTE)
# ============================================================
def load_native_bin(self, path: str, cam_meta: dict) -> np.ndarray:
"""
Decide automaticamente como carregar o .bin baseado no meta.
"""
channels = int(cam_meta.get("channels", 1))
bit_depth = int(cam_meta.get("bit_depth", 10))
shape = cam_meta.get("shape")
if channels == 1 and bit_depth == 10:
width = int(cam_meta["width"])
height = int(cam_meta["height"])
return self.load_raw10_packed_file(path, width, height)
elif channels == 3 and bit_depth == 8:
return self.load_rgb_u8_file(path, shape)
else:
raise ValueError(
f"Formato não suportado: channels={channels}, bit_depth={bit_depth}"
)
def save_native_bin(self, path: str, arr: np.ndarray, cam_meta: dict):
"""
Salva no formato correto baseado no meta.
"""
channels = int(cam_meta.get("channels", 1))
bit_depth = int(cam_meta.get("bit_depth", 10))
if channels == 1 and bit_depth == 10:
self.save_raw10_packed_file(path, arr)
elif channels == 3 and bit_depth == 8:
self.save_rgb_u8_file(path, arr)
else:
raise ValueError(
f"Formato não suportado para salvar: channels={channels}, bit_depth={bit_depth}"
)

View File

@ -0,0 +1,125 @@
import cv2
import numpy as np
from typing import Optional
class RawProcessorPreview:
def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "GBRG"):
self.sensor_width = sensor_width
self.sensor_height = sensor_height
self.bayer_pattern = bayer_pattern.upper()
def raw16_to_vis8(
self, raw16: np.ndarray,
black_level: Optional[int] = None,
white_level: Optional[int] = None,
gamma: float = 2.2,
bit_depth: int = 10
) -> np.ndarray:
"""
Conversão para visualização:
- auto-level
- gamma
"""
max_val = float((1 << bit_depth) - 1)
raw = raw16.astype(np.float32)
if black_level is None:
black_level = float(raw.min())
if white_level is None:
white_level = float(raw.max())
if white_level <= black_level:
norm = raw / max_val
else:
norm = (raw - black_level) / (white_level - black_level)
norm = np.clip(norm, 0.0, 1.0)
if gamma is not None and gamma > 0:
norm = np.power(norm, 1.0 / gamma)
return (norm * 255.0).clip(0, 255).astype(np.uint8)
def _debayer_code(self):
mapping = {
# Troque de BayerGB para BayerGR para inverter R e B
"GBRG": cv2.COLOR_BayerGR2BGR,
"GRBG": cv2.COLOR_BayerGB2BGR,
"RGGB": cv2.COLOR_BayerBG2BGR,
"BGGR": cv2.COLOR_BayerRG2BGR,
}
if self.bayer_pattern not in mapping:
raise ValueError(f"Padrão Bayer não suportado: {self.bayer_pattern}")
return mapping[self.bayer_pattern]
def apply_preview_white_balance(self, bgr: np.ndarray, strength: float = 1.0) -> np.ndarray:
"""
Gray-world simples para deixar o preview mais agradável.
Não usar no raw de treino.
"""
img = bgr.astype(np.float32)
mean_b = float(img[:, :, 0].mean())
mean_g = float(img[:, :, 1].mean())
mean_r = float(img[:, :, 2].mean())
mean_gray = (mean_b + mean_g + mean_r) / 3.0
eps = 1e-6
gain_b = mean_gray / max(mean_b, eps)
gain_g = mean_gray / max(mean_g, eps)
gain_r = mean_gray / max(mean_r, eps)
# strength=1 aplica total, strength=0 não aplica
gain_b = 1.0 + (gain_b - 1.0) * strength
gain_g = 1.0 + (gain_g - 1.0) * strength
gain_r = 1.0 + (gain_r - 1.0) * strength
img[:, :, 0] *= gain_b
img[:, :, 1] *= gain_g
img[:, :, 2] *= gain_r
return np.clip(img, 0, 255).astype(np.uint8)
def apply_preview_contrast(self, bgr: np.ndarray, alpha: float = 1.08, beta: float = 0.0) -> np.ndarray:
"""
Ajuste leve de contraste/brilho para preview.
"""
out = cv2.convertScaleAbs(bgr, alpha=alpha, beta=beta)
return out
def raw16_to_preview_bgr(
self,
raw16: np.ndarray,
gamma: float = 2.2,
wb_strength: float = 0.8,
apply_wb: bool = True,
apply_contrast: bool = True,
bit_depth: int = 10,
) -> np.ndarray:
"""
Pipeline de preview bonito:
1. auto-level + gamma no mosaico
2. demosaic
3. white balance simples
4. leve contraste final
"""
vis8 = self.raw16_to_vis8(raw16, gamma=gamma, bit_depth=bit_depth)
bgr = cv2.cvtColor(vis8, self._debayer_code())
if apply_wb:
bgr = self.apply_preview_white_balance(bgr, strength=wb_strength)
if apply_contrast:
bgr = self.apply_preview_contrast(bgr, alpha=1.08, beta=0.0)
return bgr
def raw16_to_preview_jpg_bytes(self, raw16: np.ndarray, jpeg_quality: int = 95) -> bytes:
bgr = self.raw16_to_preview_bgr(raw16)
ok, enc = cv2.imencode(".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_quality)])
if not ok:
raise RuntimeError("Falha ao codificar preview JPG")
return enc.tobytes()

View File

@ -0,0 +1,324 @@
import json
import socket
import threading
import time
from numcodecs import Blosc
import numpy as np
class StreamReceiver:
def __init__(self, host="0.0.0.0", port=6001):
self.host = host
self.port = port
self._server_sock = None
self._client_sock = None
self._thread = None
self._running = False
self.last_frame = None
self.last_meta = None
self.last_receive_ts = None
self._codec = None
self._codec_signature = None
@property
def is_running(self):
return self._running
# =========================================================
# Lifecycle
# =========================================================
def start(self):
if self._running:
return
self._running = True
self._thread = threading.Thread(target=self._worker, daemon=True)
self._thread.start()
def stop(self):
self._running = False
try:
if self._client_sock:
self._client_sock.close()
except Exception:
pass
try:
if self._server_sock:
self._server_sock.close()
except Exception:
pass
self._client_sock = None
self._server_sock = None
# =========================================================
# Socket helpers
# =========================================================
def _recv_exact(self, sock: socket.socket, n: int) -> bytes:
chunks = []
remaining = n
while remaining > 0:
chunk = sock.recv(remaining)
if not chunk:
raise ConnectionError("Conexão encerrada durante recv")
chunks.append(chunk)
remaining -= len(chunk)
return b"".join(chunks)
# =========================================================
# Worker principal
# =========================================================
def _worker(self):
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind((self.host, self.port))
server.listen(1)
server.settimeout(1.0)
self._server_sock = server
print(f"[INFO] StreamReceiver ouvindo em {self.host}:{self.port}")
while self._running:
try:
client, addr = server.accept()
except socket.timeout:
continue
print(f"[INFO] StreamReceiver conectado por {addr}")
self._client_sock = client
with client:
while self._running:
header_len = int.from_bytes(self._recv_exact(client, 4), "big")
header_bytes = self._recv_exact(client, header_len)
header = json.loads(header_bytes.decode("utf-8"))
payload_len = int.from_bytes(self._recv_exact(client, 4), "big")
payload_comp = self._recv_exact(client, payload_len)
self._ensure_codec(header)
if header.get("codec_family") == "none":
payload = payload_comp
else:
payload = self._codec.decode(payload_comp)
expected = int(header["payload_size_raw"])
if len(payload) != expected:
raise ValueError(
f"Tamanho descomprimido inválido: {len(payload)} != {expected}"
)
payload_kind = header.get("payload_kind")
if header.get("multi_payload", False) or payload_kind == "multi_array":
frame = self._decode_multi_payload(payload, header)
else:
frame = self._decode_single_payload(payload, header)
self.last_frame = frame
self.last_meta = header
self.last_receive_ts = time.perf_counter()
print("[INFO] StreamReceiver cliente desconectado")
self._client_sock = None
except Exception as e:
print(f"[WARN] StreamReceiver encerrado com erro: {e}")
finally:
self._running = False
self._client_sock = None
self._server_sock = None
# =========================================================
# Decode helpers
# =========================================================
def _numpy_dtype_from_string(self, dtype_str: str):
mapping = {
"uint8": np.uint8,
"uint16": np.uint16,
"float32": np.float32,
}
if dtype_str not in mapping:
raise RuntimeError(f"dtype não suportado: {dtype_str}")
return mapping[dtype_str]
def _numpy_dtype_from_header(self, header: dict):
dtype_str = header.get("dtype") or header.get("output_dtype") or "uint8"
if dtype_str == "multi":
# fallback conservador para protocolos mais antigos
return np.uint16
return self._numpy_dtype_from_string(dtype_str)
def _reshape_from_shape(self, payload: bytes, dtype_str: str, shape):
dtype = self._numpy_dtype_from_string(dtype_str)
arr = np.frombuffer(payload, dtype=dtype)
return arr.reshape(tuple(shape))
def _reshape_from_layout(self, payload: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
dtype = self._numpy_dtype_from_string(dtype_str)
arr = np.frombuffer(payload, dtype=dtype)
if layout == "HW":
return arr.reshape(height, width)
if layout == "CHW":
return arr.reshape(channels, height, width)
if layout == "HWC":
return arr.reshape(height, width, channels)
raise RuntimeError(f"Layout não suportado: {layout}")
def _decode_single_payload(self, payload: bytes, header: dict):
payload_parts = header.get("payload_parts", []) or []
first_part = payload_parts[0] if payload_parts else {}
dtype_str = first_part.get("dtype") or header.get("dtype") or header.get("output_dtype") or "uint8"
shape = first_part.get("shape")
if shape:
return self._reshape_from_shape(payload, dtype_str, shape)
height = int(header.get("output_height", header.get("height")))
width = int(header.get("output_width", header.get("width")))
channels = int(header.get("output_channels", header.get("channels", 1)))
layout = header.get("output_layout", "HWC")
return self._reshape_from_layout(
payload=payload,
dtype_str=dtype_str,
layout=layout,
width=width,
height=height,
channels=channels,
)
def _decode_multi_payload(self, payload: bytes, header: dict):
payload_parts = header.get("payload_parts", []) or []
camera_frames = header.get("camera_frames", {}) or {}
frames = {}
offset = 0
for part in payload_parts:
cam_id = part.get("camera_id")
if not cam_id:
raise RuntimeError("payload_parts sem camera_id")
if offset + 4 > len(payload):
raise RuntimeError("Payload multi truncado ao ler tamanho da parte")
part_size = int.from_bytes(payload[offset:offset + 4], "big")
offset += 4
if offset + part_size > len(payload):
raise RuntimeError(f"Payload multi truncado ao ler dados de {cam_id}")
part_bytes = payload[offset:offset + part_size]
offset += part_size
cam_meta = camera_frames.get(cam_id, {})
dtype_str = part.get("dtype") or header.get("dtype") or header.get("output_dtype")
shape = part.get("shape")
if dtype_str == "multi" or dtype_str is None:
channels = int(part.get("channels", cam_meta.get("channels", 1)))
dtype_str = "uint8" if channels > 1 else "uint16"
if shape:
frame = self._reshape_from_shape(part_bytes, dtype_str, shape)
else:
width = int(part.get("width", cam_meta.get("width", 0)))
height = int(part.get("height", cam_meta.get("height", 0)))
channels = int(part.get("channels", cam_meta.get("channels", 1)))
layout = "HWC" if channels > 1 else "HW"
frame = self._reshape_from_layout(
payload=part_bytes,
dtype_str=dtype_str,
layout=layout,
width=width,
height=height,
channels=channels,
)
frames[cam_id] = frame
if offset != len(payload):
raise RuntimeError(
f"Payload multi com bytes sobrando: consumidos={offset}, total={len(payload)}"
)
return frames
# =========================================================
# Codec
# =========================================================
def _normalize_shuffle(self, shuffle_value):
if isinstance(shuffle_value, int):
return shuffle_value
mapping = {
"NOSHUFFLE": Blosc.NOSHUFFLE,
"SHUFFLE": Blosc.SHUFFLE,
"BITSHUFFLE": Blosc.BITSHUFFLE,
}
key = str(shuffle_value).upper()
if key not in mapping:
raise ValueError(f"shuffle inválido: {shuffle_value}")
return mapping[key]
def _build_codec_from_header(self, header: dict):
family = header.get("codec_family")
name = header.get("codec_name")
params = dict(header.get("codec_params", {}))
if family != "numcodecs":
raise ValueError(f"Família de codec não suportada: {family}")
if name == "blosc":
params["shuffle"] = self._normalize_shuffle(params.get("shuffle", "SHUFFLE"))
return Blosc(**params)
raise ValueError(f"Codec numcodecs não suportado: {name}")
def _get_codec_signature_from_header(self, header: dict):
return (
header.get("codec_family"),
header.get("codec_name"),
tuple(sorted(dict(header.get("codec_params", {})).items()))
)
def _ensure_codec(self, header: dict):
family = header.get("codec_family")
if family == "none":
self._codec = None
self._codec_signature = ("none", None, ())
return
sig = self._get_codec_signature_from_header(header)
if self._codec is None or self._codec_signature != sig:
self._codec = self._build_codec_from_header(header)
self._codec_signature = sig

View File

@ -6,7 +6,7 @@ import cv2
import numpy as np import numpy as np
import torch import torch
from gal5000.gal_service import Gal5000Camera from gal5000.gal_service import Gal5000Camera
from raw_segformer_service import RawSegformerService from Python.OAK.datasets.gal5000.raw_segformer_service import RawSegformerService
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}") print(f"Device: {device}")

View File

@ -113,6 +113,17 @@ def desenhar_legenda_horizontal(colormap_rgb, classes, altura=30, largura_por_cl
return legenda return legenda
def _infer_ignore_id(ignore_rgb, default_id=255):
import numpy as _np
if isinstance(ignore_rgb, (list, tuple)):
if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, _np.integer)):
return int(ignore_rgb[0])
if len(ignore_rgb) == 3:
return default_id
if isinstance(ignore_rgb, (int, _np.integer)):
return int(ignore_rgb)
return default_id
# ---------------------------- # ----------------------------
# Helpers ROI # Helpers ROI
# ---------------------------- # ----------------------------

View File

@ -13,7 +13,6 @@ class CameraManager:
self.initialized = False self.initialized = False
self.camera_lock = RLock() self.camera_lock = RLock()
self.frame_lock = Lock()
self.cameras_runtime = {} self.cameras_runtime = {}
self._reconfigure_needed = False self._reconfigure_needed = False
@ -45,43 +44,34 @@ class CameraManager:
"stop_event": threading.Event(), "stop_event": threading.Event(),
"thread": None, "thread": None,
"last_signature": None,
"last_read_ts": None,
"last_new_frame_ts": None,
"frame_lock": Lock(),
} }
# ========================================================= def _make_frame_signature(self, frame: np.ndarray):
# Seleção de câmeras necessárias if frame is None:
# ========================================================= return None
def _get_required_camera_ids(self): h, w = frame.shape[:2]
if hasattr(self.state, "get_required_camera_ids_for_frame_type"):
return self.state.get_required_camera_ids_for_frame_type()
frame_type = self.state.frame_type # pega uma grade simples e barata
resolved_mode = self.state.resolve_capture_mode() step_y = max(1, h // 16)
step_x = max(1, w // 16)
if frame_type == "RGB": sample = frame[::step_y, ::step_x]
cam = getattr(self.state, "rgb_camera_id", None)
return [cam] if cam else []
if frame_type == "MULTISPEC": # reduz para assinatura curtinha
ids = [] return (
for attr in ("rgb_camera_id", "re_camera_id", "nir_camera_id"): sample.shape,
cam_id = getattr(self.state, attr, None) int(sample.mean()),
if cam_id: int(sample.std()),
ids.append(cam_id) int(sample[0, 0, 0]) if sample.ndim == 3 else int(sample[0, 0]),
return ids int(sample[-1, -1, 1]) if sample.ndim == 3 and sample.shape[2] > 1 else 0,
)
if frame_type == "RAW_BRUTO":
if resolved_mode == "TRIPLE":
return [cam.id for cam in self.state.cameras if cam.enabled]
if resolved_mode == "SINGLE":
for role in ("rgb", "re", "nir"):
if hasattr(self.state, "get_active_camera_by_role"):
cam = self.state.get_active_camera_by_role(role)
if cam:
return [cam.id]
return []
return []
# ========================================================= # =========================================================
# Inicialização # Inicialização
@ -91,7 +81,7 @@ class CameraManager:
with self.camera_lock: with self.camera_lock:
self.stop() self.stop()
required_ids = set(self._get_required_camera_ids()) required_ids = set(self.state.get_required_camera_ids_for_frame_type())
if not required_ids: if not required_ids:
print("[WARN] Nenhuma câmera requerida para o frame_type/capture_mode atual") print("[WARN] Nenhuma câmera requerida para o frame_type/capture_mode atual")
@ -149,7 +139,7 @@ class CameraManager:
picam2 = Picamera2(camera_num=cam.index) picam2 = Picamera2(camera_num=cam.index)
config = picam2.create_video_configuration( config = picam2.create_video_configuration(
main={"size": (640, 480), "format": "RGB888"}, #main={"size": (640, 480), "format": "RGB888"},
raw={"size": (cam.width, cam.height)}, raw={"size": (cam.width, cam.height)},
buffer_count=6 buffer_count=6
) )
@ -305,7 +295,7 @@ class CameraManager:
raw = request.make_array("raw") raw = request.make_array("raw")
with self.frame_lock: with runtime["frame_lock"]:
if ( if (
runtime["buffer"] is None or runtime["buffer"] is None or
runtime["buffer"].shape != raw.shape or runtime["buffer"].shape != raw.shape or
@ -329,11 +319,22 @@ class CameraManager:
def _update_loop_opencv(self, runtime): def _update_loop_opencv(self, runtime):
cap = runtime["cap"] cap = runtime["cap"]
target_fps = max(1.0, float(self.state.fps or 10))
min_period = 1.0 / target_fps
t_start = time.perf_counter()
ok, frame = cap.read() ok, frame = cap.read()
read_ts = time.perf_counter()
if not ok or frame is None: if not ok or frame is None:
raise RuntimeError("Falha ao ler frame da câmera USB") raise RuntimeError("Falha ao ler frame da câmera USB")
with self.frame_lock: signature = self._make_frame_signature(frame)
with runtime["frame_lock"]:
runtime["last_read_ts"] = read_ts
if signature != runtime.get("last_signature"):
if ( if (
runtime["buffer"] is None or runtime["buffer"] is None or
runtime["buffer"].shape != frame.shape or runtime["buffer"].shape != frame.shape or
@ -345,7 +346,14 @@ class CameraManager:
runtime["last_frame"] = runtime["buffer"] runtime["last_frame"] = runtime["buffer"]
runtime["frame_id"] += 1 runtime["frame_id"] += 1
runtime["frame_ts"] = time.perf_counter() runtime["frame_ts"] = read_ts
runtime["last_new_frame_ts"] = read_ts
runtime["last_signature"] = signature
dt = time.perf_counter() - t_start
sleep_s = min_period - dt
if sleep_s > 0:
time.sleep(sleep_s)
# ========================================================= # =========================================================
# Leitura consolidada # Leitura consolidada
@ -354,12 +362,12 @@ class CameraManager:
def capture_raw_frames(self): def capture_raw_frames(self):
result = {} result = {}
with self.frame_lock:
for cam_id, runtime in self.cameras_runtime.items(): for cam_id, runtime in self.cameras_runtime.items():
with runtime["frame_lock"]:
if runtime["last_frame"] is None: if runtime["last_frame"] is None:
continue continue
frame = runtime["last_frame"] frame = runtime["last_frame"].copy()
h, w = frame.shape[:2] h, w = frame.shape[:2]
channels = frame.shape[2] if frame.ndim == 3 else 1 channels = frame.shape[2] if frame.ndim == 3 else 1
@ -369,7 +377,8 @@ class CameraManager:
h, h,
channels, channels,
runtime["frame_id"], runtime["frame_id"],
runtime["frame_ts"] runtime["frame_ts"],
runtime["last_read_ts"]
) )
return result return result

View File

@ -42,7 +42,7 @@ class FrameService:
return rp return rp
def _capture_with_retry(self, required_sources, max_attempts=10, retry_delay_s=0.02): def _capture_with_retry(self, required_sources, previous_ids=None, max_attempts=10, retry_delay_s=0.02):
last_frames = None last_frames = None
for _ in range(max_attempts): for _ in range(max_attempts):
@ -55,10 +55,15 @@ class FrameService:
ok = False ok = False
break break
frame, width, height, channels, frame_id, frame_ts = info frame, width, height, channels, frame_id, frame_ts, last_read_ts = info
if frame is None or width <= 0 or height <= 0 or channels <= 0: if frame is None or width <= 0 or height <= 0 or channels <= 0:
ok = False ok = False
break break
if previous_ids is not None:
prev_id = previous_ids.get(cam_id)
if prev_id is not None and frame_id <= prev_id:
ok = False
break
if ok: if ok:
return frames return frames
@ -67,7 +72,7 @@ class FrameService:
time.sleep(retry_delay_s) time.sleep(retry_delay_s)
raise RuntimeError( raise RuntimeError(
f"Capture retornou frames insuficientes após {max_attempts} tentativas. " f"Capture não atingiu frames válidos/novos após {max_attempts} tentativas... "
f"required_sources={required_sources}, received_sources={list((last_frames or {}).keys())}" f"required_sources={required_sources}, received_sources={list((last_frames or {}).keys())}"
) )
@ -170,6 +175,13 @@ class FrameService:
has_csi_source = any(getattr(cam, "interface", "").upper() == "CSI" for cam in required_cams) has_csi_source = any(getattr(cam, "interface", "").upper() == "CSI" for cam in required_cams)
try: try:
previous_ids = {}
frames_before = self.camera.capture_raw_frames()
for cam_id in required_sources:
info = frames_before.get(cam_id)
if info is not None:
previous_ids[cam_id] = info[4] # frame_id
# Trigger só vale para as CSI RE/NIR. # Trigger só vale para as CSI RE/NIR.
# A USB RGB não responde a esse trigger. # A USB RGB não responde a esse trigger.
if self.state.trigger_enabled and has_csi_source: if self.state.trigger_enabled and has_csi_source:
@ -186,7 +198,7 @@ class FrameService:
dt_settle = settle_end - settle_start dt_settle = settle_end - settle_start
cap_start = time.perf_counter() cap_start = time.perf_counter()
raw_frames = self._capture_with_retry(required_sources) raw_frames = self._capture_with_retry(required_sources, previous_ids=previous_ids)
cap_end = time.perf_counter() cap_end = time.perf_counter()
dt_capture = cap_end - cap_start dt_capture = cap_end - cap_start
@ -208,12 +220,13 @@ class FrameService:
camera_frames_meta = {} camera_frames_meta = {}
for cam_id, info in raw_frames.items(): for cam_id, info in raw_frames.items():
frame, width, height, channels, cam_frame_id, cam_frame_ts = info frame, width, height, channels, cam_frame_id, cam_frame_ts, cam_last_read_ts = info
cam = self._get_camera_spec(cam_id) cam = self._get_camera_spec(cam_id)
camera_frames_meta[cam_id] = { camera_frames_meta[cam_id] = {
"camera_frame_id": int(cam_frame_id), "camera_frame_id": int(cam_frame_id),
"camera_frame_ts": cam_frame_ts, "camera_frame_ts": cam_frame_ts,
"camera_last_read_ts": cam_last_read_ts,
"width": int(width), "width": int(width),
"height": int(height), "height": int(height),
"channels": int(channels), "channels": int(channels),
@ -295,7 +308,7 @@ class FrameService:
if len(sources) == 1: if len(sources) == 1:
cam_id = sources[0] cam_id = sources[0]
frame, width, height, channels, cam_frame_id, cam_frame_ts = raw_frames[cam_id] frame, width, height, channels, cam_frame_id, cam_frame_ts, cam_last_read_ts = raw_frames[cam_id]
cam = self._get_camera_spec(cam_id) cam = self._get_camera_spec(cam_id)
if channels == 1: if channels == 1:
@ -333,7 +346,7 @@ class FrameService:
sources_meta = [] sources_meta = []
for cam_id in sources: for cam_id in sources:
frame, width, height, channels, cam_frame_id, cam_frame_ts = raw_frames[cam_id] frame, width, height, channels, cam_frame_id, cam_frame_ts, cam_frame_last_read_ts = raw_frames[cam_id]
cam = self._get_camera_spec(cam_id) cam = self._get_camera_spec(cam_id)
frames_out[cam_id] = frame frames_out[cam_id] = frame
@ -373,7 +386,7 @@ class FrameService:
def _process_rgb(self, raw_frames): def _process_rgb(self, raw_frames):
cam_id = self.state.payload.sources[0] cam_id = self.state.payload.sources[0]
frame, width, height, channels, _, _ = raw_frames[cam_id] frame, width, height, channels, _, _, _ = raw_frames[cam_id]
cam = self._get_camera_spec(cam_id) cam = self._get_camera_spec(cam_id)
if not self._is_usb_rgb_camera(cam): if not self._is_usb_rgb_camera(cam):
@ -429,7 +442,7 @@ class FrameService:
raise RuntimeError("MULTISPEC requer uma câmera RGB ativa") raise RuntimeError("MULTISPEC requer uma câmera RGB ativa")
# RGB USB # RGB USB
rgb_frame, rgb_w, rgb_h, rgb_channels, _, _ = raw_frames[rgb_cam.id] rgb_frame, rgb_w, rgb_h, rgb_channels, _, _, _ = raw_frames[rgb_cam.id]
if rgb_channels != 3: if rgb_channels != 3:
raise RuntimeError(f"Frame RGB USB inválido: channels={rgb_channels}, shape={rgb_frame.shape}") raise RuntimeError(f"Frame RGB USB inválido: channels={rgb_channels}, shape={rgb_frame.shape}")

View File

@ -259,12 +259,23 @@ class ModuleServer:
output_dtype = parse_output_dtype(msg.get("output_dtype", self.state.output_dtype)) output_dtype = parse_output_dtype(msg.get("output_dtype", self.state.output_dtype))
capture_mode = parse_capture_mode(msg.get("capture_mode", self.state.capture_mode)) capture_mode = parse_capture_mode(msg.get("capture_mode", self.state.capture_mode))
old_frame_type = self.state.frame_type
old_output_dtype = self.state.output_dtype
old_capture_mode = self.state.capture_mode
self.state.frame_type = frame_type self.state.frame_type = frame_type
self.state.output_dtype = output_dtype self.state.output_dtype = output_dtype
self.state.capture_mode = capture_mode self.state.capture_mode = capture_mode
self.state.update_payload_spec() self.state.update_payload_spec()
if self.state.initialized: if self.state.initialized:
if (
old_frame_type != self.state.frame_type or
old_output_dtype != self.state.output_dtype or
old_capture_mode != self.state.capture_mode
):
self._mark_reconfigure_needed()
self._restart_module_if_needed() self._restart_module_if_needed()
return { return {
"ok": True, "ok": True,
@ -652,3 +663,7 @@ class ModuleServer:
daemon=True daemon=True
) )
t.start() t.start()
def stop(self):
self.running = False

View File

@ -205,56 +205,53 @@ class ModuleState:
return None return None
def get_required_camera_ids_for_frame_type(self) -> List[str]: def get_required_camera_ids_for_frame_type(self) -> List[str]:
def enabled(role: str):
return self.get_enabled_camera_by_role(role)
def get_all_enabled_camera_ids(): rgb_cam = enabled("rgb")
return [cam.id for cam in self.cameras if cam.enabled] re_cam = enabled("re")
nir_cam = enabled("nir")
if self.frame_type == "RGB":
cam = get_enabled_camera_by_role("rgb")
return [cam.id] if cam else []
if self.frame_type == "MULTISPEC":
ids = []
re_cam = get_enabled_camera_by_role("re")
nir_cam = get_enabled_camera_by_role("nir")
rgb_cam = get_enabled_camera_by_role("rgb")
if re_cam:
ids.append(re_cam.id)
if nir_cam:
ids.append(nir_cam.id)
if rgb_cam:
ids.append(rgb_cam.id)
return ids
if self.frame_type == "RAW_BRUTO":
resolved = self.resolve_capture_mode() resolved = self.resolve_capture_mode()
if self.frame_type == "RGB":
return [rgb_cam.id] if rgb_cam else []
if self.frame_type == "MULTISPEC":
if resolved == "TRIPLE":
if rgb_cam and re_cam and nir_cam:
return [rgb_cam.id, re_cam.id, nir_cam.id]
return []
if resolved == "DOUBLE":
if rgb_cam and re_cam:
return [rgb_cam.id, re_cam.id]
if rgb_cam and nir_cam:
return [rgb_cam.id, nir_cam.id]
return []
# MULTISPEC não faz sentido com SINGLE
return []
if self.frame_type == "RAW_BRUTO":
if resolved == "TRIPLE": if resolved == "TRIPLE":
return [cam.id for cam in self.cameras if cam.enabled] return [cam.id for cam in self.cameras if cam.enabled]
if resolved == "DOUBLE": if resolved == "DOUBLE":
ids = [] ids = []
rgb_cam = self.get_enabled_camera_by_role("rgb")
re_cam = self.get_enabled_camera_by_role("re")
nir_cam = self.get_enabled_camera_by_role("nir")
if rgb_cam: if rgb_cam:
ids.append(rgb_cam.id) ids.append(rgb_cam.id)
if re_cam: if re_cam:
ids.append(re_cam.id) ids.append(re_cam.id)
elif nir_cam: elif nir_cam:
ids.append(nir_cam.id) ids.append(nir_cam.id)
return ids return ids
if resolved == "SINGLE": if resolved == "SINGLE":
for role in ("rgb", "re", "nir"): for cam in (rgb_cam, re_cam, nir_cam):
cam = self.get_enabled_camera_by_role(role)
if cam: if cam:
return [cam.id] return [cam.id]
return []
return [] return []
@ -356,27 +353,36 @@ class ModuleState:
self.detected_mode = "NONE" self.detected_mode = "NONE"
def resolve_capture_mode(self) -> str: def resolve_capture_mode(self) -> str:
enabled_re = self.get_enabled_camera_by_role("re") is not None connected_re = self.get_enabled_camera_by_role("re") is not None
enabled_nir = self.get_enabled_camera_by_role("nir") is not None connected_nir = self.get_enabled_camera_by_role("nir") is not None
enabled_rgb = self.get_enabled_camera_by_role("rgb") is not None connected_rgb = self.get_enabled_camera_by_role("rgb") is not None
if self.capture_mode == "AUTO": if self.capture_mode == "AUTO":
if enabled_rgb and enabled_re and enabled_nir and self.multi_camera_enabled: # MULTISPEC e RAW_BRUTO: preferir DOUBLE no automático
return "TRIPLE" if self.frame_type in ("MULTISPEC", "RAW_BRUTO"):
if enabled_rgb and (enabled_re or enabled_nir): if connected_rgb and (connected_re or connected_nir):
return "DOUBLE" return "DOUBLE"
if enabled_rgb or enabled_re or enabled_nir: if self.frame_type == "RAW_BRUTO":
if connected_rgb or connected_re or connected_nir:
return "SINGLE"
return "NONE"
if connected_rgb and connected_re and connected_nir and self.multi_camera_enabled:
return "TRIPLE"
if connected_rgb and (connected_re or connected_nir):
return "DOUBLE"
if connected_rgb or connected_re or connected_nir:
return "SINGLE" return "SINGLE"
return "NONE" return "NONE"
if self.capture_mode == "TRIPLE": if self.capture_mode == "TRIPLE":
return "TRIPLE" if (enabled_rgb and enabled_re and enabled_nir and self.multi_camera_enabled) else "NONE" return "TRIPLE" if (connected_rgb and connected_re and connected_nir and self.multi_camera_enabled) else "NONE"
if self.capture_mode == "DOUBLE": if self.capture_mode == "DOUBLE":
return "DOUBLE" if (enabled_rgb and (enabled_re or enabled_nir)) else "NONE" return "DOUBLE" if (connected_rgb and (connected_re or connected_nir)) else "NONE"
if self.capture_mode == "SINGLE": if self.capture_mode == "SINGLE":
return "SINGLE" if (enabled_rgb or enabled_re or enabled_nir) else "NONE" return "SINGLE" if (connected_rgb or connected_re or connected_nir) else "NONE"
return "NONE" return "NONE"
@ -474,8 +480,8 @@ class ModuleState:
layout="CHW", layout="CHW",
channels=3, channels=3,
channel_names=["R", "G", "B"], channel_names=["R", "G", "B"],
width=rgb_cam.width // 2, width=rgb_cam.width,
height=rgb_cam.height // 2, height=rgb_cam.height,
sources=[rgb_cam.id], sources=[rgb_cam.id],
complete=True, complete=True,
) )

View File

@ -32,6 +32,7 @@ class StreamSender:
self._frames_dropped = 0 self._frames_dropped = 0
self._capture_errors = 0 self._capture_errors = 0
self._send_errors = 0 self._send_errors = 0
self._last_queued_capture_signature = None
self._last_sent_capture_signature = None self._last_sent_capture_signature = None
@property @property
@ -66,6 +67,7 @@ class StreamSender:
self.state.stream_host = host self.state.stream_host = host
self.state.stream_port = port self.state.stream_port = port
self.state.stream_fps = fps self.state.stream_fps = fps
self._last_queued_capture_signature = None
self._last_sent_capture_signature = None self._last_sent_capture_signature = None
self._thread_capture = threading.Thread( self._thread_capture = threading.Thread(
@ -125,6 +127,7 @@ class StreamSender:
self.state.stream_host = None self.state.stream_host = None
self.state.stream_port = None self.state.stream_port = None
self.state.stream_fps = None self.state.stream_fps = None
self.state.stream_frame_id_sent = 0
def _clear_queue(self): def _clear_queue(self):
while not self._queue.empty(): while not self._queue.empty():
@ -369,7 +372,7 @@ class StreamSender:
continue continue
capture_signature = self._build_capture_signature(meta) capture_signature = self._build_capture_signature(meta)
if capture_signature is not None and capture_signature == self._last_sent_capture_signature: if capture_signature is not None and capture_signature == self._last_queued_capture_signature:
time.sleep(0.001) time.sleep(0.001)
continue continue
@ -412,7 +415,7 @@ class StreamSender:
self._frames_dropped += 1 self._frames_dropped += 1
if queued: if queued:
self._last_sent_capture_signature = capture_signature self._last_queued_capture_signature = capture_signature
if frame_interval > 0: if frame_interval > 0:
next_deadline += frame_interval next_deadline += frame_interval
@ -451,7 +454,9 @@ class StreamSender:
header["send_errors"] = self._send_errors header["send_errors"] = self._send_errors
self._send_packet(sock, header, payload) self._send_packet(sock, header, payload)
self.state.stream_frame_id = header["frame_id"] self.state.stream_frame_id_sent = header["frame_id"]
capture_signature = self._build_capture_signature(header)
self._last_sent_capture_signature = capture_signature
except Exception as e: except Exception as e:
self._send_errors += 1 self._send_errors += 1