ajustes hierarquia de arquivos para treinamento de IA

This commit is contained in:
Diego Freitas 2026-04-20 15:53:28 -03:00
parent 00142d2215
commit 7bcd259d9f
40 changed files with 5528 additions and 774 deletions

View File

@ -1,327 +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()
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)
apply_ir_comp = True
ir_k_r = 0.4
ir_k_g = 0.1
ir_k_b = 0.5
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)
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

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

View File

@ -30,7 +30,7 @@ from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ====================== Configurações base ======================
def carregar_config_e_paths():
with open("config.json", "r", encoding="utf-8") as f:
with open("config_oak.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)

View File

@ -29,7 +29,7 @@ import albumentations as A
import argparse
# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
with open("config_oak.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera", ".")
USE_MASKS2 = config.get("dual_head", False)

View File

@ -27,7 +27,7 @@ import numpy as np
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
with open("config_oak.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
@ -351,7 +351,7 @@ def main(args):
# Se quiser, você pode nomear os canais explicitamente
# dependendo da convenção do raw4:
channel_names = ["R", "G", "IR", "B"]
channel_names = ["R", "G", "B"]
stats = {
"channels": channel_names[:len(mean_list)],

View File

@ -33,7 +33,7 @@ import random
import argparse
# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
with open("config_oak.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)

View File

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

View File

@ -595,7 +595,7 @@ def apply_warmup(optimizer, base_lr: float, epoch: int, warmup_epochs: int):
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", default="config.json")
ap.add_argument("--config", default="config_oak.json")
ap.add_argument("--epochs", type=int, default=120)
ap.add_argument("--batch", type=int, default=4)
ap.add_argument("--num_workers", type=int, default=4)

View File

@ -215,7 +215,7 @@ def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Lê config do projeto (mesmo padrão do fastscnn)
with open("config.json", "r") as f:
with open("config_oak.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]

View File

@ -457,7 +457,7 @@ def main():
device = find_device()
print("Device:", device)
with open("config.json", "r") as f:
with open("config_oak.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"] # ex: "segformer_b3"

View File

@ -1,7 +1,7 @@
{
"camera": "oak-d",
"modelo": "segformer_b0",
"model_name": "nav_big",
"model_name": "nav_mit",
"dual_head": false,
"main_class_name": "navegavel",
"es_classes": "",
@ -13,5 +13,5 @@
"shaves": 3,
"channels": 3,
"use_ndvi": false,
"backbone": "nvidia/segformer-b0-finetuned-ade-512-512"
"backbone": "nvidia/mit-b0"
}

View File

@ -0,0 +1,422 @@
import os
import time
import json
import math
import ctypes as C
from ctypes import wintypes as W
from datetime import datetime
import numpy as np
import cv2
# =========================
# CONFIG
# =========================
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
# =========================
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):
return lo if v < lo else hi if v > hi else v
def norm8(x, p_lo=2, p_hi=98):
lo = np.percentile(x, p_lo)
hi = np.percentile(x, p_hi)
if hi <= lo + 1:
return x.astype(np.uint8)
y = (x.astype(np.float32) - lo) * (255.0 / (hi - lo))
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)
bgr = np.dstack([Bn, Gn, Rn]) # OpenCV usa BGR
if upscale and upscale != 1:
bgr = cv2.resize(bgr, (bgr.shape[1]*upscale, bgr.shape[0]*upscale), interpolation=cv2.INTER_NEAREST)
return bgr
def measure_raw_g_metrics(raw: np.ndarray):
"""
Mede brilho no canal G cru usando uma ROI na base (mais parecido com chão).
Retorna p90/p95 e fração saturada.
"""
G = raw[0::2, 1::2] # H/2 x W/2
h2, w2 = G.shape
# ROI: base da imagem, cortando laterais
y0, y1 = int(h2 * 0.55), int(h2 * 0.95)
x0, x1 = int(w2 * 0.15), int(w2 * 0.85)
roi = G[y0:y1, x0:x1]
p90 = float(np.percentile(roi, 90))
p95 = float(np.percentile(roi, 95))
sat = float(np.mean(roi >= 250))
return p90, p95, sat
class RobustAE:
"""
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):
p90, p95, sat = measure_raw_g_metrics(raw)
# EMA do p95 (estabiliza)
if self.p95_ema is None:
self.p95_ema = p95
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:
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 3, cv2.LINE_AA)
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255,255,255), 2, cv2.LINE_AA)
y += 28
def save_sample(raw: np.ndarray, bgr_preview: np.ndarray, meta: dict):
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)
cv2.imwrite(png_path, bgr_preview)
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
def main():
# scan
n = C.c_ubyte(0)
ck(dll.VT_DeviceScan(C.byref(n), DEVICE_UDEF), "VT_DeviceScan")
if n.value == 0:
raise RuntimeError("Nenhuma câmera encontrada.")
# open
idx = C.c_ubyte(0)
h = W.HANDLE()
ck(dll.VT_DeviceOpen(C.byref(idx), C.byref(h), DEVICE_INDEX, DEVICE_UDEF), "VT_DeviceOpen")
print("DeviceOpen OK, handle=", h.value)
print("Saving to:", SESSION_DIR)
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
# Estado local (não dependemos de GET)
exp_raw = 1500
gain_a = 0
gain_d = 0
# Aplica estado inicial
try:
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
param_set_int(h, PARAM_ID_SENSOR_GAINANALOGRAW, gain_a)
param_set_int(h, PARAM_ID_SENSOR_GAINDIGITRAW, gain_d)
except Exception as e:
print("[WARN] Falhou set inicial:", e)
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
last_auto_t = 0.0
# FPS
t0 = time.time()
frames = 0
fps = 0.0
last_msg = ""
last_msg_t = 0.0
try:
while True:
raw = capture_raw8(h)
# soft AE
ae_dbg = {}
if ae_on:
new_exp, ae_dbg = ae.step(raw, exp_raw)
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
rgb_clean = make_rgb_preview(raw, upscale=UPSCALE)
bgr = rgb_clean.copy()
# FPS
frames += 1
dt = time.time() - t0
if dt >= 1.0:
fps = frames / dt
frames = 0
t0 = time.time()
# HUD
lines = [
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Interval: {CAPTURE_INTERVAL_S:.1f}s",
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)}",
"Keys: C/SPACE=save | A=toggle autosave | E=toggle AE | +/- exp | Q/ESC quit",
]
overlay_hud(bgr, lines)
# msg pós-save
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)
# autosave
now = time.time()
if auto_save and (now - last_auto_t) >= CAPTURE_INTERVAL_S:
meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"raw_w": RAW_W, "raw_h": RAW_H,
"exp_raw": int(exp_raw),
"gain_a": int(gain_a),
"gain_d": int(gain_d),
"ae_on": bool(ae_on),
"note": "autosave",
}
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
last_msg = f"SAVED: {os.path.basename(raw_path)}"
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('e'), ord('E')):
ae_on = not ae_on
last_msg = f"AE -> {'ON' if ae_on else 'OFF'}"
last_msg_t = time.time()
elif k in (ord('c'), ord('C'), 32): # C ou SPACE
meta = {
"ts": datetime.now().isoformat(timespec="milliseconds"),
"raw_w": RAW_W, "raw_h": RAW_H,
"exp_raw": int(exp_raw),
"gain_a": int(gain_a),
"gain_d": int(gain_d),
"ae_on": bool(ae_on),
"note": "manual",
}
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
last_msg = f"SAVED: {os.path.basename(raw_path)}"
last_msg_t = time.time()
elif k in (ord('+'), ord('=')):
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('_')):
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)
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()
print("Fim.")
if __name__ == "__main__":
main()

View File

@ -22,9 +22,9 @@ RAW_WIDTH = RESOLUCAO[0] # largura real do raw (ex: 1296)
RAW_HEIGHT = RESOLUCAO[1] # altura real do raw (ex: 1028)
RAW_EXTS = [".raw"] # se tiver outro, adiciona aqui
PASTA_FINAL_PREVIEWS = os.path.join(MODELO, "dataset", "original", "previews")
PASTA_FINAL_MASKS = os.path.join(MODELO, "dataset", "original", "masks")
PASTA_FINAL_RAWS = os.path.join(MODELO, "dataset", "original", "raws")
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
# Cor da classe será obtida via labelmap
# COR_CLASSE_RGB = (128, 0, 0)
@ -200,7 +200,7 @@ def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None:
return None
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False):
source = os.path.join(MODELO, "dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join(MODELO, "dataset", "original")
source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original")
pasta_new_previews = os.path.join(source, "group", grupo, "previews")
pasta_new_masks = os.path.join(source, "group", grupo, "masks")
pasta_new_raws = os.path.join(source, "group", grupo, "raws")
@ -296,7 +296,7 @@ if __name__ == "__main__":
args = ap.parse_args()
# Caminho padrão pro labelmap, caso não seja informado
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
labelmap_path = os.path.join("dataset", "labelmap.txt")
if not os.path.exists(labelmap_path):
raise SystemExit(f"Labelmap não encontrado em: {labelmap_path}")

View File

@ -12,15 +12,15 @@ with open("config.json", "r") as f:
MODELO = config["camera"]
# Raiz das brutas agrupadas
PASTA_BRUTAS_GROUP_ROOT = os.path.join(MODELO, "dataset", "brutas")
PASTA_BRUTAS_GROUP_ROOT = os.path.join("dataset", "brutas")
# Onde as máscaras novas (rotuladas externamente) são colocadas
PASTA_NEW_MASKS = os.path.join(MODELO, "dataset", "new_masks")
PASTA_NEW_MASKS = os.path.join("dataset", "new_masks")
# Destino final do dataset consolidado
PASTA_FINAL_PREVIEWS = os.path.join(MODELO, "dataset", "original", "previews")
PASTA_FINAL_RAWS = os.path.join(MODELO, "dataset", "original", "raws")
PASTA_FINAL_MASKS = os.path.join(MODELO, "dataset", "original", "masks")
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
EXT_MASKS = (".png",)

View File

@ -12,14 +12,14 @@ with open("config.json", "r") as f:
MODELO = config["camera"]
# Raiz das brutas (todas as canas/horários/grupos)
PASTA_BRUTAS_ROOT = os.path.join(MODELO, "dataset", "brutas")
PASTA_BRUTAS_ROOT = os.path.join("dataset", "brutas")
# Onde você coloca os previews selecionados (tudo misturado)
PASTA_SELECTED_PREVIEWS = os.path.join(MODELO, "dataset", "selected_previews")
PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews")
# Destino final, organizado por grupo:
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
PASTA_ORIGINAL_GROUP_ROOT = os.path.join(MODELO, "dataset", "original", "group")
PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group")
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
EXT_RAWS = (".raw",)

View File

@ -309,8 +309,8 @@ def build_cli():
if __name__ == "__main__":
ap = build_cli()
args = ap.parse_args()
originals_dir = os.path.join(MODELO, "dataset", "original")
labelmap_dir = os.path.join(MODELO, "dataset", "labelmap.txt")
originals_dir = os.path.join("dataset", "original")
labelmap_dir = os.path.join("dataset", "labelmap.txt")
processar(
originals_dir=originals_dir,
labelmap_path=labelmap_dir,

View File

@ -6,15 +6,15 @@ Augmenta imagens/máscaras *por grupo* usando RAW4 (float32, 0..1, R G IR B).
Novo fluxo (GAL5000):
Entrada:
MODELO/dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
MODELO/dataset/original/group/<grupo>/masks (.png/.jpg)
(opcional) MODELO/dataset/original/group/<grupo>/masks2
dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
dataset/original/group/<grupo>/masks (.png/.jpg)
(opcional) dataset/original/group/<grupo>/masks2
Saída:
MODELO/dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
MODELO/dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
MODELO/dataset/augmented/group/<grupo>/masks
MODELO/dataset/augmented/group/<grupo>/masks2
dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
dataset/augmented/group/<grupo>/masks
dataset/augmented/group/<grupo>/masks2
Regras:
- GEOMETRIA (flip, shift, rotate) aplicada em RAW + masks (+ masks2).
@ -44,7 +44,7 @@ RESOLUCAO = config.get("raw_size")
USE_MASKS2 = config.get("dual_head", False)
# Pastas base
DATASET_BASE = os.path.join(MODELO, "dataset")
DATASET_BASE = os.path.join("dataset")
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")

View File

@ -4,11 +4,11 @@
Normaliza/redimensiona PREVIEW + RAW + MASK (+MASK2), mantendo ESTRUTURA POR GRUPO.
Entradas:
MODELO/dataset/original/group/<grupo>/{previews,raws,masks,(masks2)}
MODELO/dataset/augmented/group/<grupo>/{previews,raws,masks,(masks2)}
dataset/original/group/<grupo>/{previews,raws,masks,(masks2)}
dataset/augmented/group/<grupo>/{previews,raws,masks,(masks2)}
Saídas (por resolução):
MODELO/dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
Conversão de máscara:
- máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids
@ -37,9 +37,10 @@ USE_MASKS2 = config["dual_head"]
RESOLUCAO = tuple(config["resolucao"]) # [W,H]
MODEL_NAME = config["model_name"]
CHANNELS = int(config.get("channels", 4))
pasta_base = os.path.join(MODELO, "dataset")
pasta_base = os.path.join("dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
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])}
FONTES = ["original", "augmented"]

View File

@ -5,10 +5,10 @@ Split estratificado por GRUPO com **val/test só do ORIGINAL** e
garantia de NÃO VAZAMENTO entre splits, agora para estrutura RAW:
de:
MODELO/dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
Escreve em:
MODELO/dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}
dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}
Definições:
- "Família" = todas as variações da MESMA base original:
@ -33,8 +33,8 @@ USE_MASKS2 = config.get("dual_head", False)
RESOLUCAO = tuple(config.get("resolucao"))
# Pastas (ajustadas para PREVIEWS/RAWS)
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
pasta_destino = os.path.join(MODELO, "dataset", "split")
pasta_origem = os.path.join("dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
pasta_destino = os.path.join("dataset", "split")
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXT = ".png" # máscaras normalizadas em PNG
@ -376,7 +376,6 @@ def main():
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
ap.add_argument("--modelo", type=str, default=None, help="Sobrescreve MODELO do config.json.")
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")
ap.add_argument("--cap-train-families", type=str, default="",
@ -384,7 +383,6 @@ def main():
args = ap.parse_args()
modelo = args.modelo or MODELO
if args.resolucao:
try:
w, h = args.resolucao.lower().split("x")
@ -406,8 +404,8 @@ def main():
caps_map = parse_cap_map(args.cap_train_families)
global pasta_origem, pasta_destino
pasta_origem = os.path.join(modelo, "dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
pasta_destino = os.path.join(modelo, "dataset", "split")
pasta_origem = os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
pasta_destino = os.path.join("dataset", "split")
soma = args.train + args.val + args.test
if soma <= 0:

View File

@ -358,10 +358,10 @@ def main():
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, experiment_tag)
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, experiment_tag)
#save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
dataset_path = os.path.join(MODELO, "dataset")
#save_path = os.path.join("backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
dataset_path = os.path.join("dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
os.makedirs(save_path, exist_ok=True)
@ -378,7 +378,7 @@ def main():
# 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(MODELO, "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:
norm_stats_path = args.norm_stats
@ -547,7 +547,6 @@ def main():
strict=False,
)
patch_segformer_input_channels(model, in_ch=CHANNELS)
elif FUSION_MODE == "dual_branch":
model = build_dual_branch_segformer_model(
num_classes=num_classes,

View File

@ -243,7 +243,7 @@ def main():
print(f"[cfg] resolucao nominal RAW: {W}x{H}")
print(f"[cfg] channels={CHANNELS} use_ndvi={USE_NDVI}")
dataset_path = os.path.join(MODELO, "dataset")
dataset_path = os.path.join("dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
# ==========================
@ -254,7 +254,7 @@ def main():
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
norm_stats_path = os.path.join(MODELO, "backup", modelo_folder, MODEL_NAME, experiment_tag, "norm_stats.json")
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):
@ -307,7 +307,7 @@ def main():
if args.ckpt is not None:
ckpt_path = args.ckpt
else:
save_path = os.path.join(MODELO, "backup", modelo_folder, MODEL_NAME, experiment_tag)
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):

View File

@ -1,19 +1,19 @@
{
"camera": "gal5000",
"modelo": "segformer_b1",
"model_name": "pulv_new",
"model_name": "pulv_mit",
"dual_head": false,
"main_class_name": "cana",
"es_classes": "",
"model_to_use": "geral",
"raw_size": [1296, 1028],
"resolucao": [1008, 800],
"resolucao": [1024, 800],
"roi_inicio": 0.0,
"roi_tamanho": 1.0,
"shaves": 3,
"channels": 4,
"use_ndvi": false,
"backbone": "nvidia/segformer-b1-finetuned-ade-512-512",
"backbone": "nvidia/mit-b1",
"fusion_mode": "stacked",
"stats_source_tag": "stacked_raw4"
}

View File

@ -0,0 +1,390 @@
# -*- coding: utf-8 -*-
"""
Regenera previews a partir dos RAWs antigos da GAL5000,
aplicando compensação de IR no preview.
Estrutura esperada:
dataset/brutas/group/{GRUPO}/
masks/
previews/
raws/
metas/
Uso exemplo:
python regen_previews_from_raws.py ^
--root dataset/brutas/group ^
--raw-h 1028 ^
--raw-w 1296 ^
--layout rgirb ^
--ir-k-r 0.8 ^
--ir-k-g 0.4 ^
--ir-k-b 0.9
Observações:
- Para .raw, informe --raw-h e --raw-w
- Para .npy/.npz, o shape é lido automaticamente
- O script recria os previews em cada pasta previews do grupo
- Por padrão sobrescreve os previews existentes
"""
from __future__ import annotations
import argparse
import os
from pathlib import Path
from typing import Optional, Tuple
import numpy as np
try:
import cv2
except ImportError:
cv2 = None
from PIL import Image
# ---------------------------------------------------------
# Preview
# ---------------------------------------------------------
def make_bgr_preview_from_raw(
raw_np: np.ndarray,
rgirb: bool,
preview_fast: bool,
preview_scale: int = 2,
apply_ir_comp: bool = True,
ir_k_r: float = 0.40,
ir_k_g: float = 0.10,
ir_k_b: float = 0.50,
) -> np.ndarray:
"""
Gera preview BGR (OpenCV) a partir do tensor raw_np (C,H,W) float32 em 0..1.
Casos esperados:
- rgirb=False:
raw_np = [R, G, B] ou [R, G, B, NDVI]
- rgirb=True:
raw_np = [R, G, IR, B] ou [R, G, IR, B, NDVI]
"""
assert raw_np.ndim == 3, "raw_np deve ser (C,H,W)"
C = raw_np.shape[0]
if C not in (3, 4, 5):
raise RuntimeError(f"Esperado C=3, 4 ou 5, veio {C}")
raw_np = raw_np.astype(np.float32, copy=False)
r = raw_np[0]
g = raw_np[1]
b = raw_np[3] if rgirb else raw_np[2]
if preview_fast:
if preview_scale > 1:
r = r[::preview_scale, ::preview_scale]
g = g[::preview_scale, ::preview_scale]
b = b[::preview_scale, ::preview_scale]
if rgirb and apply_ir_comp and C >= 4:
ir = raw_np[2]
if preview_scale > 1:
ir = ir[::preview_scale, ::preview_scale]
r = np.clip(r - ir_k_r * ir, 0.0, 1.0)
g = np.clip(g - ir_k_g * ir, 0.0, 1.0)
b = np.clip(b - ir_k_b * ir, 0.0, 1.0)
bgr = np.stack([b, g, r], axis=0)
bgr = np.power(np.clip(bgr, 0.0, 1.0), 1 / 1.8)
bgr8 = (bgr * 255.0).clip(0, 255).astype(np.uint8)
return np.transpose(bgr8, (1, 2, 0)).copy()
if rgirb and apply_ir_comp and C >= 4:
ir = raw_np[2]
r = np.clip(r - ir_k_r * ir, 0.0, 1.0)
g = np.clip(g - ir_k_g * ir, 0.0, 1.0)
b = np.clip(b - ir_k_b * ir, 0.0, 1.0)
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)
r = stretch_channel(r)
g = stretch_channel(g)
b = stretch_channel(b)
bgr = np.stack([b, g, r], axis=0).astype(np.float32)
bgr = np.power(np.clip(bgr, 0.0, 1.0), 1 / 2.0)
bgr8 = (bgr * 255.0).clip(0, 255).astype(np.uint8)
return np.transpose(bgr8, (1, 2, 0)).copy()
# ---------------------------------------------------------
# IO helpers
# ---------------------------------------------------------
def save_bgr_image(path: Path, bgr: np.ndarray) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
if cv2 is not None:
ok = cv2.imwrite(str(path), bgr)
if not ok:
raise RuntimeError(f"Falha ao salvar imagem: {path}")
else:
rgb = bgr[..., ::-1]
Image.fromarray(rgb).save(path)
def to_chw_float01(arr: np.ndarray, layout: str) -> np.ndarray:
"""
Converte entrada para (C,H,W) float32 0..1.
layout:
- rgirb
- rgbir
- rgb
"""
arr = np.asarray(arr)
if arr.ndim != 3:
raise RuntimeError(f"Esperava array 3D, veio shape={arr.shape}")
# HWC -> CHW
if arr.shape[-1] in (3, 4, 5) and arr.shape[0] not in (3, 4, 5):
arr = np.transpose(arr, (2, 0, 1))
if arr.shape[0] not in (3, 4, 5):
raise RuntimeError(f"Não consegui interpretar canais em shape={arr.shape}")
arr = arr.astype(np.float32, copy=False)
# Normalização para 0..1
if arr.dtype == np.uint8:
arr = arr / 255.0
elif arr.dtype == np.uint16:
arr = arr / 65535.0
else:
# Se já vier float mas fora de 0..1, tenta ajustar
maxv = float(np.nanmax(arr)) if arr.size else 1.0
if maxv > 1.0:
arr = arr / maxv
arr = np.clip(arr, 0.0, 1.0)
# Reorganiza para o contrato do preview
# Queremos:
# rgirb=True -> [R,G,IR,B]
# rgirb=False -> [R,G,B]
if layout == "rgirb":
# já assume [R,G,IR,B] ou [R,G,IR,B,NDVI]
return arr
elif layout == "rgbir":
# [R,G,B,IR] -> [R,G,IR,B]
if arr.shape[0] < 4:
raise RuntimeError("layout=rgbir exige pelo menos 4 canais")
if arr.shape[0] == 4:
arr = arr[[0, 1, 3, 2], :, :]
else:
# [R,G,B,IR,NDVI] -> [R,G,IR,B,NDVI]
arr = arr[[0, 1, 3, 2, 4], :, :]
return arr
elif layout == "rgb":
return arr[:3]
else:
raise ValueError(f"layout inválido: {layout}")
def load_raw_file(raw_path: Path, raw_hw: Optional[Tuple[int, int]], layout: str) -> np.ndarray:
"""
Retorna (C,H,W) float32 em 0..1.
Suporta:
- .npy
- .npz
- .raw
Para .raw:
- mosaico uint8 HxW em padrão 2x2 R,G / IR,B
- ou RAW4 float32 (4,H,W) salvo em [R,G,IR,B]
"""
ext = raw_path.suffix.lower()
if ext == ".npy":
arr = np.load(raw_path)
return to_chw_float01(arr, layout)
if ext == ".npz":
z = np.load(raw_path)
key = list(z.keys())[0]
arr = z[key]
return to_chw_float01(arr, layout)
if ext == ".raw":
if raw_hw is None:
raise RuntimeError(f"{raw_path.name}: para .raw informe --raw-h e --raw-w")
H, W = raw_hw
size_bytes = raw_path.stat().st_size
mosa_bytes = H * W
raw4_bytes = 4 * H * W * 4 # float32
if size_bytes == mosa_bytes:
# mosaico uint8 cru
arr = np.fromfile(raw_path, dtype=np.uint8).reshape(H, W)
if (H % 2) != 0 or (W % 2) != 0:
raise RuntimeError(f"{raw_path.name}: H e W precisam ser pares para mosaico 2x2")
r_sub = arr[0::2, 0::2]
g_sub = arr[0::2, 1::2]
ir_sub = arr[1::2, 0::2]
b_sub = arr[1::2, 1::2]
if cv2 is not None:
r = cv2.resize(r_sub, (W, H), interpolation=cv2.INTER_LINEAR)
g = cv2.resize(g_sub, (W, H), interpolation=cv2.INTER_LINEAR)
ir = cv2.resize(ir_sub, (W, H), interpolation=cv2.INTER_LINEAR)
b = cv2.resize(b_sub, (W, H), interpolation=cv2.INTER_LINEAR)
else:
r = np.array(Image.fromarray(r_sub).resize((W, H), resample=Image.BILINEAR))
g = np.array(Image.fromarray(g_sub).resize((W, H), resample=Image.BILINEAR))
ir = np.array(Image.fromarray(ir_sub).resize((W, H), resample=Image.BILINEAR))
b = np.array(Image.fromarray(b_sub).resize((W, H), resample=Image.BILINEAR))
chw = np.stack([r, g, ir, b], axis=0).astype(np.float32) / 255.0
return chw
if size_bytes == raw4_bytes:
# RAW4 float32 salvo como (4,H,W) em [R,G,IR,B]
arr = np.fromfile(raw_path, dtype=np.float32).reshape(4, H, W)
arr = np.clip(arr, 0.0, 1.0)
return to_chw_float01(arr, layout="rgirb")
raise RuntimeError(
f"{raw_path.name}: tamanho inesperado {size_bytes} bytes "
f"(esperado mosaico={mosa_bytes} ou raw4 float32={raw4_bytes})"
)
raise RuntimeError(f"Extensão não suportada: {raw_path}")
# ---------------------------------------------------------
# Processamento
# ---------------------------------------------------------
def process_group(
group_dir: Path,
raw_hw: Optional[Tuple[int, int]],
layout: str,
preview_fast: bool,
preview_scale: int,
apply_ir_comp: bool,
ir_k_r: float,
ir_k_g: float,
ir_k_b: float,
overwrite: bool,
) -> tuple[int, int]:
raws_dir = group_dir / "raws"
previews_dir = group_dir / "previews"
if not raws_dir.is_dir():
return 0, 0
previews_dir.mkdir(parents=True, exist_ok=True)
raw_files = []
for ext in ("*.raw", "*.npy", "*.npz"):
raw_files.extend(sorted(raws_dir.glob(ext)))
done = 0
failed = 0
for raw_path in raw_files:
out_path = previews_dir / f"{raw_path.stem}.jpg"
if out_path.exists() and not overwrite:
continue
try:
raw_np = load_raw_file(raw_path, raw_hw=raw_hw, layout=layout)
bgr = make_bgr_preview_from_raw(
raw_np=raw_np,
rgirb=(layout == "rgirb" or layout == "rgbir"),
preview_fast=preview_fast,
preview_scale=preview_scale,
apply_ir_comp=apply_ir_comp,
ir_k_r=ir_k_r,
ir_k_g=ir_k_g,
ir_k_b=ir_k_b,
)
save_bgr_image(out_path, bgr)
done += 1
print(f"[OK] {group_dir.name}/{raw_path.name} -> {out_path.name}")
except Exception as e:
failed += 1
print(f"[ERRO] {group_dir.name}/{raw_path.name}: {e}")
return done, failed
def main():
parser = argparse.ArgumentParser(description="Regenera previews a partir dos RAWs antigos.")
parser.add_argument("--root", type=str, required=True, help="Pasta group, ex: dataset/brutas/group")
parser.add_argument("--raw-h", type=int, default=None, help="Altura do RAW para arquivos .raw")
parser.add_argument("--raw-w", type=int, default=None, help="Largura do RAW para arquivos .raw")
parser.add_argument("--layout", type=str, default="rgirb", choices=["rgirb", "rgbir", "rgb"], help="Layout dos canais dos RAWs")
parser.add_argument("--preview-fast", action="store_true", help="Usa modo rápido")
parser.add_argument("--preview-scale", type=int, default=2, help="Escala no preview_fast")
parser.add_argument("--no-ir-comp", action="store_true", help="Desliga compensação de IR")
parser.add_argument("--ir-k-r", type=float, default=0.8)
parser.add_argument("--ir-k-g", type=float, default=0.4)
parser.add_argument("--ir-k-b", type=float, default=0.9)
parser.add_argument("--no-overwrite", action="store_true", help="Não sobrescreve previews existentes")
args = parser.parse_args()
root = Path(args.root)
if not root.is_dir():
raise RuntimeError(f"Pasta root não encontrada: {root}")
raw_hw = None
if args.raw_h is not None and args.raw_w is not None:
raw_hw = (args.raw_h, args.raw_w)
total_done = 0
total_failed = 0
group_dirs = [p for p in sorted(root.iterdir()) if p.is_dir()]
if not group_dirs:
raise RuntimeError(f"Nenhum grupo encontrado em: {root}")
for group_dir in group_dirs:
done, failed = process_group(
group_dir=group_dir,
raw_hw=raw_hw,
layout=args.layout,
preview_fast=args.preview_fast,
preview_scale=args.preview_scale,
apply_ir_comp=not args.no_ir_comp,
ir_k_r=args.ir_k_r,
ir_k_g=args.ir_k_g,
ir_k_b=args.ir_k_b,
overwrite=not args.no_overwrite,
)
total_done += done
total_failed += failed
print("\n============================================")
print("Regeneração concluída")
print(f"Previews gerados : {total_done}")
print(f"Falhas : {total_failed}")
print("============================================")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,669 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
raw_triplet_sorter.py
-------------------------------------------------
Ferramenta para classificar manualmente amostras RAW do dataset
(arquivos .raw + .png + .json) em pastas de labels,
usando atalhos de teclado.
Fluxo:
- Entrada: uma ou mais pastas de sessão (cada uma contendo N arquivos):
- NOME.raw
- NOME.png
- 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:
out_root/
<label>/
raws/
previews/
metas/
masks/
Teclas:
- 1..9, 0 -> envia para a label correspondente
- Espaço / n / seta direita -> pular (skip)
- p / seta esquerda -> voltar
- b -> undo (desfaz última ação)
- q / Esc -> sair
Depedências:
- Python 3.8+
- Pillow (PIL): pip install pillow
"""
import argparse
import csv
import json
import os
import shutil
import sys
from datetime import datetime
from pathlib import Path
from typing import List, Tuple, Dict, Optional
import tkinter as tk
from tkinter import messagebox, filedialog
try:
from PIL import Image, ImageTk
except Exception as e:
print("ERROR: Pillow (PIL) is required. Install with: pip install pillow", file=sys.stderr)
raise
# Sufixos/padrões esperados:
# - preview: NOME.png
# - raw: NOME.raw
# - meta: NOME.json
PREVIEW_EXT = ".png"
RAW_EXT = ".raw"
META_EXT = ".json"
def find_triplets_in_folder(folder: Path) -> List[Tuple[Path, Path, Path, str]]:
"""
Encontra tripletas (raw, preview, json) em uma pasta:
NOME.png
NOME.raw
NOME.json
Retorna lista de tuplas (preview_path, raw_path, json_path, base_name).
"""
if not folder.is_dir():
return []
previews: Dict[str, Path] = {}
raws: Dict[str, Path] = {}
metas: Dict[str, Path] = {}
for p in folder.rglob("*"):
if not p.is_file():
continue
stem = p.stem
suffix = p.suffix.lower()
# Preview
if suffix == PREVIEW_EXT:
base = stem
previews[base] = p
continue
# RAW
if suffix == RAW_EXT:
base = stem
raws[base] = p
continue
# JSON
if suffix == META_EXT:
base = stem
metas[base] = p
continue
triplets = []
for base, prev_path in previews.items():
raw_path = raws.get(base)
meta_path = metas.get(base)
if raw_path is not None and meta_path is not None:
triplets.append((prev_path, raw_path, meta_path, base))
# Ordena por base (se for número, ordena numérico)
def sort_key(t):
b = t[3]
return (0, int(b)) if b.isdigit() else (1, b)
triplets.sort(key=sort_key)
return triplets
def collect_all_triplets(folders: List[Path]) -> List[Tuple[Path, Path, Path, str]]:
all_tr = []
for f in folders:
all_tr.extend(find_triplets_in_folder(f))
# já vem ordenado por pasta, mas garantimos ordenação global
def sort_key(t):
b = t[3]
return (0, int(b)) if b.isdigit() else (1, b)
all_tr.sort(key=sort_key)
return all_tr
class ActionLogger:
"""
Loga ações em:
- CSV: sorting_log.csv
- JSONL: sorting_log.jsonl
Serve também para suportar "resume" (não reprocessar o que foi classificado/pulado).
"""
def __init__(self, out_root: Path):
self.out_root = out_root
self.log_csv = out_root / "sorting_log.csv"
self.log_json = out_root / "sorting_log.jsonl"
self.seen = set()
if self.log_csv.exists():
try:
with self.log_csv.open("r", newline="", encoding="utf-8") as f:
reader = csv.DictReader(f)
for row in reader:
key = row.get("preview_path", "")
if key:
self.seen.add(key)
except Exception:
pass
def already_logged(self, preview_path: Path) -> bool:
return str(preview_path) in self.seen
def log(
self,
action: str,
preview_path: Path,
raw_path: Path,
meta_path: Path,
label: Optional[str] = None,
dest_preview: Optional[Path] = None,
dest_raw: Optional[Path] = None,
dest_meta: Optional[Path] = None,
):
self.out_root.mkdir(parents=True, exist_ok=True)
# CSV
new_file = not self.log_csv.exists()
with self.log_csv.open("a", newline="", encoding="utf-8") as f:
writer = csv.writer(f)
if new_file:
writer.writerow([
"timestamp",
"action",
"label",
"preview_path",
"raw_path",
"meta_path",
"dest_preview",
"dest_raw",
"dest_meta",
])
writer.writerow([
datetime.now().isoformat(timespec="seconds"),
action,
label or "",
str(preview_path),
str(raw_path),
str(meta_path),
str(dest_preview or ""),
str(dest_raw or ""),
str(dest_meta or ""),
])
# JSONL
with self.log_json.open("a", encoding="utf-8") as f:
rec = {
"ts": datetime.now().isoformat(timespec="seconds"),
"action": action,
"label": label,
"preview_path": str(preview_path),
"raw_path": str(raw_path),
"meta_path": str(meta_path),
"dest_preview": str(dest_preview) if dest_preview else None,
"dest_raw": str(dest_raw) if dest_raw else None,
"dest_meta": str(dest_meta) if dest_meta else None,
}
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
self.seen.add(str(preview_path))
class RawTripletSorterApp:
def __init__(
self,
triplets: List[Tuple[Path, Path, Path, str]],
labels: List[str],
out_root: Path,
move: bool,
resume: bool,
display_height: int = 512,
):
self.all_triplets = triplets
self.labels = labels
self.out_root = out_root
self.move = move
self.display_height = display_height
self.logger = ActionLogger(out_root)
if resume:
self.all_triplets = [
t for t in self.all_triplets if not self.logger.already_logged(t[0])
]
# Estado de sessão
self.idx = 0
self.history: List[dict] = [] # para undo
# Prepara pastas destino
for label in self.labels:
(out_root / label / "raws").mkdir(parents=True, exist_ok=True)
(out_root / label / "previews").mkdir(parents=True, exist_ok=True)
(out_root / label / "metas").mkdir(parents=True, exist_ok=True)
(out_root / label / "masks").mkdir(parents=True, exist_ok=True)
# UI
self.root = tk.Tk()
self.root.title("Agrobotics RAW Triplet Sorter (preview | raw | json)")
self.root.geometry("1100x750")
self.root.bind("<Key>", self.on_key)
# Topo: info + legenda
self.top_frame = tk.Frame(self.root)
self.top_frame.pack(side=tk.TOP, fill=tk.X)
self.info_label = tk.Label(self.top_frame, text="", font=("Segoe UI", 11))
self.info_label.pack(side=tk.LEFT, padx=10, pady=6)
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)
# Área da imagem
self.img_frame = tk.Frame(self.root)
self.img_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
self.preview_label = tk.Label(self.img_frame)
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_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)
# Rodapé: ajuda
self.footer = tk.Label(
self.root,
text="1..9/0 = labels | espaço/n/→ = próxima (skip) | p/← = anterior | b = undo | q/Esc = sair",
font=("Segoe UI", 10),
)
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
# Primeira renderização
self.render()
def build_legend_text(self) -> str:
parts = []
for i, label in enumerate(self.labels, start=1):
key = i if i <= 9 else 0 # 0 = 10ª label
parts.append(f"[{key}] {label}")
return " | ".join(parts)
def pil_load_preview(self, preview_path: Path) -> ImageTk.PhotoImage:
img = Image.open(preview_path).convert("RGB")
# escala pra altura desejada mantendo proporção
h_target = self.display_height
w, h = img.size
new_w = int(w * (h_target / h))
img = img.resize((new_w, h_target), Image.BILINEAR)
return ImageTk.PhotoImage(img)
def render(self):
if not self.all_triplets:
messagebox.showinfo("Fim", "Não há amostras para exibir (talvez tudo já foi classificado?).")
self.root.destroy()
return
self.idx = max(0, min(self.idx, len(self.all_triplets) - 1))
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx]
try:
tk_img = self.pil_load_preview(preview_path)
self.preview_tk = tk_img
self.preview_label.configure(image=self.preview_tk)
except Exception as e:
self.info_label.configure(text=f"Erro ao abrir preview: {e}")
return
self.info_label.configure(
text=f"{self.idx+1}/{len(self.all_triplets)} | base='{base}' | PREVIEW: {preview_path.name}"
)
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):
if label_index < 0 or label_index >= len(self.labels):
return
label = self.labels[label_index]
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx]
dst_preview = self.out_root / label / "previews" / preview_path.name
dst_raw = self.out_root / label / "raws" / raw_path.name
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:
self.do_copy_or_move(preview_path, dst_preview)
self.do_copy_or_move(raw_path, dst_raw)
self.do_copy_or_move(meta_path, dst_meta)
self.logger.log(
"assign",
preview_path=preview_path,
raw_path=raw_path,
meta_path=meta_path,
label=label,
dest_preview=dst_preview,
dest_raw=dst_raw,
dest_meta=dst_meta,
)
self.history.append({
"action": "assign",
"label": label,
"preview_src": preview_path,
"raw_src": raw_path,
"meta_src": meta_path,
"preview_dst": dst_preview,
"raw_dst": dst_raw,
"meta_dst": dst_meta,
"moved": self.move,
"index": self.idx,
})
self.status_var.set(
f"{'Movido' if self.move else 'Copiado'}'{label}': {preview_path.name} (+raw+json)"
)
self.idx += 1
if self.idx >= len(self.all_triplets):
messagebox.showinfo("Concluído", "Você chegou ao final da fila!")
self.root.destroy()
return
self.render()
except Exception as e:
messagebox.showerror("Erro", f"Falha ao copiar/mover: {e}")
self.status_var.set(f"ERRO: {e}")
def undo(self):
if not self.history:
return
last = self.history.pop()
if last["action"] != "assign":
return
try:
if last["moved"]:
# mover de volta pro source
shutil.move(str(last["preview_dst"]), str(last["preview_src"]))
shutil.move(str(last["raw_dst"]), str(last["raw_src"]))
shutil.move(str(last["meta_dst"]), str(last["meta_src"]))
else:
# apagar os arquivos de destino
for p in [last["preview_dst"], last["raw_dst"], last["meta_dst"]]:
if p.exists():
p.unlink()
self.logger.log(
"undo",
preview_path=last["preview_src"],
raw_path=last["raw_src"],
meta_path=last["meta_src"],
label=last["label"],
dest_preview=last["preview_dst"],
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.status_var.set(f"Desfeito ← '{last['label']}': {last['preview_src'].name}")
self.render()
except Exception as e:
messagebox.showerror("Erro", f"Falha no undo: {e}")
self.status_var.set(f"ERRO: {e}")
def on_key(self, event):
ch = event.keysym.lower()
# 1..9 / 0 -> labels
if ch in [str(i) for i in range(1, 10)] or ch == "0":
label_idx = 9 if ch == "0" else int(ch) - 1
self.send_to_label(label_idx)
return
# espaço / n / seta direita = skip / próxima
if ch in ("space", "s", "n", "right"):
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx]
self.logger.log("skip", preview_path, raw_path, meta_path, label=None)
self.status_var.set(f"Pulada → {preview_path.name}")
self.idx += 1
if self.idx >= len(self.all_triplets):
messagebox.showinfo("Concluído", "Você chegou ao final da fila!")
self.root.destroy()
return
self.render()
return
# p / seta esquerda = voltar
if ch in ("p", "left"):
self.idx = max(0, self.idx - 1)
self.status_var.set("Voltou uma imagem.")
self.render()
return
# b = undo
if ch in ("b",):
self.undo()
return
# q / Esc = sair
if ch in ("q", "escape"):
self.root.destroy()
return
def run(self):
self.root.mainloop()
class SetupWindow:
"""
Janela de setup (igual ao sorter anterior, mas voltada para RAW triplets).
"""
def __init__(self):
self.root = tk.Tk()
self.root.title("Configurar - RAW Triplet Sorter")
self.root.geometry("720x520")
# Pastas de entrada
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)
self.inputs_listbox = tk.Listbox(frm_in, height=6)
self.inputs_listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=6, pady=6)
btns_in = tk.Frame(frm_in)
btns_in.pack(side=tk.RIGHT, fill=tk.Y, padx=6, pady=6)
tk.Button(btns_in, text="Adicionar pasta...", command=self.add_input).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)
# Pasta de saída
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)
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.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.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")
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.pack(fill=tk.X, expand=False, padx=10, pady=8)
self.move_var = tk.BooleanVar(value=False)
self.resume_var = tk.BooleanVar(value=True)
self.height_var = tk.IntVar(value=512)
tk.Checkbutton(frm_opts, text="Mover em vez de copiar", variable=self.move_var).pack(anchor="w", padx=6)
tk.Checkbutton(frm_opts, text="Retomar (pular itens já logados)", variable=self.resume_var).pack(anchor="w", padx=6)
frm_height = tk.Frame(frm_opts)
frm_height.pack(fill=tk.X, padx=6, pady=6)
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)
# Botão iniciar
tk.Button(self.root, text="Iniciar classificação", command=self.start).pack(pady=10)
# Hint
tk.Label(
self.root,
text="Teclas: 1..9 (0=10ª), espaço/n/→=pular, p/←=anterior, b=undo, q/Esc=sair",
fg="#555",
).pack(pady=4)
self.result = None
def add_input(self):
p = filedialog.askdirectory(title="Selecione a pasta de sessão (contendo .raw, .png, .json)")
if p:
self.inputs_listbox.insert(tk.END, p)
def remove_selected(self):
sel = list(self.inputs_listbox.curselection())
for idx in reversed(sel):
self.inputs_listbox.delete(idx)
def clear_inputs(self):
self.inputs_listbox.delete(0, tk.END)
def choose_out_root(self):
p = filedialog.askdirectory(title="Selecione a pasta raiz de saída")
if p:
self.out_root_var.set(p)
def start(self):
inputs = [self.inputs_listbox.get(i) for i in range(self.inputs_listbox.size())]
out_root = self.out_root_var.get().strip()
labels_raw = self.labels_var.get().strip()
labels = [s.strip() for s in labels_raw.split(",") if s.strip()]
if not inputs:
messagebox.showwarning("Faltando entradas", "Adicione pelo menos uma pasta de entrada.")
return
if not out_root:
messagebox.showwarning("Faltando saída", "Escolha a pasta raiz de saída.")
return
if not labels:
messagebox.showwarning("Faltando labels", "Informe pelo menos uma label (classe).")
return
self.result = (
[Path(p) for p in inputs],
labels,
Path(out_root),
self.move_var.get(),
self.resume_var.get(),
self.height_var.get(),
)
self.root.destroy()
def run(self):
self.root.mainloop()
return self.result
def run_with_gui_setup():
setup = SetupWindow()
res = setup.run()
if not res:
return
input_folders, labels, out_root, move, resume, height = res
triplets = collect_all_triplets(input_folders)
if not triplets:
messagebox.showinfo(
"Sem amostras",
"Nenhuma tripleta encontrada.\nCertifique-se de que existam arquivos .raw, .png, .json.",
)
return
app = RawTripletSorterApp(
triplets=triplets,
labels=labels,
out_root=out_root,
move=move,
resume=resume,
display_height=height,
)
app.run()
def main():
parser = argparse.ArgumentParser(
description="Sortear manualmente amostras (raw + preview + json) em pastas de labels, usando hotkeys."
)
parser.add_argument("--inputs", nargs="+", help="Pastas de sessão contendo .raw, .png e .json")
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("--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("--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)")
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:
run_with_gui_setup()
return
# Modo apenas CLI
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).")
sys.exit(1)
input_folders = [Path(p) for p in args.inputs]
labels = args.labels
out_root = Path(args.out_root)
triplets = collect_all_triplets(input_folders)
if not triplets:
print("Nenhuma tripleta encontrada nas pastas informadas.")
sys.exit(1)
app = RawTripletSorterApp(
triplets=triplets,
labels=labels,
out_root=out_root,
move=args.move,
resume=args.resume,
display_height=args.display_height,
)
app.run()
if __name__ == "__main__":
main()

View File

@ -0,0 +1,318 @@
import json
import os
import cv2
import csv
import shutil
import argparse
import numpy as np
from utils import carregar_labelmap_completo # importa da sua utils.py
# ===================================================
# ⚙️ Configurações base vindas do config.json
# ===================================================
with open("config.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]
RESOLUCAO = config["raw_size"]
RAW_WIDTH = RESOLUCAO[0] # largura real do raw (ex: 1296)
RAW_HEIGHT = RESOLUCAO[1] # altura real do raw (ex: 1028)
RAW_EXTS = [".raw"] # se tiver outro, adiciona aqui
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
# Cor da classe será obtida via labelmap
# COR_CLASSE_RGB = (128, 0, 0)
COPIAR_IMAGENS = True
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
MASK_EXT_OUT = ".png" # saída das máscaras sempre PNG
FORCAR_SOBRESCRITA_NEW_MASK = False # sobrescrever máscara em new_masks se já existir
VALIDAR_DIM_MASK_EXISTENTE = True # se existir, validar dimensões
MANIFESTO = "manifest.csv"
# ===================================================
def garantir_pasta(p):
os.makedirs(p, exist_ok=True)
def criar_mask_solida(dim_h, dim_w, cor_rgb):
r, g, b = cor_rgb
mask_bgr = np.zeros((dim_h, dim_w, 3), dtype=np.uint8)
mask_bgr[:] = (b, g, r) # OpenCV usa 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]
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
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)
if img is None:
raise RuntimeError(f"Erro ao abrir: {caminho_img}")
h, w = img.shape[:2]
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)
caminho_mask = caminho_mask_new_para_img(pasta_new_masks, caminho_img)
if os.path.exists(caminho_mask) and not forcar:
if validar_dim:
m = cv2.imread(caminho_mask, cv2.IMREAD_COLOR)
if m is None or m.shape[0] != h or m.shape[1] != w:
print(f"[ALERTA] Máscara existente com dimensão diferente, recriando: {caminho_mask}")
else:
print(f"[OK] Já existe (mantida): {os.path.basename(caminho_mask)}")
return caminho_mask
else:
print(f"[OK] Já existe (mantida): {os.path.basename(caminho_mask)}")
return caminho_mask
mask_bgr = criar_mask_solida(h, w, cor_rgb)
garantir_pasta(os.path.dirname(caminho_mask))
cv2.imwrite(caminho_mask, mask_bgr)
print(f"[CRIADA] {os.path.basename(caminho_mask)} ({w}x{h})")
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):
"""
Verifica se algum dos arquivos (img/mask/raw) com esse base existe.
"""
if os.path.exists(os.path.join(dir_img, base + img_ext)):
return True
if os.path.exists(os.path.join(dir_mask, base + mask_ext)):
return True
if raw_ext and dir_raw:
if os.path.exists(os.path.join(dir_raw, base + raw_ext)):
return True
return False
def gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
"""
Gera um 'base' comum livre para img/mask/raw.
"""
if not _base_esta_ocupado(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
return base
i = 1
while True:
candidate = f"{base}_{i:03d}"
if not _base_esta_ocupado(candidate, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
return candidate
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).
"""
garantir_pasta(dest_img_dir)
garantir_pasta(dest_mask_dir)
if dest_raw_dir:
garantir_pasta(dest_raw_dir)
base = os.path.splitext(os.path.basename(caminho_img_src))[0]
img_ext = os.path.splitext(caminho_img_src)[1].lower()
mask_ext = ".png"
raw_ext = os.path.splitext(caminho_raw_src)[1].lower() if caminho_raw_src else None
new_base = gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dest_img_dir, dest_mask_dir, dest_raw_dir)
dst_img_path = os.path.join(dest_img_dir, new_base + img_ext)
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)
if caminho_raw_src and dst_raw_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
# ===================================================
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)
alvo = nome_classe.strip().lower()
target_id = None
for cid, cname in id_para_nome.items():
if cname.strip().lower() == alvo:
target_id = cid
break
if target_id is None:
disponiveis = ", ".join(sorted(id_para_nome.values()))
raise RuntimeError(
f"Classe '{nome_classe}' não encontrada no labelmap.\n"
f"Classes disponíveis: {disponiveis}"
)
try:
cor_rgb = colormap_rgb[target_id]
except IndexError:
raise RuntimeError(
f"ID {target_id} da classe '{nome_classe}' não possui cor correspondente em colormap_rgb."
)
print(f"[LABELMAP] Classe '{nome_classe}' -> ID={target_id} -> cor RGB={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):
source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original")
pasta_new_previews = os.path.join(source, "group", grupo, "previews")
pasta_new_masks = os.path.join(source, "group", grupo, "masks")
pasta_new_raws = os.path.join(source, "group", grupo, "raws")
garantir_pasta(pasta_new_previews)
garantir_pasta(pasta_new_masks)
garantir_pasta(pasta_new_raws)
garantir_pasta(PASTA_FINAL_PREVIEWS)
garantir_pasta(PASTA_FINAL_MASKS)
garantir_pasta(PASTA_FINAL_RAWS)
registros = []
total, criadas_mask, copiados, puladas, erros = 0, 0, 0, 0, 0
for nome in os.listdir(pasta_new_previews):
if not nome.lower().endswith(EXT_IMAGENS):
continue
total += 1
caminho_img = os.path.join(pasta_new_previews, nome)
try:
# 1) criar máscara sólida em new_masks
antes = os.path.exists(caminho_mask_new_para_img(pasta_new_masks, caminho_img))
caminho_mask_new = salvar_mask_solidaria(
pasta_new_masks,
caminho_img,
cor_classe_rgb,
forcar=FORCAR_SOBRESCRITA_NEW_MASK,
validar_dim=VALIDAR_DIM_MASK_EXISTENTE
)
if caminho_mask_new and not antes:
criadas_mask += 1
# 2) copiar imagem + máscara + RAW para as pastas finais (com renome se necessário)
if fazer_copia_final:
raw_src = localizar_raw_correspondente(pasta_new_raws, nome)
if raw_src is None:
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,
raw_src,
PASTA_FINAL_PREVIEWS,
PASTA_FINAL_MASKS,
PASTA_FINAL_RAWS,
)
copiados += 1
registros.append([
caminho_img,
caminho_mask_new,
raw_src or "",
dst_img,
dst_mask,
dst_raw or "",
])
else:
puladas += 1
except Exception as e:
erros += 1
print(f"[ERRO] {nome}: {e}")
if manifesto_csv and registros:
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow([
"src_image",
"src_mask",
"src_raw",
"dst_image",
"dst_mask",
"dst_raw",
])
w.writerows(registros)
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}")
def build_cli():
ap = argparse.ArgumentParser(
description="Gera máscaras sólidas para novas imagens de UMA classe (via labelmap) e copia para dataset final com dedup."
)
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("--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("--no-copy", action="store_true", help="Não copia para as pastas finais (só cria masks em new_masks).")
ap.add_argument("--from-originals", action="store_true", help="Faz o procedimento na pasta em originais")
ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a gerar (ou vazio para não gerar).")
return ap
if __name__ == "__main__":
ap = build_cli()
args = ap.parse_args()
# Caminho padrão pro labelmap, caso não seja informado
labelmap_path = os.path.join("dataset", "labelmap.txt")
if not os.path.exists(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)
fazer_copia = not args.no_copy
manifesto_csv = args.manifest if args.manifest else None
processar_novas_imagens(
cana=args.cana,
horario=args.horario,
grupo=args.grupo,
cor_classe_rgb=cor_rgb,
fazer_copia_final=fazer_copia,
manifesto_csv=manifesto_csv,
orignais=args.from_originals
)

View File

@ -0,0 +1,205 @@
import json
import os
import shutil
import argparse
import csv
# ================= CONFIG =================
with open("config.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]
# Raiz das brutas agrupadas
PASTA_BRUTAS_GROUP_ROOT = os.path.join("dataset", "brutas")
# Onde as máscaras novas (rotuladas externamente) são colocadas
PASTA_NEW_MASKS = os.path.join("dataset", "new_masks")
# Destino final do dataset consolidado
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
EXT_MASKS = (".png",)
EXT_RAWS = (".raw",)
# =========================================
def garantir_pasta(p: str):
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:
"""
Devolve um caminho disponível em dest_dir com base + ext.
Se existir, adiciona _001, _002, ...
"""
p = os.path.join(dest_dir, base + ext)
if not os.path.exists(p):
return p
i = 1
while True:
p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}")
if not os.path.exists(p):
return p
i += 1
def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
garantir_pasta(PASTA_FINAL_PREVIEWS)
garantir_pasta(PASTA_FINAL_RAWS)
garantir_pasta(PASTA_FINAL_MASKS)
if not os.path.isdir(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, raws_map = indexar_brutas_group(PASTA_BRUTAS_GROUP_ROOT)
registros = []
total, copiados, ignorados = 0, 0, 0
for nome in os.listdir(PASTA_NEW_MASKS):
if not nome.lower().endswith(EXT_MASKS):
continue
total += 1
base = os.path.splitext(nome)[0]
caminho_mask_src = os.path.join(PASTA_NEW_MASKS, nome)
preview_src = previews_map.get(base)
raw_src = raws_map.get(base)
if preview_src is None or raw_src is None:
ignorados += 1
motivo = []
if preview_src is None:
motivo.append("preview")
if raw_src is None:
motivo.append("raw")
print(f"[SKIP] {base} -> faltando: {', '.join(motivo)}")
# NÃO move a máscara, ela fica em new_masks pra você analisar depois
continue
prev_ext = os.path.splitext(preview_src)[1].lower()
mask_ext = os.path.splitext(nome)[1].lower()
raw_ext = os.path.splitext(raw_src)[1].lower()
# Gera um base final único usando o preview como referência
dst_preview = nome_disponivel(PASTA_FINAL_PREVIEWS, base, prev_ext)
new_base = os.path.splitext(os.path.basename(dst_preview))[0]
dst_raw = os.path.join(PASTA_FINAL_RAWS, new_base + raw_ext)
dst_mask = os.path.join(PASTA_FINAL_MASKS, new_base + mask_ext)
# Copia preview + raw, move ou copia a máscara
shutil.copy2(preview_src, dst_preview)
shutil.copy2(raw_src, dst_raw)
if mover:
shutil.move(caminho_mask_src, dst_mask)
else:
shutil.copy2(caminho_mask_src, dst_mask)
copiados += 1
registros.append([
preview_src,
raw_src,
caminho_mask_src,
dst_preview,
dst_raw,
dst_mask,
])
print(f"[OK] {new_base}")
if manifesto_csv and registros:
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow([
"src_preview",
"src_raw",
"src_mask",
"dst_preview",
"dst_raw",
"dst_mask",
])
w.writerows(registros)
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
print(f"\nResumo: total_masks={total} | ingestas={copiados} | ignoradas={ignorados}")
# ================= CLI =================
def build_cli():
ap = argparse.ArgumentParser(
description=(
"Ingere máscaras novas (dataset/new_masks) procurando previews e raws "
"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("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
ap.add_argument("--copy", action="store_true", help="Copia as máscaras em vez de mover (por padrão, move e esvazia new_masks conforme ingere).")
ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).")
return ap
if __name__ == "__main__":
ap = build_cli()
args = ap.parse_args()
manifesto = args.manifest if args.manifest.strip() else None
mover = not args.copy
PASTA_BRUTAS_GROUP_ROOT = os.path.join(PASTA_BRUTAS_GROUP_ROOT, f"cana_{args.cana}", args.horario, "group")
print(f"[INFO] MODELO : {MODELO}")
print(f"[INFO] Brutas (group root): {PASTA_BRUTAS_GROUP_ROOT}")
print(f"[INFO] New masks : {PASTA_NEW_MASKS}")
print(f"[INFO] Mover máscaras? : {mover}")
print(f"[INFO] Manifesto : {manifesto or '(nenhum)'}")
print()
processar_new_masks(mover=mover, manifesto_csv=manifesto)

View File

@ -0,0 +1,328 @@
import json
import os
import shutil
import argparse
import csv
# ================= CONFIG =================
with open("config.json", "r") as f:
config = json.load(f)
MODELO = config["camera"]
# Raiz das brutas (todas as canas/horários/grupos)
PASTA_BRUTAS_ROOT = os.path.join("dataset", "brutas")
# Onde você coloca os previews selecionados (tudo misturado)
PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews")
# Destino final, organizado por grupo:
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group")
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
EXT_RAWS = (".raw",)
EXT_MASKS = (".png",)
# Ajusta se suas metas tiverem outra extensão
EXT_METAS = (".json", ".yml", ".yaml", ".txt", ".csv")
# ================= HELPERS =================
def garantir_pasta(p: str):
os.makedirs(p, exist_ok=True)
def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
"""
Devolve um caminho disponível em dest_dir com base + ext.
Se existir, adiciona _001, _002, ...
"""
p = os.path.join(dest_dir, base + ext)
if not os.path.exists(p):
return p
i = 1
while True:
p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}")
if not os.path.exists(p):
return p
i += 1
def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None:
"""
Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'.
Retorna o caminho completo ou None.
"""
if not os.path.isdir(pasta):
return None
for nome in os.listdir(pasta):
nome_lower = nome.lower()
root, ext = os.path.splitext(nome_lower)
if root == base.lower() and ext in exts:
return os.path.join(pasta, nome)
return None
def indexar_brutas_por_raw(root_brutas: str):
"""
Varre dataset/brutas recursivamente, olhando pastas 'raws' e montando índice:
base -> {
"raw": caminho_raw,
"cana": "cana_alta" | "cana_baixa" | ... (se conseguir inferir),
"horario": "meio_dia" | "cedo" | ... (se conseguir inferir),
"grupo": "<GRUPO>",
"group_dir": caminho_da_pasta_do_grupo
}
Estrutura esperada (relativa a root_brutas):
cana_alta/meio_dia/group/<GRUPO>/raws/*.raw
"""
index = {}
if not os.path.isdir(root_brutas):
print(f"[AVISO] Pasta de brutas não existe: {root_brutas}")
return index
for dirpath, dirnames, filenames in os.walk(root_brutas):
base_dir = os.path.basename(dirpath).lower()
if base_dir != "raws":
continue
# dirpath = .../cana_x/horario/group/<GRUPO>/raws
group_dir = os.path.dirname(dirpath) # .../cana_x/horario/group/<GRUPO>
rel = os.path.relpath(dirpath, root_brutas)
parts = rel.split(os.sep)
# Defaults
cana = None
horario = None
grupo = None
if len(parts) >= 5:
# [0] = cana_<...>
# [1] = horario
# [2] = "group"
# [3] = <GRUPO>
cana = parts[0]
horario = parts[1]
# parts[2] deve ser "group"
grupo = parts[3]
else:
# fallback bem genérico
if "group" in parts:
i = parts.index("group")
if i + 1 < len(parts):
grupo = parts[i + 1]
for nome in filenames:
if not nome.lower().endswith(EXT_RAWS):
continue
base = os.path.splitext(nome)[0]
if base in index:
# Conflito (mesmo base em dois lugares) -> loga e mantém o primeiro
print(f"[CONFLITO] base repetida em raws: {base}")
continue
raw_path = os.path.join(dirpath, nome)
index[base] = {
"raw": raw_path,
"cana": cana,
"horario": horario,
"grupo": grupo,
"group_dir": group_dir,
}
print(f"[INDEX] Entradas indexadas por RAW: {len(index)}")
return index
def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | None = None):
if not os.path.isdir(PASTA_SELECTED_PREVIEWS):
raise SystemExit(f"[ERRO] Pasta selected_previews não existe: {PASTA_SELECTED_PREVIEWS}")
# Indexa todas as brutas a partir dos RAWs
index_raws = indexar_brutas_por_raw(PASTA_BRUTAS_ROOT)
registros = []
total, movidos, ignorados = 0, 0, 0
for nome in os.listdir(PASTA_SELECTED_PREVIEWS):
caminho_preview_sel = os.path.join(PASTA_SELECTED_PREVIEWS, nome)
if not os.path.isfile(caminho_preview_sel):
continue
if not nome.lower().endswith(EXT_PREVIEWS):
continue
total += 1
base = os.path.splitext(nome)[0]
info = index_raws.get(base)
if info is None:
ignorados += 1
print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via RAW)")
continue
grupo = info.get("grupo") or "unknown"
cana = info.get("cana") or "unknown"
horario = info.get("horario") or "unknown"
group_dir = info["group_dir"]
raw_src = info["raw"]
# Pastas irmãs em brutas
metas_src_dir = os.path.join(group_dir, "metas")
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)
# Destino: dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
dest_group_root = os.path.join(PASTA_ORIGINAL_GROUP_ROOT, grupo)
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_mask_dir = os.path.join(dest_group_root, "masks")
garantir_pasta(dest_prev_dir)
garantir_pasta(dest_raw_dir)
garantir_pasta(dest_meta_dir)
garantir_pasta(dest_mask_dir)
# Define extensões
prev_ext_sel = os.path.splitext(nome)[1].lower()
raw_ext = os.path.splitext(raw_src)[1].lower()
# Gera nome final único com base no preview
dst_preview = nome_disponivel(dest_prev_dir, base, prev_ext_sel)
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)
if mask_src is not None:
mask_ext = os.path.splitext(mask_src)[1].lower()
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
# Copia/move preview selecionado
if copy_only:
shutil.copy2(caminho_preview_sel, dst_preview)
else:
shutil.move(caminho_preview_sel, dst_preview)
# Copia RAW e metas/masks (se existirem)
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:
shutil.copy2(meta_src, dst_meta)
else:
shutil.move(meta_src, dst_meta)
if mask_src is not None:
if copy_only:
shutil.copy2(mask_src, dst_mask)
else:
shutil.move(mask_src, dst_mask)
movidos += 1
registros.append([
base,
grupo,
cana,
horario,
caminho_preview_sel,
raw_src,
meta_src or "",
mask_src or "",
dst_preview,
dst_raw,
dst_meta or "",
dst_mask or "",
])
print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario}")
if manifesto_csv and registros:
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow([
"base",
"grupo",
"cana",
"horario",
"src_preview_selected",
"src_raw",
"src_meta",
"src_mask",
"dst_preview",
"dst_raw",
"dst_meta",
"dst_mask",
])
w.writerows(registros)
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
print(f"\nResumo: total_selected={total} | organizadas={movidos} | ignoradas={ignorados}")
# ================= CLI =================
def build_cli():
ap = argparse.ArgumentParser(
description=(
"Organiza previews selecionadas (dataset/selected_previews) "
"descobrindo cana/horário/grupo em dataset/brutas e copiando/movendo "
"preview + raw + meta (+ mask se existir) para dataset/original/group/{GRUPO}."
)
)
ap.add_argument(
"--copy",
action="store_true",
help="Copia os previews em vez de mover (por padrão, move e esvazia selected_previews conforme organiza)."
)
ap.add_argument(
"--manifest",
default="",
help="Caminho para CSV de manifesto (opcional)."
)
return ap
if __name__ == "__main__":
ap = build_cli()
args = ap.parse_args()
manifesto = args.manifest if args.manifest.strip() else None
copy_only = args.copy
print(f"[INFO] MODELO : {MODELO}")
print(f"[INFO] Brutas root : {PASTA_BRUTAS_ROOT}")
print(f"[INFO] Selected previews : {PASTA_SELECTED_PREVIEWS}")
print(f"[INFO] Original group root : {PASTA_ORIGINAL_GROUP_ROOT}")
print(f"[INFO] Copy only? : {copy_only}")
print(f"[INFO] Manifesto : {manifesto or '(nenhum)'}")
print()
organizar_selected_previews(copy_only=copy_only, manifesto_csv=manifesto)

View File

@ -0,0 +1,322 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Versão OTIMIZADA Agrupamento ultra-rápido por classes presentes na máscara.
Substitui extrair_ids_presentes por uma versão 50× mais rápida.
Entrada:
originals/
previews/
raws/
masks/
(opcional) masks2/
Saída:
originals/group/<grupo>/previews
/raws
/masks
/masks2
"""
import os
import cv2
import csv
import json
import shutil
import argparse
import numpy as np
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
USE_MASKS2 = config["dual_head"]
EXT_PREVIEWS = (".jpg", ".jpeg", ".png")
EXT_MASKS = (".png", ".jpg", ".jpeg")
EXT_MASKS2 = (".png", ".jpg", ".jpeg")
EXT_RAW = (".raw",)
MANIFESTO_DEFAULT = "manifest.csv"
# ============================================================
# 🔥 EXTRAÇÃO DE CLASSES — VERSÃO ULTRA OTIMIZADA
# ============================================================
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
"""
Extração ultra-rápida de IDs presentes na máscara.
Estratégia:
1) Máscara 1 canal np.unique (instantâneo).
2) Máscara 3 canais:
- se assume_rgb=True: labelmap está em RGB,
convertemos a imagem BGR -> RGB para casar com cor_para_id.
- se assume_rgb=False: labelmap está em BGR,
mantemos a imagem em BGR.
3) Amostragem em grid + converter_mask_rgb_para_ids.
4) Fallback full-scan se necessário.
"""
m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if m is None:
raise RuntimeError(f"Falha ao abrir máscara: {mask_path}")
# -------------------------------------------------------
# CASO 1: máscara indexada (1 canal) — instantâneo
# -------------------------------------------------------
if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1):
ids = np.unique(m)
return set(int(v) for v in ids)
# -------------------------------------------------------
# CASO 2: máscara RGB
# -------------------------------------------------------
if assume_rgb:
# Labelmap em RGB, OpenCV em BGR -> converte
img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB)
else:
# Labelmap em BGR, OpenCV já em BGR -> usa direto
img = m
# Agora cor_para_id e img estão no MESMO espaço de cor
mapa_rgb = cor_para_id
max_classes = len(cor_para_id)
# ---------- AMOSTRAGEM ----------
step = 8 # pode virar 4 se quiser mais precisão, 16 se quiser mais velocidade
amostra = img[::step, ::step]
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)
if len(ids) >= max_classes:
return ids
# ---------- FULL-SCAN OTIMIZADO (último caso) ----------
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)
return ids
# ============================================================
# 🔧 Helpers
# ============================================================
def garantir_pasta(p):
os.makedirs(p, exist_ok=True)
def nome_disponivel(dest_dir, base_name, ext):
cand = os.path.join(dest_dir, base_name + ext)
if not os.path.exists(cand):
return cand
i = 1
while True:
cand = os.path.join(dest_dir, f"{base_name}_{i:03d}{ext}")
if not os.path.exists(cand):
return cand
i += 1
def mapear_por_base_priorizando_png(pasta, exts):
if not pasta or not os.path.isdir(pasta):
return {}
mapa = {}
for nome in os.listdir(pasta):
lower = nome.lower()
if not lower.endswith(exts):
continue
base, ext = os.path.splitext(nome)
cam = os.path.join(pasta, nome)
if base not in mapa:
mapa[base] = cam
else:
atual_ext = os.path.splitext(mapa[base])[1].lower()
if atual_ext != ".png" and ext.lower() == ".png":
mapa[base] = cam
return mapa
def inferir_ignore_id(ignore_rgb, cor_para_id):
if ignore_rgb is None:
return None
if isinstance(ignore_rgb, (list, tuple)) and len(ignore_rgb) == 3:
return cor_para_id.get(tuple(ignore_rgb))
if isinstance(ignore_rgb, int):
return ignore_rgb
return None
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:
return "sem_classe"
nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)]
return "_".join(nomes)
# ============================================================
# 📦 PIPELINE PRINCIPAL
# ============================================================
def copiar_ou_mover_tripla(preview_src, raw_src, mask_src,
dest_prev_dir, dest_raw_dir, dest_mask_dir,
mover=False,
mask2_src=None, dest_mask2_dir=None):
garantir_pasta(dest_prev_dir)
garantir_pasta(dest_raw_dir)
garantir_pasta(dest_mask_dir)
base_prev = os.path.splitext(os.path.basename(preview_src))[0]
prev_ext = os.path.splitext(preview_src)[1].lower()
raw_ext = os.path.splitext(raw_src)[1].lower()
mask_ext = os.path.splitext(mask_src)[1].lower()
dst_prev = nome_disponivel(dest_prev_dir, base_prev, prev_ext)
new_base = os.path.splitext(os.path.basename(dst_prev))[0]
dst_raw = os.path.join(dest_raw_dir, new_base + raw_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_mask = nome_disponivel(dest_mask_dir, new_base, mask_ext)
new_base = os.path.splitext(os.path.basename(dst_mask))[0]
dst_prev = nome_disponivel(dest_prev_dir, new_base, prev_ext)
dst_raw = nome_disponivel(dest_raw_dir, new_base, raw_ext)
if mover:
shutil.move(preview_src, dst_prev)
shutil.move(raw_src, dst_raw)
shutil.move(mask_src, dst_mask)
if mask2_src:
shutil.move(mask2_src, os.path.join(dest_mask2_dir, new_base + os.path.splitext(mask2_src)[1]))
else:
shutil.copy2(preview_src, dst_prev)
shutil.copy2(raw_src, dst_raw)
shutil.copy2(mask_src, dst_mask)
if mask2_src:
garantir_pasta(dest_mask2_dir)
shutil.copy2(mask2_src, os.path.join(dest_mask2_dir, new_base + os.path.splitext(mask2_src)[1]))
return dst_prev, dst_raw, dst_mask, None
# ============================================================
# 🚀 PROCESSAMENTO
# ============================================================
def processar(originals_dir, labelmap_path, mover=False,
manifesto=MANIFESTO_DEFAULT, validar_dim=True,
estrito=False, labelmap_bgr=False):
originals_dir = os.path.abspath(originals_dir)
previews_dir = os.path.join(originals_dir, "previews")
raws_dir = os.path.join(originals_dir, "raws")
masks_dir = os.path.join(originals_dir, "masks")
masks2_dir = os.path.join(originals_dir, "masks2")
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):
raise RuntimeError("Estrutura inválida em originals/")
usar_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir)
cor_para_id, colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = inferir_ignore_id(ignore_rgb, cor_para_id)
garantir_pasta(group_dir)
mapa_previews = mapear_por_base_priorizando_png(previews_dir, EXT_PREVIEWS)
mapa_masks = mapear_por_base_priorizando_png(masks_dir, EXT_MASKS)
mapa_raws = {os.path.splitext(n)[0]: os.path.join(raws_dir, n)
for n in os.listdir(raws_dir) if n.lower().endswith(EXT_RAW)}
mapa_masks2 = mapear_por_base_priorizando_png(masks2_dir, EXT_MASKS2) if usar_masks2 else {}
registros = []
totais = {"total_masks":0, "processados":0, "pulados":0}
for base, mask_path in sorted(mapa_masks.items()):
totais["total_masks"] += 1
prev_path = mapa_previews.get(base)
raw_path = mapa_raws.get(base)
if not prev_path or not raw_path:
totais["pulados"] += 1
continue
# ------------------------------
# EXTRAÇÃO ULTRA RÁPIDA
# ------------------------------
ids_presentes = extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=not labelmap_bgr)
if ignore_id in ids_presentes:
ids_presentes.discard(ignore_id)
grupo = montar_nome_grupo(ids_presentes, id_para_nome)
# ------------------------------
# CRIA PASTAS
# ------------------------------
dest_prev_dir = os.path.join(group_dir, grupo, "previews")
dest_raw_dir = os.path.join(group_dir, grupo, "raws")
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
dst_prev, dst_raw, dst_mask, _ = copiar_ou_mover_tripla(
prev_path, raw_path, mask_path,
dest_prev_dir, dest_raw_dir, dest_mask_dir,
mover=mover,
mask2_src=mapa_masks2.get(base),
dest_mask2_dir=dest_mask2_dir
)
registros.append([grupo, prev_path, raw_path, mask_path, dst_prev, dst_raw, dst_mask])
totais["processados"] += 1
# ------------------------------
# MANIFESTO
# ------------------------------
if manifesto:
with open(manifesto, "w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow(["grupo","src_preview","src_raw","src_mask","dst_preview","dst_raw","dst_mask"])
for r in registros:
w.writerow(r)
print("\nFinalizado!")
print("Total masks:", totais["total_masks"])
print("Processados:", totais["processados"])
print("Pulados :", totais["pulados"])
# ============================================================
# CLI
# ============================================================
def build_cli():
ap = argparse.ArgumentParser()
ap.add_argument("--move", action="store_true", help="Mover em vez de copiar")
ap.add_argument("--manifest", default=MANIFESTO_DEFAULT)
ap.add_argument("--no-validate", action="store_true")
ap.add_argument("--strict", action="store_true")
ap.add_argument("--labels-bgr", action="store_true")
return ap
if __name__ == "__main__":
ap = build_cli()
args = ap.parse_args()
originals_dir = os.path.join("dataset", "original")
labelmap_dir = os.path.join("dataset", "labelmap.txt")
processar(
originals_dir=originals_dir,
labelmap_path=labelmap_dir,
mover=args.move,
manifesto=args.manifest,
validar_dim=not args.no_validate,
estrito=args.strict,
labelmap_bgr=args.labels_bgr
)

View File

@ -0,0 +1,708 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Augmenta imagens/máscaras *por grupo* usando RAW4 (float32, 0..1, R G IR B).
Novo fluxo (GAL5000):
Entrada:
dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
dataset/original/group/<grupo>/masks (.png/.jpg)
(opcional) dataset/original/group/<grupo>/masks2
Saída:
dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
dataset/augmented/group/<grupo>/masks
dataset/augmented/group/<grupo>/masks2
Regras:
- GEOMETRIA (flip, shift, rotate) aplicada em RAW + masks (+ masks2).
- BLUR / MOTION BLUR aplicado no RAW, canal por canal, com o mesmo kernel.
- Preview é SEMPRE derivado do RAW transformado (geom + blur).
- Máscaras não sofrem blur, apenas geom.
Uso:
python _5_augmentation_raw.py --copies 5
python _5_augmentation_raw.py --copies 5 --groups chao,chao_erva,cana
"""
import os
import json
import cv2
from PIL import Image
import albumentations as A
from albumentations import ReplayCompose
import argparse
import numpy as np
# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera", ".")
RESOLUCAO = config.get("raw_size")
USE_MASKS2 = config.get("dual_head", False)
# Pastas base
DATASET_BASE = os.path.join("dataset")
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
# Fallback (modo antigo, sem grupos) mantido por compat
ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images")
ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks")
ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2")
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
# Extensões aceitas
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg")
MSK2_EXTS = (".png", ".jpg", ".jpeg")
RAW_EXTS = (".raw",)
def garantir_dir(p):
os.makedirs(p, exist_ok=True)
# ===============================
# Pipelines de augmentations
# ===============================
# 1) Geométricas: aplicam em preview “temporário”, mask, mask2 e RAW.
# RAW é passado como 'raw' em additional_targets para receber só geom.
train_geo = ReplayCompose([
A.HorizontalFlip(p=0.5),
A.ShiftScaleRotate(
shift_limit=0.01,
scale_limit=0.10,
rotate_limit=5,
border_mode=cv2.BORDER_REFLECT_101,
interpolation=cv2.INTER_LINEAR,
p=0.30
),
], additional_targets={
'mask2': 'mask',
'raw': 'mask', # trata raw como “mask” para receber apenas geom
})
# ===============================
# Funções de blur físico no RAW
# ===============================
def motion_blur_kernel(ksize=5, angle=0.0):
"""
Gera um kernel de motion blur 2D com tamanho ksize e ângulo em graus.
Usado tanto para RAW quanto para preview.
"""
ksize = int(ksize)
if ksize < 3:
ksize = 3
if ksize % 2 == 0:
ksize += 1
kernel = np.zeros((ksize, ksize), dtype=np.float32)
# linha horizontal no centro
kernel[ksize // 2, :] = 1.0
# rotaciona a linha para o ângulo desejado
center = (ksize / 2.0 - 0.5, ksize / 2.0 - 0.5)
M = cv2.getRotationMatrix2D(center, angle, 1.0)
kernel = cv2.warpAffine(kernel, M, (ksize, ksize))
s = kernel.sum()
if s > 0:
kernel /= s
return kernel
def apply_coherent_blur_raw(raw, p_motion=0.20, p_gauss=0.20):
"""
Aplica blur/motion blur de forma COERENTE no RAW, canal por canal.
- Mesmo kernel em TODOS os canais.
- Máscaras não passam aqui.
- Se nada for sorteado, retorna o RAW original.
"""
if raw is None:
return None
r = np.random.rand()
if r >= (p_motion + p_gauss):
# sem blur nesta amostra
return raw
# escolhe tamanho de kernel
k_choices = [3, 5, 7]
ksize = int(np.random.choice(k_choices))
if r < p_motion:
# Motion blur com ângulo pequeno (vibração / movimento leve)
angle = float(np.random.uniform(-20.0, 20.0))
kernel = motion_blur_kernel(ksize=ksize, angle=angle)
else:
# Gaussian blur (leve desfoque de foco)
if ksize % 2 == 0:
ksize += 1
g1d = cv2.getGaussianKernel(ksize, 0)
kernel = g1d @ g1d.T # vira kernel 2D
# Aplica kernel canal a canal
if raw.ndim == 2:
raw_blur = cv2.filter2D(raw, ddepth=-1, kernel=kernel)
else:
raw_blur = np.empty_like(raw)
h, w, c = raw.shape
for ch in range(c):
raw_blur[..., ch] = cv2.filter2D(
raw[..., ch],
ddepth=-1,
kernel=kernel,
borderType=cv2.BORDER_REFLECT_101,
)
# Garante que continua em [0,1] para float
if np.issubdtype(raw_blur.dtype, np.floating):
raw_blur = np.clip(raw_blur, 0.0, 1.0)
return raw_blur
# ===============================
# IO de imagens / RAW
# ===============================
def raw_to_preview_rgb(raw_np: np.ndarray) -> np.ndarray:
"""
raw_np: (H,W,C) float32 em 0..1, C=4 ou 5.
Layout esperado (após _load_raw):
- C=4 -> [R,G,B,IR]
- C=5 -> [R,G,B,IR,NDVI] (se um dia tiver)
Usamos sempre R,G,B para o preview e devolvemos (H,W,3) uint8.
"""
assert raw_np.ndim == 3, "raw_np deve ser (H,W,C)"
H, W, C = raw_np.shape
if C not in (4, 5):
raise RuntimeError(f"Esperado C=4 ou 5, veio {C} em raw_to_preview_rgb")
# RAW já vem como [R,G,B,IR] do _load_raw
r = raw_np[..., 0].astype(np.float32)
g = raw_np[..., 1].astype(np.float32)
b = raw_np[..., 2].astype(np.float32)
# Monta RGB em (H,W,3) para trabalhar
rgb = np.stack([r, g, b], axis=-1) # (H,W,3)
# Auto white balance simples (referenciado no G)
means = rgb.reshape(-1, 3).mean(axis=0) # (3,)
ref = means[1] # canal G
gains = ref / (means + 1e-6)
gains = np.clip(gains, 0.5, 2.0)
rgb_wb = rgb * gains # broadcasting em (3,)
# Força do WB
alpha = 1.0 # teste 0.4..0.8 se quiser misturar com original
if alpha >= 0.999:
rgb = rgb_wb
else:
rgb_orig = rgb.copy()
rgb = (1 - alpha) * rgb_orig + alpha * rgb_wb
rgb = np.clip(rgb, 0.0, 1.0)
# Gamma mais suave
rgb = np.power(rgb, 1/1.9, out=rgb)
# Converte pra uint8 OpenCV-friendly
rgb8 = (rgb * 255.0).clip(0, 255).astype(np.uint8) # (H,W,3)
return rgb8
def load_rgb(path):
im = cv2.imread(path, cv2.IMREAD_COLOR)
if im is None:
raise FileNotFoundError(path)
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
def save_rgb(path, arr_rgb):
Image.fromarray(arr_rgb).save(path)
def _load_raw(path: str, raw_hw) -> np.ndarray:
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":
if raw_hw is None:
raise RuntimeError(
f"Para ler .raw preciso de raw_hw=(H,W). "
f"Passe raw_hw ao criar o dataset."
)
H, W = raw_hw
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 uint8
arr = np.fromfile(path, dtype=np.uint8)
raw2d = arr.reshape(H, W)
if (H % 2) != 0 or (W % 2) != 0:
raise RuntimeError(f"raw_hw deve ser par em H e W p/ mosaico 2x2 (R,G,IR,B), 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,IR,B]
arr = np.stack([r_full, g_full, ir_full, b_full], axis=-1).astype(np.uint8)
# REORGANIZA PARA CONTRATO INTERNO [R,G,B,IR]
arr = arr[..., [0, 1, 3, 2]]
elif size_bytes == raw4_bytes:
# NOVO MODO: RAW4 float32 salvo pelo normalize
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 PARA 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})")
# Até aqui queremos (H,W,4) base
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)
def _scale_to_float01(raw: np.ndarray, max_value = None) -> np.ndarray:
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:
# float já vem normalizado ou não; se passar max_value, aplica
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 = raw_f / mv
# clip pra não explodir
raw_f = np.clip(raw_f, 0.0, 1.0)
return raw_f
def save_raw_any(path, arr):
arr = np.asarray(arr)
if arr.ndim == 3 and arr.shape[-1] == 4:
# arr está no contrato interno [R,G,B,IR]
arr_cam = internal_rgbir_to_camera_raw4(arr) # [R,G,IR,B]
arr_cam = np.transpose(arr_cam, (2, 0, 1)) # (4,H,W)
arr_cam.astype(np.float32).tofile(path)
else:
arr.tofile(path)
# ===============================
# Helpers de diretório / mapeamento
# ===============================
def list_groups(root):
"""
Lista grupos válidos para o NOVO fluxo:
- group/<g>/raws
- group/<g>/masks
(previews/images não são mais obrigatórios)
"""
if not os.path.isdir(root):
return []
grupos = []
for name in sorted(os.listdir(root)):
gdir = os.path.join(root, name)
if not os.path.isdir(gdir):
continue
has_masks = os.path.isdir(os.path.join(gdir, "masks"))
has_raws = os.path.isdir(os.path.join(gdir, "raws"))
if has_masks and has_raws:
grupos.append(name)
return grupos
def map_by_base_priorizando_png(msk_dir, exts):
"""Mapeia arquivos por base (prioriza .png)."""
by_base = {}
if not os.path.isdir(msk_dir):
return by_base
for fname in os.listdir(msk_dir):
f_lower = fname.lower()
if not f_lower.endswith(exts):
continue
base, ext = os.path.splitext(fname)
cand = os.path.join(msk_dir, fname)
if base not in by_base:
by_base[base] = cand
else:
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):
by_base = {}
if not os.path.isdir(raw_dir):
return by_base
for fname in os.listdir(raw_dir):
f_lower = fname.lower()
if not f_lower.endswith(RAW_EXTS):
continue
base, _ = os.path.splitext(fname)
by_base[base] = os.path.join(raw_dir, fname)
return by_base
def ensure_aug_dirs(group_name=None, use_masks2=False, use_raw=False, is_preview_mode=True):
"""
Cria diretórios de saída para o grupo ou modo antigo.
is_preview_mode:
True -> saídas em 'previews' + 'raws'
False -> saídas em 'images' + (sem raws)
"""
if group_name:
base = os.path.join(AUG_GROUP_ROOT, group_name)
if is_preview_mode:
img_out = os.path.join(base, "previews")
else:
img_out = os.path.join(base, "images")
msk_out = os.path.join(base, "masks")
msk2_out = os.path.join(base, "masks2") if use_masks2 else None
raw_out = os.path.join(base, "raws") if use_raw else None
else:
# modo legacy
img_out = AUG_OLD_IMG
msk_out = AUG_OLD_MSK
msk2_out = AUG_OLD_MSK2 if use_masks2 else None
raw_out = None
garantir_dir(img_out)
garantir_dir(msk_out)
if use_masks2 and msk2_out:
garantir_dir(msk2_out)
if use_raw and raw_out:
garantir_dir(raw_out)
return img_out, msk_out, msk2_out, raw_out
# ===============================
# Núcleo: augmentação de um sample
# ===============================
def augment_sample(
raw_path,
msk_path,
img_out_dir,
msk_out_dir,
copies,
msk2_path=None,
msk2_out_dir=None,
raw_out_dir=None,
aug_suffix="aug",
):
"""
Faz a augmentação a partir de UM RAW4 + máscara (e opcionalmente máscara2).
Passos:
1) Carrega máscara (para saber HxW).
2) Carrega RAW4 float32 usando HxW da máscara.
3) Gera preview RGB a partir do RAW4 (para servir de "image" no Albumentations).
4) Aplica GEOMETRIA (mesma em preview, mask, mask2 e raw).
5) Aplica BLUR/MOTION BLUR no RAW geométrico, canal a canal, com o mesmo kernel.
6) Gera preview final a partir do RAW blurado.
7) Salva:
- preview_n em img_out_dir
- mask_n em msk_out_dir
- raw_n em raw_out_dir
- mask2_n em msk2_out_dir (se houver)
"""
base_raw, raw_ext = os.path.splitext(os.path.basename(raw_path))
base = base_raw
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
if base_msk != base:
print(f"[WARN] RAW '{base}' e máscara '{base_msk}' com bases diferentes, mas seguindo mesmo assim.")
msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
# 1) Máscaras
msk = load_rgb(msk_path)
msk2 = load_rgb(msk2_path) if msk2_path else None
# 2) RAW4
raw = _load_raw(raw_path, raw_hw=(RESOLUCAO[1], RESOLUCAO[0])) # (H,W,4) uint8
raw = _scale_to_float01(raw) # float32 0..1
# 3) Preview inicial derivado do RAW (antes de geom/blur, só referência)
img = raw_to_preview_rgb(raw)
gen = 0
for i in range(copies):
# 4) GEOMETRIA: mesma transformação em img, mask, mask2 e raw
if msk2 is not None and msk2_out_dir:
aug_geo = train_geo(image=img, mask=msk, mask2=msk2, raw=raw)
else:
aug_geo = train_geo(image=img, mask=msk, raw=raw)
msk_g = aug_geo["mask"]
raw_g = aug_geo.get("raw", None)
msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None
# 5) BLUR FÍSICO: aplica (ou não) o MESMO blur no RAW geométrico
raw_b = apply_coherent_blur_raw(raw_g)
# 6) Preview final derivado do RAW blurado
img_b = raw_to_preview_rgb(raw_b)
# 7) Salva cópias
out_img = os.path.join(img_out_dir, f"{base}_{aug_suffix}_{i:02d}.png")
out_msk = os.path.join(msk_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk_ext}")
save_rgb(out_img, img_b)
save_rgb(out_msk, msk_g)
if msk2_g is not None and msk2_out_dir:
out_msk2 = os.path.join(msk2_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk2_ext}")
save_rgb(out_msk2, msk2_g)
if raw_b is not None and raw_out_dir:
out_raw = os.path.join(raw_out_dir, f"{base}_{aug_suffix}_{i:02d}{raw_ext}")
save_raw_any(out_raw, raw_b)
gen += 1
return gen
# ===============================
# Processamento por grupo
# ===============================
def process_group(group_name, copies, limit=None, seed=42, aug_suffix="aug"):
"""
Processa um grupo único usando APENAS:
- group/<g>/raws
- group/<g>/masks (+ masks2 opcional)
Previews de entrada são ignoradas.
As de saída são geradas do RAW4 (geom + blur).
"""
gdir = os.path.join(ORIG_GROUP_ROOT, group_name)
raw_dir = os.path.join(gdir, "raws")
msk_dir = os.path.join(gdir, "masks")
msk2_dir = os.path.join(gdir, "masks2")
if not (os.path.isdir(raw_dir) and os.path.isdir(msk_dir)):
print(f"[WARN] Grupo '{group_name}' inválido (precisa de 'raws' e 'masks'). Pulando.")
return 0
use_raw = True
use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
raw_map = map_raws_by_base(raw_dir)
items = sorted(raw_map.items())
if limit is not None and limit > 0:
rng = np.random.default_rng(seed)
if limit < len(items):
idx = rng.choice(len(items), size=limit, replace=False)
idx = sorted(idx.tolist())
items = [items[i] for i in idx]
msk_map = map_by_base_priorizando_png(msk_dir, MSK_EXTS)
msk2_map = map_by_base_priorizando_png(msk2_dir, MSK2_EXTS) if use_masks2 else {}
img_out_dir, msk_out_dir, msk2_out_dir, raw_out_dir = ensure_aug_dirs(
group_name,
use_masks2=use_masks2,
use_raw=use_raw,
is_preview_mode=True
)
count = 0
for base, raw_file in items:
msk_file = msk_map.get(base)
if not msk_file:
print(f"[WARN] [{group_name}] Máscara não encontrada para RAW {base}, pulando.")
continue
msk2_file = msk2_map.get(base) if use_masks2 else None
if use_masks2 and not msk2_file:
print(f"[WARN] [{group_name}] mask2 não encontrada para {base}, gerando só RAW+preview+mask.")
try:
count += augment_sample(
raw_path=raw_file,
msk_path=msk_file,
img_out_dir=img_out_dir,
msk_out_dir=msk_out_dir,
copies=copies,
msk2_path=msk2_file,
msk2_out_dir=msk2_out_dir,
raw_out_dir=raw_out_dir,
aug_suffix=aug_suffix,
)
except Exception as e:
print(f"[ERRO] [{group_name}] {base}: {e}")
print(f"[OK] Grupo '{group_name}'{count} amostras geradas.")
return count
# ===============================
# Fallback legacy (sem RAW)
# ===============================
def process_legacy(copies):
"""
Fallback: modo sem grupos (original/images e original/masks).
Mantido por compat, mas não usa RAW.
"""
if not (os.path.isdir(ORIG_OLD_IMG) and os.path.isdir(ORIG_OLD_MSK)):
print("[WARN] Modo legacy não encontrado. Nada a fazer.")
return 0
imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS]
msk_map = map_by_base_priorizando_png(ORIG_OLD_MSK, MSK_EXTS)
use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
msk2_map = map_by_base_priorizando_png(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}
img_out_dir, msk_out_dir, msk2_out_dir, _ = ensure_aug_dirs(
group_name=None,
use_masks2=use_masks2,
use_raw=False,
is_preview_mode=False
)
count = 0
for img_file in sorted(imgs):
base, img_ext = os.path.splitext(img_file)
msk_file = msk_map.get(base)
if not msk_file:
print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
continue
msk2_file = msk2_map.get(base) if use_masks2 else None
if use_masks2 and not msk2_file:
print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.")
try:
img_path = os.path.join(ORIG_OLD_IMG, img_file)
img = load_rgb(img_path)
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_file))
msk = load_rgb(msk_file)
msk2 = load_rgb(msk2_file) if msk2_file else None
for i in range(copies):
if msk2 is not None and msk2_out_dir:
aug_geo = train_geo(image=img, mask=msk, mask2=msk2)
else:
aug_geo = train_geo(image=img, mask=msk)
img_g = aug_geo["image"]
msk_g = aug_geo["mask"]
msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None
out_img = os.path.join(img_out_dir, f"{base}_aug_{i:02d}{img_ext}")
out_msk = os.path.join(msk_out_dir, f"{base_msk}_aug_{i:02d}{msk_ext}")
save_rgb(out_img, img_g)
save_rgb(out_msk, msk_g)
if msk2_g is not None and msk2_out_dir:
msk2_ext = os.path.splitext(msk2_file)[1]
out_msk2 = os.path.join(msk2_out_dir, f"{base_msk}_aug_{i:02d}{msk2_ext}")
save_rgb(out_msk2, msk2_g)
count += 1
except Exception as e:
print(f"[ERRO] (legacy) {img_file}: {e}")
print(f"[OK] Legacy → {count} amostras geradas.")
return count
def internal_rgbir_to_camera_raw4(arr: np.ndarray) -> np.ndarray:
"""
Converte do contrato interno (H,W,4) [R,G,B,IR]
para o contrato bruto da câmera (H,W,4) [R,G,IR,B].
"""
if arr is None:
return None
if arr.ndim != 3 or arr.shape[-1] != 4:
raise RuntimeError(f"Esperado array (H,W,4), veio {arr.shape}")
return arr[..., [0, 1, 3, 2]]
# ===============================
# main
# ===============================
def main(copies=5, groups_csv=None, limit=None, seed=42, suffix="aug"):
total = 0
if os.path.isdir(ORIG_GROUP_ROOT):
grupos = list_groups(ORIG_GROUP_ROOT)
if groups_csv:
want = {g.strip() for g in groups_csv.split(",") if g.strip()}
grupos = [g for g in grupos if g in want]
if not grupos:
print("[WARN] Nenhum grupo válido encontrado após filtro.")
if not grupos:
print("[WARN] Nenhum grupo encontrado em original/group. Tentando modo legacy...")
total += process_legacy(copies)
else:
print(f"Grupos encontrados: {', '.join(grupos)}")
for g in grupos:
total += process_group(g, copies, limit=limit, seed=seed, aug_suffix=suffix)
else:
total += process_legacy(copies)
print(f"\nAugmentation completed! Total: {total} amostras geradas.")
if __name__ == "__main__":
ap = argparse.ArgumentParser(description="Augmentação por grupos usando RAW4 (preview gerado do RAW).")
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem (default=5).")
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula (ex: chao,erva_cana).")
ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de imagens originais do grupo a augmentar.")
ap.add_argument("--seed", type=int, default=42, help="Seed para seleção reproduzível quando usar --limit.")
ap.add_argument("--suffix", type=str, default="aug", help="Sufixo usado no nome dos arquivos gerados.")
args = ap.parse_args()
main(
copies=args.copies,
groups_csv=args.groups,
limit=args.limit,
seed=args.seed,
suffix=args.suffix,
)

View File

@ -0,0 +1,325 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Normaliza/redimensiona PREVIEW + RAW + MASK (+MASK2), mantendo ESTRUTURA POR GRUPO.
Entradas:
dataset/original/group/<grupo>/{previews,raws,masks,(masks2)}
dataset/augmented/group/<grupo>/{previews,raws,masks,(masks2)}
Saídas (por resolução):
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
Conversão de máscara:
- máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids
- ignore_id conforme labelmap (default 255)
RAW:
- Detecta dtype (uint8/uint16) pelo tamanho do arquivo
- Carrega como (H,W), redimensiona, salva em .raw
"""
import argparse
import os
import json
import cv2
import numpy as np
from typing import Dict, List, Tuple
from gal5000.gal_service import mosaic_to_raw4_resized_buf
from raw_segformer_service import _infer_ignore_id
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ===== config =====
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config["camera"]
USE_MASKS2 = config["dual_head"]
RESOLUCAO = tuple(config["resolucao"]) # [W,H]
MODEL_NAME = config["model_name"]
CHANNELS = int(config.get("channels", 4))
pasta_base = os.path.join("dataset")
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])}
FONTES = ["original", "augmented"]
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):
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]:
by_base = {}
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:
h, w = src_hw
npx = h * w
fsize = os.path.getsize(path)
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)
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
grupos = list_groups_raw(fonte_root)
if not grupos:
return 0
not_want = {g.strip() for g in groups_except.split(",") if g.strip()}
total = 0
for nome_res, dim in RESOLUCOES.items():
out_root = os.path.join(pasta_base, nome_res, "group")
for grupo in grupos:
if grupo in not_want:
print(f"[WARN] Grupo desconsiderado: {grupo}")
continue
in_prev = os.path.join(fonte_root, grupo, "previews")
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"[WARN] Grupo inválido (sem previews/masks): {grupo}")
continue
usar_raw = os.path.isdir(in_raw)
usar_msk2 = USE_MASKS2 and os.path.isdir(in_msk2)
out_prev = os.path.join(out_root, grupo, "previews")
out_raw = os.path.join(out_root, grupo, "raws") if usar_raw else None
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)
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]
msk_map = map_by_base_priorizando_png(in_msk, MSK_EXTS)
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)
for i, fname in enumerate(sorted(prev_files), 1):
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:
print(f"[WARN] [{fonte_nome} | {grupo}] Sem máscara p/ {fname}, pulando.")
continue
# --- preview ---
prev_bgr = cv2.imread(prev_path, cv2.IMREAD_COLOR)
if prev_bgr is None:
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)
out_name_prev = f"{fonte_nome}_{fname}"
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)
# --- mask ids ---
ids_res = normalize_mask_ids(msk_path, cor_para_id, ignore_id, dim)
cv2.imwrite(os.path.join(out_msk, out_name_base + ".png"), ids_res)
# --- mask2 ---
if usar_msk2 and out_msk2:
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 ---
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 = []
for k in range(raw4.shape[0]):
ch = raw4[k]
ch_res = cv2.resize(ch, (out_w, out_h), interpolation=cv2.INTER_AREA)
chans.append(ch_res.astype(np.float32))
raw4 = np.stack(chans, axis=0) # (4, out_h, out_w)
# Atualiza acumuladores de stats
# raw4: (C,H,W) -> (C,N)
c, hh, ww = raw4.shape
if GLOBAL_SUM is None:
GLOBAL_SUM = np.zeros(c, dtype=np.float64)
GLOBAL_SUMSQ = np.zeros(c, dtype=np.float64)
flat = raw4.reshape(c, -1).astype(np.float64)
GLOBAL_SUM += flat.sum(axis=1)
GLOBAL_SUMSQ += (flat ** 2).sum(axis=1)
GLOBAL_PIXELS += hh * ww # por canal é o mesmo H*W
# Salva como float32 "linearzão" (4 * H * W floats)
save_raw(os.path.join(out_raw, out_name_base + ".raw"), raw4.astype(np.float32))
else:
print(f"[WARN] [{fonte_nome} | {grupo}] Sem RAW p/ {fname} (seguindo só preview+mask).")
total += 1
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n}{fname}")
return total
def main(args):
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
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 = {
"channels": channel_names[:len(mean_list)],
"mean": mean_list,
"std": std_list,
"pixels_per_channel": int(GLOBAL_PIXELS),
}
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}")
print(f" mean: {mean_list}")
print(f" std : {std_list}")
else:
print("⚠️ Nenhum RAW processado, não há stats para salvar.")
if __name__ == "__main__":
ap = argparse.ArgumentParser(description="Normalize por grupos (RAW: previews/raws/masks)")
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

@ -0,0 +1,451 @@
#!/usr/bin/env python3
# -*- 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 re
import json
import shutil
import random
import argparse
# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
RESOLUCAO = tuple(config.get("resolucao"))
# Pastas (ajustadas para PREVIEWS/RAWS)
pasta_origem = os.path.join("dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
pasta_destino = os.path.join("dataset", "split")
IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXT = ".png" # máscaras normalizadas em PNG
RAW_EXT = ".raw"
# Regex para identificar famílias
RE_ORIGINAL_PREFIX = re.compile(r'^original_(.+)$', re.IGNORECASE)
RE_AUGMENTED_FAMILY = re.compile(r'^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$', re.IGNORECASE)
RE_AUG_SUFFIX = re.compile(r'_aug[a-zA-Z0-9]*_\d+$', re.IGNORECASE)
def garantir(p):
os.makedirs(p, exist_ok=True)
def lista_grupos(root):
"""Lista grupos que têm previews + masks (raws opcional)."""
if not os.path.isdir(root):
return []
out = []
for g in sorted(os.listdir(root)):
gdir = os.path.join(root, g)
if not os.path.isdir(gdir):
continue
if os.path.isdir(os.path.join(gdir, "previews")) and os.path.isdir(os.path.join(gdir, "masks")):
out.append(g)
return out
def listar_previews(prev_dir):
if not os.path.isdir(prev_dir):
return []
fs = []
for f in os.listdir(prev_dir):
ext = os.path.splitext(f.lower())[1]
if ext in IMG_EXTS:
fs.append(f)
return sorted(fs)
def mask_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return base + MSK_EXT
def mask2_from_image_name(img_name):
base, _ = os.path.splitext(img_name)
return base + MSK_EXT # masks2 também normalizadas em PNG
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):
"""
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)
if m:
return "original", m.group(1)
m = RE_AUGMENTED_FAMILY.match(filename_no_ext)
if m:
return "augmented", m.group(1)
# legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
if RE_AUG_SUFFIX.search(filename_no_ext):
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
return "augmented", fam
return "unknown", filename_no_ext
def build_family_index(prev_dir, msk_dir):
"""
Constroi índice de famílias a partir de previews/masks.
Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]}
(strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe)
"""
familias = {}
imgs = listar_previews(prev_dir)
for img_name in imgs:
base_no_ext, ext = os.path.splitext(img_name)
mask_name = mask_from_image_name(img_name)
if not os.path.exists(os.path.join(msk_dir, mask_name)):
continue # garante pareamento preview/mask
source, fam = classify_source_and_family(base_no_ext)
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
d["all"].append(img_name)
if source == "original":
d["original"] = img_name
elif source == "augmented":
d["augmented"].append(img_name)
else:
# trata como original desconhecido para não perder dado
if d["original"] is None:
d["original"] = img_name
else:
d["augmented"].append(img_name)
return familias
def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
"""
Mesmo esquema do script original: calcula quantas FAMÍLIAS vão para train/val/test.
"""
n_train = int(round(n * p_train))
n_val = int(round(n * p_val))
n_test = n - n_train - n_val
if n_test < 0:
excesso = -n_test
take_train = min(excesso, max(0, n_train))
n_train -= take_train
excesso -= take_train
if excesso > 0:
take_val = min(excesso, max(0, n_val))
n_val -= take_val
excesso -= take_val
n_test = 0
min_sum = min_train + min_val + min_test
if n >= min_sum:
n_train = max(n_train, min_train)
n_val = max(n_val, min_val)
n_test = max(n_test, min_test)
total = n_train + n_val + n_test
while total > n:
if n_test > min_test:
n_test -= 1
elif n_val > min_val:
n_val -= 1
elif n_train > min_train:
n_train -= 1
else:
break
total = n_train + n_val + n_test
while total < n:
if n_train - min_train <= n_val - min_val:
n_train += 1
else:
n_val += 1
total = n_train + n_val + n_test
else:
n_train = min(n, max(1, min_train))
resto = n - n_train
n_val = max(0, min(resto, min_val))
n_test = max(0, resto - n_val)
# ajuste final
diff = n - (n_train + n_val + n_test)
if diff != 0:
if diff > 0:
# adiciona em train, depois val
take = min(diff, n - n_train)
n_train += take
diff -= take
if diff > 0:
n_val += diff
else:
diff = -diff
# tira de test, depois val
take = min(diff, n_test)
n_test -= take
diff -= take
if diff > 0:
n_val -= diff
return n_train, n_val, n_test
def copiar(
nomes,
src_prev_dir,
src_msk_dir,
dst_prev_dir,
dst_msk_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.
"""
garantir(dst_prev_dir)
garantir(dst_msk_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
for nome in nomes:
mask_name = mask_from_image_name(nome)
src_prev = os.path.join(src_prev_dir, nome)
src_msk = os.path.join(src_msk_dir, mask_name)
if not (os.path.exists(src_prev) and os.path.exists(src_msk)):
continue
shutil.copy2(src_prev, os.path.join(dst_prev_dir, nome))
shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name))
if use_msk2:
m2_name = mask2_from_image_name(nome)
src_m2 = os.path.join(src_msk2_dir, m2_name)
if os.path.exists(src_m2):
shutil.copy2(src_m2, os.path.join(dst_msk2_dir, m2_name))
if use_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
return moved
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_msk_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)
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]
total_familias = len(familias_originais)
if total_familias == 0:
print(f"[{group_name}] 0 famílias com original, pulando.")
return {"train": 0, "val": 0, "test": 0, "familias": 0}
rng = random.Random(seed)
rng.shuffle(familias_originais)
n_tr, n_va, n_te = allocate_counts(
total_familias, p_train, p_val, p_test,
mins["train"], mins["val"], mins["test"]
)
fam_train = set(familias_originais[:n_tr])
fam_val = set(familias_originais[n_tr:n_tr+n_va])
fam_test = set(familias_originais[n_tr+n_va: n_tr+n_va+n_te])
# CAP por grupo (apenas no TRAIN)
if caps_map and group_name in caps_map:
cap = caps_map[group_name]
if len(fam_train) > cap:
fam_list = list(fam_train)
rng.shuffle(fam_list)
kept = set(fam_list[:cap])
dropped = set(fam_list[cap:])
fam_train = kept
print(f"[{group_name}] cap-train-families={cap} → mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN")
nomes_train, nomes_val, nomes_test = [], [], []
for fam, d in familias.items():
if fam in fam_train:
if d["original"]:
nomes_train.append(d["original"])
if d["augmented"]:
nomes_train.extend(d["augmented"])
elif fam in fam_val:
if d["original"]:
nomes_val.append(d["original"])
elif fam in fam_test:
if d["original"]:
nomes_test.append(d["original"])
else:
pass
# dest dirs
dest_train_prev = os.path.join(pasta_destino, "train", "group", group_name, "previews")
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
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
dest_train_raw = os.path.join(pasta_destino, "train", "group", group_name, "raws") if use_raw else None
dest_val_raw = os.path.join(pasta_destino, "val", "group", group_name, "raws") if use_raw else None
dest_test_raw = os.path.join(pasta_destino, "test", "group", group_name, "raws") if use_raw else None
m_train = copiar(
nomes_train,
src_prev_dir, src_msk_dir,
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}")
return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}
def main():
ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (RAW: previews/raws/masks).")
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).")
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação (default=0.29).")
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste (default=0.01).")
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento (default=42).")
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")
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'")
args = ap.parse_args()
if args.resolucao:
try:
w, h = args.resolucao.lower().split("x")
resolucao = (int(w), int(h))
except Exception:
resolucao = RESOLUCAO
else:
resolucao = RESOLUCAO
def parse_cap_map(s):
caps = {}
if not s:
return caps
for item in s.split(","):
k, v = item.strip().split(":")
caps[k.strip()] = int(v)
return caps
caps_map = parse_cap_map(args.cap_train_families)
global pasta_origem, pasta_destino
pasta_origem = os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
pasta_destino = os.path.join("dataset", "split")
soma = args.train + args.val + args.test
if soma <= 0:
raise ValueError("Soma de proporções deve ser > 0.")
p_train = args.train / soma
p_val = args.val / soma
p_test = args.test / soma
mins = {
"train": max(0, args.min_train),
"val": max(0, args.min_val),
"test": max(0, args.min_test),
}
garantir(pasta_destino)
grupos = lista_grupos(pasta_origem)
if not grupos:
print(f"[WARN] Nenhum grupo encontrado em: {pasta_origem}")
return
random.seed(args.seed)
total_global = {"train": 0, "val": 0, "test": 0, "familias": 0}
print(f"Grupos: {', '.join(grupos)}")
print(f"Proporções normalizadas: train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
print(f"Mínimos por grupo (famílias): train={mins['train']} val={mins['val']} test={mins['test']}")
for g in grupos:
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map)
for k in total_global.keys():
total_global[k] += res.get(k, 0)
print("\nResumo global (imagens copiadas):")
print(f" train: {total_global['train']}")
print(f" val: {total_global['val']}")
print(f" test: {total_global['test']}")
print(f" famílias (total): {total_global['familias']}")
print("\n✅ Split RAW sem vazamento concluído!")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,779 @@
#python _8_train_segformer_b3_raw.py --config config.json --epochs 20 --batch 4 --lr 3e-5 --wd 0.01 --num_workers 4 --amp --amp_val --grad_accum 2 --class_weights auto
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128")
import json
import time
import argparse
from typing import Dict, List, Optional, Tuple, Any
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.amp import autocast, GradScaler
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)
# ----------------------------
# Utils
# ----------------------------
def set_seed(seed: int = 42):
import random
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
@torch.no_grad()
def update_confusion_matrix(
cm: torch.Tensor,
preds: torch.Tensor,
labels: torch.Tensor,
num_classes: int,
ignore_index: int = 255,
):
preds = preds.view(-1)
labels = labels.view(-1)
valid = labels != ignore_index
preds = preds[valid]
labels = labels[valid]
idx = labels * num_classes + preds
bins = torch.bincount(idx, minlength=num_classes * num_classes)
cm += bins.view(num_classes, num_classes)
@torch.no_grad()
def compute_iou_from_cm(
cm: torch.Tensor, eps: float = 1e-6
) -> Tuple[float, List[float]]:
cm = cm.float()
tp = torch.diag(cm)
fp = cm.sum(0) - tp
fn = cm.sum(1) - tp
denom = tp + fp + fn + eps
iou = (tp / denom).cpu().tolist()
miou = float(np.mean(iou))
return miou, iou
@torch.no_grad()
def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float:
cm = cm.float()
return (torch.diag(cm).sum() / (cm.sum() + eps)).item()
def save_checkpoint(
path: str,
model: nn.Module,
optimizer: torch.optim.Optimizer,
scaler: Optional[GradScaler],
epoch: int,
best_miou: float,
best_main_iou: float,
extra: Optional[Dict[str, Any]] = None,
):
ckpt = {
"epoch": epoch,
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"best_miou": best_miou,
"best_main_iou": best_main_iou,
}
if scaler is not None:
ckpt["scaler"] = scaler.state_dict()
if extra:
ckpt["extra"] = extra
torch.save(ckpt, path)
def load_checkpoint(
path: str,
model: nn.Module,
optimizer: Optional[torch.optim.Optimizer] = None,
scaler: Optional[GradScaler] = None,
map_location: str = "cpu",
) -> Dict[str, Any]:
ckpt = torch.load(path, map_location=map_location, weights_only=False)
model.load_state_dict(ckpt["model"], strict=True)
if optimizer is not None and "optimizer" in ckpt:
optimizer.load_state_dict(ckpt["optimizer"])
if scaler is not None and "scaler" in ckpt:
scaler.load_state_dict(ckpt["scaler"])
return ckpt
def default_collate(batch):
imgs, masks = [], []
for item in batch:
img = item["image"] if isinstance(item, dict) else item[0]
mask = item["mask"] if isinstance(item, dict) else item[1]
if isinstance(img, np.ndarray):
img = torch.from_numpy(img)
if isinstance(mask, np.ndarray):
mask = torch.from_numpy(mask)
imgs.append(img.float())
masks.append(mask.long())
return torch.stack(imgs, 0), torch.stack(masks, 0)
# ----------------------------
# Train/Val loop
# ----------------------------
def dice_loss(
logits: torch.Tensor, # (B, C, H, W)
target: torch.Tensor, # (B, H, W) com ids
num_classes: int,
ignore_index: int = -100,
smooth: float = 1.0
) -> torch.Tensor:
"""
Dice loss multi-classe (macro) com suporte a ignore_index.
Retorna um escalar.
"""
# Softmax -> probabilidades
probs = torch.softmax(logits, dim=1) # (B,C,H,W)
# Máscara de pixels válidos
valid = (target != ignore_index) # (B,H,W)
# Se por algum motivo não tem pixel válido no batch
if valid.sum() == 0:
return logits.new_tensor(0.0)
# One-hot do target (clamp só pra evitar index inválido)
target_clamped = target.clone()
target_clamped[~valid] = 0 # coloca qualquer classe nos ignore, pois vamos mascarar depois
# GARANTIR tipo inteiro pra one_hot
target_clamped = target_clamped.long()
target_1h = F.one_hot(target_clamped, num_classes=num_classes) # (B,H,W,C)
target_1h = target_1h.permute(0, 3, 1, 2).float() # (B,C,H,W)
# Aplica máscara de validade em probs e target
valid_f = valid.unsqueeze(1).float() # (B,1,H,W)
probs = probs * valid_f
target_1h = target_1h * valid_f
# Dice por classe (macro)
dims = (0, 2, 3) # soma em batch e espaço
inter = (probs * target_1h).sum(dims)
den = probs.sum(dims) + target_1h.sum(dims)
dice_per_class = (2.0 * inter + smooth) / (den + smooth) # (C,)
loss = 1.0 - dice_per_class.mean()
return loss
def run_one_epoch(
model: nn.Module,
loader: DataLoader,
optimizer: Optional[torch.optim.Optimizer],
device: torch.device,
num_classes: int,
ignore_index: int,
criterion: nn.Module,
amp: bool,
scaler: Optional[GradScaler],
train: bool,
grad_accum: int = 1,
normalizer=None
) -> Dict[str, Any]:
model.train(train)
total_loss = 0.0
cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)
t0 = time.time()
n_batches = 0
with torch.set_grad_enabled(train):
if train and optimizer is not None:
optimizer.zero_grad(set_to_none=True)
for step, (imgs, masks) in enumerate(loader):
# imgs: (B,C,H,W), C=4 ou 5
imgs = imgs.to(device, non_blocking=True)
masks = masks.to(device, non_blocking=True) # (B,H,W)
if normalizer is not None:
imgs = normalizer(imgs)
else:
imgs = normalize_raw(imgs)
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
out = model(pixel_values=imgs)
logits = out.logits
if logits.shape[-2:] != masks.shape[-2:]:
logits = torch.nn.functional.interpolate(
logits,
size=masks.shape[-2:],
mode="bilinear",
align_corners=False,
)
ce = criterion(logits, masks) # CrossEntropyLoss com weight/ignore_index
dice = dice_loss(logits=logits, target=masks, num_classes=num_classes, ignore_index=ignore_index, smooth=1.0)
loss = 0.7 * ce + 0.3 * dice
#loss = criterion(logits, masks)
if train and grad_accum > 1:
loss = loss / grad_accum
if train and optimizer is not None:
if amp and scaler is not None and device.type == "cuda":
scaler.scale(loss).backward()
if ((step + 1) % grad_accum) == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
else:
loss.backward()
if ((step + 1) % grad_accum) == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
total_loss += float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0)
n_batches += 1
preds = torch.argmax(logits, dim=1)
update_confusion_matrix(
cm,
preds,
masks,
num_classes=num_classes,
ignore_index=ignore_index,
)
dt = time.time() - t0
avg_loss = total_loss / max(n_batches, 1)
miou, iou_per_class = compute_iou_from_cm(cm)
acc = compute_pixel_acc_from_cm(cm)
return {
"loss": avg_loss,
"miou": miou,
"iou_per_class": iou_per_class,
"acc": acc,
"time_s": dt,
}
def estimate_class_weights(
ds: Dataset,
num_classes: int,
ignore_index: int = 255,
max_samples: int = 800,
) -> torch.Tensor:
n = min(len(ds), max_samples)
idxs = np.random.choice(len(ds), size=n, replace=False)
counts = np.zeros(num_classes, dtype=np.float64)
for i in idxs:
item = ds[i]
mask = item["mask"] if isinstance(item, dict) else item[1]
m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
m = m.reshape(-1)
m = m[m != ignore_index]
if m.size == 0:
continue
counts += np.bincount(m, minlength=num_classes)[:num_classes]
freq = counts / (counts.sum() + 1e-12)
freq = np.clip(freq, 1e-12, 1.0)
weights = 1.0 / np.log(1.02 + freq)
weights = weights / weights.mean()
return torch.tensor(weights, dtype=torch.float32)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="config.json")
parser.add_argument("--epochs", type=int, default=120)
parser.add_argument("--batch", type=int, default=2)
parser.add_argument("--lr", type=float, default=6e-5)
parser.add_argument("--wd", type=float, default=0.01)
parser.add_argument("--num_workers", type=int, default=4)
parser.add_argument("--save_every", type=int, default=10)
parser.add_argument("--resume", action="store_true")
parser.add_argument("--resume_ckpt", type=str, default=None)
parser.add_argument("--amp", action="store_true")
parser.add_argument("--amp_val", action="store_true")
parser.add_argument("--grad_accum", type=int, default=1)
parser.add_argument("--grad_ckpt", action="store_true")
parser.add_argument("--ignore_index", type=int, default=255)
parser.add_argument("--class_weights", type=str, default="auto")
parser.add_argument("--main_class", type=str, default=None)
parser.add_argument("--seed", type=int, default=42)
# RAW specifics
parser.add_argument("--raw_max", type=float, default=None, help="Ex: 1023, 4095, 16383, 65535. Se None: auto por dtype.")
parser.add_argument("--resize_h", type=int, default=None)
parser.add_argument("--resize_w", type=int, default=None)
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("--fusion_mode", type=str, default=None)
args = parser.parse_args()
set_seed(args.seed)
with open(args.config, "r") as f:
config = json.load(f)
MODELO = config["camera"]
MODEL_NAME = config["model_name"]
RESOLUCAO = config["resolucao"]
W, H = RESOLUCAO[0], RESOLUCAO[1]
MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower()
# novos parâmetros do config
USE_NDVI = bool(config.get("use_ndvi", False))
CHANNELS = int(config.get("channels", 4))
BACKBONE = config.get("backbone")
if args.main_class is not None:
MAIN_CLASS_NAME = args.main_class.lower()
FUSION_MODE = config.get("fusion_mode", "stacked")
if args.fusion_mode is not None:
FUSION_MODE = args.fusion_mode
if FUSION_MODE == "dual_branch" and CHANNELS not in (4, 5):
raise ValueError("dual_branch requer channels=4 ou 5")
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, experiment_tag)
#save_path = os.path.join("backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
dataset_path = os.path.join("dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
os.makedirs(save_path, exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
print(f"Config: channels={CHANNELS} use_ndvi={USE_NDVI}")
# ==========================
# Normalizador (fixo ou dinâmico)
# ==========================
norm_stats = 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")
if args.norm_stats is not None:
norm_stats_path = args.norm_stats
if norm_stats_path is not None and os.path.exists(norm_stats_path):
with open(norm_stats_path, "r", encoding="utf-8") as f:
norm_stats = json.load(f)
stats_channels = norm_stats.get("channels", [])
stats_mean = norm_stats.get("mean", [])
stats_std = norm_stats.get("std", [])
print(f"[NORM] usando stats fixos de: {norm_stats_path}")
print(f"[NORM] channels={stats_channels}")
print(f"[NORM] mean={stats_mean}")
print(f"[NORM] std ={stats_std}")
# 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:
# ordem interna bruta do JSON: R,G,IR,B
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
if args.resize_h is not None and args.resize_w is not None:
resize_hw = (args.resize_h, args.resize_w)
else:
# Se quiser, pode forçar pra RESOLUCAO do config (H,W)
# resize_hw = (RESOLUCAO[1], RESOLUCAO[0])
pass
# Datasets (RAW com 4 ou 5 canais)
ds_train = RawSegDataset(
os.path.join(dataset_path, "split", "train"),
labelmap_path=labelmap_path,
max_value=args.raw_max,
resize_hw=resize_hw,
raw_hw=(H, W),
use_ndvi=USE_NDVI,
channels=CHANNELS,
)
ds_val = RawSegDataset(
os.path.join(dataset_path, "split", "val"),
labelmap_path=labelmap_path,
max_value=args.raw_max,
resize_hw=resize_hw,
raw_hw=(H, W),
use_ndvi=USE_NDVI,
channels=CHANNELS,
)
# classes
if ds_train.classes is None:
raise RuntimeError("Crie labelmap.txt (um nome por linha) para definirmos num_classes com segurança.")
num_classes = len(ds_train.classes)
class_id_by_name = {
name.lower(): cid for cid, name in ds_train.classes.items()
}
class_name_by_id = {cid: name for cid, name in ds_train.classes.items()}
main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None)
if main_class_id is None:
print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.")
else:
print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}")
dl_train = DataLoader(
ds_train,
batch_size=args.batch,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
collate_fn=default_collate,
drop_last=True,
)
dl_val = DataLoader(
ds_val,
batch_size=1,
shuffle=False,
num_workers=max(2, args.num_workers // 2),
pin_memory=True,
collate_fn=default_collate,
drop_last=False,
)
if FUSION_MODE == "stacked":
model = build_raw_segformer_model(
num_classes=num_classes,
channels=CHANNELS,
backbone=BACKBONE,
device=device,
ckpt_path=None,
strict=False,
)
patch_segformer_input_channels(model, in_ch=CHANNELS)
elif FUSION_MODE == "dual_branch":
model = build_dual_branch_segformer_model(
num_classes=num_classes,
channels=CHANNELS,
backbone=BACKBONE,
device=device,
ckpt_path=None,
strict=False,
)
else:
raise ValueError(f"fusion_mode inválido: {FUSION_MODE}")
if args.grad_ckpt:
try:
model.gradient_checkpointing_enable()
print("[OK] gradient checkpointing enabled")
except Exception as e:
print(f"[WARN] grad_ckpt não suportado ({type(e).__name__}: {e}). Seguindo sem.")
model.to(device)
# Loss weights
if args.class_weights.lower() == "none":
weights = None
elif args.class_weights.lower() == "auto":
w = estimate_class_weights(ds_train, num_classes=num_classes, ignore_index=args.ignore_index)
weights = w.to(device)
print("Class weights (auto):", w.cpu().numpy().round(3).tolist())
else:
parts = [float(x) for x in args.class_weights.split(",")]
if len(parts) != num_classes:
raise ValueError(f"class_weights manual precisa ter {num_classes} valores, recebeu {len(parts)}.")
weights = torch.tensor(parts, dtype=torch.float32, device=device)
criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=args.ignore_index)
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=6, threshold=1e-4)
scaler = GradScaler(enabled=args.amp and device.type == "cuda")
last_ckpt_path = os.path.join(save_path, "last.pt")
best_miou_path = os.path.join(save_path, "best_miou.pt")
best_main_path = os.path.join(save_path, "best_main.pt")
start_epoch = 1
best_miou = -1.0
best_main_iou = -1.0
resume_path = args.resume_ckpt if args.resume_ckpt is not None else last_ckpt_path
if args.resume and os.path.exists(resume_path):
ckpt = load_checkpoint(
resume_path,
model,
optimizer,
scaler=scaler,
map_location="cpu",
)
start_epoch = int(ckpt["epoch"]) + 1
best_miou = float(ckpt.get("best_miou", -1.0))
best_main_iou = float(ckpt.get("best_main_iou", -1.0))
print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f} {resume_path}")
def pretty_iou(iou_list):
return " | ".join(
[
f"{class_name_by_id.get(cid, cid)}:{v:.3f}"
for cid, v in enumerate(iou_list)
]
)
# Debug opcional de layout de canais/máscara
if False:
sample = ds_train[0]
img = sample["image"] # tensor (C,H,W)
mask = sample["mask"] # tensor (H,W)
import cv2
img_np = img.cpu().numpy() # (C,H,W)
mask_np = mask.cpu().numpy()
use_ndvi_dbg = getattr(ds_train, "use_ndvi", False)
channels_dbg = getattr(ds_train, "channels", img_np.shape[0])
if use_ndvi_dbg:
if channels_dbg == 4:
# [R,G,B,NDVI]
r = (img_np[0] * 255).astype("uint8")
g = (img_np[1] * 255).astype("uint8")
b = (img_np[2] * 255).astype("uint8")
ndvi = (img_np[3] * 255).astype("uint8")
cv2.imwrite(os.path.join(save_path, "dbg_r.png"), r)
cv2.imwrite(os.path.join(save_path, "dbg_g.png"), g)
cv2.imwrite(os.path.join(save_path, "dbg_b.png"), b)
cv2.imwrite(os.path.join(save_path, "dbg_ndvi.png"), ndvi)
elif channels_dbg == 5:
# [R,G,B,IR,NDVI]
r = (img_np[0] * 255).astype("uint8")
g = (img_np[1] * 255).astype("uint8")
b = (img_np[2] * 255).astype("uint8")
ir = (img_np[3] * 255).astype("uint8")
ndvi = (img_np[4] * 255).astype("uint8")
cv2.imwrite(os.path.join(save_path, "dbg_r.png"), r)
cv2.imwrite(os.path.join(save_path, "dbg_g.png"), g)
cv2.imwrite(os.path.join(save_path, "dbg_b.png"), b)
cv2.imwrite(os.path.join(save_path, "dbg_ir.png"), ir)
cv2.imwrite(os.path.join(save_path, "dbg_ndvi.png"), ndvi)
else:
# layout clássico [R,G,IR,B] (C deve ser 4)
r = (img_np[0] * 255).astype("uint8")
g = (img_np[1] * 255).astype("uint8")
ir = (img_np[2] * 255).astype("uint8")
b = (img_np[3] * 255).astype("uint8")
cv2.imwrite(os.path.join(save_path, "dbg_r.png"), r)
cv2.imwrite(os.path.join(save_path, "dbg_g.png"), g)
cv2.imwrite(os.path.join(save_path, "dbg_ir.png"), ir)
cv2.imwrite(os.path.join(save_path, "dbg_b.png"), b)
cv2.imwrite(
os.path.join(save_path, "dbg_mask.png"), mask_np.astype("uint8")
)
for epoch in range(start_epoch, args.epochs + 1):
lr_now = optimizer.param_groups[0]["lr"]
print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====")
if device.type == "cuda":
torch.cuda.empty_cache()
tr = run_one_epoch(
model=model,
loader=dl_train,
optimizer=optimizer,
device=device,
num_classes=num_classes,
ignore_index=args.ignore_index,
criterion=criterion,
amp=args.amp,
scaler=scaler,
train=True,
grad_accum=max(1, args.grad_accum),
normalizer=normalizer,
)
if device.type == "cuda":
torch.cuda.empty_cache()
va = run_one_epoch(
model=model,
loader=dl_val,
optimizer=None,
device=device,
num_classes=num_classes,
ignore_index=args.ignore_index,
criterion=criterion,
amp=args.amp_val,
scaler=None,
train=False,
grad_accum=1,
normalizer=normalizer,
)
scheduler.step(va["loss"])
main_iou = (
va["miou"]
if main_class_id is None
else va["iou_per_class"][main_class_id]
)
print(f"TRAIN: loss={tr['loss']:.4f} acc={tr['acc']:.4f} miou={tr['miou']:.4f} (t={tr['time_s']:.1f}s)")
print(f"VAL : loss={va['loss']:.4f} acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={float(main_iou):.4f} (t={va['time_s']:.1f}s)")
print("IoU per class:", pretty_iou(va["iou_per_class"]))
save_checkpoint(
last_ckpt_path,
model,
optimizer,
scaler=scaler,
epoch=epoch,
best_miou=best_miou,
best_main_iou=best_main_iou,
extra={
"val_loss": va["loss"],
"val_miou": va["miou"],
"val_main_iou": float(main_iou),
},
)
if args.save_every > 0 and (epoch % args.save_every == 0):
save_checkpoint(
os.path.join(save_path, f"epoch_{epoch:04d}.pt"),
model,
optimizer,
scaler=scaler,
epoch=epoch,
best_miou=best_miou,
best_main_iou=best_main_iou,
)
if va["miou"] > best_miou:
best_miou = va["miou"]
save_checkpoint(
best_miou_path,
model,
optimizer,
scaler=scaler,
epoch=epoch,
best_miou=best_miou,
best_main_iou=best_main_iou,
)
print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}")
if float(main_iou) > best_main_iou:
best_main_iou = float(main_iou)
save_checkpoint(
best_main_path,
model,
optimizer,
scaler=scaler,
epoch=epoch,
best_miou=best_miou,
best_main_iou=best_main_iou,
)
print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,535 @@
#!/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

@ -0,0 +1,19 @@
{
"camera": "multispec_module",
"modelo": "segformer_b1",
"model_name": "pulv_mit",
"dual_head": false,
"main_class_name": "cana",
"es_classes": "",
"model_to_use": "geral",
"raw_size": [1296, 1028],
"resolucao": [1024, 800],
"roi_inicio": 0.0,
"roi_tamanho": 1.0,
"shaves": 3,
"channels": 5,
"use_ndvi": false,
"backbone": "nvidia/mit-b1",
"fusion_mode": "stacked",
"stats_source_tag": "stacked_raw5"
}

View File

@ -1,146 +0,0 @@
import socket
import json
import numpy as np
import base64
import time
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()
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:
pass
try:
if self.sock:
self.sock.close()
except:
pass
self.file = None
self.sock = None
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())
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):
return self._send_command({"cmd": "begin"})
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_resolution(self, width: int, height: int):
return self._send_command({
"cmd": "set_resolution",
"width": width,
"height": height
})
def capture_jpg_disk(self):
return self._send_command({"cmd": "capture_jpg_disk"})
def capture_jpg_bytes(self):
return self._send_command({"cmd": "capture_jpg_bytes"})
def capture_jpg_base64(self):
resp = self._send_command({"cmd": "capture_jpg_base64"})
if not resp.get("ok"):
raise RuntimeError(resp.get("error", "Falha ao capturar JPG"))
return base64.b64decode(resp["data"])
def capture_frame_array(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"))
raw = base64.b64decode(resp["data"])
width = resp["width"]
height = resp["height"]
channels = resp["channels"]
arr = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, channels)
t1 = time.perf_counter()
meta = {
"frame_type": resp.get("frame_type"),
"width": width,
"height": height,
"channels": channels,
"dtype": resp.get("dtype"),
"size": resp.get("size"),
# 👇 TELEMETRIA PI
"ts_pi": resp.get("ts_pi"),
"dt_trigger": resp.get("dt_trigger"),
"dt_settle": resp.get("dt_settle"),
"dt_capture": resp.get("dt_capture"),
"dt_total_pi": resp.get("dt_total_pi"),
# 👇 TELEMETRIA PC
"dt_total_pc": t1 - t0,
}
return arr, meta
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"})

View File

@ -1,139 +0,0 @@
import json
import socket
import threading
import time
import lz4.frame
import zstandard as zstd
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._zstd_d = zstd.ZstdDecompressor()
self._codec = Blosc(cname="lz4", clevel=1, shuffle=Blosc.SHUFFLE)
@property
def is_running(self):
return self._running
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:
pass
try:
if self._server_sock:
self._server_sock.close()
except:
pass
self._client_sock = None
self._server_sock = None
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)
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)
codec = header["codec"]
if codec == "lz4":
payload = lz4.frame.decompress(payload_comp)
elif codec == "zstd":
payload = self._zstd_d.decompress(payload_comp)
elif codec == "numcodecs":
payload = self._codec.decode(payload_comp)
else:
raise ValueError(f"Codec não suportado: {codec}")
expected = header["payload_size_raw"]
if len(payload) != expected:
raise ValueError(
f"Tamanho descomprimido inválido: {len(payload)} != {expected}"
)
height = header["height"]
width = header["width"]
channels = header["channels"]
dtype = np.uint8 if header["dtype"] == "uint8" else None
if dtype is None:
raise RuntimeError(f"dtype não suportado: {header['dtype']}")
frame = np.frombuffer(payload, dtype=dtype).reshape(height, width, channels)
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

View File

@ -1,19 +0,0 @@
import time
from multispectral_service import MultiSpectralService
svc = MultiSpectralService(host="192.168.105.6", port=5000)
svc.connect()
print("BEGIN:", svc.begin())
tempos_pi = []
tempos_pc = []
for i in range(10):
frame, meta = svc.capture_frame_array()
tempos_pi.append(meta["dt_total_pi"])
tempos_pc.append(meta["dt_total_pc"])
print(i, meta["dt_total_pi"], meta["dt_total_pc"])
print("STOP:", svc.stop())
svc.disconnect()

View File

@ -1,29 +0,0 @@
import time
from multispectral_service import MultiSpectralService
svc = MultiSpectralService(host="192.168.105.6", port=5000)
svc.connect()
print("PING:", svc.ping())
print("STATUS:", svc.get_status())
print("SET FPS:", svc.set_fps(15))
print("SET JPG:", svc.set_jpeg_quality(85))
print("SET RES:", svc.set_resolution(1280, 720))
print("BEGIN:", svc.begin())
jpg = None
for i in range(1, 6): # Começa em 1 e vai até 5
t0 = time.time()
jpg = svc.capture_jpg_base64()
status = "OK" if jpg is not None else "Falha"
tempo = time.time() - t0
print(f"CAPTURE {i} JPG: {status}, Tempo: {tempo:.4f}")
print("STOP:", svc.stop())
print("CONFIG:", svc.get_config())
print("STATUS FINAL:", svc.get_status())
svc.disconnect()
if jpg is not None:
with open("capture.jpg", "wb") as f:
f.write(jpg)

View File

@ -1,55 +0,0 @@
import time
from multispectral_service import MultiSpectralService
from stream_receiver import StreamReceiver
STREAM_PORT = 6001
PI_HOST = "192.168.105.6"
PC_HOST = "192.168.105.5"
def main():
receiver = StreamReceiver(host="0.0.0.0", port=STREAM_PORT)
receiver.start()
time.sleep(0.5)
svc = MultiSpectralService(host=PI_HOST, port=5000, timeout=10)
svc.connect()
try:
_fps = 15
print("SET RES:", svc.set_resolution(640, 480))
print("SET FPS:", svc.set_fps(_fps))
print("BEGIN:", svc.begin())
print("START STREAM:", svc.start_stream(PC_HOST, STREAM_PORT, fps=_fps))
t0 = time.perf_counter()
last_frame_id = -1
while time.perf_counter() - t0 < 5:
meta = receiver.last_meta
frame = receiver.last_frame
if meta is not None and meta["frame_id"] != last_frame_id:
last_frame_id = meta["frame_id"]
print(meta)
#print(
# f"frame_id={meta['frame_id']} "
# f"shape={frame.shape if frame is not None else None} "
# f"dt_total_pi={meta.get('dt_total_pi'):.4f}"
#)
time.sleep(0.02)
print("STOP STREAM:", svc.stop_stream())
print("STOP:", svc.stop())
finally:
svc.disconnect()
receiver.stop()
if __name__ == "__main__":
main()