Limpeza e refinamento na firmware do modulo multiespectral
This commit is contained in:
parent
fc1289f177
commit
5c26d66378
|
|
@ -65,6 +65,8 @@ Python/OAK/datasets/oak-d/dataset/
|
|||
Python/OAK/datasets/oak-d/backup/
|
||||
Python/OAK/datasets/gal5000/dataset/
|
||||
Python/OAK/datasets/gal5000/backup/
|
||||
Python/OAK/datasets/multiespec_module/dataset/
|
||||
Python/OAK/datasets/multiespec_module/backup/
|
||||
|
||||
Python/yolov8-seg/venv/
|
||||
Python/yolov8-seg/images/
|
||||
|
|
|
|||
|
|
@ -1,328 +0,0 @@
|
|||
import os
|
||||
import time
|
||||
import json
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from raw_segformer_service import make_bgr_preview_from_raw
|
||||
from gal5000.gal_service import Gal5000Camera # ajuste o nome do módulo se estiver diferente
|
||||
|
||||
|
||||
# =========================
|
||||
# Helpers gerais
|
||||
# =========================
|
||||
|
||||
def clamp(v, lo, hi):
|
||||
return lo if v < lo else hi if v > hi else v
|
||||
|
||||
|
||||
def ts_name() -> str:
|
||||
"""Timestamp legível e único para nome de arquivo."""
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||||
|
||||
|
||||
def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
|
||||
"""
|
||||
Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
|
||||
Pensado pra deixar o preview bonitinho sem estourar tudo.
|
||||
"""
|
||||
x = np.asarray(x)
|
||||
if x.dtype != np.float32 and x.dtype != np.float64:
|
||||
x = x.astype(np.float32)
|
||||
|
||||
# Se o canal já está em 0..1, escala pra 0..255 antes de cortar
|
||||
if x.max() <= 1.5:
|
||||
x = x * 255.0
|
||||
|
||||
lo = np.percentile(x, p_lo)
|
||||
hi = np.percentile(x, p_hi)
|
||||
|
||||
if hi <= lo + 1e-3:
|
||||
y = x
|
||||
else:
|
||||
y = (x - lo) * (255.0 / (hi - lo))
|
||||
|
||||
return np.clip(y, 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def overlay_hud(
|
||||
img_bgr: np.ndarray,
|
||||
lines: list[str],
|
||||
base_h: int = 720,
|
||||
base_font_scale: float = 0.75,
|
||||
base_line_step: int = 28,
|
||||
):
|
||||
"""
|
||||
Escreve textos empilhados no canto superior esquerdo,
|
||||
ajustando o tamanho do texto de acordo com a altura da imagem.
|
||||
|
||||
base_h: altura de referência (ex: 720 ou a RAW_H original).
|
||||
"""
|
||||
h, w = img_bgr.shape[:2]
|
||||
|
||||
# Fator de escala com base na altura atual
|
||||
scale = h / float(base_h)
|
||||
|
||||
# Evita ficar microscópico em resoluções muito baixas
|
||||
scale = max(scale, 0.4)
|
||||
|
||||
font_scale = base_font_scale * scale
|
||||
line_step = int(base_line_step * scale)
|
||||
|
||||
# Espessuras proporcionais
|
||||
thick_outline = max(1, int(3 * scale))
|
||||
thick_text = max(1, int(2 * scale))
|
||||
|
||||
# Margem superior / esquerda também escaladas
|
||||
y = int(24 * scale)
|
||||
x = int(12 * scale)
|
||||
|
||||
for s in lines:
|
||||
# contorno preto
|
||||
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
|
||||
# texto branco
|
||||
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
|
||||
y += line_step
|
||||
|
||||
|
||||
def save_sample_raw4(
|
||||
base_dir: str,
|
||||
raw4: np.ndarray,
|
||||
preview_bgr: np.ndarray,
|
||||
meta: dict,
|
||||
):
|
||||
"""
|
||||
Salva:
|
||||
- RAW4 como .raw float32 (4,H,W)
|
||||
- preview RGB como .png
|
||||
- metadados como .json
|
||||
dentro de base_dir.
|
||||
"""
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
name = ts_name()
|
||||
|
||||
raw_path = os.path.join(base_dir, f"{name}.raw")
|
||||
png_path = os.path.join(base_dir, f"{name}.png")
|
||||
json_path = os.path.join(base_dir, f"{name}.json")
|
||||
|
||||
# RAW4
|
||||
#np.save(raw_path, raw4.astype(np.float32))
|
||||
raw4.astype(np.float32).tofile(raw_path)
|
||||
|
||||
# Preview
|
||||
cv2.imwrite(png_path, preview_bgr)
|
||||
|
||||
# Metadados
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return raw_path, png_path, json_path
|
||||
|
||||
|
||||
# =========================
|
||||
# MAIN
|
||||
# =========================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
||||
parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
||||
parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
|
||||
parser.add_argument("--dll_dir", default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls", help="Pasta onde está a VT_SDK64.dll (usada pelo Gal5000Camera).")
|
||||
parser.add_argument("--dll_name", default="VT_SDK64.dll", help="Nome da DLL da câmera.")
|
||||
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
|
||||
parser.add_argument("--no_ae", action="store_true", help="Desliga o AutoExposure do service (por padrão ele vem ligado).")
|
||||
parser.add_argument("--upscale", type=int, default=2, help="Fator de upscale visual do preview.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# ===== config =====
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
RAW_W = config["raw_size"][0]
|
||||
RAW_H = config["raw_size"][1]
|
||||
|
||||
# Define diretório de sessão:
|
||||
# dataset/cana_<estado>/<horario>/<YYYYMMDD>/
|
||||
session_dir = os.path.join(MODELO, args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
|
||||
os.makedirs(session_dir, exist_ok=True)
|
||||
|
||||
print("============================================")
|
||||
print("Coleta de dataset RAW4 - SegFormer B0")
|
||||
print(f"Cana : {args.cana}")
|
||||
print(f"Horário : {args.horario}")
|
||||
print(f"Saída : {session_dir}")
|
||||
print("============================================")
|
||||
|
||||
window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
|
||||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
||||
|
||||
auto_save = False
|
||||
last_auto_t = 0.0
|
||||
upscale = args.upscale
|
||||
|
||||
# Estatísticas simples de FPS
|
||||
t_fps = time.time()
|
||||
frames = 0
|
||||
fps = 0.0
|
||||
|
||||
last_msg = ""
|
||||
last_msg_t = 0.0
|
||||
|
||||
try:
|
||||
cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
|
||||
with cam:
|
||||
print("[CAM] Status inicial:", cam.get_status())
|
||||
|
||||
# opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
|
||||
cam.configure_fps(20)
|
||||
cam.start_streaming()
|
||||
|
||||
apply_ir_comp = True
|
||||
ir_k_r = 0.8
|
||||
ir_k_g = 0.4
|
||||
ir_k_b = 0.9
|
||||
|
||||
while True:
|
||||
t0 = time.time()
|
||||
raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
|
||||
t1 = time.time()
|
||||
|
||||
ae_dbg = dbg.get("ae", {}) or {}
|
||||
exp_raw = dbg.get("exp_raw", None)
|
||||
gain_a = dbg.get("gain_a", None)
|
||||
gain_d = dbg.get("gain_d", None)
|
||||
|
||||
bgr = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=upscale > 0, preview_scale=upscale, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
||||
|
||||
# FPS
|
||||
frames += 1
|
||||
dt_fps = time.time() - t_fps
|
||||
if dt_fps >= 1.0:
|
||||
fps = frames / dt_fps
|
||||
frames = 0
|
||||
t_fps = time.time()
|
||||
|
||||
# AE info
|
||||
p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
|
||||
sat_disp = ae_dbg.get("sat", 0.0)
|
||||
hold = ae_dbg.get("hold", False)
|
||||
|
||||
ae_on = cam.is_auto_exposure_enabled()
|
||||
|
||||
# HUD principal
|
||||
lines = [
|
||||
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
|
||||
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
|
||||
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
|
||||
f"AEdbg: p95={p95_disp:.1f} sat={sat_disp:.3f} hold={hold}",
|
||||
"Keys: C/SPACE=save | A=auto-save | E=AE toggle | M=preview scale | Q/Esc=quit",
|
||||
]
|
||||
overlay_hud(bgr, lines, base_h=RAW_H)
|
||||
|
||||
# Mensagem rápida (ex: arquivo salvo)
|
||||
if last_msg and (time.time() - last_msg_t) < 2.0:
|
||||
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
|
||||
|
||||
cv2.imshow(window_name, bgr)
|
||||
|
||||
# Auto-save
|
||||
now = time.time()
|
||||
if auto_save and (now - last_auto_t) >= args.interval:
|
||||
meta = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"cana": args.cana,
|
||||
"horario": args.horario,
|
||||
"raw4_shape": list(raw4_base.shape),
|
||||
"out_h": RAW_H,
|
||||
"out_w": RAW_W,
|
||||
"ae_enabled": bool(ae_on),
|
||||
"exp_raw": int(exp_raw) if exp_raw is not None else None,
|
||||
"gain_a": int(gain_a) if gain_a is not None else None,
|
||||
"gain_d": int(gain_d) if gain_d is not None else None,
|
||||
"apply_ir_comp": apply_ir_comp,
|
||||
"ir_k_r": ir_k_r,
|
||||
"ir_k_g": ir_k_g,
|
||||
"ir_k_b": ir_k_b,
|
||||
"ae_dbg": {
|
||||
k: (float(v) if isinstance(v, (int, float, np.floating)) else v)
|
||||
for k, v in ae_dbg.items()
|
||||
},
|
||||
"note": "autosave",
|
||||
}
|
||||
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
||||
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
||||
last_msg = f"SALVO (auto): {os.path.basename(raw_path)}"
|
||||
last_msg_t = now
|
||||
last_auto_t = now
|
||||
|
||||
# Teclado
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
if k in (ord("q"), ord("Q"), 27): # Q ou ESC
|
||||
break
|
||||
|
||||
elif k in (ord("a"), ord("A")):
|
||||
auto_save = not auto_save
|
||||
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("e"), ord("E")):
|
||||
cam.enable_auto_exposure(not ae_on)
|
||||
last_msg = f"AE -> {'ON' if cam.is_auto_exposure_enabled() else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("m"), ord("M")):
|
||||
upscale = 0 if upscale else args.upscale
|
||||
last_msg = f"Preview UPSCALE -> {upscale}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("c"), ord("C"), 32): # C ou SPACE
|
||||
meta = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"cana": args.cana,
|
||||
"horario": args.horario,
|
||||
"raw4_shape": list(raw4_base.shape),
|
||||
"out_h": RAW_H,
|
||||
"out_w": RAW_W,
|
||||
"ae_enabled": bool(ae_on),
|
||||
"exp_raw": int(exp_raw) if exp_raw is not None else None,
|
||||
"gain_a": int(gain_a) if gain_a is not None else None,
|
||||
"gain_d": int(gain_d) if gain_d is not None else None,
|
||||
"apply_ir_comp": apply_ir_comp,
|
||||
"ir_k_r": ir_k_r,
|
||||
"ir_k_g": ir_k_g,
|
||||
"ir_k_b": ir_k_b,
|
||||
"ae_dbg": {
|
||||
k2: (float(v2) if isinstance(v2, (int, float, np.floating)) else v2)
|
||||
for k2, v2 in ae_dbg.items()
|
||||
},
|
||||
"note": "manual",
|
||||
}
|
||||
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
||||
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
||||
last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
# você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)
|
||||
|
||||
# Só pra não ficar rodando a 1000 FPS na UI
|
||||
# mas sem travar muito a captura
|
||||
dt_loop = time.time() - t0
|
||||
if dt_loop < 0.001:
|
||||
time.sleep(0.001)
|
||||
|
||||
finally:
|
||||
cv2.destroyAllWindows()
|
||||
print("Fim da captura.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -1,302 +1,176 @@
|
|||
import os
|
||||
import time
|
||||
import json
|
||||
import math
|
||||
import ctypes as C
|
||||
from ctypes import wintypes as W
|
||||
import argparse
|
||||
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"
|
||||
from raw_segformer_service import make_bgr_preview_from_raw
|
||||
from gal_service import Gal5000Camera # ajuste o nome do módulo se estiver diferente
|
||||
|
||||
# Onde salvar o dataset
|
||||
OUT_ROOT = os.path.join(os.path.dirname(__file__), "dataset")
|
||||
SESSION_DIR = os.path.join(OUT_ROOT, datetime.now().strftime("%Y%m%d"))
|
||||
os.makedirs(SESSION_DIR, exist_ok=True)
|
||||
|
||||
# Camera scan/open
|
||||
DEVICE_UDEF = 0
|
||||
DEVICE_INDEX = 0
|
||||
DATA_RAW = 0
|
||||
|
||||
# RAW geometry (se mudar no futuro, ajuste)
|
||||
RAW_W = 2592
|
||||
RAW_H = 2056
|
||||
|
||||
TIMEOUT_MS = 2000
|
||||
WINDOW_NAME = "GAL5000 Dataset Capture (C/SPACE=save | A=auto-save | E=AE toggle | Q=quit)"
|
||||
|
||||
# Preview
|
||||
UPSCALE = 2
|
||||
|
||||
# Auto-save
|
||||
CAPTURE_INTERVAL_S = 1.0
|
||||
|
||||
# Param IDs (VT_Param.h)
|
||||
BUF_SIZE = 256
|
||||
PARAM_ID_SENSOR_EXPOSURETIMERAW = 0x00003010
|
||||
PARAM_ID_SENSOR_GAINANALOGRAW = 0x00003020
|
||||
PARAM_ID_SENSOR_GAINDIGITRAW = 0x0000302A
|
||||
|
||||
# PARAM_VALUETYPE
|
||||
VALUE_INT = 0
|
||||
VALUE_FLOAT = 1
|
||||
VALUE_STR = 2
|
||||
|
||||
# Exposure/Gain limits (ajuste depois conforme o sensor aceitar)
|
||||
EXP_MIN = 1
|
||||
EXP_MAX = 20000
|
||||
|
||||
GAIN_A_MIN, GAIN_A_MAX = 0, 255
|
||||
GAIN_D_MIN, GAIN_D_MAX = 0, 255
|
||||
|
||||
# =========================
|
||||
# Helpers
|
||||
# Helpers gerais
|
||||
# =========================
|
||||
def ck(ret: int, name: str):
|
||||
if ret != 0:
|
||||
raise RuntimeError(f"{name} falhou, ret={ret}")
|
||||
|
||||
def ts_name() -> str:
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||||
|
||||
def clamp(v, lo, hi):
|
||||
return lo if v < lo else hi if v > hi else v
|
||||
|
||||
def norm8(x, p_lo=2, p_hi=98):
|
||||
|
||||
def ts_name() -> str:
|
||||
"""Timestamp legível e único para nome de arquivo."""
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||||
|
||||
|
||||
def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
|
||||
"""
|
||||
Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
|
||||
Pensado pra deixar o preview bonitinho sem estourar tudo.
|
||||
"""
|
||||
x = np.asarray(x)
|
||||
if x.dtype != np.float32 and x.dtype != np.float64:
|
||||
x = x.astype(np.float32)
|
||||
|
||||
# Se o canal já está em 0..1, escala pra 0..255 antes de cortar
|
||||
if x.max() <= 1.5:
|
||||
x = x * 255.0
|
||||
|
||||
lo = np.percentile(x, p_lo)
|
||||
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))
|
||||
|
||||
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 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):
|
||||
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,
|
||||
):
|
||||
"""
|
||||
Mede brilho no canal G cru usando uma ROI na base (mais parecido com chão).
|
||||
Retorna p90/p95 e fração saturada.
|
||||
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).
|
||||
"""
|
||||
G = raw[0::2, 1::2] # H/2 x W/2
|
||||
h2, w2 = G.shape
|
||||
h, w = img_bgr.shape[:2]
|
||||
|
||||
# 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]
|
||||
# Fator de escala com base na altura atual
|
||||
scale = h / float(base_h)
|
||||
|
||||
p90 = float(np.percentile(roi, 90))
|
||||
p95 = float(np.percentile(roi, 95))
|
||||
sat = float(np.mean(roi >= 250))
|
||||
return p90, p95, sat
|
||||
# Evita ficar microscópico em resoluções muito baixas
|
||||
scale = max(scale, 0.4)
|
||||
|
||||
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
|
||||
font_scale = base_font_scale * scale
|
||||
line_step = int(base_line_step * scale)
|
||||
|
||||
def step(self, raw, exp_raw):
|
||||
p90, p95, sat = measure_raw_g_metrics(raw)
|
||||
# Espessuras proporcionais
|
||||
thick_outline = max(1, int(3 * scale))
|
||||
thick_text = max(1, int(2 * scale))
|
||||
|
||||
# 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
|
||||
# Margem superior / esquerda também escaladas
|
||||
y = int(24 * scale)
|
||||
x = int(12 * scale)
|
||||
|
||||
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
|
||||
# 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(raw: np.ndarray, bgr_preview: np.ndarray, meta: dict):
|
||||
|
||||
def save_sample_raw4(
|
||||
base_dir: str,
|
||||
raw4: np.ndarray,
|
||||
preview_bgr: np.ndarray,
|
||||
meta: dict,
|
||||
):
|
||||
"""
|
||||
Salva:
|
||||
- RAW4 como .raw float32 (4,H,W)
|
||||
- preview RGB como .png
|
||||
- metadados como .json
|
||||
dentro de base_dir.
|
||||
"""
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
name = ts_name()
|
||||
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)
|
||||
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():
|
||||
# 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.")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
|
||||
# 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)
|
||||
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.")
|
||||
|
||||
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Estado local (não dependemos de GET)
|
||||
exp_raw = 1500
|
||||
gain_a = 0
|
||||
gain_d = 0
|
||||
# ===== 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]
|
||||
|
||||
# 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)
|
||||
# Define diretório de sessão:
|
||||
# dataset/cana_<estado>/<horario>/<YYYYMMDD>/
|
||||
session_dir = os.path.join(args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
|
||||
os.makedirs(session_dir, exist_ok=True)
|
||||
|
||||
print("============================================")
|
||||
print("Coleta de dataset RAW4 - SegFormer B0")
|
||||
print(f"Cana : {args.cana}")
|
||||
print(f"Horário : {args.horario}")
|
||||
print(f"Saída : {session_dir}")
|
||||
print("============================================")
|
||||
|
||||
window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
|
||||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
||||
|
||||
ae = RobustAE(target_p95=140.0, deadband=6.0, k=0.12, max_step=0.10, ema_alpha=0.20, sat_limit=0.01)
|
||||
ae_on = True
|
||||
auto_save = False
|
||||
last_auto_t = 0.0
|
||||
upscale = args.upscale
|
||||
|
||||
# FPS
|
||||
t0 = time.time()
|
||||
# Estatísticas simples de FPS
|
||||
t_fps = time.time()
|
||||
frames = 0
|
||||
fps = 0.0
|
||||
|
||||
|
|
@ -304,119 +178,151 @@ def main():
|
|||
last_msg_t = 0.0
|
||||
|
||||
try:
|
||||
cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
|
||||
with cam:
|
||||
print("[CAM] Status inicial:", cam.get_status())
|
||||
|
||||
# opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
|
||||
cam.configure_fps(20)
|
||||
cam.start_streaming()
|
||||
|
||||
apply_ir_comp = True
|
||||
ir_k_r = 0.8
|
||||
ir_k_g = 0.4
|
||||
ir_k_b = 0.9
|
||||
|
||||
while True:
|
||||
raw = capture_raw8(h)
|
||||
t0 = time.time()
|
||||
raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
|
||||
t1 = time.time()
|
||||
|
||||
# soft AE
|
||||
ae_dbg = {}
|
||||
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
|
||||
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)
|
||||
|
||||
# preview RGB bonitão
|
||||
rgb_clean = make_rgb_preview(raw, upscale=UPSCALE)
|
||||
bgr = rgb_clean.copy()
|
||||
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 = time.time() - t0
|
||||
if dt >= 1.0:
|
||||
fps = frames / dt
|
||||
dt_fps = time.time() - t_fps
|
||||
if dt_fps >= 1.0:
|
||||
fps = frames / dt_fps
|
||||
frames = 0
|
||||
t0 = time.time()
|
||||
t_fps = time.time()
|
||||
|
||||
# HUD
|
||||
# AE info
|
||||
p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
|
||||
sat_disp = ae_dbg.get("sat", 0.0)
|
||||
hold = ae_dbg.get("hold", False)
|
||||
|
||||
ae_on = cam.is_auto_exposure_enabled()
|
||||
|
||||
# HUD principal
|
||||
lines = [
|
||||
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Interval: {CAPTURE_INTERVAL_S:.1f}s",
|
||||
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
|
||||
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
|
||||
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
|
||||
f"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",
|
||||
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)
|
||||
overlay_hud(bgr, lines, base_h=RAW_H)
|
||||
|
||||
# msg pós-save
|
||||
# Mensagem rápida (ex: arquivo salvo)
|
||||
if last_msg and (time.time() - last_msg_t) < 2.0:
|
||||
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2, cv2.LINE_AA)
|
||||
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)
|
||||
cv2.imshow(window_name, bgr)
|
||||
|
||||
# autosave
|
||||
# Auto-save
|
||||
now = time.time()
|
||||
if auto_save and (now - last_auto_t) >= CAPTURE_INTERVAL_S:
|
||||
if auto_save and (now - last_auto_t) >= args.interval:
|
||||
meta = {
|
||||
"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),
|
||||
"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",
|
||||
}
|
||||
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
|
||||
last_msg = f"SAVED: {os.path.basename(raw_path)}"
|
||||
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):
|
||||
if k in (ord("q"), ord("Q"), 27): # Q ou ESC
|
||||
break
|
||||
|
||||
elif k in (ord('a'), ord('A')):
|
||||
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'}"
|
||||
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('c'), ord('C'), 32): # C ou SPACE
|
||||
elif k in (ord("m"), ord("M")):
|
||||
upscale = 0 if upscale else args.upscale
|
||||
last_msg = f"Preview UPSCALE -> {upscale}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("c"), ord("C"), 32): # C ou SPACE
|
||||
meta = {
|
||||
"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),
|
||||
"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",
|
||||
}
|
||||
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
|
||||
last_msg = f"SAVED: {os.path.basename(raw_path)}"
|
||||
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
||||
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
||||
last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
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)
|
||||
# você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)
|
||||
|
||||
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)
|
||||
# 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:
|
||||
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 da captura.")
|
||||
|
||||
print("Fim.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -1,11 +1,13 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import csv
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo # importa da sua utils.py
|
||||
|
||||
# ===================================================
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ Saída:
|
|||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import csv
|
||||
import json
|
||||
|
|
@ -28,6 +29,7 @@ import shutil
|
|||
import argparse
|
||||
import numpy as np
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
|
|
|
|||
|
|
@ -21,11 +21,14 @@ RAW:
|
|||
import argparse
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
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
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
# ===== config =====
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ Modos:
|
|||
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import argparse
|
||||
from collections import deque
|
||||
|
|
@ -29,11 +30,12 @@ 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
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import converter_mask_ids_para_bgr, desenhar_legenda_horizontal
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Helpers para modo "sem máscara"
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@
|
|||
import os
|
||||
import json
|
||||
from contextlib import nullcontext
|
||||
import sys
|
||||
from typing import List, Optional, Tuple
|
||||
import threading
|
||||
import time
|
||||
|
|
@ -23,6 +24,8 @@ from torch.utils.data import Dataset
|
|||
import torch.nn.functional as F
|
||||
from transformers import SegformerForSemanticSegmentation
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from utils import (carregar_labelmap_completo, converter_mask_ids_para_bgr)
|
||||
|
||||
def build_raw_segformer_model(
|
||||
|
|
@ -0,0 +1,822 @@
|
|||
import os
|
||||
import time
|
||||
import json
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from multispectral_service import MultiSpectralService
|
||||
from stream_receiver import StreamReceiver
|
||||
from pi.raw_processor_core import RawProcessorCore
|
||||
from pi.raw_processor_preview import RawProcessorPreview
|
||||
|
||||
|
||||
STREAM_PORT = 6001
|
||||
PI_HOST = "192.168.105.6"
|
||||
PC_HOST = "192.168.105.5"
|
||||
|
||||
|
||||
# =========================
|
||||
# Helpers gerais
|
||||
# =========================
|
||||
|
||||
def ts_name() -> str:
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||||
|
||||
|
||||
def overlay_hud(
|
||||
img_bgr: np.ndarray,
|
||||
lines: list[str],
|
||||
base_h: int = 720,
|
||||
base_font_scale: float = 0.75,
|
||||
base_line_step: int = 28,
|
||||
):
|
||||
h, w = img_bgr.shape[:2]
|
||||
|
||||
scale = h / float(base_h)
|
||||
scale = max(scale, 0.4)
|
||||
|
||||
font_scale = base_font_scale * scale
|
||||
line_step = int(base_line_step * scale)
|
||||
|
||||
thick_outline = max(1, int(3 * scale))
|
||||
thick_text = max(1, int(2 * scale))
|
||||
|
||||
y = int(24 * scale)
|
||||
x = int(12 * scale)
|
||||
|
||||
for s in lines:
|
||||
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
|
||||
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
|
||||
y += line_step
|
||||
|
||||
|
||||
def save_sample(
|
||||
base_dir: str,
|
||||
frame_type: str,
|
||||
preview_bgr: np.ndarray,
|
||||
meta: dict,
|
||||
raw_payload: np.ndarray | None = None,
|
||||
packed_raw: np.ndarray | None = None,
|
||||
packed_raw_by_camera: dict | None = None,
|
||||
):
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
name = ts_name()
|
||||
|
||||
png_path = os.path.join(base_dir, f"{name}.png")
|
||||
json_path = os.path.join(base_dir, f"{name}.json")
|
||||
|
||||
if frame_type in ("RGB", "MULTISPEC"):
|
||||
if raw_payload is None:
|
||||
raise ValueError(f"raw_payload não pode ser None quando frame_type='{frame_type}'")
|
||||
|
||||
payload_path = os.path.join(base_dir, f"{name}.raw")
|
||||
raw_payload.astype(np.float32).tofile(payload_path)
|
||||
|
||||
meta["saved_payload_type"] = frame_type.lower()
|
||||
meta["saved_payload_path"] = os.path.basename(payload_path)
|
||||
meta["saved_payload_dtype"] = "float32"
|
||||
meta["saved_payload_shape"] = list(raw_payload.shape)
|
||||
|
||||
elif frame_type == "RAW_BRUTO":
|
||||
if packed_raw_by_camera is not None:
|
||||
payload_files = {}
|
||||
payload_shapes = {}
|
||||
payload_dtypes = {}
|
||||
|
||||
for cam_id, arr in packed_raw_by_camera.items():
|
||||
path = os.path.join(base_dir, f"{name}_{cam_id}.bin")
|
||||
arr.tofile(path)
|
||||
payload_files[cam_id] = os.path.basename(path)
|
||||
payload_shapes[cam_id] = list(arr.shape)
|
||||
payload_dtypes[cam_id] = str(arr.dtype)
|
||||
|
||||
meta["saved_payload_type"] = "raw_native_multi"
|
||||
meta["saved_payload_paths"] = payload_files
|
||||
meta["saved_payload_shapes"] = payload_shapes
|
||||
meta["saved_payload_dtypes"] = payload_dtypes
|
||||
|
||||
else:
|
||||
if packed_raw is None:
|
||||
raise ValueError("packed_raw não pode ser None quando frame_type='RAW_BRUTO'")
|
||||
|
||||
payload_path = os.path.join(base_dir, f"{name}.bin")
|
||||
packed_raw.tofile(payload_path)
|
||||
|
||||
meta["saved_payload_type"] = "raw_native_single"
|
||||
meta["saved_payload_path"] = os.path.basename(payload_path)
|
||||
meta["saved_payload_dtype"] = str(packed_raw.dtype)
|
||||
meta["saved_payload_shape"] = list(packed_raw.shape)
|
||||
|
||||
else:
|
||||
raise ValueError(f"frame_type não suportado para save: {frame_type}")
|
||||
|
||||
cv2.imwrite(png_path, preview_bgr)
|
||||
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return png_path, json_path
|
||||
|
||||
|
||||
def get_camera_map_from_status(status: dict) -> dict:
|
||||
result = {}
|
||||
for cam in status.get("cameras", []):
|
||||
result[cam.get("id")] = cam
|
||||
return result
|
||||
|
||||
|
||||
def validate_module_ready(status: dict, frame_type: str, raw_policy: str, capture_mode: str):
|
||||
if not status.get("ok", True):
|
||||
raise RuntimeError(f"Status inválido retornado pelo módulo: {status}")
|
||||
|
||||
active_ids = list(status.get("active_camera_ids", []))
|
||||
active_count = int(status.get("camera_count_active", 0))
|
||||
|
||||
if frame_type == "RGB":
|
||||
if "cam2" not in active_ids:
|
||||
raise RuntimeError(
|
||||
"Modo RGB requer cam2 ativa (USB RGB), mas o módulo não reportou cam2 como ativa."
|
||||
)
|
||||
return
|
||||
|
||||
if frame_type == "MULTISPEC":
|
||||
if capture_mode == "TRIPLE":
|
||||
missing = [cid for cid in ("cam0", "cam1", "cam2") if cid not in active_ids]
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
f"Modo MULTISPEC/TRIPLE requer cam0, cam1 e cam2 ativas. "
|
||||
f"Faltando: {missing}. Ativas atuais: {active_ids}"
|
||||
)
|
||||
return
|
||||
|
||||
if capture_mode == "DOUBLE":
|
||||
has_rgb = "cam2" in active_ids
|
||||
has_spec = ("cam0" in active_ids) or ("cam1" in active_ids)
|
||||
|
||||
if not has_rgb or not has_spec:
|
||||
raise RuntimeError(
|
||||
f"Modo MULTISPEC/DOUBLE requer cam2 + (cam0 ou cam1). "
|
||||
f"Ativas atuais: {active_ids}"
|
||||
)
|
||||
return
|
||||
|
||||
# AUTO ou outros casos
|
||||
has_rgb = "cam2" in active_ids
|
||||
has_spec = ("cam0" in active_ids) or ("cam1" in active_ids)
|
||||
|
||||
if not (has_rgb and has_spec):
|
||||
raise RuntimeError(
|
||||
f"Modo MULTISPEC requer pelo menos RGB + 1 canal espectral. "
|
||||
f"Ativas atuais: {active_ids}"
|
||||
)
|
||||
return
|
||||
|
||||
if frame_type == "RAW_BRUTO":
|
||||
if raw_policy == "require_triple":
|
||||
missing = [cid for cid in ("cam0", "cam1", "cam2") if cid not in active_ids]
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
f"RAW_BRUTO com política require_triple exige três câmeras ativas. "
|
||||
f"Faltando: {missing}. Ativas atuais: {active_ids}"
|
||||
)
|
||||
else:
|
||||
if active_count < 1:
|
||||
raise RuntimeError("RAW_BRUTO requer ao menos uma câmera ativa, mas nenhuma foi detectada.")
|
||||
return
|
||||
|
||||
raise RuntimeError(f"frame_type desconhecido para validação: {frame_type}")
|
||||
|
||||
|
||||
def build_preview_from_raw_payload(
|
||||
frame,
|
||||
meta: dict,
|
||||
processor_core: RawProcessorCore,
|
||||
processor_preview: RawProcessorPreview,
|
||||
):
|
||||
"""
|
||||
Gera preview priorizando a câmera RGB (cam2).
|
||||
Se cam2 não estiver presente, cai para fallback usando a primeira câmera mono disponível.
|
||||
Retorna:
|
||||
preview_bgr
|
||||
payload_float_preview
|
||||
preview_source_id
|
||||
"""
|
||||
payload_sources = meta.get("payload_sources", []) or []
|
||||
|
||||
# Caso multi-payload: tenta usar cam2 primeiro
|
||||
if isinstance(frame, dict):
|
||||
if "cam2" in frame:
|
||||
rgb_frame = frame["cam2"]
|
||||
|
||||
if rgb_frame.ndim != 3 or rgb_frame.shape[2] != 3:
|
||||
raise RuntimeError(f"cam2 recebida mas inválida para preview RGB: shape={rgb_frame.shape}")
|
||||
|
||||
preview_bgr = rgb_frame.copy()
|
||||
payload_float = rgb_frame[:, :, ::-1].astype(np.float32) / 255.0
|
||||
payload_float = np.transpose(payload_float, (2, 0, 1))
|
||||
|
||||
return preview_bgr, payload_float, "cam2"
|
||||
|
||||
# fallback: usa a primeira câmera mono disponível
|
||||
fallback_id = None
|
||||
for cid in ("cam0", "cam1"):
|
||||
if cid in frame:
|
||||
fallback_id = cid
|
||||
break
|
||||
|
||||
if fallback_id is None:
|
||||
raise RuntimeError("Nenhuma câmera disponível no payload para gerar preview")
|
||||
|
||||
packed = frame[fallback_id]
|
||||
if packed.ndim == 3 and packed.shape[2] == 1:
|
||||
packed = packed[:, :, 0]
|
||||
|
||||
cam_frames = meta.get("camera_frames", {}) or {}
|
||||
cam_meta = cam_frames.get(fallback_id, {})
|
||||
bit_depth = int(cam_meta.get("bit_depth", 10))
|
||||
|
||||
raw16 = processor_core.unpack_raw10_packed(packed)
|
||||
preview_bgr = processor_preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)
|
||||
|
||||
payload_float = processor_core.build_training_rgb(
|
||||
raw16,
|
||||
output_dtype="float32",
|
||||
bit_depth=bit_depth,
|
||||
)
|
||||
|
||||
return preview_bgr, payload_float, fallback_id
|
||||
|
||||
# Caso single-payload
|
||||
if isinstance(frame, np.ndarray):
|
||||
# Se vier HWC/3ch, tratamos como RGB USB
|
||||
if frame.ndim == 3 and frame.shape[2] == 3:
|
||||
preview_bgr = frame.copy()
|
||||
payload_float = frame[:, :, ::-1].astype(np.float32) / 255.0
|
||||
payload_float = np.transpose(payload_float, (2, 0, 1))
|
||||
return preview_bgr, payload_float, "cam2"
|
||||
|
||||
# Se vier mono packed, fallback antigo
|
||||
packed = frame
|
||||
if packed.ndim == 3 and packed.shape[2] == 1:
|
||||
packed = packed[:, :, 0]
|
||||
|
||||
source_camera = meta.get("source_camera") or {}
|
||||
bit_depth = int(source_camera.get("bit_depth", meta.get("source_bit_depth", 10)))
|
||||
|
||||
raw16 = processor_core.unpack_raw10_packed(packed)
|
||||
preview_bgr = processor_preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)
|
||||
|
||||
payload_float = processor_core.build_training_rgb(
|
||||
raw16,
|
||||
output_dtype="float32",
|
||||
bit_depth=bit_depth,
|
||||
)
|
||||
|
||||
return preview_bgr, payload_float, source_camera.get("id", "unknown")
|
||||
|
||||
raise RuntimeError(f"Tipo de frame não suportado para preview: {type(frame)}")
|
||||
|
||||
|
||||
def resolve_effective_capture_mode(frame_type: str, raw_policy: str, requested_mode: str) -> str:
|
||||
"""
|
||||
Decide o modo real que será pedido ao módulo.
|
||||
|
||||
Nova regra:
|
||||
- Se o usuário pediu explicitamente SINGLE/DOUBLE/TRIPLE, respeitamos.
|
||||
- Se pediu AUTO, deixamos AUTO ir para o módulo.
|
||||
- A única exceção opcional é MULTISPEC + require_triple implícito,
|
||||
mas mesmo assim podemos deixar o módulo resolver se preferirmos.
|
||||
"""
|
||||
if requested_mode in ("SINGLE", "DOUBLE", "TRIPLE"):
|
||||
return requested_mode
|
||||
|
||||
# requested_mode == AUTO
|
||||
return "AUTO"
|
||||
|
||||
# =========================
|
||||
# MAIN
|
||||
# =========================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Captura de dataset usando módulo multispectral Pi + StreamReceiver.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
||||
parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
||||
parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
|
||||
parser.add_argument("--pi_host", default=PI_HOST, help="IP do servidor no Raspberry Pi.")
|
||||
parser.add_argument("--pc_host", default=PC_HOST, help="IP local do notebook/PC que receberá o stream.")
|
||||
parser.add_argument("--stream_port", type=int, default=STREAM_PORT, help="Porta TCP do receiver de stream.")
|
||||
parser.add_argument("--server_port", type=int, default=5000, help="Porta TCP do servidor de comandos no Pi.")
|
||||
parser.add_argument("--fps", type=int, default=20, help="FPS desejado.")
|
||||
parser.add_argument("--width", type=int, default=640, help="Largura óptica da câmera.")
|
||||
parser.add_argument("--height", type=int, default=480, help="Altura óptica da câmera.")
|
||||
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
|
||||
parser.add_argument("--preview_upscale", type=int, default=2, help="Fator de upscale visual do preview.")
|
||||
parser.add_argument("--bayer", default="GBRG", choices=["GBRG", "GRBG", "RGGB", "BGGR"], help="Padrão Bayer das câmeras.")
|
||||
parser.add_argument("--output_dtype", default="float32", choices=["uint8", "uint16", "float32"], help="Dtype do payload processado no Pi.")
|
||||
parser.add_argument("--frame_type", default="RAW_BRUTO", choices=["RAW_BRUTO", "RGB", "MULTISPEC"], help="Tipo de payload pedido ao Pi.")
|
||||
parser.add_argument("--capture_mode", default="AUTO", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"], help="Modo de captura desejado no módulo.")
|
||||
parser.add_argument("--raw_policy", default="allow_single", choices=["allow_single", "require_triple"], help="Quando frame_type=RAW_BRUTO, define se o script aceita 1 câmera ou exige 3.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
effective_capture_mode = resolve_effective_capture_mode(
|
||||
frame_type=args.frame_type,
|
||||
raw_policy=args.raw_policy,
|
||||
requested_mode=args.capture_mode,
|
||||
)
|
||||
|
||||
raw_w = args.width
|
||||
raw_h = args.height
|
||||
|
||||
session_dir = os.path.join(
|
||||
args.out_root,
|
||||
"brutas",
|
||||
f"cana_{args.cana}",
|
||||
args.horario,
|
||||
datetime.now().strftime("%Y%m%d"),
|
||||
)
|
||||
os.makedirs(session_dir, exist_ok=True)
|
||||
|
||||
print("============================================")
|
||||
print("Coleta de dataset - Módulo Multiespectral")
|
||||
print(f"Cana : {args.cana}")
|
||||
print(f"Horário : {args.horario}")
|
||||
print(f"Saída : {session_dir}")
|
||||
print(f"Sensor : {raw_w}x{raw_h} | Bayer={args.bayer}")
|
||||
print(f"FrameType : {args.frame_type}")
|
||||
print(f"CaptureMode : {args.capture_mode} -> efetivo={effective_capture_mode}")
|
||||
print(f"RAW policy : {args.raw_policy}")
|
||||
print("============================================")
|
||||
|
||||
auto_save = False
|
||||
last_auto_t = 0.0
|
||||
preview_upscale = args.preview_upscale
|
||||
|
||||
t_view_fps = time.time()
|
||||
view_frames = 0
|
||||
fps_view = 0.0
|
||||
|
||||
t_stream_fps = time.time()
|
||||
last_stream_frame_id = None
|
||||
stream_frames_accum = 0
|
||||
fps_stream = 0.0
|
||||
|
||||
last_msg = ""
|
||||
last_msg_t = 0.0
|
||||
|
||||
receiver = StreamReceiver(host="0.0.0.0", port=args.stream_port)
|
||||
svc = MultiSpectralService(host=args.pi_host, port=args.server_port, timeout=10)
|
||||
|
||||
print(f"[INFO] Verificando conexão com o módulo em {args.pi_host}:{args.server_port}...")
|
||||
|
||||
svc.ensure_alive()
|
||||
|
||||
if not svc.is_alive():
|
||||
raise RuntimeError("Módulo não encontrado ou não respondeu ao ping.")
|
||||
|
||||
print("[OK] Módulo conectado e respondendo.")
|
||||
|
||||
window_name = "Dataset Capture (C/SPACE=save | A=auto-save | M=preview scale | Q=quit)"
|
||||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
||||
|
||||
processor_core_cam0 = RawProcessorCore(
|
||||
sensor_width=raw_w,
|
||||
sensor_height=raw_h,
|
||||
bayer_pattern=args.bayer,
|
||||
)
|
||||
processor_preview_cam0 = RawProcessorPreview(
|
||||
sensor_width=raw_w,
|
||||
sensor_height=raw_h,
|
||||
bayer_pattern=args.bayer,
|
||||
)
|
||||
|
||||
last_frame_id = -1
|
||||
last_payload_float = None
|
||||
last_packed_raw = None
|
||||
last_packed_raw_by_camera = None
|
||||
last_preview_bgr = None
|
||||
last_meta_stream = None
|
||||
|
||||
try:
|
||||
receiver.start()
|
||||
time.sleep(0.5)
|
||||
|
||||
svc.connect()
|
||||
|
||||
if args.frame_type in ("RAW_BRUTO", "MULTISPEC"):
|
||||
modes_resp = svc.get_sensor_modes()
|
||||
if not modes_resp.get("ok"):
|
||||
print(f"[WARN] Falha ao obter sensor_modes: {modes_resp}")
|
||||
else:
|
||||
for mode in modes_resp.get("sensor_modes", []):
|
||||
print(
|
||||
f"[cam={mode.get('camera_id')} mode={mode.get('mode_index')}] "
|
||||
f"size={mode.get('size')} "
|
||||
f"format={mode.get('format')} "
|
||||
f"bit_depth={mode.get('bit_depth')} "
|
||||
f"fps={mode.get('fps')}"
|
||||
)
|
||||
else:
|
||||
print("[INFO] get_sensor_modes pulado para frame_type=RGB")
|
||||
|
||||
print("SET CAM0 RES:", svc.set_camera_resolution(0, raw_w, raw_h))
|
||||
print("SET CAM1 RES:", svc.set_camera_resolution(1, raw_w, raw_h))
|
||||
print("SET CAM2 RES:", svc.set_camera_resolution(2, raw_w, raw_h))
|
||||
print("SET CAM0 BAYER:", svc.set_camera_bayer(0, args.bayer))
|
||||
print("SET CAM1 BAYER:", svc.set_camera_bayer(1, args.bayer))
|
||||
print("SET FPS:", svc.set_fps(args.fps))
|
||||
print("SET CAPTURE MODE:", svc.set_capture_mode(effective_capture_mode))
|
||||
print("SET FRAME TYPE:", svc.set_frame_type(args.frame_type))
|
||||
print("SET OUTPUT DTYPE:", svc.set_output_dtype(args.output_dtype))
|
||||
|
||||
begin_resp = svc.begin(
|
||||
frame_type=args.frame_type,
|
||||
output_dtype=args.output_dtype,
|
||||
capture_mode=effective_capture_mode,
|
||||
)
|
||||
print("BEGIN:", begin_resp)
|
||||
|
||||
status = svc.get_status()
|
||||
print("STATUS:", json.dumps({
|
||||
"status": status.get("status"),
|
||||
"detected_mode": status.get("detected_mode"),
|
||||
"camera_count_active": status.get("camera_count_active"),
|
||||
"active_camera_ids": status.get("active_camera_ids"),
|
||||
}, ensure_ascii=False))
|
||||
|
||||
validate_module_ready(status, args.frame_type, args.raw_policy, effective_capture_mode)
|
||||
|
||||
print("START STREAM:", svc.start_stream(args.pc_host, args.stream_port, fps=args.fps))
|
||||
|
||||
camera_ctrl = svc.get_camera_controls()
|
||||
|
||||
ae_enabled = bool(camera_ctrl.get("ae_enable", True))
|
||||
awb_enabled = bool(camera_ctrl.get("awb_enable", True))
|
||||
manual_exposure_us = camera_ctrl.get("exposure_time_us", None)
|
||||
manual_gain = camera_ctrl.get("analogue_gain", None)
|
||||
manual_colour_gains = camera_ctrl.get("colour_gains", None)
|
||||
|
||||
while True:
|
||||
t0 = time.time()
|
||||
|
||||
meta = receiver.last_meta
|
||||
frame = receiver.last_frame
|
||||
|
||||
if meta is not None and frame is not None and meta.get("frame_id") != last_frame_id:
|
||||
last_frame_id = meta["frame_id"]
|
||||
|
||||
try:
|
||||
frame_type = meta.get("frame_type", "RAW_BRUTO")
|
||||
dtype_str = meta.get("dtype") or meta.get("output_dtype", "uint8")
|
||||
preview_source_id = "cam2"
|
||||
|
||||
if frame_type == "RAW_BRUTO":
|
||||
if isinstance(frame, dict):
|
||||
packed_by_camera = frame
|
||||
|
||||
preview_bgr, raw3_preview, preview_source_id = build_preview_from_raw_payload(
|
||||
frame=frame,
|
||||
meta=meta,
|
||||
processor_core=processor_core_cam0,
|
||||
processor_preview=processor_preview_cam0,
|
||||
)
|
||||
|
||||
last_packed_raw = None
|
||||
last_packed_raw_by_camera = {cam_id: arr.copy() for cam_id, arr in packed_by_camera.items()}
|
||||
last_payload_float = raw3_preview.copy()
|
||||
|
||||
else:
|
||||
preview_bgr, raw3_preview, preview_source_id = build_preview_from_raw_payload(
|
||||
frame=frame,
|
||||
meta=meta,
|
||||
processor_core=processor_core_cam0,
|
||||
processor_preview=processor_preview_cam0,
|
||||
)
|
||||
|
||||
last_packed_raw = frame.copy()
|
||||
last_packed_raw_by_camera = None
|
||||
last_payload_float = raw3_preview.copy()
|
||||
|
||||
elif frame_type == "RGB":
|
||||
rgb_chw = frame
|
||||
if not isinstance(rgb_chw, np.ndarray) or rgb_chw.ndim != 3:
|
||||
raise RuntimeError(f"Frame RGB inválido: type={type(rgb_chw)}")
|
||||
|
||||
if dtype_str == "uint8":
|
||||
payload_float = rgb_chw.astype(np.float32) / 255.0
|
||||
elif dtype_str == "float32":
|
||||
payload_float = rgb_chw.astype(np.float32)
|
||||
elif dtype_str == "uint16":
|
||||
payload_float = rgb_chw.astype(np.float32) / 65535.0
|
||||
else:
|
||||
raise RuntimeError(f"dtype RGB não suportado: {dtype_str}")
|
||||
|
||||
preview_rgb = np.transpose(payload_float, (1, 2, 0))
|
||||
preview_bgr = cv2.cvtColor(
|
||||
np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8),
|
||||
cv2.COLOR_RGB2BGR
|
||||
)
|
||||
|
||||
last_payload_float = payload_float.copy()
|
||||
last_packed_raw = None
|
||||
last_packed_raw_by_camera = None
|
||||
|
||||
elif frame_type == "MULTISPEC":
|
||||
multispec_chw = frame
|
||||
if not isinstance(multispec_chw, np.ndarray) or multispec_chw.ndim != 3 or multispec_chw.shape[0] not in (4, 5):
|
||||
raise RuntimeError(f"Frame MULTISPEC inválido: shape={getattr(multispec_chw, 'shape', None)}")
|
||||
|
||||
if dtype_str == "uint8":
|
||||
payload_float = multispec_chw.astype(np.float32) / 255.0
|
||||
elif dtype_str == "float32":
|
||||
payload_float = multispec_chw.astype(np.float32)
|
||||
elif dtype_str == "uint16":
|
||||
payload_float = multispec_chw.astype(np.float32) / 65535.0
|
||||
else:
|
||||
raise RuntimeError(f"dtype MULTISPEC não suportado: {dtype_str}")
|
||||
|
||||
preview_rgb = np.transpose(payload_float[:3], (1, 2, 0))
|
||||
preview_bgr = cv2.cvtColor(
|
||||
np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8),
|
||||
cv2.COLOR_RGB2BGR
|
||||
)
|
||||
|
||||
last_payload_float = payload_float.copy()
|
||||
last_packed_raw = None
|
||||
last_packed_raw_by_camera = None
|
||||
|
||||
else:
|
||||
raise RuntimeError(f"frame_type não suportado neste script: {frame_type}")
|
||||
|
||||
if preview_upscale and preview_upscale > 1:
|
||||
preview_show = cv2.resize(
|
||||
preview_bgr,
|
||||
(preview_bgr.shape[1] * preview_upscale, preview_bgr.shape[0] * preview_upscale),
|
||||
interpolation=cv2.INTER_NEAREST,
|
||||
)
|
||||
else:
|
||||
preview_show = preview_bgr.copy()
|
||||
|
||||
curr_frame_id = meta.get("frame_id")
|
||||
|
||||
if curr_frame_id is not None:
|
||||
if last_stream_frame_id != curr_frame_id:
|
||||
stream_frames_accum += 1
|
||||
|
||||
last_stream_frame_id = curr_frame_id
|
||||
|
||||
dt_stream = time.time() - t_stream_fps
|
||||
if dt_stream >= 1.0:
|
||||
fps_stream = stream_frames_accum / dt_stream
|
||||
stream_frames_accum = 0
|
||||
t_stream_fps = time.time()
|
||||
|
||||
view_frames += 1
|
||||
dt_view = time.time() - t_view_fps
|
||||
if dt_view >= 1.0:
|
||||
fps_view = view_frames / dt_view
|
||||
view_frames = 0
|
||||
t_view_fps = time.time()
|
||||
|
||||
active_sources = meta.get("payload_sources")
|
||||
lines = [
|
||||
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
|
||||
f"Type={meta.get('frame_type')} | CaptureMode={effective_capture_mode} | RAW policy={args.raw_policy}",
|
||||
f"Sources={active_sources} | FPS_STREAM={fps_stream:.1f} | FPS_VIEW={fps_view:.1f}",
|
||||
f"frame_id={meta.get('frame_id')} | layout={meta.get('output_layout')} | dtype={meta.get('dtype') or meta.get('output_dtype')}",
|
||||
f"codec={meta.get('codec_name', meta.get('codec_family', '-'))} | comp={meta.get('dt_comp', 0):.4f}s | send={meta.get('dt_send_payload_prev', 0):.4f}s",
|
||||
f"AE={'ON' if ae_enabled else 'OFF'} | AWB={'ON' if awb_enabled else 'OFF'} | EXP={manual_exposure_us} | GAIN={manual_gain}",
|
||||
"Keys: C/SPACE=save | A=auto-save | E=AE | W=AWB | I/K=exp | O/L=gain | R=reset | M=preview | Q/Esc=quit",
|
||||
]
|
||||
overlay_hud(preview_show, lines, base_h=raw_h)
|
||||
|
||||
if last_msg and (time.time() - last_msg_t) < 2.0:
|
||||
cv2.putText(preview_show, last_msg, (12, preview_show.shape[0] - 18),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
|
||||
|
||||
cv2.imshow(window_name, preview_show)
|
||||
|
||||
last_preview_bgr = preview_bgr.copy()
|
||||
last_meta_stream = dict(meta)
|
||||
|
||||
except Exception as e:
|
||||
err = np.zeros((500, 1200, 3), dtype=np.uint8)
|
||||
cv2.putText(err, f"Erro ao processar frame: {e}", (20, 60),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2, cv2.LINE_AA)
|
||||
cv2.imshow(window_name, err)
|
||||
print(f"[ERRO FRAME] {e}")
|
||||
|
||||
now = time.time()
|
||||
can_save = (
|
||||
last_meta_stream is not None and
|
||||
last_preview_bgr is not None and
|
||||
(
|
||||
(last_meta_stream.get("frame_type") in ("RGB", "MULTISPEC") and last_payload_float is not None) or
|
||||
(last_meta_stream.get("frame_type") == "RAW_BRUTO" and (last_packed_raw is not None or last_packed_raw_by_camera is not None))
|
||||
)
|
||||
)
|
||||
|
||||
if auto_save and can_save and (now - last_auto_t) >= args.interval:
|
||||
frame_type_save = last_meta_stream.get("frame_type")
|
||||
|
||||
meta_save = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"cana": args.cana,
|
||||
"horario": args.horario,
|
||||
"sensor_width": raw_w,
|
||||
"sensor_height": raw_h,
|
||||
"bayer_pattern": args.bayer,
|
||||
"fps_target": args.fps,
|
||||
"frame_type": frame_type_save,
|
||||
"capture_mode_requested": args.capture_mode,
|
||||
"capture_mode_effective": effective_capture_mode,
|
||||
"raw_policy": args.raw_policy,
|
||||
"stream_meta": last_meta_stream,
|
||||
"note": "autosave",
|
||||
"raw_preview_reference_camera": preview_source_id,
|
||||
}
|
||||
|
||||
save_sample(
|
||||
session_dir,
|
||||
frame_type=frame_type_save,
|
||||
preview_bgr=last_preview_bgr,
|
||||
meta=meta_save,
|
||||
raw_payload=last_payload_float,
|
||||
packed_raw=last_packed_raw,
|
||||
packed_raw_by_camera=last_packed_raw_by_camera,
|
||||
)
|
||||
|
||||
last_msg = "SALVO (auto)"
|
||||
last_msg_t = now
|
||||
last_auto_t = now
|
||||
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
if k in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
|
||||
elif k in (ord("a"), ord("A")):
|
||||
auto_save = not auto_save
|
||||
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("m"), ord("M")):
|
||||
preview_upscale = 0 if preview_upscale else args.preview_upscale
|
||||
last_msg = f"Preview UPSCALE -> {preview_upscale}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("e"), ord("E")):
|
||||
ae_enabled = not ae_enabled
|
||||
resp = svc.set_ae_enable(ae_enabled)
|
||||
ae_enabled = bool(resp.get("ae_enable", ae_enabled))
|
||||
last_msg = f"AE -> {'ON' if ae_enabled else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("w"), ord("W")):
|
||||
awb_enabled = not awb_enabled
|
||||
resp = svc.set_awb_enable(awb_enabled)
|
||||
awb_enabled = bool(resp.get("awb_enable", awb_enabled))
|
||||
last_msg = f"AWB -> {'ON' if awb_enabled else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("i"), ord("I")):
|
||||
if manual_exposure_us is None:
|
||||
manual_exposure_us = 15000
|
||||
else:
|
||||
manual_exposure_us = min(int(manual_exposure_us * 1.15), 200000)
|
||||
|
||||
if ae_enabled:
|
||||
ae_enabled = False
|
||||
svc.set_ae_enable(False)
|
||||
|
||||
resp = svc.set_exposure_time(int(manual_exposure_us))
|
||||
manual_exposure_us = resp.get("exposure_time_us", manual_exposure_us)
|
||||
last_msg = f"ExposureTime -> {manual_exposure_us} us"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("k"), ord("K")):
|
||||
if manual_exposure_us is None:
|
||||
manual_exposure_us = 15000
|
||||
else:
|
||||
manual_exposure_us = max(int(manual_exposure_us / 1.15), 100)
|
||||
|
||||
if ae_enabled:
|
||||
ae_enabled = False
|
||||
svc.set_ae_enable(False)
|
||||
|
||||
resp = svc.set_exposure_time(int(manual_exposure_us))
|
||||
manual_exposure_us = resp.get("exposure_time_us", manual_exposure_us)
|
||||
last_msg = f"ExposureTime -> {manual_exposure_us} us"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("o"), ord("O")):
|
||||
if manual_gain is None:
|
||||
manual_gain = 1.0
|
||||
else:
|
||||
manual_gain = min(float(manual_gain) * 1.10, 32.0)
|
||||
|
||||
if ae_enabled:
|
||||
ae_enabled = False
|
||||
svc.set_ae_enable(False)
|
||||
|
||||
resp = svc.set_analogue_gain(float(manual_gain))
|
||||
manual_gain = resp.get("analogue_gain", manual_gain)
|
||||
last_msg = f"AnalogueGain -> {manual_gain:.2f}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("l"), ord("L")):
|
||||
if manual_gain is None:
|
||||
manual_gain = 1.0
|
||||
else:
|
||||
manual_gain = max(float(manual_gain) / 1.10, 1.0)
|
||||
|
||||
if ae_enabled:
|
||||
ae_enabled = False
|
||||
svc.set_ae_enable(False)
|
||||
|
||||
resp = svc.set_analogue_gain(float(manual_gain))
|
||||
manual_gain = resp.get("analogue_gain", manual_gain)
|
||||
last_msg = f"AnalogueGain -> {manual_gain:.2f}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("r"), ord("R")):
|
||||
svc.clear_exposure_time()
|
||||
svc.clear_analogue_gain()
|
||||
svc.clear_colour_gains()
|
||||
|
||||
manual_exposure_us = None
|
||||
manual_gain = None
|
||||
manual_colour_gains = None
|
||||
|
||||
last_msg = "Manual controls resetados"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("c"), ord("C"), 32):
|
||||
if can_save:
|
||||
frame_type_save = last_meta_stream.get("frame_type")
|
||||
meta_save = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"cana": args.cana,
|
||||
"horario": args.horario,
|
||||
"sensor_width": raw_w,
|
||||
"sensor_height": raw_h,
|
||||
"bayer_pattern": args.bayer,
|
||||
"fps_target": args.fps,
|
||||
"frame_type": frame_type_save,
|
||||
"capture_mode_requested": args.capture_mode,
|
||||
"capture_mode_effective": effective_capture_mode,
|
||||
"raw_policy": args.raw_policy,
|
||||
"stream_meta": last_meta_stream,
|
||||
"camera_controls": {
|
||||
"ae_enable": ae_enabled,
|
||||
"awb_enable": awb_enabled,
|
||||
"exposure_time_us": manual_exposure_us,
|
||||
"analogue_gain": manual_gain,
|
||||
"colour_gains": manual_colour_gains,
|
||||
},
|
||||
"note": "manual",
|
||||
"raw_preview_reference_camera": preview_source_id,
|
||||
}
|
||||
|
||||
save_sample(
|
||||
session_dir,
|
||||
frame_type=frame_type_save,
|
||||
preview_bgr=last_preview_bgr,
|
||||
meta=meta_save,
|
||||
raw_payload=last_payload_float,
|
||||
packed_raw=last_packed_raw,
|
||||
packed_raw_by_camera=last_packed_raw_by_camera,
|
||||
)
|
||||
|
||||
last_msg = "SALVO (manual)"
|
||||
last_msg_t = time.time()
|
||||
|
||||
dt_loop = time.time() - t0
|
||||
if dt_loop < 0.001:
|
||||
time.sleep(0.001)
|
||||
|
||||
finally:
|
||||
try:
|
||||
print("STOP STREAM:", svc.stop_stream())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
print("STOP:", svc.stop())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
svc.disconnect()
|
||||
receiver.stop()
|
||||
cv2.destroyAllWindows()
|
||||
print("Fim da captura.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,226 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
Conta a porcentagem de pixels por classe dentro da ROI, por grupo e por split,
|
||||
na estrutura nova baseada em tensors/masks .npy.
|
||||
|
||||
Estrutura esperada:
|
||||
dataset/split/<split>/group/<grupo>/masks/
|
||||
original_xxx.npy
|
||||
original_xxx.png # opcional, debug
|
||||
augmented_xxx.npy
|
||||
augmented_xxx.png # opcional, debug
|
||||
|
||||
Regras:
|
||||
- Prioriza máscaras .npy
|
||||
- Se não encontrar .npy, pode usar .png como fallback
|
||||
- Ignora pixels com ignore_id
|
||||
- Normaliza a porcentagem somente sobre classes válidas
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
import numpy as np
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo, _infer_ignore_id
|
||||
|
||||
|
||||
# ================= CONFIG =================
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
W, H = config["resolucao"][0], config["resolucao"][1]
|
||||
ROI_INICIO = float(config["roi_inicio"])
|
||||
ROI_TAMANHO = float(config["roi_tamanho"])
|
||||
|
||||
pasta_base = "dataset"
|
||||
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
|
||||
|
||||
MAX_SAMPLES_PER_GROUP = 1000
|
||||
|
||||
|
||||
# ================= HELPERS =================
|
||||
def roi_slice(h: int):
|
||||
"""
|
||||
Replica a lógica antiga de ROI vertical.
|
||||
"""
|
||||
y_fim = int((1.0 - ROI_TAMANHO) * h)
|
||||
y_ini = int(ROI_INICIO * h)
|
||||
|
||||
if y_ini <= y_fim:
|
||||
y_fim = max(0, h - int(ROI_TAMANHO * h))
|
||||
y_ini = h
|
||||
|
||||
return slice(y_fim, y_ini)
|
||||
|
||||
|
||||
def load_mask(path: str) -> np.ndarray:
|
||||
"""
|
||||
Carrega máscara .npy ou .png.
|
||||
A máscara oficial agora é .npy (IDs inteiros).
|
||||
"""
|
||||
ext = os.path.splitext(path)[1].lower()
|
||||
|
||||
if ext == ".npy":
|
||||
arr = np.load(path)
|
||||
if arr.ndim != 2:
|
||||
raise RuntimeError(f"Máscara .npy inválida (esperado 2D): {path} shape={arr.shape}")
|
||||
return arr.astype(np.int32)
|
||||
|
||||
if ext == ".png":
|
||||
import cv2
|
||||
m = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
if m is None:
|
||||
raise RuntimeError(f"Falha ao ler PNG: {path}")
|
||||
if m.ndim == 3:
|
||||
# PNG de debug pode ter 3 canais; nesse caso ele não serve como máscara de IDs
|
||||
raise RuntimeError(
|
||||
f"PNG colorido encontrado em {path}. "
|
||||
f"Para estatística, use a máscara .npy correspondente."
|
||||
)
|
||||
return m.astype(np.int32)
|
||||
|
||||
raise RuntimeError(f"Extensão de máscara não suportada: {path}")
|
||||
|
||||
|
||||
def list_mask_files(mask_dir: str, prefer_npy: bool = True):
|
||||
"""
|
||||
Lista máscaras evitando contar duas vezes a mesma amostra
|
||||
quando existem .npy e .png com o mesmo base.
|
||||
"""
|
||||
if not os.path.isdir(mask_dir):
|
||||
return []
|
||||
|
||||
files = os.listdir(mask_dir)
|
||||
by_base = {}
|
||||
|
||||
for fname in files:
|
||||
low = fname.lower()
|
||||
if not (low.endswith(".npy") or low.endswith(".png")):
|
||||
continue
|
||||
|
||||
base, ext = os.path.splitext(fname)
|
||||
ext = ext.lower()
|
||||
full = os.path.join(mask_dir, fname)
|
||||
|
||||
if base not in by_base:
|
||||
by_base[base] = full
|
||||
else:
|
||||
cur_ext = os.path.splitext(by_base[base])[1].lower()
|
||||
|
||||
if prefer_npy:
|
||||
if cur_ext != ".npy" and ext == ".npy":
|
||||
by_base[base] = full
|
||||
else:
|
||||
if cur_ext != ".png" and ext == ".png":
|
||||
by_base[base] = full
|
||||
|
||||
return [by_base[k] for k in sorted(by_base.keys())]
|
||||
|
||||
|
||||
def summarize_group(mask_dir, class_ids_sorted, classes, ignore_id, max_samples):
|
||||
totals = {cid: 0 for cid in class_ids_sorted}
|
||||
ignored_pixels = 0
|
||||
valid_pixels = 0
|
||||
n = 0
|
||||
|
||||
for path in list_mask_files(mask_dir, prefer_npy=True):
|
||||
mask = load_mask(path)
|
||||
|
||||
rs = roi_slice(mask.shape[0])
|
||||
roi = mask[rs, :]
|
||||
|
||||
ignored_pixels += int((roi == ignore_id).sum())
|
||||
|
||||
for cid in class_ids_sorted:
|
||||
c = int((roi == cid).sum())
|
||||
totals[cid] += c
|
||||
valid_pixels += c
|
||||
|
||||
n += 1
|
||||
if n >= max_samples:
|
||||
break
|
||||
|
||||
return totals, valid_pixels, ignored_pixels, n
|
||||
|
||||
|
||||
# ================= MAIN =================
|
||||
def main(args):
|
||||
_, _colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
|
||||
|
||||
class_ids_sorted = sorted(classes.keys())
|
||||
|
||||
root = os.path.join(pasta_base, "split", args.split, "group")
|
||||
if not os.path.isdir(root):
|
||||
raise SystemExit(f"Nenhum diretório encontrado em: {root}")
|
||||
|
||||
grupos = [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))]
|
||||
|
||||
global_totals = {cid: 0 for cid in class_ids_sorted}
|
||||
global_valid_pixels = 0
|
||||
global_ignored_pixels = 0
|
||||
global_samples = 0
|
||||
|
||||
print(f"Split: {args.split}")
|
||||
print(f"ROI: inicio={ROI_INICIO:.3f} tamanho={ROI_TAMANHO:.3f}")
|
||||
print(f"Ignore ID: {ignore_id}")
|
||||
print()
|
||||
|
||||
for g in sorted(grupos):
|
||||
mdir = os.path.join(root, g, "masks")
|
||||
if not os.path.isdir(mdir):
|
||||
continue
|
||||
|
||||
totals, valid_pixels, ignored_pixels, n = summarize_group(
|
||||
mask_dir=mdir,
|
||||
class_ids_sorted=class_ids_sorted,
|
||||
classes=classes,
|
||||
ignore_id=ignore_id,
|
||||
max_samples=args.max_samples,
|
||||
)
|
||||
|
||||
denom = valid_pixels if valid_pixels > 0 else 1
|
||||
|
||||
parts = []
|
||||
for cid in class_ids_sorted:
|
||||
name = classes[cid]
|
||||
perc = totals[cid] / denom
|
||||
parts.append(f"{name}={perc:6.2%}")
|
||||
|
||||
ignore_ratio = ignored_pixels / max(1, (valid_pixels + ignored_pixels))
|
||||
|
||||
print(
|
||||
f"{g:18s} "
|
||||
+ " ".join(parts)
|
||||
+ f" | ignore={ignore_ratio:6.2%} | amostras={n}"
|
||||
)
|
||||
|
||||
for cid in class_ids_sorted:
|
||||
global_totals[cid] += totals[cid]
|
||||
global_valid_pixels += valid_pixels
|
||||
global_ignored_pixels += ignored_pixels
|
||||
global_samples += n
|
||||
|
||||
print("\nResumo global:")
|
||||
denom = global_valid_pixels if global_valid_pixels > 0 else 1
|
||||
parts = []
|
||||
for cid in class_ids_sorted:
|
||||
name = classes[cid]
|
||||
perc = global_totals[cid] / denom
|
||||
parts.append(f"{name}={perc:6.2%}")
|
||||
|
||||
ignore_ratio = global_ignored_pixels / max(1, (global_valid_pixels + global_ignored_pixels))
|
||||
print(" " + " ".join(parts))
|
||||
print(f" ignore={ignore_ratio:6.2%} | amostras={global_samples}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser(description="Verifica densidade de classes no dataset splitado.")
|
||||
ap.add_argument("--split", type=str, default="train", help="Split para verificar: train, val ou test.")
|
||||
ap.add_argument("--max-samples", type=int, default=MAX_SAMPLES_PER_GROUP, help="Máximo de amostras por grupo.")
|
||||
args = ap.parse_args()
|
||||
main(args)
|
||||
|
|
@ -1,134 +1,152 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
raw_triplet_sorter.py
|
||||
sample_bundle_sorter.py
|
||||
-------------------------------------------------
|
||||
Ferramenta para classificar manualmente amostras RAW do dataset
|
||||
(arquivos .raw + .png + .json) em pastas de labels,
|
||||
usando atalhos de teclado.
|
||||
Classificador manual para o formato de dataset do módulo multiespectral.
|
||||
|
||||
Cada amostra é formada por:
|
||||
<base>.png
|
||||
<base>.json
|
||||
<base>_cam0.bin
|
||||
<base>_cam1.bin
|
||||
<base>_cam2.bin # opcional
|
||||
|
||||
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.
|
||||
- Entrada: uma ou mais pastas-raiz contendo amostras em subpastas.
|
||||
- O script mostra o preview (.png).
|
||||
- Você usa teclas 1..9/0 para enviar o conjunto da amostra para uma label.
|
||||
|
||||
Estrutura de saída:
|
||||
out_root/
|
||||
<label>/
|
||||
raws/
|
||||
bins/
|
||||
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
|
||||
Regras:
|
||||
- PNG vai para previews/
|
||||
- JSON vai para metas/
|
||||
- BINs vão para bins/
|
||||
- masks/ é criada vazia
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Dict, Optional
|
||||
from typing import List, Optional
|
||||
|
||||
import tkinter as tk
|
||||
from tkinter import messagebox, filedialog
|
||||
|
||||
try:
|
||||
from PIL import Image, ImageTk
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
print("ERROR: Pillow (PIL) is required. Install with: pip install pillow", file=sys.stderr)
|
||||
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
|
||||
@dataclass
|
||||
class SampleBundle:
|
||||
sample_dir: Path
|
||||
preview_path: Path
|
||||
meta_path: Path
|
||||
bin_paths: List[Path]
|
||||
sample_id: str
|
||||
|
||||
Retorna lista de tuplas (preview_path, raw_path, json_path, base_name).
|
||||
|
||||
def find_sample_bundles_in_folder(folder: Path) -> List[SampleBundle]:
|
||||
"""
|
||||
Procura amostras no formato:
|
||||
<base>.png
|
||||
<base>.json
|
||||
<base>_cam0.bin
|
||||
<base>_cam1.bin
|
||||
<base>_cam2.bin
|
||||
|
||||
dentro da mesma pasta.
|
||||
"""
|
||||
import re
|
||||
|
||||
if not folder.is_dir():
|
||||
return []
|
||||
|
||||
previews: Dict[str, Path] = {}
|
||||
raws: Dict[str, Path] = {}
|
||||
metas: Dict[str, Path] = {}
|
||||
bundles: List[SampleBundle] = []
|
||||
candidate_dirs = [folder]
|
||||
candidate_dirs.extend([p for p in folder.rglob("*") if p.is_dir()])
|
||||
|
||||
for p in folder.rglob("*"):
|
||||
bin_re = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
|
||||
|
||||
for d in candidate_dirs:
|
||||
pngs = {}
|
||||
jsons = {}
|
||||
bins_by_base = {}
|
||||
|
||||
for p in d.iterdir():
|
||||
if not p.is_file():
|
||||
continue
|
||||
stem = p.stem
|
||||
|
||||
suffix = p.suffix.lower()
|
||||
|
||||
# Preview
|
||||
if suffix == PREVIEW_EXT:
|
||||
base = stem
|
||||
previews[base] = p
|
||||
pngs[p.stem] = p
|
||||
continue
|
||||
|
||||
# RAW
|
||||
if suffix == RAW_EXT:
|
||||
base = stem
|
||||
raws[base] = p
|
||||
continue
|
||||
|
||||
# JSON
|
||||
if suffix == META_EXT:
|
||||
base = stem
|
||||
metas[base] = p
|
||||
jsons[p.stem] = 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))
|
||||
if suffix == ".bin":
|
||||
m = bin_re.match(p.name)
|
||||
if m:
|
||||
base = m.group("base")
|
||||
cam = int(m.group("cam"))
|
||||
bins_by_base.setdefault(base, []).append((cam, p))
|
||||
|
||||
# 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)
|
||||
valid_bases = sorted(set(pngs.keys()) & set(jsons.keys()) & set(bins_by_base.keys()))
|
||||
|
||||
triplets.sort(key=sort_key)
|
||||
return triplets
|
||||
for base in valid_bases:
|
||||
sorted_bins = [p for cam, p in sorted(bins_by_base[base], key=lambda x: x[0])]
|
||||
if not sorted_bins:
|
||||
continue
|
||||
|
||||
bundles.append(
|
||||
SampleBundle(
|
||||
sample_dir=d,
|
||||
preview_path=pngs[base],
|
||||
meta_path=jsons[base],
|
||||
bin_paths=sorted_bins,
|
||||
sample_id=base,
|
||||
)
|
||||
)
|
||||
|
||||
bundles.sort(key=lambda x: (str(x.sample_dir), x.sample_id))
|
||||
return bundles
|
||||
|
||||
|
||||
def collect_all_triplets(folders: List[Path]) -> List[Tuple[Path, Path, Path, str]]:
|
||||
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
|
||||
def collect_all_bundles(folders: List[Path]) -> List[SampleBundle]:
|
||||
all_bundles: List[SampleBundle] = []
|
||||
seen_dirs = set()
|
||||
|
||||
for folder in folders:
|
||||
for bundle in find_sample_bundles_in_folder(folder):
|
||||
key = str((bundle.sample_dir / bundle.sample_id).resolve())
|
||||
if key in seen_dirs:
|
||||
continue
|
||||
seen_dirs.add(key)
|
||||
all_bundles.append(bundle)
|
||||
|
||||
all_bundles.sort(key=lambda x: (str(x.sample_dir), x.sample_id))
|
||||
return all_bundles
|
||||
|
||||
|
||||
class ActionLogger:
|
||||
|
|
@ -136,7 +154,8 @@ 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 já foi classificado/pulado).
|
||||
|
||||
O resume usa sample_dir como chave.
|
||||
"""
|
||||
|
||||
def __init__(self, out_root: Path):
|
||||
|
|
@ -150,29 +169,29 @@ class ActionLogger:
|
|||
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", "")
|
||||
key = row.get("meta_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 already_logged(self, sample_key: Path) -> bool:
|
||||
return str(sample_key.resolve()) in self.seen
|
||||
|
||||
def log(
|
||||
self,
|
||||
action: str,
|
||||
sample_dir: Path,
|
||||
preview_path: Path,
|
||||
raw_path: Path,
|
||||
meta_path: Path,
|
||||
bin_paths: List[Path],
|
||||
label: Optional[str] = None,
|
||||
dest_preview: Optional[Path] = None,
|
||||
dest_raw: Optional[Path] = None,
|
||||
dest_meta: Optional[Path] = None,
|
||||
dest_sample_dir: Optional[Path] = None,
|
||||
):
|
||||
self.out_root.mkdir(parents=True, exist_ok=True)
|
||||
sample_dir_str = str(sample_dir.resolve())
|
||||
meta_key_str = str(meta_path.resolve())
|
||||
bin_paths_str = [str(p.resolve()) for p in bin_paths]
|
||||
|
||||
# CSV
|
||||
new_file = not self.log_csv.exists()
|
||||
with self.log_csv.open("a", newline="", encoding="utf-8") as f:
|
||||
writer = csv.writer(f)
|
||||
|
|
@ -181,54 +200,50 @@ class ActionLogger:
|
|||
"timestamp",
|
||||
"action",
|
||||
"label",
|
||||
"sample_dir",
|
||||
"preview_path",
|
||||
"raw_path",
|
||||
"meta_path",
|
||||
"dest_preview",
|
||||
"dest_raw",
|
||||
"dest_meta",
|
||||
"bin_paths_json",
|
||||
"dest_sample_dir",
|
||||
])
|
||||
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 ""),
|
||||
sample_dir_str,
|
||||
str(preview_path.resolve()),
|
||||
str(meta_path.resolve()),
|
||||
json.dumps(bin_paths_str, ensure_ascii=False),
|
||||
str(dest_sample_dir.resolve()) if dest_sample_dir else "",
|
||||
])
|
||||
|
||||
# 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,
|
||||
"sample_dir": sample_dir_str,
|
||||
"preview_path": str(preview_path.resolve()),
|
||||
"meta_path": str(meta_path.resolve()),
|
||||
"bin_paths": bin_paths_str,
|
||||
"dest_sample_dir": str(dest_sample_dir.resolve()) if dest_sample_dir else None,
|
||||
}
|
||||
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
|
||||
self.seen.add(str(preview_path))
|
||||
self.seen.add(str(meta_path.resolve()))
|
||||
|
||||
|
||||
class RawTripletSorterApp:
|
||||
class SampleBundleSorterApp:
|
||||
def __init__(
|
||||
self,
|
||||
triplets: List[Tuple[Path, Path, Path, str]],
|
||||
bundles: List[SampleBundle],
|
||||
labels: List[str],
|
||||
out_root: Path,
|
||||
move: bool,
|
||||
resume: bool,
|
||||
display_height: int = 512,
|
||||
):
|
||||
self.all_triplets = triplets
|
||||
self.all_bundles = bundles
|
||||
self.labels = labels
|
||||
self.out_root = out_root
|
||||
self.move = move
|
||||
|
|
@ -236,28 +251,25 @@ class RawTripletSorterApp:
|
|||
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])
|
||||
self.all_bundles = [
|
||||
b for b in self.all_bundles if not self.logger.already_logged(b.meta_path)
|
||||
]
|
||||
|
||||
# Estado de sessão
|
||||
self.idx = 0
|
||||
self.history: List[dict] = [] # para undo
|
||||
self.history = []
|
||||
|
||||
# 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)
|
||||
label_root = out_root / label
|
||||
(label_root / "bins").mkdir(parents=True, exist_ok=True)
|
||||
(label_root / "previews").mkdir(parents=True, exist_ok=True)
|
||||
(label_root / "metas").mkdir(parents=True, exist_ok=True)
|
||||
(label_root / "masks").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# UI
|
||||
self.root = tk.Tk()
|
||||
self.root.title("Agrobotics RAW Triplet Sorter (preview | raw | json)")
|
||||
self.root.title("Agrobotics Sample Bundle Sorter")
|
||||
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)
|
||||
|
||||
|
|
@ -267,19 +279,16 @@ class RawTripletSorterApp:
|
|||
self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10))
|
||||
self.legend_label.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",
|
||||
|
|
@ -287,19 +296,31 @@ class RawTripletSorterApp:
|
|||
)
|
||||
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
|
||||
|
||||
# Primeira renderização
|
||||
self.render()
|
||||
|
||||
def unique_file(self, p: Path) -> Path:
|
||||
if not p.exists():
|
||||
return p
|
||||
|
||||
stem = p.stem
|
||||
suffix = p.suffix
|
||||
k = 1
|
||||
|
||||
while True:
|
||||
cand = p.with_name(f"{stem}__{k}{suffix}")
|
||||
if not cand.exists():
|
||||
return cand
|
||||
k += 1
|
||||
|
||||
def build_legend_text(self) -> str:
|
||||
parts = []
|
||||
for i, label in enumerate(self.labels, start=1):
|
||||
key = i if i <= 9 else 0 # 0 = 10ª label
|
||||
key = i if i <= 9 else 0
|
||||
parts.append(f"[{key}] {label}")
|
||||
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))
|
||||
|
|
@ -307,135 +328,140 @@ class RawTripletSorterApp:
|
|||
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?).")
|
||||
if not self.all_bundles:
|
||||
messagebox.showinfo("Fim", "Não há amostras para exibir.")
|
||||
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]
|
||||
self.idx = max(0, min(self.idx, len(self.all_bundles) - 1))
|
||||
bundle = self.all_bundles[self.idx]
|
||||
|
||||
try:
|
||||
tk_img = self.pil_load_preview(preview_path)
|
||||
tk_img = self.pil_load_preview(bundle.preview_path)
|
||||
self.preview_tk = tk_img
|
||||
self.preview_label.configure(image=self.preview_tk)
|
||||
except Exception as e:
|
||||
self.info_label.configure(text=f"Erro ao abrir preview: {e}")
|
||||
return
|
||||
|
||||
bins_text = ", ".join(p.name for p in bundle.bin_paths)
|
||||
self.info_label.configure(
|
||||
text=f"{self.idx+1}/{len(self.all_triplets)} | base='{base}' | PREVIEW: {preview_path.name}"
|
||||
text=(
|
||||
f"{self.idx+1}/{len(self.all_bundles)} | "
|
||||
f"sample='{bundle.sample_id}' | "
|
||||
f"bins={len(bundle.bin_paths)} [{bins_text}]"
|
||||
)
|
||||
)
|
||||
|
||||
def do_copy_or_move(self, src: Path, dst: Path):
|
||||
if self.move:
|
||||
shutil.move(str(src), str(dst))
|
||||
else:
|
||||
shutil.copy2(str(src), str(dst))
|
||||
|
||||
def unique_path(self, p: Path) -> Path:
|
||||
"""Se o caminho já existe, gera um novo com sufixo __k."""
|
||||
if not p.exists():
|
||||
return p
|
||||
stem, ext = p.stem, p.suffix
|
||||
k = 1
|
||||
while True:
|
||||
cand = p.with_name(f"{stem}__{k}{ext}")
|
||||
if not cand.exists():
|
||||
return cand
|
||||
k += 1
|
||||
|
||||
def send_to_label(self, label_index: int):
|
||||
if label_index < 0 or label_index >= len(self.labels):
|
||||
return
|
||||
|
||||
label = self.labels[label_index]
|
||||
bundle = self.all_bundles[self.idx]
|
||||
|
||||
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)
|
||||
label_root = self.out_root / label
|
||||
dst_preview = self.unique_file(label_root / "previews" / bundle.preview_path.name)
|
||||
dst_meta = self.unique_file(label_root / "metas" / bundle.meta_path.name)
|
||||
dst_bins = [self.unique_file(label_root / "bins" / p.name) for p in bundle.bin_paths]
|
||||
|
||||
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)
|
||||
if self.move:
|
||||
shutil.move(str(bundle.preview_path), str(dst_preview))
|
||||
shutil.move(str(bundle.meta_path), str(dst_meta))
|
||||
for src_bin, dst_bin in zip(bundle.bin_paths, dst_bins):
|
||||
shutil.move(str(src_bin), str(dst_bin))
|
||||
else:
|
||||
shutil.copy2(bundle.preview_path, dst_preview)
|
||||
shutil.copy2(bundle.meta_path, dst_meta)
|
||||
for src_bin, dst_bin in zip(bundle.bin_paths, dst_bins):
|
||||
shutil.copy2(src_bin, dst_bin)
|
||||
|
||||
self.logger.log(
|
||||
"assign",
|
||||
preview_path=preview_path,
|
||||
raw_path=raw_path,
|
||||
meta_path=meta_path,
|
||||
action="assign",
|
||||
sample_dir=bundle.sample_dir,
|
||||
preview_path=bundle.preview_path,
|
||||
meta_path=bundle.meta_path,
|
||||
bin_paths=bundle.bin_paths,
|
||||
label=label,
|
||||
dest_preview=dst_preview,
|
||||
dest_raw=dst_raw,
|
||||
dest_meta=dst_meta,
|
||||
dest_sample_dir=label_root,
|
||||
)
|
||||
|
||||
self.history.append({
|
||||
"action": "assign",
|
||||
"label": label,
|
||||
"preview_src": preview_path,
|
||||
"raw_src": raw_path,
|
||||
"meta_src": meta_path,
|
||||
"sample_src": bundle.sample_dir,
|
||||
"preview_src": bundle.preview_path,
|
||||
"meta_src": bundle.meta_path,
|
||||
"bin_srcs": list(bundle.bin_paths),
|
||||
"preview_dst": dst_preview,
|
||||
"raw_dst": dst_raw,
|
||||
"meta_dst": dst_meta,
|
||||
"bin_dsts": dst_bins,
|
||||
"moved": self.move,
|
||||
"index": self.idx,
|
||||
})
|
||||
|
||||
self.status_var.set(
|
||||
f"{'Movido' if self.move else 'Copiado'} → '{label}': {preview_path.name} (+raw+json)"
|
||||
f"{'Movido' if self.move else 'Copiado'} → '{label}': {bundle.sample_id}"
|
||||
)
|
||||
|
||||
self.idx += 1
|
||||
if self.idx >= len(self.all_triplets):
|
||||
if self.idx >= len(self.all_bundles):
|
||||
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}")
|
||||
messagebox.showerror("Erro", f"Falha ao copiar/mover amostra: {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:
|
||||
preview_src = Path(last["preview_src"])
|
||||
meta_src = Path(last["meta_src"])
|
||||
bin_srcs = [Path(p) for p in last["bin_srcs"]]
|
||||
|
||||
preview_dst = Path(last["preview_dst"])
|
||||
meta_dst = Path(last["meta_dst"])
|
||||
bin_dsts = [Path(p) for p in last["bin_dsts"]]
|
||||
|
||||
if last["moved"]:
|
||||
# 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"]))
|
||||
shutil.move(str(preview_dst), str(preview_src))
|
||||
shutil.move(str(meta_dst), str(meta_src))
|
||||
for dst_bin, src_bin in zip(bin_dsts, bin_srcs):
|
||||
shutil.move(str(dst_bin), str(src_bin))
|
||||
else:
|
||||
# apagar os arquivos de destino
|
||||
for p in [last["preview_dst"], last["raw_dst"], last["meta_dst"]]:
|
||||
if p.exists():
|
||||
p.unlink()
|
||||
if preview_dst.exists():
|
||||
preview_dst.unlink()
|
||||
if meta_dst.exists():
|
||||
meta_dst.unlink()
|
||||
for dst_bin in bin_dsts:
|
||||
if dst_bin.exists():
|
||||
dst_bin.unlink()
|
||||
|
||||
self.logger.log(
|
||||
"undo",
|
||||
preview_path=last["preview_src"],
|
||||
raw_path=last["raw_src"],
|
||||
meta_path=last["meta_src"],
|
||||
action="undo",
|
||||
sample_dir=Path(last["sample_src"]),
|
||||
preview_path=preview_src,
|
||||
meta_path=meta_src,
|
||||
bin_paths=bin_srcs,
|
||||
label=last["label"],
|
||||
dest_preview=last["preview_dst"],
|
||||
dest_raw=last["raw_dst"],
|
||||
dest_meta=last["meta_dst"],
|
||||
dest_sample_dir=self.out_root / last["label"],
|
||||
)
|
||||
|
||||
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.idx = max(0, min(last.get("index", self.idx), len(self.all_bundles) - 1))
|
||||
self.status_var.set(f"Desfeito ← '{last['label']}': {preview_src.stem}")
|
||||
self.render()
|
||||
|
||||
except Exception as e:
|
||||
messagebox.showerror("Erro", f"Falha no undo: {e}")
|
||||
self.status_var.set(f"ERRO: {e}")
|
||||
|
|
@ -443,38 +469,41 @@ class RawTripletSorterApp:
|
|||
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}")
|
||||
bundle = self.all_bundles[self.idx]
|
||||
self.logger.log(
|
||||
action="skip",
|
||||
sample_dir=bundle.sample_dir,
|
||||
preview_path=bundle.preview_path,
|
||||
meta_path=bundle.meta_path,
|
||||
bin_paths=bundle.bin_paths,
|
||||
label=None,
|
||||
dest_sample_dir=None,
|
||||
)
|
||||
self.status_var.set(f"Pulada → {bundle.sample_id}")
|
||||
self.idx += 1
|
||||
if self.idx >= len(self.all_triplets):
|
||||
if self.idx >= len(self.all_bundles):
|
||||
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.status_var.set("Voltou uma amostra.")
|
||||
self.render()
|
||||
return
|
||||
|
||||
# b = undo
|
||||
if ch in ("b",):
|
||||
if ch == "b":
|
||||
self.undo()
|
||||
return
|
||||
|
||||
# q / Esc = sair
|
||||
if ch in ("q", "escape"):
|
||||
self.root.destroy()
|
||||
return
|
||||
|
|
@ -484,17 +513,12 @@ class RawTripletSorterApp:
|
|||
|
||||
|
||||
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.title("Configurar - Sample Bundle 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 = tk.LabelFrame(self.root, text="Pastas de entrada (raízes com subpastas de amostras)")
|
||||
frm_in.pack(fill=tk.BOTH, expand=False, padx=10, pady=8)
|
||||
|
||||
self.inputs_listbox = tk.Listbox(frm_in, height=6)
|
||||
|
|
@ -506,20 +530,17 @@ class SetupWindow:
|
|||
tk.Button(btns_in, text="Remover selecionada", command=self.remove_selected).pack(fill=tk.X, pady=2)
|
||||
tk.Button(btns_in, text="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)
|
||||
|
|
@ -534,10 +555,8 @@ class SetupWindow:
|
|||
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",
|
||||
|
|
@ -547,7 +566,7 @@ class SetupWindow:
|
|||
self.result = None
|
||||
|
||||
def add_input(self):
|
||||
p = filedialog.askdirectory(title="Selecione a pasta de sessão (contendo .raw, .png, .json)")
|
||||
p = filedialog.askdirectory(title="Selecione a pasta raiz contendo as amostras")
|
||||
if p:
|
||||
self.inputs_listbox.insert(tk.END, p)
|
||||
|
||||
|
|
@ -600,18 +619,20 @@ def run_with_gui_setup():
|
|||
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:
|
||||
bundles = collect_all_bundles(input_folders)
|
||||
if not bundles:
|
||||
messagebox.showinfo(
|
||||
"Sem amostras",
|
||||
"Nenhuma tripleta encontrada.\nCertifique-se de que existam arquivos .raw, .png, .json.",
|
||||
"Nenhuma amostra válida encontrada.\n"
|
||||
"Cada amostra precisa ter <base>.png, <base>.json e um ou mais <base>_camX.bin.",
|
||||
)
|
||||
return
|
||||
|
||||
app = RawTripletSorterApp(
|
||||
triplets=triplets,
|
||||
app = SampleBundleSorterApp(
|
||||
bundles=bundles,
|
||||
labels=labels,
|
||||
out_root=out_root,
|
||||
move=move,
|
||||
|
|
@ -623,39 +644,37 @@ def run_with_gui_setup():
|
|||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Sortear manualmente amostras (raw + preview + json) em pastas de labels, usando hotkeys."
|
||||
description="Classificar manualmente amostras no novo formato do módulo multiespectral."
|
||||
)
|
||||
parser.add_argument("--inputs", nargs="+", help="Pastas de sessão contendo .raw, .png e .json")
|
||||
parser.add_argument("--inputs", nargs="+", help="Pastas-raiz contendo subpastas de amostras")
|
||||
parser.add_argument("--labels", nargs="+", help="Labels (classes) mapeadas para teclas 1..9/0")
|
||||
parser.add_argument("--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)")
|
||||
parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir a GUI de setup")
|
||||
|
||||
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).")
|
||||
print("ERRO: informe --inputs, --labels e --out-root ou use a GUI.")
|
||||
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.")
|
||||
bundles = collect_all_bundles(input_folders)
|
||||
if not bundles:
|
||||
print("Nenhuma amostra válida encontrada nas pastas informadas.")
|
||||
sys.exit(1)
|
||||
|
||||
app = RawTripletSorterApp(
|
||||
triplets=triplets,
|
||||
app = SampleBundleSorterApp(
|
||||
bundles=bundles,
|
||||
labels=labels,
|
||||
out_root=out_root,
|
||||
move=args.move,
|
||||
|
|
|
|||
|
|
@ -3,10 +3,11 @@ import os
|
|||
import shutil
|
||||
import argparse
|
||||
import csv
|
||||
import re
|
||||
|
||||
# ================= CONFIG =================
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
|
|
@ -14,18 +15,17 @@ 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)
|
||||
# Onde você coloca os previews selecionados manualmente
|
||||
PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews")
|
||||
|
||||
# Destino final, organizado por grupo:
|
||||
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
|
||||
# dataset/original/group/{GRUPO}/{previews,metas,bins,masks}
|
||||
PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group")
|
||||
|
||||
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")
|
||||
EXT_METAS = (".json",)
|
||||
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
|
||||
|
||||
|
||||
# ================= HELPERS =================
|
||||
|
|
@ -51,7 +51,7 @@ def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
|
|||
i += 1
|
||||
|
||||
|
||||
def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None:
|
||||
def buscar_por_base(pasta: str, base: str, exts):
|
||||
"""
|
||||
Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'.
|
||||
Retorna o caminho completo ou None.
|
||||
|
|
@ -60,28 +60,54 @@ def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None:
|
|||
return None
|
||||
|
||||
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:
|
||||
root, ext = os.path.splitext(nome)
|
||||
if root.lower() == base.lower() and ext.lower() in tuple(e.lower() for e in exts):
|
||||
return os.path.join(pasta, nome)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def indexar_brutas_por_raw(root_brutas: str):
|
||||
def buscar_bins_por_base(pasta_bins: str, base: str):
|
||||
"""
|
||||
Varre dataset/brutas recursivamente, olhando pastas 'raws' e montando índice:
|
||||
Procura todos os bins do tipo:
|
||||
<base>_cam0.bin
|
||||
<base>_cam1.bin
|
||||
<base>_cam2.bin
|
||||
Retorna lista ordenada por cam.
|
||||
"""
|
||||
encontrados = []
|
||||
|
||||
if not os.path.isdir(pasta_bins):
|
||||
return encontrados
|
||||
|
||||
for nome in os.listdir(pasta_bins):
|
||||
m = BIN_RE.match(nome)
|
||||
if not m:
|
||||
continue
|
||||
if m.group("base").lower() != base.lower():
|
||||
continue
|
||||
cam = int(m.group("cam"))
|
||||
encontrados.append((cam, os.path.join(pasta_bins, nome)))
|
||||
|
||||
encontrados.sort(key=lambda x: x[0])
|
||||
return [p for _, p in encontrados]
|
||||
|
||||
|
||||
def indexar_brutas_por_meta(root_brutas: str):
|
||||
"""
|
||||
Varre dataset/brutas recursivamente, olhando pastas 'metas' e montando índice:
|
||||
|
||||
base -> {
|
||||
"raw": caminho_raw,
|
||||
"cana": "cana_alta" | "cana_baixa" | ... (se conseguir inferir),
|
||||
"horario": "meio_dia" | "cedo" | ... (se conseguir inferir),
|
||||
"meta": caminho_meta,
|
||||
"bins": [cam0, cam1, ...],
|
||||
"cana": "cana_alta" | "cana_baixa" | ...,
|
||||
"horario": "meio_dia" | "cedo" | ...,
|
||||
"grupo": "<GRUPO>",
|
||||
"group_dir": caminho_da_pasta_do_grupo
|
||||
}
|
||||
|
||||
Estrutura esperada (relativa a root_brutas):
|
||||
cana_alta/meio_dia/group/<GRUPO>/raws/*.raw
|
||||
Estrutura esperada:
|
||||
cana_alta/meio_dia/group/<GRUPO>/{previews,metas,bins,masks}
|
||||
"""
|
||||
index = {}
|
||||
|
||||
|
|
@ -91,59 +117,56 @@ def indexar_brutas_por_raw(root_brutas: str):
|
|||
|
||||
for dirpath, dirnames, filenames in os.walk(root_brutas):
|
||||
base_dir = os.path.basename(dirpath).lower()
|
||||
|
||||
if base_dir != "raws":
|
||||
if base_dir != "metas":
|
||||
continue
|
||||
|
||||
# dirpath = .../cana_x/horario/group/<GRUPO>/raws
|
||||
group_dir = os.path.dirname(dirpath) # .../cana_x/horario/group/<GRUPO>
|
||||
# dirpath = .../cana_x/horario/group/<GRUPO>/metas
|
||||
group_dir = os.path.dirname(dirpath)
|
||||
bins_dir = os.path.join(group_dir, "bins")
|
||||
|
||||
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):
|
||||
if not nome.lower().endswith(EXT_METAS):
|
||||
continue
|
||||
|
||||
base = os.path.splitext(nome)[0]
|
||||
meta_path = os.path.join(dirpath, nome)
|
||||
bin_paths = buscar_bins_por_base(bins_dir, base)
|
||||
|
||||
if base in index:
|
||||
# Conflito (mesmo base em dois lugares) -> loga e mantém o primeiro
|
||||
print(f"[CONFLITO] base repetida em raws: {base}")
|
||||
if not bin_paths:
|
||||
print(f"[SKIP] Meta sem bins correspondentes: {meta_path}")
|
||||
continue
|
||||
|
||||
raw_path = os.path.join(dirpath, nome)
|
||||
if base in index:
|
||||
print(f"[CONFLITO] base repetida em metas: {base}")
|
||||
continue
|
||||
|
||||
index[base] = {
|
||||
"raw": raw_path,
|
||||
"meta": meta_path,
|
||||
"bins": bin_paths,
|
||||
"cana": cana,
|
||||
"horario": horario,
|
||||
"grupo": grupo,
|
||||
"group_dir": group_dir,
|
||||
}
|
||||
|
||||
print(f"[INDEX] Entradas indexadas por RAW: {len(index)}")
|
||||
print(f"[INDEX] Entradas indexadas por META: {len(index)}")
|
||||
return index
|
||||
|
||||
|
||||
|
|
@ -151,8 +174,8 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
|
|||
if not os.path.isdir(PASTA_SELECTED_PREVIEWS):
|
||||
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)
|
||||
# Indexa todas as brutas a partir das METAS + BINS
|
||||
index_samples = indexar_brutas_por_meta(PASTA_BRUTAS_ROOT)
|
||||
|
||||
registros = []
|
||||
total, movidos, ignorados = 0, 0, 0
|
||||
|
|
@ -167,59 +190,54 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
|
|||
continue
|
||||
|
||||
total += 1
|
||||
|
||||
base = os.path.splitext(nome)[0]
|
||||
|
||||
info = index_raws.get(base)
|
||||
info = index_samples.get(base)
|
||||
if info is None:
|
||||
ignorados += 1
|
||||
print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via RAW)")
|
||||
print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via META+BINS)")
|
||||
continue
|
||||
|
||||
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"]
|
||||
meta_src = info["meta"]
|
||||
bins_src = info["bins"]
|
||||
|
||||
# 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}
|
||||
# Destino: dataset/original/group/{GRUPO}/{previews,metas,bins,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_bins_dir = os.path.join(dest_group_root, "bins")
|
||||
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_bins_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()
|
||||
meta_ext = os.path.splitext(meta_src)[1].lower()
|
||||
|
||||
# Gera nome final único com base no preview
|
||||
# Nome-base final único definido a partir do preview
|
||||
dst_preview = nome_disponivel(dest_prev_dir, base, prev_ext_sel)
|
||||
new_base = os.path.splitext(os.path.basename(dst_preview))[0]
|
||||
|
||||
dst_raw = os.path.join(dest_raw_dir, new_base + raw_ext)
|
||||
dst_meta = None
|
||||
dst_mask = None
|
||||
|
||||
if meta_src is not None:
|
||||
meta_ext = os.path.splitext(meta_src)[1].lower()
|
||||
dst_meta = os.path.join(dest_meta_dir, new_base + meta_ext)
|
||||
dst_mask = os.path.join(dest_mask_dir, new_base + ".png") if mask_src else None
|
||||
|
||||
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)
|
||||
dst_bins = []
|
||||
for src_bin in bins_src:
|
||||
m = BIN_RE.match(os.path.basename(src_bin))
|
||||
if not m:
|
||||
continue
|
||||
cam = m.group("cam")
|
||||
dst_bins.append(os.path.join(dest_bins_dir, f"{new_base}_cam{cam}.bin"))
|
||||
|
||||
# Copia/move preview selecionado
|
||||
if copy_only:
|
||||
|
|
@ -227,18 +245,20 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
|
|||
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:
|
||||
# Copia/move meta
|
||||
if copy_only:
|
||||
shutil.copy2(meta_src, dst_meta)
|
||||
else:
|
||||
shutil.move(meta_src, dst_meta)
|
||||
|
||||
# Copia/move bins
|
||||
for src_bin, dst_bin in zip(bins_src, dst_bins):
|
||||
if copy_only:
|
||||
shutil.copy2(src_bin, dst_bin)
|
||||
else:
|
||||
shutil.move(src_bin, dst_bin)
|
||||
|
||||
# Copia/move mask se existir
|
||||
if mask_src is not None:
|
||||
if copy_only:
|
||||
shutil.copy2(mask_src, dst_mask)
|
||||
|
|
@ -253,16 +273,16 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
|
|||
cana,
|
||||
horario,
|
||||
caminho_preview_sel,
|
||||
raw_src,
|
||||
meta_src or "",
|
||||
meta_src,
|
||||
json.dumps(bins_src, ensure_ascii=False),
|
||||
mask_src or "",
|
||||
dst_preview,
|
||||
dst_raw,
|
||||
dst_meta or "",
|
||||
dst_meta,
|
||||
json.dumps(dst_bins, ensure_ascii=False),
|
||||
dst_mask or "",
|
||||
])
|
||||
|
||||
print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario}")
|
||||
print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario} | bins={len(dst_bins)}")
|
||||
|
||||
if manifesto_csv and registros:
|
||||
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
|
||||
|
|
@ -273,12 +293,12 @@ def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | No
|
|||
"cana",
|
||||
"horario",
|
||||
"src_preview_selected",
|
||||
"src_raw",
|
||||
"src_meta",
|
||||
"src_bins_json",
|
||||
"src_mask",
|
||||
"dst_preview",
|
||||
"dst_raw",
|
||||
"dst_meta",
|
||||
"dst_bins_json",
|
||||
"dst_mask",
|
||||
])
|
||||
w.writerows(registros)
|
||||
|
|
@ -294,7 +314,7 @@ def build_cli():
|
|||
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}."
|
||||
"preview + meta + bins (+ mask se existir) para dataset/original/group/{GRUPO}."
|
||||
)
|
||||
)
|
||||
ap.add_argument(
|
||||
|
|
@ -1,76 +1,117 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
_2_create_full_mask_multispec.py
|
||||
-------------------------------------------------
|
||||
Gera máscaras sólidas para grupos de classe única no novo formato
|
||||
multiespectral.
|
||||
|
||||
Estrutura esperada de entrada:
|
||||
dataset/
|
||||
brutas/
|
||||
cana_<estado>/
|
||||
<horario>/
|
||||
group/
|
||||
<grupo>/
|
||||
previews/
|
||||
metas/
|
||||
bins/
|
||||
masks/
|
||||
|
||||
Ou, usando --from-originals:
|
||||
dataset/
|
||||
original/
|
||||
previews/
|
||||
metas/
|
||||
bins/
|
||||
masks/
|
||||
|
||||
Cada amostra é formada por:
|
||||
<base>.png
|
||||
<base>.json
|
||||
<base>_cam0.bin
|
||||
<base>_cam1.bin
|
||||
<base>_cam2.bin # opcional
|
||||
|
||||
O script:
|
||||
1. percorre os previews do grupo
|
||||
2. localiza o meta correspondente
|
||||
3. localiza todos os bins correspondentes
|
||||
4. cria uma máscara sólida com a cor da classe
|
||||
5. opcionalmente copia preview + meta + bins + mask para dataset/original
|
||||
com deduplicação de nome-base
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import csv
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
import re
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from utils import carregar_labelmap_completo # importa da sua utils.py
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo
|
||||
|
||||
# ===================================================
|
||||
# ⚙️ Configurações base vindas do config.json
|
||||
# ===================================================
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
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
|
||||
RAW_WIDTH = RESOLUCAO[0]
|
||||
RAW_HEIGHT = RESOLUCAO[1]
|
||||
|
||||
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)
|
||||
PASTA_FINAL_METAS = os.path.join("dataset", "original", "metas")
|
||||
PASTA_FINAL_BINS = os.path.join("dataset", "original", "bins")
|
||||
|
||||
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
|
||||
|
||||
MASK_EXT_OUT = ".png"
|
||||
FORCAR_SOBRESCRITA_NEW_MASK = False
|
||||
VALIDAR_DIM_MASK_EXISTENTE = True
|
||||
MANIFESTO = "manifest.csv"
|
||||
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
|
||||
|
||||
|
||||
# ===================================================
|
||||
# Utilidades gerais
|
||||
# ===================================================
|
||||
|
||||
def garantir_pasta(p):
|
||||
def garantir_pasta(p: str):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
def criar_mask_solida(dim_h, dim_w, cor_rgb):
|
||||
|
||||
def criar_mask_solida(dim_h: int, dim_w: int, cor_rgb: Tuple[int, int, int]):
|
||||
r, g, b = cor_rgb
|
||||
mask_bgr = np.zeros((dim_h, dim_w, 3), dtype=np.uint8)
|
||||
mask_bgr[:] = (b, g, r) # OpenCV usa BGR
|
||||
mask_bgr[:] = (b, g, r)
|
||||
return mask_bgr
|
||||
|
||||
def caminho_mask_new_para_img(pasta_new_masks, caminho_img):
|
||||
base = os.path.splitext(os.path.basename(caminho_img))[0]
|
||||
|
||||
def caminho_mask_new_para_base(pasta_new_masks: str, base: str):
|
||||
return os.path.join(pasta_new_masks, base + MASK_EXT_OUT)
|
||||
|
||||
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
|
||||
def ler_dim(caminho_img: str):
|
||||
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)
|
||||
|
||||
def salvar_mask_solidaria_por_base(pasta_new_masks: str, base: str, caminho_preview: str,
|
||||
cor_rgb: Tuple[int, int, int], forcar: bool = False,
|
||||
validar_dim: bool = True):
|
||||
(h, w), _ = ler_dim(caminho_preview)
|
||||
caminho_mask = caminho_mask_new_para_base(pasta_new_masks, base)
|
||||
|
||||
if os.path.exists(caminho_mask) and not forcar:
|
||||
if validar_dim:
|
||||
|
|
@ -90,76 +131,130 @@ def salvar_mask_solidaria(pasta_new_masks, caminho_img, cor_rgb, forcar=False, v
|
|||
print(f"[CRIADA] {os.path.basename(caminho_mask)} ({w}x{h})")
|
||||
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 já existe.
|
||||
"""
|
||||
if os.path.exists(os.path.join(dir_img, base + img_ext)):
|
||||
# ===================================================
|
||||
# Descoberta de arquivos da amostra
|
||||
# ===================================================
|
||||
|
||||
def localizar_meta_correspondente(pasta_metas: str, base: str) -> Optional[str]:
|
||||
candidate = os.path.join(pasta_metas, base + ".json")
|
||||
return candidate if os.path.exists(candidate) else None
|
||||
|
||||
|
||||
def localizar_bins_correspondentes(pasta_bins: str, base: str) -> List[str]:
|
||||
encontrados = []
|
||||
if not os.path.isdir(pasta_bins):
|
||||
return encontrados
|
||||
|
||||
for nome in os.listdir(pasta_bins):
|
||||
m = BIN_RE.match(nome)
|
||||
if not m:
|
||||
continue
|
||||
if m.group("base") != base:
|
||||
continue
|
||||
cam = int(m.group("cam"))
|
||||
encontrados.append((cam, os.path.join(pasta_bins, nome)))
|
||||
|
||||
encontrados.sort(key=lambda x: x[0])
|
||||
return [p for _, p in encontrados]
|
||||
|
||||
|
||||
# ===================================================
|
||||
# Deduplicação para cópia ao dataset final
|
||||
# ===================================================
|
||||
|
||||
def _base_esta_ocupado_multispec(base: str, preview_ext: str, dir_preview: str,
|
||||
dir_mask: str, dir_meta: str, dir_bins: str,
|
||||
bin_src_paths: List[str]) -> bool:
|
||||
if os.path.exists(os.path.join(dir_preview, base + preview_ext)):
|
||||
return True
|
||||
if os.path.exists(os.path.join(dir_mask, base + mask_ext)):
|
||||
if os.path.exists(os.path.join(dir_mask, base + ".png")):
|
||||
return True
|
||||
if raw_ext and dir_raw:
|
||||
if os.path.exists(os.path.join(dir_raw, base + raw_ext)):
|
||||
if os.path.exists(os.path.join(dir_meta, base + ".json")):
|
||||
return True
|
||||
|
||||
for src_bin in bin_src_paths:
|
||||
nome = os.path.basename(src_bin)
|
||||
m = BIN_RE.match(nome)
|
||||
if not m:
|
||||
continue
|
||||
cam = m.group("cam")
|
||||
bin_dest = os.path.join(dir_bins, f"{base}_cam{cam}.bin")
|
||||
if os.path.exists(bin_dest):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
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):
|
||||
|
||||
def gerar_base_disponivel_multispec(base: str, preview_ext: str, dir_preview: str,
|
||||
dir_mask: str, dir_meta: str, dir_bins: str,
|
||||
bin_src_paths: List[str]) -> str:
|
||||
if not _base_esta_ocupado_multispec(base, preview_ext, dir_preview, dir_mask, dir_meta, dir_bins, bin_src_paths):
|
||||
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):
|
||||
if not _base_esta_ocupado_multispec(candidate, preview_ext, dir_preview, dir_mask, dir_meta, dir_bins, bin_src_paths):
|
||||
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)
|
||||
def copiar_conjunto_multispec(caminho_preview_src: str, caminho_mask_src: str,
|
||||
caminho_meta_src: str, caminhos_bins_src: List[str],
|
||||
dest_preview_dir: str, dest_mask_dir: str,
|
||||
dest_meta_dir: str, dest_bins_dir: str):
|
||||
garantir_pasta(dest_preview_dir)
|
||||
garantir_pasta(dest_mask_dir)
|
||||
if dest_raw_dir:
|
||||
garantir_pasta(dest_raw_dir)
|
||||
garantir_pasta(dest_meta_dir)
|
||||
garantir_pasta(dest_bins_dir)
|
||||
|
||||
base = os.path.splitext(os.path.basename(caminho_img_src))[0]
|
||||
base = os.path.splitext(os.path.basename(caminho_preview_src))[0]
|
||||
preview_ext = os.path.splitext(caminho_preview_src)[1].lower()
|
||||
|
||||
img_ext = os.path.splitext(caminho_img_src)[1].lower()
|
||||
mask_ext = ".png"
|
||||
raw_ext = os.path.splitext(caminho_raw_src)[1].lower() if caminho_raw_src else None
|
||||
new_base = gerar_base_disponivel_multispec(
|
||||
base,
|
||||
preview_ext,
|
||||
dest_preview_dir,
|
||||
dest_mask_dir,
|
||||
dest_meta_dir,
|
||||
dest_bins_dir,
|
||||
caminhos_bins_src,
|
||||
)
|
||||
|
||||
new_base = gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dest_img_dir, dest_mask_dir, dest_raw_dir)
|
||||
dst_preview_path = os.path.join(dest_preview_dir, new_base + preview_ext)
|
||||
dst_mask_path = os.path.join(dest_mask_dir, new_base + ".png")
|
||||
dst_meta_path = os.path.join(dest_meta_dir, new_base + ".json")
|
||||
|
||||
dst_img_path = os.path.join(dest_img_dir, new_base + img_ext)
|
||||
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_preview_src, dst_preview_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)")
|
||||
shutil.copy2(caminho_meta_src, dst_meta_path)
|
||||
|
||||
return dst_img_path, dst_mask_path, dst_raw_path
|
||||
dst_bins_paths = []
|
||||
for src_bin in caminhos_bins_src:
|
||||
nome = os.path.basename(src_bin)
|
||||
m = BIN_RE.match(nome)
|
||||
if not m:
|
||||
continue
|
||||
cam = m.group("cam")
|
||||
dst_bin = os.path.join(dest_bins_dir, f"{new_base}_cam{cam}.bin")
|
||||
shutil.copy2(src_bin, dst_bin)
|
||||
dst_bins_paths.append(dst_bin)
|
||||
|
||||
print(
|
||||
f"[COPIADO] {os.path.basename(dst_preview_path)} | "
|
||||
f"{os.path.basename(dst_mask_path)} | "
|
||||
f"{os.path.basename(dst_meta_path)} | bins={len(dst_bins_paths)}"
|
||||
)
|
||||
|
||||
return dst_preview_path, dst_mask_path, dst_meta_path, dst_bins_paths
|
||||
|
||||
|
||||
# ===================================================
|
||||
# Labelmap
|
||||
# ===================================================
|
||||
|
||||
def obter_cor_da_classe(nome_classe: str, caminho_labelmap: str):
|
||||
"""
|
||||
Lê 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()
|
||||
|
|
@ -186,31 +281,36 @@ def obter_cor_da_classe(nome_classe: str, caminho_labelmap: str):
|
|||
print(f"[LABELMAP] Classe '{nome_classe}' -> ID={target_id} -> cor RGB={cor_rgb}")
|
||||
return cor_rgb
|
||||
|
||||
def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None:
|
||||
"""
|
||||
Dado o nome do preview (ex: 20260122_091833_277.png),
|
||||
tenta achar o RAW correspondente em PASTA_NEW_RAWS:
|
||||
20260122_091833_277.raw (ou outras extensões de RAW_EXTS)
|
||||
"""
|
||||
base = os.path.splitext(nome_img)[0]
|
||||
for ext in RAW_EXTS:
|
||||
candidate = os.path.join(pasta_new_raws, base + ext)
|
||||
if os.path.exists(candidate):
|
||||
return candidate
|
||||
return None
|
||||
|
||||
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False):
|
||||
# ===================================================
|
||||
# Processamento principal
|
||||
# ===================================================
|
||||
|
||||
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb,
|
||||
fazer_copia_final=True, manifesto_csv=None,
|
||||
orignais=False):
|
||||
source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original")
|
||||
|
||||
if not orignais:
|
||||
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")
|
||||
pasta_new_metas = os.path.join(source, "group", grupo, "metas")
|
||||
pasta_new_bins = os.path.join(source, "group", grupo, "bins")
|
||||
else:
|
||||
pasta_new_previews = os.path.join(source, "previews")
|
||||
pasta_new_masks = os.path.join(source, "masks")
|
||||
pasta_new_metas = os.path.join(source, "metas")
|
||||
pasta_new_bins = os.path.join(source, "bins")
|
||||
|
||||
garantir_pasta(pasta_new_previews)
|
||||
garantir_pasta(pasta_new_masks)
|
||||
garantir_pasta(pasta_new_raws)
|
||||
garantir_pasta(pasta_new_metas)
|
||||
garantir_pasta(pasta_new_bins)
|
||||
|
||||
garantir_pasta(PASTA_FINAL_PREVIEWS)
|
||||
garantir_pasta(PASTA_FINAL_MASKS)
|
||||
garantir_pasta(PASTA_FINAL_RAWS)
|
||||
garantir_pasta(PASTA_FINAL_METAS)
|
||||
garantir_pasta(PASTA_FINAL_BINS)
|
||||
|
||||
registros = []
|
||||
total, criadas_mask, copiados, puladas, erros = 0, 0, 0, 0, 0
|
||||
|
|
@ -218,42 +318,53 @@ def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_fi
|
|||
for nome in os.listdir(pasta_new_previews):
|
||||
if not nome.lower().endswith(EXT_IMAGENS):
|
||||
continue
|
||||
|
||||
total += 1
|
||||
caminho_img = os.path.join(pasta_new_previews, nome)
|
||||
caminho_preview = os.path.join(pasta_new_previews, nome)
|
||||
base = os.path.splitext(nome)[0]
|
||||
|
||||
try:
|
||||
# 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(
|
||||
caminho_meta = localizar_meta_correspondente(pasta_new_metas, base)
|
||||
if caminho_meta is None:
|
||||
raise RuntimeError(f"Meta não encontrado para {nome} em {pasta_new_metas}")
|
||||
|
||||
caminhos_bins = localizar_bins_correspondentes(pasta_new_bins, base)
|
||||
if not caminhos_bins:
|
||||
raise RuntimeError(f"Nenhum bin encontrado para {nome} em {pasta_new_bins}")
|
||||
|
||||
antes = os.path.exists(caminho_mask_new_para_base(pasta_new_masks, base))
|
||||
caminho_mask_new = salvar_mask_solidaria_por_base(
|
||||
pasta_new_masks,
|
||||
caminho_img,
|
||||
base,
|
||||
caminho_preview,
|
||||
cor_classe_rgb,
|
||||
forcar=FORCAR_SOBRESCRITA_NEW_MASK,
|
||||
validar_dim=VALIDAR_DIM_MASK_EXISTENTE
|
||||
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,
|
||||
dst_preview, dst_mask, dst_meta, dst_bins = copiar_conjunto_multispec(
|
||||
caminho_preview,
|
||||
caminho_mask_new,
|
||||
raw_src,
|
||||
caminho_meta,
|
||||
caminhos_bins,
|
||||
PASTA_FINAL_PREVIEWS,
|
||||
PASTA_FINAL_MASKS,
|
||||
PASTA_FINAL_RAWS,
|
||||
PASTA_FINAL_METAS,
|
||||
PASTA_FINAL_BINS,
|
||||
)
|
||||
copiados += 1
|
||||
registros.append([
|
||||
caminho_img,
|
||||
caminho_preview,
|
||||
caminho_mask_new,
|
||||
raw_src or "",
|
||||
dst_img,
|
||||
caminho_meta,
|
||||
json.dumps(caminhos_bins, ensure_ascii=False),
|
||||
dst_preview,
|
||||
dst_mask,
|
||||
dst_raw or "",
|
||||
dst_meta,
|
||||
json.dumps(dst_bins, ensure_ascii=False),
|
||||
])
|
||||
else:
|
||||
puladas += 1
|
||||
|
|
@ -266,44 +377,51 @@ def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_fi
|
|||
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
|
||||
w = csv.writer(f)
|
||||
w.writerow([
|
||||
"src_image",
|
||||
"src_preview",
|
||||
"src_mask",
|
||||
"src_raw",
|
||||
"dst_image",
|
||||
"src_meta",
|
||||
"src_bins_json",
|
||||
"dst_preview",
|
||||
"dst_mask",
|
||||
"dst_raw",
|
||||
"dst_meta",
|
||||
"dst_bins_json",
|
||||
])
|
||||
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}")
|
||||
print(
|
||||
f"\nResumo: total_imgs={total} | masks_criadas={criadas_mask} | "
|
||||
f"copiados={copiados} | puladas={puladas} | erros={erros}"
|
||||
)
|
||||
|
||||
|
||||
# ===================================================
|
||||
# CLI
|
||||
# ===================================================
|
||||
|
||||
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."
|
||||
description="Gera máscaras sólidas para grupos de classe única no novo formato multiespectral."
|
||||
)
|
||||
ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
||||
ap.add_argument("--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).")
|
||||
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")
|
||||
ap.add_argument("--no-copy", action="store_true", help="Não copia para as pastas finais, só cria masks no grupo.")
|
||||
ap.add_argument("--from-originals", action="store_true", help="Executa na pasta dataset/original em vez de dataset/brutas/.../group/... .")
|
||||
ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a 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
|
||||
|
||||
|
|
@ -314,5 +432,5 @@ if __name__ == "__main__":
|
|||
cor_classe_rgb=cor_rgb,
|
||||
fazer_copia_final=fazer_copia,
|
||||
manifesto_csv=manifesto_csv,
|
||||
orignais=args.from_originals
|
||||
orignais=args.from_originals,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,10 +3,11 @@ import os
|
|||
import shutil
|
||||
import argparse
|
||||
import csv
|
||||
import re
|
||||
|
||||
# ================= CONFIG =================
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
|
|
@ -19,58 +20,20 @@ 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_METAS = os.path.join("dataset", "original", "metas")
|
||||
PASTA_FINAL_BINS = os.path.join("dataset", "original", "bins")
|
||||
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
|
||||
|
||||
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
||||
EXT_MASKS = (".png",)
|
||||
EXT_RAWS = (".raw",)
|
||||
EXT_METAS = (".json",)
|
||||
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
|
||||
|
||||
# =========================================
|
||||
|
||||
def garantir_pasta(p: str):
|
||||
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:
|
||||
"""
|
||||
|
|
@ -88,16 +51,90 @@ def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
|
|||
return p
|
||||
i += 1
|
||||
|
||||
|
||||
def buscar_bins_por_base(pasta_bins: str, base: str):
|
||||
encontrados = []
|
||||
|
||||
if not os.path.isdir(pasta_bins):
|
||||
return encontrados
|
||||
|
||||
for nome in os.listdir(pasta_bins):
|
||||
m = BIN_RE.match(nome)
|
||||
if not m:
|
||||
continue
|
||||
if m.group("base").lower() != base.lower():
|
||||
continue
|
||||
cam = int(m.group("cam"))
|
||||
encontrados.append((cam, os.path.join(pasta_bins, nome)))
|
||||
|
||||
encontrados.sort(key=lambda x: x[0])
|
||||
return [p for _, p in encontrados]
|
||||
|
||||
|
||||
def indexar_brutas_group(root_group: str):
|
||||
"""
|
||||
Varre dataset/brutas/<cana>/<horario>/group recursivamente e monta dois índices:
|
||||
base -> caminho_preview
|
||||
base -> {meta, bins}
|
||||
|
||||
Considera a estrutura:
|
||||
brutas/cana_x/horario/group/<grupo>/previews/*.png|jpg
|
||||
brutas/cana_x/horario/group/<grupo>/metas/*.json
|
||||
brutas/cana_x/horario/group/<grupo>/bins/*_camX.bin
|
||||
"""
|
||||
previews_map = {}
|
||||
amostras_map = {}
|
||||
|
||||
if not os.path.isdir(root_group):
|
||||
print(f"[AVISO] Pasta de grupos não existe: {root_group}")
|
||||
return previews_map, amostras_map
|
||||
|
||||
for dirpath, dirnames, filenames in os.walk(root_group):
|
||||
base_dir = os.path.basename(dirpath).lower()
|
||||
parent_group_dir = os.path.dirname(dirpath)
|
||||
|
||||
if base_dir == "previews":
|
||||
for nome in filenames:
|
||||
if not nome.lower().endswith(EXT_PREVIEWS):
|
||||
continue
|
||||
base = os.path.splitext(nome)[0]
|
||||
if base not in previews_map:
|
||||
previews_map[base] = os.path.join(dirpath, nome)
|
||||
|
||||
elif base_dir == "metas":
|
||||
bins_dir = os.path.join(parent_group_dir, "bins")
|
||||
for nome in filenames:
|
||||
if not nome.lower().endswith(EXT_METAS):
|
||||
continue
|
||||
base = os.path.splitext(nome)[0]
|
||||
meta_path = os.path.join(dirpath, nome)
|
||||
bin_paths = buscar_bins_por_base(bins_dir, base)
|
||||
|
||||
if not bin_paths:
|
||||
print(f"[SKIP] Meta sem bins correspondentes: {meta_path}")
|
||||
continue
|
||||
|
||||
if base not in amostras_map:
|
||||
amostras_map[base] = {
|
||||
"meta": meta_path,
|
||||
"bins": bin_paths,
|
||||
}
|
||||
|
||||
print(f"[INDEX] Previews indexados : {len(previews_map)}")
|
||||
print(f"[INDEX] Amostras indexadas : {len(amostras_map)}")
|
||||
return previews_map, amostras_map
|
||||
|
||||
|
||||
def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
|
||||
garantir_pasta(PASTA_FINAL_PREVIEWS)
|
||||
garantir_pasta(PASTA_FINAL_RAWS)
|
||||
garantir_pasta(PASTA_FINAL_METAS)
|
||||
garantir_pasta(PASTA_FINAL_BINS)
|
||||
garantir_pasta(PASTA_FINAL_MASKS)
|
||||
|
||||
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)
|
||||
previews_map, amostras_map = indexar_brutas_group(PASTA_BRUTAS_GROUP_ROOT)
|
||||
|
||||
registros = []
|
||||
total, copiados, ignorados = 0, 0, 0
|
||||
|
|
@ -111,33 +148,45 @@ def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
|
|||
caminho_mask_src = os.path.join(PASTA_NEW_MASKS, nome)
|
||||
|
||||
preview_src = previews_map.get(base)
|
||||
raw_src = raws_map.get(base)
|
||||
amostra = amostras_map.get(base)
|
||||
|
||||
if preview_src is None or raw_src is None:
|
||||
if preview_src is None or amostra is None:
|
||||
ignorados += 1
|
||||
motivo = []
|
||||
if preview_src is None:
|
||||
motivo.append("preview")
|
||||
if raw_src is None:
|
||||
motivo.append("raw")
|
||||
if amostra is None:
|
||||
motivo.append("meta/bins")
|
||||
print(f"[SKIP] {base} -> faltando: {', '.join(motivo)}")
|
||||
# NÃO move a máscara, ela fica em new_masks pra você analisar depois
|
||||
continue
|
||||
|
||||
meta_src = amostra["meta"]
|
||||
bins_src = amostra["bins"]
|
||||
|
||||
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()
|
||||
meta_ext = os.path.splitext(meta_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_meta = os.path.join(PASTA_FINAL_METAS, new_base + meta_ext)
|
||||
dst_mask = os.path.join(PASTA_FINAL_MASKS, new_base + mask_ext)
|
||||
|
||||
# Copia preview + raw, move ou copia a máscara
|
||||
dst_bins = []
|
||||
for src_bin in bins_src:
|
||||
m = BIN_RE.match(os.path.basename(src_bin))
|
||||
if not m:
|
||||
continue
|
||||
cam = m.group("cam")
|
||||
dst_bins.append(os.path.join(PASTA_FINAL_BINS, f"{new_base}_cam{cam}.bin"))
|
||||
|
||||
shutil.copy2(preview_src, dst_preview)
|
||||
shutil.copy2(raw_src, dst_raw)
|
||||
shutil.copy2(meta_src, dst_meta)
|
||||
for src_bin, dst_bin in zip(bins_src, dst_bins):
|
||||
shutil.copy2(src_bin, dst_bin)
|
||||
|
||||
if mover:
|
||||
shutil.move(caminho_mask_src, dst_mask)
|
||||
else:
|
||||
|
|
@ -146,24 +195,28 @@ def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
|
|||
copiados += 1
|
||||
registros.append([
|
||||
preview_src,
|
||||
raw_src,
|
||||
meta_src,
|
||||
json.dumps(bins_src, ensure_ascii=False),
|
||||
caminho_mask_src,
|
||||
dst_preview,
|
||||
dst_raw,
|
||||
dst_meta,
|
||||
json.dumps(dst_bins, ensure_ascii=False),
|
||||
dst_mask,
|
||||
])
|
||||
|
||||
print(f"[OK] {new_base}")
|
||||
print(f"[OK] {new_base} | bins={len(dst_bins)}")
|
||||
|
||||
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_meta",
|
||||
"src_bins_json",
|
||||
"src_mask",
|
||||
"dst_preview",
|
||||
"dst_raw",
|
||||
"dst_meta",
|
||||
"dst_bins_json",
|
||||
"dst_mask",
|
||||
])
|
||||
w.writerows(registros)
|
||||
|
|
@ -171,13 +224,14 @@ def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
|
|||
|
||||
print(f"\nResumo: total_masks={total} | ingestas={copiados} | ignoradas={ignorados}")
|
||||
|
||||
|
||||
# ================= 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."
|
||||
"Ingere máscaras novas (dataset/new_masks) procurando previews, metas e bins "
|
||||
"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.")
|
||||
|
|
@ -186,6 +240,7 @@ def build_cli():
|
|||
ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).")
|
||||
return ap
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
|
|
@ -2,32 +2,41 @@
|
|||
# -*- 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.
|
||||
Agrupa automaticamente as amostras do dataset/original por classes presentes na máscara.
|
||||
|
||||
Entrada:
|
||||
originals/
|
||||
previews/
|
||||
raws/
|
||||
metas/
|
||||
bins/
|
||||
masks/
|
||||
(opcional) masks2/
|
||||
|
||||
Saída:
|
||||
originals/group/<grupo>/previews
|
||||
/raws
|
||||
/metas
|
||||
/bins
|
||||
/masks
|
||||
/masks2
|
||||
|
||||
Cada amostra é:
|
||||
<base>.png (preview)
|
||||
<base>.json (meta)
|
||||
<base>_cam0.bin, <base>_cam1.bin, ... (bins)
|
||||
<base>.png (mask)
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import csv
|
||||
import json
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
import re
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
|
|
@ -38,13 +47,14 @@ USE_MASKS2 = config["dual_head"]
|
|||
EXT_PREVIEWS = (".jpg", ".jpeg", ".png")
|
||||
EXT_MASKS = (".png", ".jpg", ".jpeg")
|
||||
EXT_MASKS2 = (".png", ".jpg", ".jpeg")
|
||||
EXT_RAW = (".raw",)
|
||||
EXT_METAS = (".json",)
|
||||
BIN_RE = re.compile(r"^(?P<base>.+)_cam(?P<cam>\d+)\.bin$", re.IGNORECASE)
|
||||
|
||||
MANIFESTO_DEFAULT = "manifest.csv"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 🔥 EXTRAÇÃO DE CLASSES — VERSÃO ULTRA OTIMIZADA
|
||||
# EXTRAÇÃO DE CLASSES DA MÁSCARA
|
||||
# ============================================================
|
||||
|
||||
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
|
||||
|
|
@ -52,42 +62,30 @@ def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
|
|||
Extração ultra-rápida de IDs presentes na máscara.
|
||||
|
||||
Estratégia:
|
||||
1) Máscara 1 canal → np.unique (instantâneo).
|
||||
1) Máscara 1 canal -> np.unique
|
||||
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.
|
||||
- se assume_rgb=True: converte BGR -> RGB
|
||||
- senão usa BGR direto
|
||||
3) Amostragem em grid + converter_mask_rgb_para_ids
|
||||
4) Fallback full-scan
|
||||
"""
|
||||
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
|
||||
step = 8
|
||||
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)
|
||||
|
|
@ -95,19 +93,19 @@ def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
|
|||
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
|
||||
# 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):
|
||||
|
|
@ -119,6 +117,7 @@ def nome_disponivel(dest_dir, base_name, ext):
|
|||
return cand
|
||||
i += 1
|
||||
|
||||
|
||||
def mapear_por_base_priorizando_png(pasta, exts):
|
||||
if not pasta or not os.path.isdir(pasta):
|
||||
return {}
|
||||
|
|
@ -137,6 +136,42 @@ def mapear_por_base_priorizando_png(pasta, exts):
|
|||
mapa[base] = cam
|
||||
return mapa
|
||||
|
||||
|
||||
def mapear_metas_por_base(pasta, exts):
|
||||
if not pasta or not os.path.isdir(pasta):
|
||||
return {}
|
||||
return {
|
||||
os.path.splitext(nome)[0]: os.path.join(pasta, nome)
|
||||
for nome in os.listdir(pasta)
|
||||
if nome.lower().endswith(exts)
|
||||
}
|
||||
|
||||
|
||||
def mapear_bins_por_base(pasta_bins):
|
||||
"""
|
||||
Retorna:
|
||||
{
|
||||
base: [path_cam0, path_cam1, ...]
|
||||
}
|
||||
"""
|
||||
mapa = {}
|
||||
if not pasta_bins or not os.path.isdir(pasta_bins):
|
||||
return mapa
|
||||
|
||||
for nome in os.listdir(pasta_bins):
|
||||
m = BIN_RE.match(nome)
|
||||
if not m:
|
||||
continue
|
||||
base = m.group("base")
|
||||
cam = int(m.group("cam"))
|
||||
mapa.setdefault(base, []).append((cam, os.path.join(pasta_bins, nome)))
|
||||
|
||||
for base in list(mapa.keys()):
|
||||
mapa[base] = [p for _, p in sorted(mapa[base], key=lambda x: x[0])]
|
||||
|
||||
return mapa
|
||||
|
||||
|
||||
def inferir_ignore_id(ignore_rgb, cor_para_id):
|
||||
if ignore_rgb is None:
|
||||
return None
|
||||
|
|
@ -146,11 +181,8 @@ def inferir_ignore_id(ignore_rgb, cor_para_id):
|
|||
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)]
|
||||
|
|
@ -158,54 +190,64 @@ def montar_nome_grupo(ids_presentes, id_para_nome):
|
|||
|
||||
|
||||
# ============================================================
|
||||
# 📦 PIPELINE PRINCIPAL
|
||||
# CÓPIA / MOVIMENTAÇÃO DO CONJUNTO MULTISPEC
|
||||
# ============================================================
|
||||
|
||||
def copiar_ou_mover_tripla(preview_src, raw_src, mask_src,
|
||||
dest_prev_dir, dest_raw_dir, dest_mask_dir,
|
||||
def copiar_ou_mover_conjunto(preview_src, meta_src, bins_src, mask_src,
|
||||
dest_prev_dir, dest_meta_dir, dest_bins_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_meta_dir)
|
||||
garantir_pasta(dest_bins_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()
|
||||
meta_ext = os.path.splitext(meta_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_meta = os.path.join(dest_meta_dir, new_base + meta_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)
|
||||
dst_bins = []
|
||||
for src_bin in bins_src:
|
||||
m = BIN_RE.match(os.path.basename(src_bin))
|
||||
if not m:
|
||||
continue
|
||||
cam = m.group("cam")
|
||||
dst_bins.append(os.path.join(dest_bins_dir, f"{new_base}_cam{cam}.bin"))
|
||||
|
||||
if mover:
|
||||
shutil.move(preview_src, dst_prev)
|
||||
shutil.move(raw_src, dst_raw)
|
||||
shutil.move(meta_src, dst_meta)
|
||||
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)
|
||||
for src_bin, dst_bin in zip(bins_src, dst_bins):
|
||||
shutil.move(src_bin, dst_bin)
|
||||
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]))
|
||||
mask2_ext = os.path.splitext(mask2_src)[1].lower()
|
||||
shutil.move(mask2_src, os.path.join(dest_mask2_dir, new_base + mask2_ext))
|
||||
else:
|
||||
shutil.copy2(preview_src, dst_prev)
|
||||
shutil.copy2(meta_src, dst_meta)
|
||||
shutil.copy2(mask_src, dst_mask)
|
||||
for src_bin, dst_bin in zip(bins_src, dst_bins):
|
||||
shutil.copy2(src_bin, dst_bin)
|
||||
if mask2_src:
|
||||
garantir_pasta(dest_mask2_dir)
|
||||
mask2_ext = os.path.splitext(mask2_src)[1].lower()
|
||||
shutil.copy2(mask2_src, os.path.join(dest_mask2_dir, new_base + mask2_ext))
|
||||
|
||||
return dst_prev, dst_raw, dst_mask, None
|
||||
return dst_prev, dst_meta, dst_bins, dst_mask
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 🚀 PROCESSAMENTO
|
||||
# PIPELINE PRINCIPAL
|
||||
# ============================================================
|
||||
|
||||
def processar(originals_dir, labelmap_path, mover=False,
|
||||
|
|
@ -214,12 +256,13 @@ def processar(originals_dir, labelmap_path, mover=False,
|
|||
|
||||
originals_dir = os.path.abspath(originals_dir)
|
||||
previews_dir = os.path.join(originals_dir, "previews")
|
||||
raws_dir = os.path.join(originals_dir, "raws")
|
||||
metas_dir = os.path.join(originals_dir, "metas")
|
||||
bins_dir = os.path.join(originals_dir, "bins")
|
||||
masks_dir = os.path.join(originals_dir, "masks")
|
||||
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):
|
||||
if not os.path.isdir(previews_dir) or not os.path.isdir(metas_dir) or not os.path.isdir(masks_dir) or not os.path.isdir(bins_dir):
|
||||
raise RuntimeError("Estrutura inválida em originals/")
|
||||
|
||||
usar_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir)
|
||||
|
|
@ -231,8 +274,8 @@ def processar(originals_dir, labelmap_path, mover=False,
|
|||
|
||||
mapa_previews = mapear_por_base_priorizando_png(previews_dir, EXT_PREVIEWS)
|
||||
mapa_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_metas = mapear_metas_por_base(metas_dir, EXT_METAS)
|
||||
mapa_bins = mapear_bins_por_base(bins_dir)
|
||||
mapa_masks2 = mapear_por_base_priorizando_png(masks2_dir, EXT_MASKS2) if usar_masks2 else {}
|
||||
|
||||
registros = []
|
||||
|
|
@ -242,15 +285,13 @@ def processar(originals_dir, labelmap_path, mover=False,
|
|||
totais["total_masks"] += 1
|
||||
|
||||
prev_path = mapa_previews.get(base)
|
||||
raw_path = mapa_raws.get(base)
|
||||
meta_path = mapa_metas.get(base)
|
||||
bins_paths = mapa_bins.get(base)
|
||||
|
||||
if not prev_path or not raw_path:
|
||||
if not prev_path or not meta_path or not bins_paths:
|
||||
totais["pulados"] += 1
|
||||
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:
|
||||
|
|
@ -258,32 +299,47 @@ def processar(originals_dir, labelmap_path, mover=False,
|
|||
|
||||
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_meta_dir = os.path.join(group_dir, grupo, "metas")
|
||||
dest_bins_dir = os.path.join(group_dir, grupo, "bins")
|
||||
dest_mask_dir = os.path.join(group_dir, grupo, "masks")
|
||||
dest_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,
|
||||
dst_prev, dst_meta, dst_bins, dst_mask = copiar_ou_mover_conjunto(
|
||||
prev_path, meta_path, bins_paths, mask_path,
|
||||
dest_prev_dir, dest_meta_dir, dest_bins_dir, dest_mask_dir,
|
||||
mover=mover,
|
||||
mask2_src=mapa_masks2.get(base),
|
||||
dest_mask2_dir=dest_mask2_dir
|
||||
dest_mask2_dir=dest_mask2_dir,
|
||||
)
|
||||
|
||||
registros.append([grupo, prev_path, raw_path, mask_path, dst_prev, dst_raw, dst_mask])
|
||||
registros.append([
|
||||
grupo,
|
||||
prev_path,
|
||||
meta_path,
|
||||
json.dumps(bins_paths, ensure_ascii=False),
|
||||
mask_path,
|
||||
dst_prev,
|
||||
dst_meta,
|
||||
json.dumps(dst_bins, ensure_ascii=False),
|
||||
dst_mask,
|
||||
])
|
||||
totais["processados"] += 1
|
||||
|
||||
# ------------------------------
|
||||
# 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"])
|
||||
w.writerow([
|
||||
"grupo",
|
||||
"src_preview",
|
||||
"src_meta",
|
||||
"src_bins_json",
|
||||
"src_mask",
|
||||
"dst_preview",
|
||||
"dst_meta",
|
||||
"dst_bins_json",
|
||||
"dst_mask",
|
||||
])
|
||||
for r in registros:
|
||||
w.writerow(r)
|
||||
|
||||
|
|
@ -306,6 +362,7 @@ def build_cli():
|
|||
ap.add_argument("--labels-bgr", action="store_true")
|
||||
return ap
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
|
|
@ -318,5 +375,5 @@ if __name__ == "__main__":
|
|||
manifesto=args.manifest,
|
||||
validar_dim=not args.no_validate,
|
||||
estrito=args.strict,
|
||||
labelmap_bgr=args.labels_bgr
|
||||
labelmap_bgr=args.labels_bgr,
|
||||
)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -18,308 +18,187 @@ RAW:
|
|||
- Detecta dtype (uint8/uint16) pelo tamanho do arquivo
|
||||
- Carrega como (H,W), redimensiona, salva em .raw
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
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
|
||||
from pi.raw_processor_core import RawProcessorCore
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids, _infer_ignore_id
|
||||
|
||||
# ===== config =====
|
||||
# ===== 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")
|
||||
RES = tuple(config["resolucao"]) # (W,H)
|
||||
pasta_base = "dataset"
|
||||
|
||||
INPUTS = [
|
||||
("original", os.path.join(pasta_base, "original", "group")),
|
||||
("augmented", os.path.join(pasta_base, "augmented", "group")),
|
||||
]
|
||||
|
||||
OUTPUT_ROOT = os.path.join(pasta_base, f"{RES[0]}x{RES[1]}", "group")
|
||||
|
||||
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
|
||||
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"]
|
||||
# ===== GLOBAL STATS =====
|
||||
GLOBAL_SUM = None
|
||||
GLOBAL_SUMSQ = None
|
||||
GLOBAL_PIXELS = 0
|
||||
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXTS = (".png", ".jpg", ".jpeg")
|
||||
MSK2_EXTS = (".png", ".jpg", ".jpeg")
|
||||
RAW_EXTS = (".raw",)
|
||||
|
||||
# === Acumuladores globais para mean/std dos canais RAW4 ===
|
||||
GLOBAL_SUM = None # soma por canal
|
||||
GLOBAL_SUMSQ = None # soma dos quadrados por canal
|
||||
GLOBAL_PIXELS = 0 # n de pixels por canal (H*W por imagem)
|
||||
|
||||
def garantir_dir(p):
|
||||
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 list_groups(root):
|
||||
return [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))]
|
||||
|
||||
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)
|
||||
def extract_bins_for_base(bins_dir, base):
|
||||
return [
|
||||
os.path.join(bins_dir, f)
|
||||
for f in os.listdir(bins_dir)
|
||||
if f.startswith(base) and f.endswith(".bin")
|
||||
]
|
||||
|
||||
num_floats = fsize // 4
|
||||
if num_floats != 4 * npx:
|
||||
raise RuntimeError(
|
||||
f"RAW {path}: esperado 4 canais float32, "
|
||||
f"mas num_floats={num_floats}, H*W={npx}"
|
||||
)
|
||||
|
||||
data = np.fromfile(path, dtype=np.float32)
|
||||
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:
|
||||
def process():
|
||||
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
|
||||
grupos = list_groups_raw(fonte_root)
|
||||
if not grupos:
|
||||
return 0
|
||||
|
||||
not_want = {g.strip() for g in groups_except.split(",") if g.strip()}
|
||||
cor_para_id, _, _, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
ignore_id = _infer_ignore_id(ignore_rgb, 255)
|
||||
|
||||
core = RawProcessorCore(640, 480)
|
||||
|
||||
total = 0
|
||||
|
||||
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}")
|
||||
for fonte_nome, fonte_root in INPUTS:
|
||||
|
||||
if not os.path.isdir(fonte_root):
|
||||
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")
|
||||
for grupo in list_groups(fonte_root):
|
||||
|
||||
if not (os.path.isdir(in_prev) and os.path.isdir(in_msk)):
|
||||
print(f"[WARN] Grupo inválido (sem previews/masks): {grupo}")
|
||||
print(f"\n[{fonte_nome}] Grupo: {grupo}")
|
||||
|
||||
gpath = os.path.join(fonte_root, grupo)
|
||||
|
||||
bins_dir = os.path.join(gpath, "bins")
|
||||
meta_dir = os.path.join(gpath, "metas")
|
||||
mask_dir = os.path.join(gpath, "masks")
|
||||
|
||||
if not os.path.isdir(meta_dir):
|
||||
continue
|
||||
|
||||
usar_raw = os.path.isdir(in_raw)
|
||||
usar_msk2 = USE_MASKS2 and os.path.isdir(in_msk2)
|
||||
out_dir = os.path.join(OUTPUT_ROOT, grupo)
|
||||
out_tensor = os.path.join(out_dir, "tensors")
|
||||
out_mask = os.path.join(out_dir, "masks")
|
||||
|
||||
out_prev = os.path.join(out_root, grupo, "previews")
|
||||
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_tensor)
|
||||
garantir_dir(out_mask)
|
||||
|
||||
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)
|
||||
for fname in os.listdir(meta_dir):
|
||||
|
||||
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 {}
|
||||
base = fname.replace(".json", "")
|
||||
meta_path = os.path.join(meta_dir, fname)
|
||||
|
||||
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
|
||||
with open(meta_path) as f:
|
||||
meta = json.load(f)
|
||||
|
||||
if not msk_path:
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] Sem máscara p/ {fname}, pulando.")
|
||||
bins_paths = extract_bins_for_base(bins_dir, base)
|
||||
|
||||
if not bins_paths:
|
||||
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)
|
||||
# ===== LOAD BINS =====
|
||||
bins_data = []
|
||||
bins_meta = []
|
||||
|
||||
# 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
|
||||
for path in bins_paths:
|
||||
cam_id = os.path.basename(path).split("_")[-1].replace(".bin", "")
|
||||
|
||||
cv2.imwrite(os.path.join(out_prev, out_name_prev), prev_res)
|
||||
cam_meta = core.extract_camera_meta(meta, cam_id)
|
||||
data = core.load_native_bin(path, cam_meta)
|
||||
|
||||
# --- mask ids ---
|
||||
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)
|
||||
bins_data.append(data)
|
||||
bins_meta.append(cam_meta)
|
||||
|
||||
# --- 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}")
|
||||
# ===== BUILD TENSOR =====
|
||||
tensor, channel_names = core.build_multispectral_tensor(bins_data, bins_meta)
|
||||
|
||||
# --- 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
|
||||
# ===== RESIZE =====
|
||||
chans = []
|
||||
for k in range(raw4.shape[0]):
|
||||
ch = raw4[k]
|
||||
ch_res = cv2.resize(ch, (out_w, out_h), interpolation=cv2.INTER_AREA)
|
||||
for ch in tensor:
|
||||
ch_res = cv2.resize(ch, RES, 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
|
||||
tensor = np.stack(chans, axis=0)
|
||||
|
||||
# ===== STATS =====
|
||||
c, h, w = tensor.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)
|
||||
flat = tensor.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
|
||||
GLOBAL_PIXELS += 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).")
|
||||
# ===== SAVE TENSOR =====
|
||||
name = f"{fonte_nome}_{base}"
|
||||
|
||||
np.save(os.path.join(out_tensor, name + ".npy"), tensor)
|
||||
|
||||
# ===== MASK =====
|
||||
mask_path = os.path.join(mask_dir, base + ".png")
|
||||
|
||||
if os.path.exists(mask_path):
|
||||
msk = cv2.imread(mask_path)
|
||||
msk = cv2.cvtColor(msk, cv2.COLOR_BGR2RGB)
|
||||
|
||||
ids = converter_mask_rgb_para_ids(msk, cor_para_id, ignore_id)
|
||||
ids_res = cv2.resize(ids, RES, interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
np.save(os.path.join(out_mask, name + ".npy"), ids_res)
|
||||
|
||||
# debug visual
|
||||
cv2.imwrite(
|
||||
os.path.join(out_mask, name + ".png"),
|
||||
ids_res
|
||||
)
|
||||
|
||||
total += 1
|
||||
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n} → {fname}")
|
||||
print(f"OK: {name}")
|
||||
|
||||
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)
|
||||
# ===== SAVE STATS =====
|
||||
if GLOBAL_SUM is not None:
|
||||
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),
|
||||
"channels": channel_names,
|
||||
"mean": mean.tolist(),
|
||||
"std": std.tolist(),
|
||||
}
|
||||
|
||||
save_path = os.path.join("backup", config["modelo"], config["model_name"], config["stats_source_tag"])
|
||||
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.")
|
||||
with open(os.path.join(save_path, "norm_stats.json"), "w") as f:
|
||||
json.dump(stats, f, indent=2)
|
||||
|
||||
print("\n📊 STATS:")
|
||||
print(stats)
|
||||
|
||||
print(f"\n✅ FINALIZADO: {total} samples")
|
||||
|
||||
|
||||
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)
|
||||
process()
|
||||
|
|
@ -1,23 +1,6 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Split estratificado por GRUPO com **val/test só do ORIGINAL** e
|
||||
garantia de NÃO VAZAMENTO entre splits, agora para estrutura RAW:
|
||||
|
||||
Lê 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: só **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
|
||||
|
|
@ -25,25 +8,20 @@ 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)
|
||||
RESOLUCAO = tuple(config.get("resolucao"))
|
||||
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"
|
||||
TENSOR_EXT = ".npy"
|
||||
MASK_NPY_SUFFIX = ".npy"
|
||||
|
||||
# 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)
|
||||
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):
|
||||
|
|
@ -51,7 +29,6 @@ def garantir(p):
|
|||
|
||||
|
||||
def lista_grupos(root):
|
||||
"""Lista grupos que têm previews + masks (raws opcional)."""
|
||||
if not os.path.isdir(root):
|
||||
return []
|
||||
out = []
|
||||
|
|
@ -59,43 +36,27 @@ def lista_grupos(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")):
|
||||
if os.path.isdir(os.path.join(gdir, "tensors")) 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):
|
||||
def listar_tensors(tensor_dir):
|
||||
if not os.path.isdir(tensor_dir):
|
||||
return []
|
||||
fs = []
|
||||
for f in os.listdir(prev_dir):
|
||||
ext = os.path.splitext(f.lower())[1]
|
||||
if ext in IMG_EXTS:
|
||||
for f in os.listdir(tensor_dir):
|
||||
if f.lower().endswith(TENSOR_EXT):
|
||||
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 mask_npy_from_tensor_name(tensor_name):
|
||||
base, _ = os.path.splitext(tensor_name)
|
||||
return base + MASK_NPY_SUFFIX
|
||||
|
||||
|
||||
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)
|
||||
|
|
@ -104,7 +65,6 @@ def classify_source_and_family(filename_no_ext):
|
|||
if m:
|
||||
return "augmented", m.group(1)
|
||||
|
||||
# legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
|
||||
if RE_AUG_SUFFIX.search(filename_no_ext):
|
||||
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
|
||||
return "augmented", fam
|
||||
|
|
@ -112,40 +72,43 @@ def classify_source_and_family(filename_no_ext):
|
|||
return "unknown", filename_no_ext
|
||||
|
||||
|
||||
def build_family_index(prev_dir, msk_dir):
|
||||
def build_family_index(tensor_dir, mask_dir):
|
||||
"""
|
||||
Constroi índice de famílias a partir de previews/masks.
|
||||
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)
|
||||
family -> {
|
||||
"original": tensor_name or None,
|
||||
"augmented": [tensor_name, ...],
|
||||
"all": [...]
|
||||
}
|
||||
Só indexa se houver pelo menos mask .npy correspondente.
|
||||
"""
|
||||
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
|
||||
tensors = listar_tensors(tensor_dir)
|
||||
|
||||
for tensor_name in tensors:
|
||||
base_no_ext, _ = os.path.splitext(tensor_name)
|
||||
mask_npy_name = mask_npy_from_tensor_name(tensor_name)
|
||||
|
||||
if not os.path.exists(os.path.join(mask_dir, mask_npy_name)):
|
||||
continue
|
||||
|
||||
source, fam = classify_source_and_family(base_no_ext)
|
||||
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
|
||||
d["all"].append(img_name)
|
||||
d["all"].append(tensor_name)
|
||||
|
||||
if source == "original":
|
||||
d["original"] = img_name
|
||||
d["original"] = tensor_name
|
||||
elif source == "augmented":
|
||||
d["augmented"].append(img_name)
|
||||
d["augmented"].append(tensor_name)
|
||||
else:
|
||||
# trata como original desconhecido para não perder dado
|
||||
if d["original"] is None:
|
||||
d["original"] = img_name
|
||||
d["original"] = tensor_name
|
||||
else:
|
||||
d["augmented"].append(img_name)
|
||||
d["augmented"].append(tensor_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
|
||||
|
|
@ -178,6 +141,7 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
else:
|
||||
break
|
||||
total = n_train + n_val + n_test
|
||||
|
||||
while total < n:
|
||||
if n_train - min_train <= n_val - min_val:
|
||||
n_train += 1
|
||||
|
|
@ -190,11 +154,9 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
n_val = max(0, min(resto, min_val))
|
||||
n_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
|
||||
|
|
@ -202,7 +164,6 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|||
n_val += diff
|
||||
else:
|
||||
diff = -diff
|
||||
# tira de test, depois val
|
||||
take = min(diff, n_test)
|
||||
n_test -= take
|
||||
diff -= take
|
||||
|
|
@ -214,71 +175,47 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_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,
|
||||
src_tensor_dir,
|
||||
src_mask_dir,
|
||||
dst_tensor_dir,
|
||||
dst_mask_dir,
|
||||
):
|
||||
"""
|
||||
Copia preview + mask (+ mask2) (+ raw) para o split.
|
||||
Copia:
|
||||
- tensor .npy
|
||||
- mask .npy obrigatória
|
||||
- mask .png opcional (debug)
|
||||
"""
|
||||
garantir(dst_prev_dir)
|
||||
garantir(dst_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)
|
||||
garantir(dst_tensor_dir)
|
||||
garantir(dst_mask_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)
|
||||
tensor_src = os.path.join(src_tensor_dir, nome)
|
||||
mask_npy_name = mask_npy_from_tensor_name(nome)
|
||||
|
||||
if not (os.path.exists(src_prev) and os.path.exists(src_msk)):
|
||||
mask_npy_src = os.path.join(src_mask_dir, mask_npy_name)
|
||||
|
||||
if not (os.path.exists(tensor_src) and os.path.exists(mask_npy_src)):
|
||||
continue
|
||||
|
||||
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))
|
||||
shutil.copy2(tensor_src, os.path.join(dst_tensor_dir, nome))
|
||||
shutil.copy2(mask_npy_src, os.path.join(dst_mask_dir, mask_npy_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")
|
||||
src_tensor_dir = os.path.join(pasta_origem, group_name, "tensors")
|
||||
src_mask_dir = os.path.join(pasta_origem, group_name, "masks")
|
||||
|
||||
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
|
||||
use_raw = os.path.isdir(src_raw_dir)
|
||||
familias = build_family_index(src_tensor_dir, src_mask_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}
|
||||
|
|
@ -295,7 +232,6 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
|
|||
fam_val = set(familias_originais[n_tr:n_tr+n_va])
|
||||
fam_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:
|
||||
|
|
@ -320,66 +256,40 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
|
|||
elif fam in fam_test:
|
||||
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")
|
||||
dst_train_tensor = os.path.join(pasta_destino, "train", "group", group_name, "tensors")
|
||||
dst_train_mask = os.path.join(pasta_destino, "train", "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
|
||||
dst_val_tensor = os.path.join(pasta_destino, "val", "group", group_name, "tensors")
|
||||
dst_val_mask = os.path.join(pasta_destino, "val", "group", group_name, "masks")
|
||||
|
||||
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
|
||||
dst_test_tensor = os.path.join(pasta_destino, "test", "group", group_name, "tensors")
|
||||
dst_test_mask = os.path.join(pasta_destino, "test", "group", group_name, "masks")
|
||||
|
||||
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
|
||||
)
|
||||
m_train = copiar(nomes_train, src_tensor_dir, src_mask_dir, dst_train_tensor, dst_train_mask)
|
||||
m_val = copiar(nomes_val, src_tensor_dir, src_mask_dir, dst_val_tensor, dst_val_mask)
|
||||
m_test = copiar(nomes_test, src_tensor_dir, src_mask_dir, dst_test_tensor, dst_test_mask)
|
||||
|
||||
print(f"[{group_name}] famílias={total_familias} → train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}")
|
||||
print(f"[{group_name}] famílias={total_familias} → train(tensors)={m_train}, val(tensors)={m_val}, test(tensors)={m_test}")
|
||||
return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}
|
||||
|
||||
|
||||
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 = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (tensors/masks).")
|
||||
ap.add_argument("--train", type=float, default=0.70)
|
||||
ap.add_argument("--val", type=float, default=0.29)
|
||||
ap.add_argument("--test", type=float, default=0.01)
|
||||
ap.add_argument("--seed", type=int, default=42)
|
||||
|
||||
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
|
||||
ap.add_argument("--min-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("--min-train", type=int, default=1)
|
||||
ap.add_argument("--min-val", type=int, default=1)
|
||||
ap.add_argument("--min-test", type=int, default=0)
|
||||
|
||||
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")
|
||||
ap.add_argument("--resolucao", type=str, default=None,
|
||||
help="Sobrescreve resolução no formato WxH (ex: 1024x800).")
|
||||
|
||||
ap.add_argument("--cap-train-families", type=str, default="",
|
||||
help="Mapa 'grupo:cap,...' p/ limitar número de FAMÍLIAS no TRAIN. Ex.: 'chao:350'")
|
||||
help="Mapa 'grupo:cap,...' para limitar famílias no TRAIN. Ex.: 'chao:350'")
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
|
|
@ -410,6 +320,7 @@ def main():
|
|||
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
|
||||
|
|
@ -439,12 +350,12 @@ def main():
|
|||
for k in total_global.keys():
|
||||
total_global[k] += res.get(k, 0)
|
||||
|
||||
print("\nResumo global (imagens copiadas):")
|
||||
print("\nResumo global (tensors copiados):")
|
||||
print(f" train: {total_global['train']}")
|
||||
print(f" 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!")
|
||||
print("\n✅ Split sem vazamento concluído!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -16,7 +16,14 @@ 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)
|
||||
from multispec_segformer_service import (
|
||||
MultispecSegDataset,
|
||||
normalize_per_batch,
|
||||
patch_segformer_input_channels,
|
||||
build_multispec_segformer_model,
|
||||
build_dual_branch_segformer_model,
|
||||
build_fixed_normalizer,
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------
|
||||
|
|
@ -208,7 +215,7 @@ def run_one_epoch(
|
|||
if normalizer is not None:
|
||||
imgs = normalizer(imgs)
|
||||
else:
|
||||
imgs = normalize_raw(imgs)
|
||||
imgs = normalize_per_batch(imgs)
|
||||
|
||||
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
|
||||
out = model(pixel_values=imgs)
|
||||
|
|
@ -374,15 +381,11 @@ def main():
|
|||
# ==========================
|
||||
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):
|
||||
if norm_stats_path and os.path.exists(norm_stats_path):
|
||||
with open(norm_stats_path, "r", encoding="utf-8") as f:
|
||||
norm_stats = json.load(f)
|
||||
|
||||
|
|
@ -390,89 +393,22 @@ def main():
|
|||
stats_mean = norm_stats.get("mean", [])
|
||||
stats_std = norm_stats.get("std", [])
|
||||
|
||||
print(f"[NORM] usando stats fixos de: {norm_stats_path}")
|
||||
if len(stats_mean) != CHANNELS or len(stats_std) != CHANNELS:
|
||||
raise RuntimeError(
|
||||
f"norm_stats incompatível com channels={CHANNELS}: "
|
||||
f"mean={len(stats_mean)} std={len(stats_std)}"
|
||||
)
|
||||
|
||||
if len(stats_channels) != CHANNELS:
|
||||
raise RuntimeError(
|
||||
f"norm_stats.channels incompatível com channels={CHANNELS}: {stats_channels}"
|
||||
)
|
||||
|
||||
normalizer = build_fixed_normalizer(stats_mean, stats_std, device)
|
||||
print(f"[NORM] usando stats fixos: {norm_stats_path}")
|
||||
print(f"[NORM] channels={stats_channels}")
|
||||
print(f"[NORM] 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.")
|
||||
print("[NORM] sem stats fixos, usando normalize_per_batch.")
|
||||
|
||||
resize_hw = None
|
||||
if args.resize_h is not None and args.resize_w is not None:
|
||||
|
|
@ -483,23 +419,20 @@ def main():
|
|||
pass
|
||||
|
||||
# Datasets (RAW com 4 ou 5 canais)
|
||||
ds_train = RawSegDataset(
|
||||
ds_train = MultispecSegDataset(
|
||||
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,
|
||||
strict_channels=True,
|
||||
)
|
||||
ds_val = RawSegDataset(
|
||||
|
||||
ds_val = MultispecSegDataset(
|
||||
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,
|
||||
strict_channels=True,
|
||||
)
|
||||
|
||||
# classes
|
||||
|
|
@ -538,7 +471,7 @@ def main():
|
|||
)
|
||||
|
||||
if FUSION_MODE == "stacked":
|
||||
model = build_raw_segformer_model(
|
||||
model = build_multispec_segformer_model(
|
||||
num_classes=num_classes,
|
||||
channels=CHANNELS,
|
||||
backbone=BACKBONE,
|
||||
|
|
@ -0,0 +1,413 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Teste/visualização do SegFormer para o novo pipeline multiespectral.
|
||||
|
||||
Modos:
|
||||
1) Dataset (com GT):
|
||||
- Usa MultispecSegDataset a partir de split (train/val/test)
|
||||
- Mostra: [preview tensor] | [mask GT] | [overlay predito]
|
||||
|
||||
2) Pasta sem máscara (--test_folder sem masks):
|
||||
- Faz inferência só com tensors .npy (sem GT)
|
||||
- Mostra: [preview tensor] | [overlay predito]
|
||||
|
||||
3) Câmera ao vivo (--camera):
|
||||
- Mantido como placeholder para integração futura com FrameService/MULTISPEC.
|
||||
- Neste script, o foco principal é dataset + tensors finais.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import argparse
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import converter_mask_ids_para_bgr, desenhar_legenda_horizontal
|
||||
from multispec_segformer_service import (
|
||||
MultispecSegformerService,
|
||||
MultispecSegDataset,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Helpers para modo sem máscara
|
||||
# ============================================================
|
||||
|
||||
def _collect_tensor_paths(root: str) -> list[str]:
|
||||
"""
|
||||
Coleta caminhos de tensors .npy em 'root' usando a convenção nova:
|
||||
- se root/group existe: procura group/*/tensors
|
||||
- senão: procura root/tensors
|
||||
"""
|
||||
paths = []
|
||||
|
||||
dir_group = os.path.join(root, "group")
|
||||
if os.path.isdir(dir_group):
|
||||
for g in sorted(os.listdir(dir_group)):
|
||||
gdir = os.path.join(dir_group, g)
|
||||
if not os.path.isdir(gdir):
|
||||
continue
|
||||
g_tensors = os.path.join(gdir, "tensors")
|
||||
if not os.path.isdir(g_tensors):
|
||||
continue
|
||||
for fn in sorted(os.listdir(g_tensors)):
|
||||
if fn.lower().endswith(".npy"):
|
||||
paths.append(os.path.join(g_tensors, fn))
|
||||
else:
|
||||
dir_tensors = os.path.join(root, "tensors")
|
||||
if not os.path.isdir(dir_tensors):
|
||||
raise RuntimeError(f"Modo sem máscara: não achei pasta 'tensors' em {root}")
|
||||
for fn in sorted(os.listdir(dir_tensors)):
|
||||
if fn.lower().endswith(".npy"):
|
||||
paths.append(os.path.join(dir_tensors, fn))
|
||||
|
||||
return paths
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Script principal
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
# Config / treino
|
||||
parser.add_argument("--config", default="config.json", help="Caminho do config.json (mesmo do treino)")
|
||||
parser.add_argument("--split_folder", type=str, default="val", help="train, val ou test (dentro de dataset/split)")
|
||||
parser.add_argument("--root_override", type=str, default=None, help="Aponta direto para uma pasta split (sobrepõe split_folder)")
|
||||
parser.add_argument(
|
||||
"--test_folder",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Pasta de teste para visualização. "
|
||||
"Se tiver masks, usa dataset normal (mostra GT). "
|
||||
"Se tiver apenas tensors, entra em modo inferência-only."
|
||||
)
|
||||
)
|
||||
parser.add_argument("--ckpt", type=str, default=None, help="Caminho do .pt (se não informar, usa best_miou.pt da pasta rawX)")
|
||||
|
||||
parser.add_argument("--camera", action="store_true", help="Placeholder para integração futura com câmera ao vivo")
|
||||
parser.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal")
|
||||
parser.add_argument("--resize_h", type=int, default=None, help="Altura para inferência (override)")
|
||||
parser.add_argument("--resize_w", type=int, default=None, help="Largura para inferência (override)")
|
||||
parser.add_argument("--alpha", type=float, default=0.45, help="Alpha do overlay da máscara")
|
||||
parser.add_argument("--camera_frame_type", type=str, default="MULTISPEC", choices=["RAW_BRUTO", "RGB", "MULTISPEC"])
|
||||
parser.add_argument("--camera_capture_mode", type=str, default="AUTO", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"])
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Device: {device}")
|
||||
|
||||
# ---- Carrega config ----
|
||||
with open(args.config, "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
MODEL_NAME = config["model_name"]
|
||||
modelo_folder = config["modelo"]
|
||||
|
||||
CHANNELS = int(config.get("channels", 5))
|
||||
FUSION_MODE = config.get("fusion_mode", "stacked")
|
||||
RESOLUCAO = config["resolucao"]
|
||||
W, H = RESOLUCAO[0], RESOLUCAO[1]
|
||||
if args.resize_h is not None:
|
||||
H = args.resize_h
|
||||
if args.resize_w is not None:
|
||||
W = args.resize_w
|
||||
|
||||
print(f"[cfg] resolucao nominal tensor: {W}x{H}")
|
||||
print(f"[cfg] channels={CHANNELS}")
|
||||
|
||||
dataset_path = os.path.join("dataset")
|
||||
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
||||
|
||||
# ==========================
|
||||
# Normalização fixa (igual treino)
|
||||
# ==========================
|
||||
norm_mean = None
|
||||
norm_std = None
|
||||
|
||||
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
|
||||
norm_stats_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag, "norm_stats.json")
|
||||
if args.norm_stats is not None:
|
||||
norm_stats_path = args.norm_stats
|
||||
|
||||
if os.path.isfile(norm_stats_path):
|
||||
with open(norm_stats_path, "r", encoding="utf-8") as f:
|
||||
norm_stats = json.load(f)
|
||||
|
||||
stats_channels = norm_stats.get("channels", [])
|
||||
stats_mean = norm_stats.get("mean", [])
|
||||
stats_std = norm_stats.get("std", [])
|
||||
|
||||
print(f"[NORM] usando stats fixos de: {norm_stats_path}")
|
||||
print(f"[NORM] channels={stats_channels}")
|
||||
print(f"[NORM] mean={stats_mean}")
|
||||
print(f"[NORM] std ={stats_std}")
|
||||
|
||||
if len(stats_mean) != CHANNELS or len(stats_std) != CHANNELS:
|
||||
raise RuntimeError(
|
||||
f"norm_stats incompatível com channels={CHANNELS}: mean={len(stats_mean)} std={len(stats_std)}"
|
||||
)
|
||||
if len(stats_channels) != CHANNELS:
|
||||
raise RuntimeError(
|
||||
f"norm_stats.channels incompatível com channels={CHANNELS}: {stats_channels}"
|
||||
)
|
||||
|
||||
norm_mean = stats_mean
|
||||
norm_std = stats_std
|
||||
else:
|
||||
print(f"[NORM] norm_stats.json não encontrado em {norm_stats_path}. Usando normalização dinâmica por frame.")
|
||||
|
||||
# ---- Descobre checkpoint ----
|
||||
if args.ckpt is not None:
|
||||
ckpt_path = args.ckpt
|
||||
else:
|
||||
save_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag)
|
||||
ckpt_path = os.path.join(save_path, "best_miou.pt")
|
||||
|
||||
if not os.path.isfile(ckpt_path):
|
||||
raise SystemExit(f"Checkpoint não encontrado em: {ckpt_path}")
|
||||
|
||||
print(f"[model] ckpt = {ckpt_path}")
|
||||
|
||||
# ---- Instancia service do modelo ----
|
||||
model_svc = MultispecSegformerService(
|
||||
config_path=args.config,
|
||||
ckpt_path=ckpt_path,
|
||||
device=device,
|
||||
use_amp=True,
|
||||
mean=norm_mean,
|
||||
std=norm_std,
|
||||
)
|
||||
|
||||
classes = model_svc.get_classes()
|
||||
colormap_rgb = model_svc.get_colormap()
|
||||
ignore_id = model_svc.get_ignore_id()
|
||||
|
||||
print(f"[svc] classes: {classes}")
|
||||
print(f"[svc] ignore_id={ignore_id}")
|
||||
|
||||
# ======================================================================
|
||||
# MODO CÂMERA AO VIVO (placeholder)
|
||||
# ======================================================================
|
||||
if args.camera:
|
||||
print("[mode] Câmera MULTISPEC + SegFormer")
|
||||
|
||||
from multispectral_service import MultiSpectralService
|
||||
from stream_receiver import StreamReceiver
|
||||
from pi.raw_processor_core import RawProcessorCore
|
||||
|
||||
STREAM_PORT = 6001
|
||||
PI_HOST = "192.168.105.6"
|
||||
PC_HOST = "192.168.105.5"
|
||||
|
||||
receiver = StreamReceiver(host="0.0.0.0", port=STREAM_PORT)
|
||||
svc = MultiSpectralService(host=PI_HOST, port=5000, timeout=10)
|
||||
|
||||
receiver.start()
|
||||
time.sleep(0.5)
|
||||
|
||||
svc.connect()
|
||||
|
||||
core = RawProcessorCore(sensor_width=W, sensor_height=H)
|
||||
camera_frame_type = args.camera_frame_type
|
||||
camera_output_dtype = "uint8"
|
||||
camera_capture_mode = args.camera_capture_mode
|
||||
|
||||
print("SET FRAME TYPE:", svc.set_frame_type(camera_frame_type))
|
||||
print("SET OUTPUT DTYPE:", svc.set_output_dtype(camera_output_dtype))
|
||||
print("SET FPS:", svc.set_fps(15))
|
||||
|
||||
print("BEGIN:", svc.begin(
|
||||
frame_type=camera_frame_type,
|
||||
output_dtype=camera_output_dtype,
|
||||
capture_mode=camera_capture_mode,
|
||||
))
|
||||
|
||||
print("START STREAM:", svc.start_stream(PC_HOST, STREAM_PORT, fps=15))
|
||||
|
||||
win = "CAMERA MULTISPEC + SEGFORMER (Q=quit)"
|
||||
cv2.namedWindow(win, cv2.WINDOW_NORMAL)
|
||||
|
||||
last_frame_id = -1
|
||||
|
||||
try:
|
||||
while True:
|
||||
meta = receiver.last_meta
|
||||
frame = receiver.last_frame
|
||||
|
||||
if meta is None or frame is None:
|
||||
continue
|
||||
|
||||
frame_id = meta.get("frame_id")
|
||||
if frame_id == last_frame_id:
|
||||
continue
|
||||
|
||||
last_frame_id = frame_id
|
||||
|
||||
try:
|
||||
raw_np = core.build_infer_tensor_from_stream(frame, meta, channels_expected=CHANNELS)
|
||||
|
||||
# =========================================================
|
||||
# INFERÊNCIA
|
||||
# =========================================================
|
||||
pred_ids, preview_bgr, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_np, alpha=args.alpha)
|
||||
|
||||
# =========================================================
|
||||
# HUD
|
||||
# =========================================================
|
||||
txt = f"C={CHANNELS} | inf={t_inf:.1f}ms | pvw={t_pvw:.1f}ms"
|
||||
|
||||
cv2.putText(
|
||||
overlay,
|
||||
txt,
|
||||
(10, 24),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.7,
|
||||
(0, 255, 0),
|
||||
2,
|
||||
)
|
||||
|
||||
cv2.imshow(win, overlay)
|
||||
|
||||
except Exception as e:
|
||||
print("[ERRO FRAME]", e)
|
||||
|
||||
key = cv2.waitKey(1) & 0xFF
|
||||
if key in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
|
||||
finally:
|
||||
try:
|
||||
svc.stop_stream()
|
||||
except:
|
||||
pass
|
||||
|
||||
try:
|
||||
svc.stop()
|
||||
except:
|
||||
pass
|
||||
|
||||
svc.disconnect()
|
||||
receiver.stop()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
return
|
||||
|
||||
# ======================================================================
|
||||
# MODO DATASET / TEST_FOLDER
|
||||
# ======================================================================
|
||||
|
||||
if args.test_folder is not None:
|
||||
root = args.test_folder
|
||||
elif args.root_override is not None:
|
||||
root = args.root_override
|
||||
else:
|
||||
root = os.path.join(dataset_path, "split", args.split_folder)
|
||||
|
||||
print(f"[data] root = {root}")
|
||||
|
||||
has_gt = True
|
||||
ds = None
|
||||
tensor_paths = []
|
||||
|
||||
try:
|
||||
ds = MultispecSegDataset(
|
||||
root,
|
||||
labelmap_path=labelmap_path,
|
||||
resize_hw=None,
|
||||
channels=CHANNELS,
|
||||
strict_channels=True,
|
||||
)
|
||||
print("[mode] Dataset com GT (masks) detectado. Mostrando GT + overlay.")
|
||||
except RuntimeError as e:
|
||||
msg = str(e)
|
||||
if "Nenhum par tensor/mask encontrado" in msg:
|
||||
print("[mode] Nenhum par tensor/mask encontrado. Entrando em modo inferência-only (sem GT).")
|
||||
has_gt = False
|
||||
tensor_paths = _collect_tensor_paths(root)
|
||||
if not tensor_paths:
|
||||
raise RuntimeError(f"Modo inferência-only: não encontrei nenhum tensor em {root}")
|
||||
print(f"[data] Tensors encontrados: {len(tensor_paths)}")
|
||||
else:
|
||||
raise
|
||||
|
||||
n = len(ds) if has_gt else len(tensor_paths)
|
||||
|
||||
idx = 0
|
||||
print(f"[data] total de amostras: {n}")
|
||||
print("Controles: D=próxima, A=anterior, Q=sair")
|
||||
|
||||
win_name = "Tensor preview | GT | Overlay (SegFormer)" if has_gt else "Tensor preview | Overlay (SegFormer)"
|
||||
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
if has_gt:
|
||||
sample = ds[idx]
|
||||
img_tensor = sample["image"]
|
||||
mask_gt = sample["mask"]
|
||||
|
||||
raw_np = img_tensor.cpu().numpy().astype(np.float32)
|
||||
pred_ids, preview_bgr, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_np, alpha=args.alpha)
|
||||
|
||||
mask_gt_np = mask_gt.cpu().numpy().astype(np.uint8)
|
||||
gt_rgb = converter_mask_ids_para_bgr(mask_gt_np, colormap_rgb, ignore_id)
|
||||
|
||||
h, w, _ = preview_bgr.shape
|
||||
gt_resized = cv2.resize(gt_rgb, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
resultado = np.concatenate([preview_bgr, gt_resized, overlay_resized], axis=1)
|
||||
legenda = desenhar_legenda_horizontal(colormap_rgb, classes)
|
||||
legenda_resized = cv2.resize(legenda, (resultado.shape[1], legenda.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||
resultado_completo = np.concatenate([resultado, legenda_resized], axis=0)
|
||||
|
||||
header_txt = f"idx {idx + 1}/{n} | C={CHANNELS} | inf={t_inf:.1f}ms | pvw={t_pvw:.1f}ms"
|
||||
else:
|
||||
tensor_path = tensor_paths[idx]
|
||||
raw_np = np.load(tensor_path).astype(np.float32)
|
||||
pred_ids, preview_bgr, _, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_np, alpha=args.alpha)
|
||||
|
||||
h, w, _ = preview_bgr.shape
|
||||
overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
resultado = np.concatenate([preview_bgr, overlay_resized], axis=1)
|
||||
resultado_completo = resultado
|
||||
|
||||
header_txt = f"idx {idx + 1}/{n} | C={CHANNELS} | {os.path.basename(tensor_path)} | inf={t_inf:.1f}ms | pvw={t_pvw:.1f}ms"
|
||||
|
||||
cv2.putText(resultado_completo, header_txt, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||||
|
||||
try:
|
||||
_x, _y, win_w, win_h = cv2.getWindowImageRect(win_name)
|
||||
except Exception:
|
||||
win_w, win_h = 0, 0
|
||||
|
||||
if win_w > 0 and win_h > 0:
|
||||
display = cv2.resize(resultado_completo, (win_w, win_h), interpolation=cv2.INTER_NEAREST)
|
||||
else:
|
||||
display = resultado_completo
|
||||
|
||||
display_bgr = cv2.cvtColor(display, cv2.COLOR_RGB2BGR)
|
||||
cv2.imshow(win_name, display_bgr)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
elif key in (ord("d"), ord("D")):
|
||||
idx = (idx + 1) % n
|
||||
elif key in (ord("a"), ord("A")):
|
||||
idx = (idx - 1 + n) % n
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -1,535 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Teste/visualização do SegFormer RAW (4 ou 5 canais) treinado no _8_train_segformer_b3_raw.py
|
||||
|
||||
Modos:
|
||||
|
||||
1) Dataset (com GT):
|
||||
- Usa RawSegDataset (raws + masks) a partir de split (train/val/test)
|
||||
- Mostra: [preview RAW (RGB)] | [mask GT] | [overlay predito]
|
||||
|
||||
2) Pasta sem máscara (--test_folder sem masks):
|
||||
- Faz inferência só 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):
|
||||
"""
|
||||
Lê 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()
|
||||
|
|
@ -6,14 +6,14 @@
|
|||
"main_class_name": "cana",
|
||||
"es_classes": "",
|
||||
"model_to_use": "geral",
|
||||
"raw_size": [1296, 1028],
|
||||
"raw_size": [640, 480],
|
||||
"resolucao": [1024, 800],
|
||||
"roi_inicio": 0.0,
|
||||
"roi_tamanho": 1.0,
|
||||
"shaves": 3,
|
||||
"channels": 5,
|
||||
"channels": 4,
|
||||
"use_ndvi": false,
|
||||
"backbone": "nvidia/mit-b1",
|
||||
"fusion_mode": "stacked",
|
||||
"stats_source_tag": "stacked_raw5"
|
||||
"stats_source_tag": "stacked_raw4"
|
||||
}
|
||||
|
|
@ -0,0 +1,683 @@
|
|||
# multispec_segformer_service.py
|
||||
# Serviço e dataset para SegFormer no novo pipeline multiespectral.
|
||||
#
|
||||
# Contrato principal do tensor salvo:
|
||||
# (C,H,W) float32 em 0..1
|
||||
# onde C pode ser 3, 4 ou 5, conforme config/channels.
|
||||
#
|
||||
# Estruturas aceitas:
|
||||
# 1) Simples:
|
||||
# root/
|
||||
# tensors/*.npy
|
||||
# masks/*.npy
|
||||
#
|
||||
# 2) Com grupos:
|
||||
# root/
|
||||
# group/
|
||||
# <grupo>/
|
||||
# tensors/*.npy
|
||||
# masks/*.npy
|
||||
#
|
||||
# A máscara oficial é .npy (IDs inteiros). O .png fica apenas para debug.
|
||||
|
||||
import os
|
||||
import json
|
||||
from contextlib import nullcontext
|
||||
import sys
|
||||
from typing import List, Optional, Tuple
|
||||
import threading
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import amp
|
||||
from torch.utils.data import Dataset
|
||||
import torch.nn.functional as F
|
||||
from transformers import SegformerForSemanticSegmentation
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
from utils import carregar_labelmap_completo, converter_mask_ids_para_bgr
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Model builders
|
||||
# ============================================================
|
||||
|
||||
def build_multispec_segformer_model(
|
||||
num_classes: int,
|
||||
channels: int,
|
||||
backbone: str = "nvidia/segformer-b3-finetuned-ade-512-512",
|
||||
device: torch.device | None = None,
|
||||
ckpt_path: str | None = None,
|
||||
strict: bool = True,
|
||||
) -> SegformerForSemanticSegmentation:
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = SegformerForSemanticSegmentation.from_pretrained(
|
||||
backbone,
|
||||
num_labels=num_classes,
|
||||
ignore_mismatched_sizes=True,
|
||||
use_safetensors=True,
|
||||
)
|
||||
|
||||
patch_segformer_input_channels(model, in_ch=channels)
|
||||
|
||||
if ckpt_path is not None:
|
||||
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
||||
if "model" not in ckpt:
|
||||
raise RuntimeError(f"Checkpoint {ckpt_path} não contém chave 'model'.")
|
||||
model.load_state_dict(ckpt["model"], strict=strict)
|
||||
|
||||
model.to(device)
|
||||
return model
|
||||
|
||||
|
||||
def build_dual_branch_segformer_model(
|
||||
num_classes: int,
|
||||
channels: int,
|
||||
backbone: str = "nvidia/segformer-b1-finetuned-ade-512-512",
|
||||
device: torch.device | None = None,
|
||||
ckpt_path: str | None = None,
|
||||
strict: bool = True,
|
||||
) -> nn.Module:
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if channels < 4:
|
||||
raise ValueError(f"Dual branch espera channels>=4, veio {channels}")
|
||||
|
||||
spec_channels = channels - 3
|
||||
|
||||
model = DualBranchSegformerV2(
|
||||
num_classes=num_classes,
|
||||
backbone=backbone,
|
||||
spec_channels=spec_channels,
|
||||
)
|
||||
|
||||
if ckpt_path is not None:
|
||||
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
||||
if "model" not in ckpt:
|
||||
raise RuntimeError(f"Checkpoint {ckpt_path} não contém chave 'model'.")
|
||||
model.load_state_dict(ckpt["model"], strict=strict)
|
||||
|
||||
model.to(device)
|
||||
return model
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Dual-branch blocks
|
||||
# ============================================================
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, ch: int):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(ch)
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
self.conv2 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(ch)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
x = self.act(self.bn1(self.conv1(x)))
|
||||
x = self.bn2(self.conv2(x))
|
||||
x = self.act(x + identity)
|
||||
return x
|
||||
|
||||
|
||||
class SpectralEncoderV2(nn.Module):
|
||||
def __init__(self, in_ch: int, out_ch: int = 256):
|
||||
super().__init__()
|
||||
self.stem = nn.Sequential(
|
||||
nn.Conv2d(in_ch, 32, 3, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(32),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
self.stage1 = nn.Sequential(
|
||||
ResidualBlock(32),
|
||||
nn.Conv2d(32, 64, 3, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(64),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
self.stage2 = nn.Sequential(
|
||||
ResidualBlock(64),
|
||||
nn.Conv2d(64, 128, 3, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(128),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
self.stage3 = nn.Sequential(
|
||||
ResidualBlock(128),
|
||||
nn.Conv2d(128, out_ch, 1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.stem(x)
|
||||
x = self.stage1(x)
|
||||
x = self.stage2(x)
|
||||
x = self.stage3(x)
|
||||
return x
|
||||
|
||||
|
||||
class FeatureFusionBlock(nn.Module):
|
||||
def __init__(self, ch_rgb: int, ch_spec: int, ch_fused: int):
|
||||
super().__init__()
|
||||
self.rgb_proj = nn.Conv2d(ch_rgb, ch_fused, 1, bias=False)
|
||||
self.spec_proj = nn.Conv2d(ch_spec, ch_fused, 1, bias=False)
|
||||
self.fuse = nn.Sequential(
|
||||
nn.Conv2d(ch_fused * 2, ch_fused, 3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(ch_fused),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(ch_fused, ch_fused, 3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(ch_fused),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
self.gate = nn.Sequential(
|
||||
nn.Conv2d(ch_fused, ch_fused, 1),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, rgb_feat, spec_feat):
|
||||
rgb_feat = self.rgb_proj(rgb_feat)
|
||||
spec_feat = self.spec_proj(spec_feat)
|
||||
|
||||
if spec_feat.shape[-2:] != rgb_feat.shape[-2:]:
|
||||
spec_feat = F.interpolate(spec_feat, size=rgb_feat.shape[-2:], mode="bilinear", align_corners=False)
|
||||
|
||||
fused = torch.cat([rgb_feat, spec_feat], dim=1)
|
||||
fused = self.fuse(fused)
|
||||
gate = self.gate(fused)
|
||||
fused = fused * gate
|
||||
return fused
|
||||
|
||||
|
||||
class DualBranchSegformerV2(nn.Module):
|
||||
def __init__(self, num_classes: int, backbone: str, spec_channels: int, fused_channels: int = 256):
|
||||
super().__init__()
|
||||
self.rgb_model = SegformerForSemanticSegmentation.from_pretrained(
|
||||
backbone,
|
||||
num_labels=num_classes,
|
||||
ignore_mismatched_sizes=True,
|
||||
use_safetensors=True,
|
||||
)
|
||||
patch_segformer_input_channels(self.rgb_model, in_ch=3)
|
||||
|
||||
self.spec_branch = SpectralEncoderV2(in_ch=spec_channels, out_ch=fused_channels)
|
||||
self.fusion = FeatureFusionBlock(
|
||||
ch_rgb=self.rgb_model.config.hidden_sizes[-1],
|
||||
ch_spec=fused_channels,
|
||||
ch_fused=fused_channels,
|
||||
)
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Conv2d(fused_channels, fused_channels, 3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(fused_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(fused_channels, num_classes, 1, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, pixel_values: torch.Tensor):
|
||||
rgb = pixel_values[:, :3, :, :]
|
||||
spec = pixel_values[:, 3:, :, :]
|
||||
|
||||
rgb_out = self.rgb_model(pixel_values=rgb, output_hidden_states=True)
|
||||
rgb_logits = rgb_out.logits
|
||||
rgb_feat = rgb_out.hidden_states[-1]
|
||||
|
||||
spec_feat = self.spec_branch(spec)
|
||||
fused_feat = self.fusion(rgb_feat, spec_feat)
|
||||
logits = self.classifier(fused_feat)
|
||||
|
||||
if logits.shape[-2:] != rgb_logits.shape[-2:]:
|
||||
logits = F.interpolate(logits, size=rgb_logits.shape[-2:], mode="bilinear", align_corners=False)
|
||||
|
||||
return type("DualBranchOutput", (), {"logits": logits})
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Dataset
|
||||
# ============================================================
|
||||
|
||||
class MultispecSegDataset(Dataset):
|
||||
"""
|
||||
Espera tensores finais já prontos em .npy.
|
||||
|
||||
1) Simples:
|
||||
root/
|
||||
tensors/*.npy
|
||||
masks/*.npy
|
||||
|
||||
2) Com grupos:
|
||||
root/
|
||||
group/
|
||||
<grupo>/
|
||||
tensors/*.npy
|
||||
masks/*.npy
|
||||
|
||||
Saída:
|
||||
- image: tensor (C,H,W) float32
|
||||
- mask : (H,W) long
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str,
|
||||
labelmap_path: Optional[str] = None,
|
||||
resize_hw: Optional[Tuple[int, int]] = None,
|
||||
channels: int = 5,
|
||||
strict_channels: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.root = root
|
||||
self.resize_hw = resize_hw
|
||||
self.channels = int(channels)
|
||||
self.strict_channels = strict_channels
|
||||
self.items: List[Tuple[str, str]] = []
|
||||
|
||||
if self.channels not in (3, 4, 5, 6, 7):
|
||||
raise ValueError(f"[dataset] channels inválido: {self.channels}")
|
||||
|
||||
dir_group = os.path.join(root, "group")
|
||||
dir_tensor_simple = os.path.join(root, "tensors")
|
||||
dir_mask_simple = os.path.join(root, "masks")
|
||||
|
||||
if os.path.isdir(dir_group):
|
||||
for g in sorted(os.listdir(dir_group)):
|
||||
gdir = os.path.join(dir_group, g)
|
||||
if not os.path.isdir(gdir):
|
||||
continue
|
||||
g_tensors = os.path.join(gdir, "tensors")
|
||||
g_masks = os.path.join(gdir, "masks")
|
||||
if not (os.path.isdir(g_tensors) and os.path.isdir(g_masks)):
|
||||
continue
|
||||
self._append_pairs(g_tensors, g_masks)
|
||||
else:
|
||||
if not os.path.isdir(dir_tensor_simple):
|
||||
raise RuntimeError(f"Não achei pasta tensors: {dir_tensor_simple}")
|
||||
if not os.path.isdir(dir_mask_simple):
|
||||
raise RuntimeError(f"Não achei pasta masks: {dir_mask_simple}")
|
||||
self._append_pairs(dir_tensor_simple, dir_mask_simple)
|
||||
|
||||
if not self.items:
|
||||
raise RuntimeError(f"Nenhum par tensor/mask encontrado em {root} (nomes base precisam bater).")
|
||||
|
||||
_, _, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
|
||||
self.classes = classes
|
||||
self.class_ids = sorted(classes.keys())
|
||||
self.ignore_id = ignore_id
|
||||
|
||||
print(f"[dataset] root={root} items={len(self.items)} channels={self.channels}")
|
||||
|
||||
def _append_pairs(self, tensor_dir: str, mask_dir: str):
|
||||
tensors = [fn for fn in sorted(os.listdir(tensor_dir)) if fn.lower().endswith(".npy")]
|
||||
for fn in tensors:
|
||||
stem = os.path.splitext(fn)[0]
|
||||
m = os.path.join(mask_dir, stem + ".npy")
|
||||
if os.path.exists(m):
|
||||
self.items.append((os.path.join(tensor_dir, fn), m))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.items)
|
||||
|
||||
def _resize_tensor(self, x: np.ndarray, new_hw: Tuple[int, int]) -> np.ndarray:
|
||||
h, w = new_hw
|
||||
chans = []
|
||||
for ch in x:
|
||||
ch_res = F.interpolate(
|
||||
torch.from_numpy(ch[None, None, :, :]).float(),
|
||||
size=(h, w),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)[0, 0].numpy()
|
||||
chans.append(ch_res.astype(np.float32))
|
||||
return np.stack(chans, axis=0)
|
||||
|
||||
def _resize_mask(self, y: np.ndarray, new_hw: Tuple[int, int]) -> np.ndarray:
|
||||
h, w = new_hw
|
||||
out = F.interpolate(
|
||||
torch.from_numpy(y[None, None, :, :]).float(),
|
||||
size=(h, w),
|
||||
mode="nearest",
|
||||
)[0, 0].numpy()
|
||||
return out.astype(np.int64)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
tensor_path, mask_path = self.items[idx]
|
||||
|
||||
x = np.load(tensor_path).astype(np.float32) # (C,H,W)
|
||||
y = np.load(mask_path).astype(np.int64) # (H,W)
|
||||
|
||||
if x.ndim != 3:
|
||||
raise RuntimeError(f"Tensor precisa ser (C,H,W), veio {x.shape} em {tensor_path}")
|
||||
if y.ndim != 2:
|
||||
raise RuntimeError(f"Mask precisa ser (H,W), veio {y.shape} em {mask_path}")
|
||||
|
||||
if self.strict_channels and x.shape[0] != self.channels:
|
||||
raise RuntimeError(
|
||||
f"Tensor com número de canais inesperado: {x.shape[0]} em {tensor_path}. "
|
||||
f"Esperado {self.channels}."
|
||||
)
|
||||
|
||||
if self.resize_hw is not None:
|
||||
x = self._resize_tensor(x, self.resize_hw)
|
||||
y = self._resize_mask(y, self.resize_hw)
|
||||
|
||||
return {
|
||||
"image": torch.from_numpy(x).float(),
|
||||
"mask": torch.from_numpy(y).long(),
|
||||
}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Normalização
|
||||
# ============================================================
|
||||
|
||||
def normalize_per_batch(imgs: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
|
||||
mean = imgs.mean(dim=(0, 2, 3), keepdim=True)
|
||||
std = imgs.std(dim=(0, 2, 3), keepdim=True).clamp_min(eps)
|
||||
return (imgs - mean) / std
|
||||
|
||||
|
||||
def build_fixed_normalizer(mean: List[float], std: List[float], device: torch.device):
|
||||
mean_t = torch.tensor(mean, dtype=torch.float32, device=device).view(1, -1, 1, 1)
|
||||
std_t = torch.tensor(std, dtype=torch.float32, device=device).view(1, -1, 1, 1).clamp_min(1e-6)
|
||||
|
||||
def normalizer(x: torch.Tensor) -> torch.Tensor:
|
||||
return (x - mean_t) / std_t
|
||||
|
||||
return normalizer
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Preview helpers
|
||||
# ============================================================
|
||||
|
||||
def make_bgr_preview_from_tensor(raw_np: np.ndarray, preview_fast: bool = False, preview_scale: int = 2) -> np.ndarray:
|
||||
"""
|
||||
Preview sempre usando os três primeiros canais como RGB.
|
||||
Espera raw_np (C,H,W) float32 em 0..1.
|
||||
"""
|
||||
assert raw_np.ndim == 3 and raw_np.shape[0] >= 3
|
||||
|
||||
r = raw_np[0].astype(np.float32, copy=False)
|
||||
g = raw_np[1].astype(np.float32, copy=False)
|
||||
b = raw_np[2].astype(np.float32, copy=False)
|
||||
|
||||
if preview_fast and preview_scale > 1:
|
||||
r = r[::preview_scale, ::preview_scale]
|
||||
g = g[::preview_scale, ::preview_scale]
|
||||
b = b[::preview_scale, ::preview_scale]
|
||||
|
||||
def stretch_channel(x: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray:
|
||||
lo = np.percentile(x, p_low)
|
||||
hi = np.percentile(x, p_high)
|
||||
if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
|
||||
return np.clip(x, 0.0, 1.0)
|
||||
x = (x - lo) / (hi - lo)
|
||||
return np.clip(x, 0.0, 1.0)
|
||||
|
||||
if not preview_fast:
|
||||
r = stretch_channel(r)
|
||||
g = stretch_channel(g)
|
||||
b = stretch_channel(b)
|
||||
|
||||
bgr = np.stack([b, g, r], axis=0)
|
||||
gamma = 1 / (1.8 if preview_fast else 2.0)
|
||||
bgr = np.power(np.clip(bgr, 0.0, 1.0), gamma)
|
||||
bgr8 = (bgr * 255.0).clip(0, 255).astype(np.uint8)
|
||||
return np.transpose(bgr8, (1, 2, 0)).copy()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Patch input channels
|
||||
# ============================================================
|
||||
|
||||
def patch_segformer_input_channels(model: SegformerForSemanticSegmentation, in_ch: int = 4):
|
||||
enc = model.segformer.encoder
|
||||
proj = enc.patch_embeddings[0].proj
|
||||
if proj.in_channels == in_ch:
|
||||
return
|
||||
|
||||
old_w = proj.weight.data
|
||||
old_b = proj.bias.data if proj.bias is not None else None
|
||||
|
||||
new_proj = nn.Conv2d(
|
||||
in_channels=in_ch,
|
||||
out_channels=proj.out_channels,
|
||||
kernel_size=proj.kernel_size,
|
||||
stride=proj.stride,
|
||||
padding=proj.padding,
|
||||
dilation=proj.dilation,
|
||||
groups=proj.groups,
|
||||
bias=(proj.bias is not None),
|
||||
padding_mode=proj.padding_mode,
|
||||
).to(proj.weight.device)
|
||||
|
||||
with torch.no_grad():
|
||||
if old_w.shape[1] == 3 and in_ch >= 3:
|
||||
new_proj.weight[:, :3, :, :].copy_(old_w)
|
||||
if in_ch > 3:
|
||||
extra = old_w.mean(dim=1, keepdim=True)
|
||||
for c in range(3, in_ch):
|
||||
new_proj.weight[:, c:c+1, :, :].copy_(extra)
|
||||
else:
|
||||
nn.init.kaiming_normal_(new_proj.weight, mode="fan_out", nonlinearity="relu")
|
||||
|
||||
if old_b is not None:
|
||||
new_proj.bias.copy_(old_b)
|
||||
|
||||
enc.patch_embeddings[0].proj = new_proj
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Service
|
||||
# ============================================================
|
||||
|
||||
class MultispecSegformerService:
|
||||
"""
|
||||
Serviço para inferência no novo pipeline multiespectral.
|
||||
Espera tensores finais (C,H,W) já prontos, em 0..1.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config_path: str,
|
||||
ckpt_path: str | None = None,
|
||||
labelmap_path: str | None = None,
|
||||
device: torch.device | None = None,
|
||||
use_amp: bool = True,
|
||||
mean=None,
|
||||
std=None,
|
||||
load_model=True,
|
||||
):
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.device = device
|
||||
self.use_amp = use_amp and (self.device.type == "cuda")
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
if device.type == "cuda":
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
|
||||
self._norm_mean = None
|
||||
self._norm_std = None
|
||||
if mean is not None and std is not None:
|
||||
self.set_norm_stats(mean, std)
|
||||
|
||||
self._buf_shape = None
|
||||
self._cpu_pinned = None
|
||||
self._gpu_input = None
|
||||
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
self.cfg = cfg
|
||||
self.MODELO = cfg["camera"]
|
||||
self.MODEL_NAME = cfg["model_name"]
|
||||
self.modelo_folder = cfg["modelo"]
|
||||
self.res_w, self.res_h = cfg["resolucao"]
|
||||
self.backbone = cfg["backbone"]
|
||||
self.fusion_mode = cfg.get("fusion_mode", "stacked")
|
||||
self.channels = int(cfg.get("channels", 5))
|
||||
|
||||
self._preview_lock = threading.Lock()
|
||||
self._preview_busy = False
|
||||
self._last_preview = None
|
||||
self._last_preview_ts = 0.0
|
||||
self._last_preview_dt = 0.0
|
||||
|
||||
if labelmap_path is None:
|
||||
labelmap_path = os.path.join("dataset", "labelmap.txt")
|
||||
self.labelmap_path = labelmap_path
|
||||
|
||||
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
self.colormap_rgb = colormap_rgb
|
||||
self.classes = classes
|
||||
self.ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
|
||||
self.num_classes = len(classes)
|
||||
|
||||
if load_model:
|
||||
if ckpt_path is None:
|
||||
save_path = os.path.join("backup", self.modelo_folder, self.MODEL_NAME, f"{self.fusion_mode}_raw{self.channels}")
|
||||
ckpt_path = os.path.join(save_path, "best_miou.pt")
|
||||
|
||||
if not os.path.isfile(ckpt_path):
|
||||
raise FileNotFoundError(f"Checkpoint não encontrado: {ckpt_path}")
|
||||
|
||||
self.ckpt_path = ckpt_path
|
||||
self.model = self._load_model(backbone=self.backbone)
|
||||
self.model.eval()
|
||||
|
||||
def set_norm_stats(self, mean, std):
|
||||
mean = torch.tensor(mean, dtype=torch.float32, device=self.device).view(1, -1, 1, 1)
|
||||
std = torch.tensor(std, dtype=torch.float32, device=self.device).view(1, -1, 1, 1).clamp_min(1e-6)
|
||||
self._norm_mean = mean
|
||||
self._norm_std = std
|
||||
|
||||
def prepare_infer_buffers(self, C, H, W):
|
||||
shape = (1, C, H, W)
|
||||
if self._buf_shape == shape:
|
||||
return
|
||||
self._buf_shape = shape
|
||||
self._cpu_pinned = torch.empty(shape, dtype=torch.float32, pin_memory=True)
|
||||
self._gpu_input = torch.empty(shape, dtype=torch.float32, device=self.device)
|
||||
|
||||
def _load_model(self, backbone: str) -> torch.nn.Module:
|
||||
if self.fusion_mode == "dual_branch":
|
||||
model = build_dual_branch_segformer_model(
|
||||
num_classes=self.num_classes,
|
||||
channels=self.channels,
|
||||
backbone=backbone,
|
||||
device=self.device,
|
||||
ckpt_path=self.ckpt_path,
|
||||
strict=True,
|
||||
)
|
||||
elif self.fusion_mode == "stacked":
|
||||
model = build_multispec_segformer_model(
|
||||
num_classes=self.num_classes,
|
||||
channels=self.channels,
|
||||
backbone=backbone,
|
||||
device=self.device,
|
||||
ckpt_path=self.ckpt_path,
|
||||
strict=True,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"fusion_mode inválido: {self.fusion_mode}")
|
||||
return model
|
||||
|
||||
def _preprocess_tensor(self, raw_np: np.ndarray) -> torch.Tensor:
|
||||
if raw_np.dtype != np.float32:
|
||||
raw_np = raw_np.astype(np.float32)
|
||||
x = torch.from_numpy(raw_np).unsqueeze(0).to(self.device)
|
||||
if self._norm_mean is not None and self._norm_std is not None:
|
||||
x = (x - self._norm_mean) / self._norm_std
|
||||
else:
|
||||
x = normalize_per_batch(x)
|
||||
return x
|
||||
|
||||
def _build_preview(self, raw_np, pred_ids, alpha, fast):
|
||||
t0 = time.time()
|
||||
bgr = make_bgr_preview_from_tensor(raw_np, preview_fast=fast)
|
||||
pred_bgr = converter_mask_ids_para_bgr(pred_ids, self.colormap_rgb, self.ignore_id)
|
||||
a = float(np.clip(alpha, 0.0, 1.0))
|
||||
overlay = (bgr.astype(np.float32) * (1 - a) + pred_bgr.astype(np.float32) * a)
|
||||
overlay = np.clip(overlay, 0, 255).astype(np.uint8)
|
||||
dt = (time.time() - t0) * 1000.0
|
||||
return bgr, pred_bgr, overlay, dt
|
||||
|
||||
@torch.inference_mode()
|
||||
def infer_ids(self, raw_np: np.ndarray, use_amp: bool | None = None) -> np.ndarray:
|
||||
if use_amp is None:
|
||||
use_amp = self.use_amp and (self.device.type == "cuda")
|
||||
|
||||
if raw_np.dtype != np.float32:
|
||||
raw_np = raw_np.astype(np.float32, copy=False)
|
||||
|
||||
C, H, W = raw_np.shape
|
||||
self.prepare_infer_buffers(C, H, W)
|
||||
|
||||
self._cpu_pinned[0].copy_(torch.from_numpy(raw_np), non_blocking=True)
|
||||
x = self._gpu_input
|
||||
x.copy_(self._cpu_pinned, non_blocking=True)
|
||||
|
||||
if self._norm_mean is not None and self._norm_std is not None:
|
||||
x = (x - self._norm_mean) / self._norm_std
|
||||
else:
|
||||
mean = x.mean(dim=(2, 3), keepdim=True)
|
||||
std = x.std(dim=(2, 3), keepdim=True).clamp_min(1e-6)
|
||||
x = (x - mean) / std
|
||||
|
||||
ctx = amp.autocast("cuda") if use_amp else nullcontext()
|
||||
with ctx:
|
||||
out = self.model(pixel_values=x)
|
||||
logits = out.logits
|
||||
if logits.shape[-2:] != (H, W):
|
||||
logits = F.interpolate(logits, size=(H, W), mode="bilinear", align_corners=False)
|
||||
pred = torch.argmax(logits, dim=1)
|
||||
|
||||
pred_ids = pred[0].to(torch.uint8).cpu().numpy()
|
||||
return pred_ids
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_raw(self, raw_np: np.ndarray):
|
||||
return self.infer_ids(raw_np)
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_and_preview(self, raw_np: np.ndarray, alpha=0.5):
|
||||
t0 = time.time()
|
||||
pred_ids = self.infer_ids(raw_np)
|
||||
t1 = time.time()
|
||||
rgb, pred_rgb, overlay, dt = self._build_preview(raw_np, pred_ids, alpha, False)
|
||||
t2 = time.time()
|
||||
t_inf = (t1 - t0) * 1000
|
||||
t_pvw = (t2 - t1) * 1000
|
||||
return pred_ids, rgb, pred_rgb, overlay, t_inf, t_pvw
|
||||
|
||||
def get_classes(self):
|
||||
return self.classes
|
||||
|
||||
def get_colormap(self):
|
||||
return self.colormap_rgb
|
||||
|
||||
def get_ignore_id(self):
|
||||
return self.ignore_id
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Utils
|
||||
# ============================================================
|
||||
|
||||
def _infer_ignore_id(ignore_rgb, default_id=255):
|
||||
import numpy as _np
|
||||
if isinstance(ignore_rgb, (list, tuple)):
|
||||
if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, _np.integer)):
|
||||
return int(ignore_rgb[0])
|
||||
if len(ignore_rgb) == 3:
|
||||
return default_id
|
||||
if isinstance(ignore_rgb, (int, _np.integer)):
|
||||
return int(ignore_rgb)
|
||||
return default_id
|
||||
|
|
@ -0,0 +1,405 @@
|
|||
import socket
|
||||
import json
|
||||
import numpy as np
|
||||
import base64
|
||||
import time
|
||||
from typing import Optional, Any
|
||||
|
||||
|
||||
class MultiSpectralService:
|
||||
def __init__(self, host="192.168.105.6", port=5000, timeout=5):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout = timeout
|
||||
self.sock = None
|
||||
self.file = None
|
||||
|
||||
def __enter__(self):
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
self.disconnect()
|
||||
|
||||
# =========================================================
|
||||
# Conexão
|
||||
# =========================================================
|
||||
|
||||
def connect(self):
|
||||
if self.sock is not None:
|
||||
return
|
||||
|
||||
self.sock = socket.create_connection((self.host, self.port), timeout=self.timeout)
|
||||
self.sock.settimeout(self.timeout)
|
||||
self.file = self.sock.makefile("r", encoding="utf-8")
|
||||
|
||||
def disconnect(self):
|
||||
try:
|
||||
if self.file:
|
||||
self.file.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if self.sock:
|
||||
self.sock.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.file = None
|
||||
self.sock = None
|
||||
|
||||
def check_connection(self, timeout: float = None) -> bool:
|
||||
try:
|
||||
self.connect()
|
||||
resp = self._send_command({"cmd": "ping"})
|
||||
return resp.get("ok") and resp.get("reply") == "pong"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
try:
|
||||
self.connect()
|
||||
resp = self._send_command({"cmd": "ping"})
|
||||
return resp.get("ok") and resp.get("reply") == "pong"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def ensure_alive(self):
|
||||
try:
|
||||
self.connect()
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Não foi possível conectar ao módulo em {self.host}:{self.port}. "
|
||||
f"Verifique rede, IP e se o Pi está ligado. Erro: {e}"
|
||||
) from e
|
||||
|
||||
try:
|
||||
resp = self._send_command({"cmd": "ping"})
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Conectou ao endereço {self.host}:{self.port}, mas o módulo não respondeu ao ping. "
|
||||
f"Verifique se o serviço está rodando no Pi. Erro: {e}"
|
||||
) from e
|
||||
|
||||
if not resp.get("ok") or resp.get("reply") != "pong":
|
||||
raise RuntimeError(f"Resposta inválida do módulo ao ping: {resp}")
|
||||
|
||||
def _send_command(self, payload: dict) -> dict:
|
||||
if self.sock is None:
|
||||
self.connect()
|
||||
|
||||
data = (json.dumps(payload) + "\n").encode("utf-8")
|
||||
self.sock.sendall(data)
|
||||
|
||||
line = self.file.readline()
|
||||
if not line:
|
||||
self.disconnect()
|
||||
raise RuntimeError("Conexão encerrada pelo servidor")
|
||||
|
||||
return json.loads(line.strip())
|
||||
|
||||
# =========================================================
|
||||
# Helpers numpy
|
||||
# =========================================================
|
||||
|
||||
def _numpy_dtype_from_string(self, dtype_str: str):
|
||||
mapping = {
|
||||
"uint8": np.uint8,
|
||||
"uint16": np.uint16,
|
||||
"float32": np.float32,
|
||||
}
|
||||
if dtype_str not in mapping:
|
||||
raise RuntimeError(f"dtype não suportado recebido do Pi: {dtype_str}")
|
||||
return mapping[dtype_str]
|
||||
|
||||
def _reshape_array_from_shape(self, raw_bytes: bytes, dtype_str: str, shape: list | tuple):
|
||||
np_dtype = self._numpy_dtype_from_string(dtype_str)
|
||||
arr = np.frombuffer(raw_bytes, dtype=np_dtype)
|
||||
return arr.reshape(tuple(shape))
|
||||
|
||||
def _reshape_array(self, raw_bytes: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
|
||||
np_dtype = self._numpy_dtype_from_string(dtype_str)
|
||||
arr = np.frombuffer(raw_bytes, dtype=np_dtype)
|
||||
|
||||
if layout == "HW":
|
||||
return arr.reshape(height, width)
|
||||
|
||||
if layout == "CHW":
|
||||
return arr.reshape(channels, height, width)
|
||||
|
||||
if layout == "HWC":
|
||||
return arr.reshape(height, width, channels)
|
||||
|
||||
raise RuntimeError(f"Layout não suportado recebido do Pi: {layout}")
|
||||
|
||||
def _decode_single_array(self, raw: bytes, resp: dict):
|
||||
dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8"
|
||||
shape = resp.get("shape")
|
||||
|
||||
if shape:
|
||||
return self._reshape_array_from_shape(raw, dtype_str, shape)
|
||||
|
||||
output_layout = resp.get("output_layout", "HW")
|
||||
width = int(resp.get("output_width", resp.get("width", 0)))
|
||||
height = int(resp.get("output_height", resp.get("height", 0)))
|
||||
channels = int(resp.get("output_channels", resp.get("channels", 1)))
|
||||
|
||||
return self._reshape_array(
|
||||
raw_bytes=raw,
|
||||
dtype_str=dtype_str,
|
||||
layout=output_layout,
|
||||
width=width,
|
||||
height=height,
|
||||
channels=channels,
|
||||
)
|
||||
|
||||
def _decode_multi_frames_base64(self, resp: dict):
|
||||
frames_resp = resp.get("frames", {})
|
||||
payload_parts = resp.get("payload_parts", []) or []
|
||||
camera_frames = resp.get("camera_frames", {}) or {}
|
||||
|
||||
parts_by_cam = {}
|
||||
for part in payload_parts:
|
||||
cam_id = part.get("camera_id")
|
||||
if cam_id:
|
||||
parts_by_cam[cam_id] = part
|
||||
|
||||
frames = {}
|
||||
|
||||
for cam_id, item in frames_resp.items():
|
||||
raw = base64.b64decode(item["data"])
|
||||
part_meta = parts_by_cam.get(cam_id, {})
|
||||
cam_meta = camera_frames.get(cam_id, {})
|
||||
|
||||
dtype_str = part_meta.get("dtype") or resp.get("dtype") or resp.get("output_dtype")
|
||||
shape = part_meta.get("shape")
|
||||
|
||||
if dtype_str == "multi" or dtype_str is None:
|
||||
# fallback conservador
|
||||
dtype_str = "uint8" if int(cam_meta.get("channels", 1)) > 1 else "uint16"
|
||||
|
||||
if shape:
|
||||
arr = self._reshape_array_from_shape(raw, dtype_str, shape)
|
||||
else:
|
||||
width = int(part_meta.get("width", cam_meta.get("width", 0)))
|
||||
height = int(part_meta.get("height", cam_meta.get("height", 0)))
|
||||
channels = int(part_meta.get("channels", cam_meta.get("channels", 1)))
|
||||
|
||||
layout = "HWC" if channels > 1 else "HW"
|
||||
|
||||
arr = self._reshape_array(
|
||||
raw_bytes=raw,
|
||||
dtype_str=dtype_str,
|
||||
layout=layout,
|
||||
width=width,
|
||||
height=height,
|
||||
channels=channels,
|
||||
)
|
||||
|
||||
frames[cam_id] = arr
|
||||
|
||||
return frames
|
||||
|
||||
def _extract_meta(self, resp: dict, t0: float) -> dict:
|
||||
dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8"
|
||||
|
||||
return {
|
||||
"frame_id": resp.get("frame_id"),
|
||||
"frame_type": resp.get("frame_type"),
|
||||
"payload_format_version": resp.get("payload_format_version"),
|
||||
|
||||
"output_dtype": resp.get("output_dtype"),
|
||||
"dtype": dtype_str,
|
||||
"output_layout": resp.get("output_layout"),
|
||||
"output_channels": resp.get("output_channels"),
|
||||
"output_channel_names": resp.get("output_channel_names"),
|
||||
"output_width": resp.get("output_width"),
|
||||
"output_height": resp.get("output_height"),
|
||||
|
||||
"payload_sources": resp.get("payload_sources"),
|
||||
"payload_complete": resp.get("payload_complete"),
|
||||
|
||||
"source_camera": resp.get("source_camera"),
|
||||
"source_cameras": resp.get("source_cameras"),
|
||||
"camera_frames": resp.get("camera_frames"),
|
||||
|
||||
"multi_payload": resp.get("multi_payload", False),
|
||||
"payload_kind": resp.get("payload_kind"),
|
||||
"payload_parts": resp.get("payload_parts"),
|
||||
|
||||
"packed_width": resp.get("packed_width"),
|
||||
"packed_height": resp.get("packed_height"),
|
||||
"source_width": resp.get("source_width"),
|
||||
"source_height": resp.get("source_height"),
|
||||
"source_bayer_pattern": resp.get("source_bayer_pattern"),
|
||||
"source_bit_depth": resp.get("source_bit_depth"),
|
||||
|
||||
"size": resp.get("size"),
|
||||
"ts_pi": resp.get("ts_pi"),
|
||||
"ts_pi_monotonic": resp.get("ts_pi_monotonic"),
|
||||
|
||||
"dt_trigger": resp.get("dt_trigger"),
|
||||
"dt_settle": resp.get("dt_settle"),
|
||||
"dt_capture": resp.get("dt_capture"),
|
||||
"dt_process": resp.get("dt_process"),
|
||||
"dt_total_pi": resp.get("dt_total_pi"),
|
||||
|
||||
"dt_total_pc": time.perf_counter() - t0,
|
||||
}
|
||||
|
||||
# =========================================================
|
||||
# Comandos básicos
|
||||
# =========================================================
|
||||
|
||||
def ping(self):
|
||||
return self._send_command({"cmd": "ping"})
|
||||
|
||||
def get_status(self):
|
||||
return self._send_command({"cmd": "get_status"})
|
||||
|
||||
def get_config(self):
|
||||
return self._send_command({"cmd": "get_config"})
|
||||
|
||||
def begin(self, frame_type: str = "RAW_BRUTO", output_dtype: str = "uint8", capture_mode: str = "AUTO"):
|
||||
return self._send_command({
|
||||
"cmd": "begin",
|
||||
"frame_type": frame_type,
|
||||
"output_dtype": output_dtype,
|
||||
"capture_mode": capture_mode,
|
||||
})
|
||||
|
||||
def stop(self):
|
||||
return self._send_command({"cmd": "stop"})
|
||||
|
||||
def set_fps(self, fps: int):
|
||||
return self._send_command({"cmd": "set_fps", "value": fps})
|
||||
|
||||
def set_jpeg_quality(self, quality: int):
|
||||
return self._send_command({"cmd": "set_jpeg_quality", "value": quality})
|
||||
|
||||
def set_frame_type(self, frame_type: str):
|
||||
return self._send_command({"cmd": "set_frame_type", "value": frame_type})
|
||||
|
||||
def set_output_dtype(self, output_dtype: str):
|
||||
return self._send_command({"cmd": "set_output_dtype", "value": output_dtype})
|
||||
|
||||
def set_capture_mode(self, capture_mode: str):
|
||||
return self._send_command({"cmd": "set_capture_mode", "value": capture_mode})
|
||||
|
||||
def set_camera_enabled(self, index: int, enabled: bool):
|
||||
return self._send_command({
|
||||
"cmd": "set_camera_enabled",
|
||||
"index": index,
|
||||
"enabled": bool(enabled)
|
||||
})
|
||||
|
||||
def set_camera_bayer(self, index: int, bayer_pattern: str):
|
||||
return self._send_command({
|
||||
"cmd": "set_camera_bayer",
|
||||
"index": index,
|
||||
"pattern": bayer_pattern
|
||||
})
|
||||
|
||||
def set_camera_resolution(self, index: int, width: int, height: int):
|
||||
return self._send_command({
|
||||
"cmd": "set_camera_resolution",
|
||||
"index": index,
|
||||
"width": width,
|
||||
"height": height
|
||||
})
|
||||
|
||||
# =========================================================
|
||||
# Captura
|
||||
# =========================================================
|
||||
|
||||
def capture_frame(self):
|
||||
t0 = time.perf_counter()
|
||||
|
||||
resp = self._send_command({"cmd": "capture_frame"})
|
||||
if not resp.get("ok"):
|
||||
raise RuntimeError(resp.get("error", "Falha ao capturar frame"))
|
||||
|
||||
encoding = resp.get("encoding", "base64")
|
||||
meta = self._extract_meta(resp, t0)
|
||||
|
||||
if encoding == "base64":
|
||||
raw = base64.b64decode(resp["data"])
|
||||
arr = self._decode_single_array(raw, resp)
|
||||
|
||||
meta.update({
|
||||
"shape": list(arr.shape),
|
||||
"channels": int(arr.shape[0]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
|
||||
else (int(arr.shape[2]) if arr.ndim == 3 else 1),
|
||||
"width": int(arr.shape[2]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
|
||||
else (int(arr.shape[1]) if arr.ndim == 3 else int(arr.shape[1])),
|
||||
"height": int(arr.shape[1]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
|
||||
else int(arr.shape[0]),
|
||||
})
|
||||
return arr, meta
|
||||
|
||||
if encoding == "base64-multi":
|
||||
frames = self._decode_multi_frames_base64(resp)
|
||||
meta.update({
|
||||
"frames_meta": resp.get("camera_frames", {}),
|
||||
"decoded_shapes": {cam_id: list(arr.shape) for cam_id, arr in frames.items()},
|
||||
})
|
||||
return frames, meta
|
||||
|
||||
raise RuntimeError(f"encoding não suportado recebido do Pi: {encoding}")
|
||||
|
||||
def capture_frame_array(self):
|
||||
return self.capture_frame()
|
||||
|
||||
# =========================================================
|
||||
# Stream
|
||||
# =========================================================
|
||||
|
||||
def start_stream(self, host: str, port: int, fps: float):
|
||||
return self._send_command({
|
||||
"cmd": "start_stream",
|
||||
"host": host,
|
||||
"port": port,
|
||||
"fps": fps
|
||||
})
|
||||
|
||||
def stop_stream(self):
|
||||
return self._send_command({"cmd": "stop_stream"})
|
||||
|
||||
# =========================================================
|
||||
# Controles de câmera
|
||||
# =========================================================
|
||||
|
||||
def get_camera_controls(self):
|
||||
return self._send_command({"cmd": "get_camera_controls"})
|
||||
|
||||
def set_ae_enable(self, value: bool):
|
||||
return self._send_command({"cmd": "set_ae_enable", "value": bool(value)})
|
||||
|
||||
def set_awb_enable(self, value: bool):
|
||||
return self._send_command({"cmd": "set_awb_enable", "value": bool(value)})
|
||||
|
||||
def set_exposure_time(self, exposure_time_us: Optional[int] = None):
|
||||
return self._send_command({"cmd": "set_exposure_time", "value": exposure_time_us})
|
||||
|
||||
def clear_exposure_time(self):
|
||||
return self._send_command({"cmd": "clear_exposure_time"})
|
||||
|
||||
def set_analogue_gain(self, gain: float | None):
|
||||
return self._send_command({"cmd": "set_analogue_gain", "value": gain})
|
||||
|
||||
def clear_analogue_gain(self):
|
||||
return self._send_command({"cmd": "clear_analogue_gain"})
|
||||
|
||||
def set_colour_gains(self, r_gain: float, b_gain: float):
|
||||
return self._send_command({
|
||||
"cmd": "set_colour_gains",
|
||||
"r_gain": r_gain,
|
||||
"b_gain": b_gain
|
||||
})
|
||||
|
||||
def clear_colour_gains(self):
|
||||
return self._send_command({"cmd": "clear_colour_gains"})
|
||||
|
||||
def get_sensor_modes(self):
|
||||
return self._send_command({"cmd": "get_sensor_modes"})
|
||||
|
|
@ -0,0 +1,453 @@
|
|||
import os
|
||||
import numpy as np
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class RawProcessorCore:
|
||||
def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "GBRG"):
|
||||
self.sensor_width = sensor_width
|
||||
self.sensor_height = sensor_height
|
||||
self.bayer_pattern = bayer_pattern.upper()
|
||||
|
||||
def unpack_raw10_packed(
|
||||
self,
|
||||
packed_frame: np.ndarray,
|
||||
sensor_width: Optional[int] = None,
|
||||
sensor_height: Optional[int] = None
|
||||
):
|
||||
if packed_frame.ndim == 3 and packed_frame.shape[2] == 1:
|
||||
packed_frame = packed_frame[:, :, 0]
|
||||
|
||||
width = sensor_width if sensor_width is not None else self.sensor_width
|
||||
height = sensor_height if sensor_height is not None else self.sensor_height
|
||||
|
||||
if width % 4 != 0:
|
||||
raise ValueError(f"Largura {width} não é múltipla de 4 para RAW10 packed")
|
||||
|
||||
expected_packed_width = math.ceil(width * 10 / 8)
|
||||
|
||||
actual_h, actual_w = packed_frame.shape[:2]
|
||||
padding = actual_w - expected_packed_width
|
||||
|
||||
if actual_h != height:
|
||||
raise ValueError(
|
||||
f"[ERRO FRAME] Altura packed inesperada: {packed_frame.shape}, "
|
||||
f"esperado altura={height}"
|
||||
)
|
||||
|
||||
if actual_w < expected_packed_width:
|
||||
raise ValueError(
|
||||
f"[ERRO FRAME] Largura packed menor que a útil esperada: {packed_frame.shape}, "
|
||||
f"esperado pelo menos ({height}, {expected_packed_width})"
|
||||
)
|
||||
|
||||
if padding > 64:
|
||||
raise ValueError(
|
||||
f"[ERRO FRAME] Padding excessivo no packed: {packed_frame.shape}, "
|
||||
f"esperado útil ({height}, {expected_packed_width}), padding={padding}"
|
||||
)
|
||||
|
||||
packed_frame = packed_frame[:, :expected_packed_width]
|
||||
groups = packed_frame.reshape(height, width // 4, 5).astype(np.uint16)
|
||||
|
||||
b0 = groups[:, :, 0]
|
||||
b1 = groups[:, :, 1]
|
||||
b2 = groups[:, :, 2]
|
||||
b3 = groups[:, :, 3]
|
||||
b4 = groups[:, :, 4]
|
||||
|
||||
p0 = (b0 << 2) | ((b4 >> 0) & 0x03)
|
||||
p1 = (b1 << 2) | ((b4 >> 2) & 0x03)
|
||||
p2 = (b2 << 2) | ((b4 >> 4) & 0x03)
|
||||
p3 = (b3 << 2) | ((b4 >> 6) & 0x03)
|
||||
|
||||
raw16 = np.empty((height, width), dtype=np.uint16)
|
||||
raw16[:, 0::4] = p0
|
||||
raw16[:, 1::4] = p1
|
||||
raw16[:, 2::4] = p2
|
||||
raw16[:, 3::4] = p3
|
||||
|
||||
return raw16
|
||||
|
||||
def extract_bayer_channels(self, raw16: np.ndarray) -> dict:
|
||||
p = self.bayer_pattern
|
||||
|
||||
if p == "GBRG":
|
||||
g1 = raw16[0::2, 0::2]
|
||||
b = raw16[0::2, 1::2]
|
||||
r = raw16[1::2, 0::2]
|
||||
g2 = raw16[1::2, 1::2]
|
||||
elif p == "GRBG":
|
||||
g1 = raw16[0::2, 0::2]
|
||||
r = raw16[0::2, 1::2]
|
||||
b = raw16[1::2, 0::2]
|
||||
g2 = raw16[1::2, 1::2]
|
||||
elif p == "RGGB":
|
||||
r = raw16[0::2, 0::2]
|
||||
g1 = raw16[0::2, 1::2]
|
||||
g2 = raw16[1::2, 0::2]
|
||||
b = raw16[1::2, 1::2]
|
||||
elif p == "BGGR":
|
||||
b = raw16[0::2, 0::2]
|
||||
g1 = raw16[0::2, 1::2]
|
||||
g2 = raw16[1::2, 0::2]
|
||||
r = raw16[1::2, 1::2]
|
||||
else:
|
||||
raise ValueError(f"Padrão Bayer não suportado: {p}")
|
||||
|
||||
return {"R": r, "G1": g1, "G2": g2, "B": b}
|
||||
|
||||
def build_training_rgb(
|
||||
self,
|
||||
raw16: np.ndarray,
|
||||
output_dtype: str = "float32",
|
||||
bit_depth: int = 10,
|
||||
) -> np.ndarray:
|
||||
ch = self.extract_bayer_channels(raw16)
|
||||
|
||||
max_val = float((1 << bit_depth) - 1)
|
||||
|
||||
r = ch["R"].astype(np.float32) / max_val
|
||||
g = ((ch["G1"].astype(np.float32) + ch["G2"].astype(np.float32)) * 0.5) / max_val
|
||||
b = ch["B"].astype(np.float32) / max_val
|
||||
|
||||
chw = np.stack([r, g, b], axis=0).astype(np.float32)
|
||||
chw = np.clip(chw, 0.0, 1.0)
|
||||
|
||||
if output_dtype == "float32":
|
||||
return chw
|
||||
|
||||
if output_dtype == "uint8":
|
||||
return (chw * 255.0).clip(0, 255).astype(np.uint8)
|
||||
|
||||
if output_dtype == "uint16":
|
||||
return (chw * 65535.0).clip(0, 65535).astype(np.uint16)
|
||||
|
||||
raise ValueError(f"output_dtype não suportado: {output_dtype}")
|
||||
|
||||
def build_multispectral_tensor(self, bins_data, bins_meta):
|
||||
rgb = None
|
||||
re = None
|
||||
nir = None
|
||||
|
||||
for data, meta in zip(bins_data, bins_meta):
|
||||
role = meta.get("role")
|
||||
|
||||
if role == "rgb":
|
||||
rgb = data.astype(np.float32) / 255.0
|
||||
rgb = rgb.transpose(2, 0, 1)
|
||||
|
||||
elif role == "re":
|
||||
re = data.astype(np.float32) / 1023.0
|
||||
re = re[None, :, :]
|
||||
|
||||
elif role == "nir":
|
||||
nir = data.astype(np.float32) / 1023.0
|
||||
nir = nir[None, :, :]
|
||||
|
||||
if rgb is None:
|
||||
raise RuntimeError("RGB obrigatório")
|
||||
|
||||
arrays = [rgb]
|
||||
channel_names = ["R", "G", "B"]
|
||||
|
||||
if re is not None:
|
||||
arrays.append(re)
|
||||
channel_names.append("RE")
|
||||
|
||||
if nir is not None:
|
||||
arrays.append(nir)
|
||||
channel_names.append("NIR")
|
||||
|
||||
min_h = min(a.shape[1] for a in arrays)
|
||||
min_w = min(a.shape[2] for a in arrays)
|
||||
|
||||
arrays = [a[:, :min_h, :min_w] for a in arrays]
|
||||
tensor = np.concatenate(arrays, axis=0)
|
||||
|
||||
return tensor, channel_names
|
||||
|
||||
def build_infer_tensor_from_stream(self, frame, meta, channels_expected):
|
||||
"""
|
||||
Converte o frame vindo do stream do Pi em tensor (C,H,W) float32 0..1
|
||||
compatível com o modelo.
|
||||
Suporta:
|
||||
- RGB uint8/float32 já pronto
|
||||
- MULTISPEC uint8/float32 já pronto
|
||||
- RAW_BRUTO multi_payload (cam2 RGB + cam0/cam1 packed)
|
||||
"""
|
||||
frame_type = meta.get("frame_type")
|
||||
dtype_str = meta.get("dtype") or meta.get("output_dtype", "uint8")
|
||||
camera_frames = meta.get("camera_frames", {}) or {}
|
||||
|
||||
# -------------------------------------------------
|
||||
# RAW_BRUTO multi_payload
|
||||
# -------------------------------------------------
|
||||
if frame_type == "RAW_BRUTO":
|
||||
if not isinstance(frame, dict):
|
||||
raise RuntimeError("RAW_BRUTO esperado como dict de câmeras no modo multi")
|
||||
|
||||
arrays = []
|
||||
channel_names = []
|
||||
|
||||
# RGB USB
|
||||
if "cam2" in frame:
|
||||
rgb_bgr = frame["cam2"]
|
||||
if rgb_bgr.ndim != 3 or rgb_bgr.shape[2] != 3:
|
||||
raise RuntimeError(f"cam2 RGB inválida: shape={rgb_bgr.shape}")
|
||||
|
||||
rgb = rgb_bgr[:, :, ::-1].astype(np.float32) / 255.0
|
||||
rgb_chw = np.transpose(rgb, (2, 0, 1))
|
||||
arrays.append(rgb_chw)
|
||||
channel_names.extend(["R", "G", "B"])
|
||||
else:
|
||||
raise RuntimeError("RAW_BRUTO para inferência precisa incluir cam2 (RGB)")
|
||||
|
||||
# RE / NIR
|
||||
for cam_id, spec_name in (("cam0", "RE"), ("cam1", "NIR")):
|
||||
if cam_id not in frame:
|
||||
continue
|
||||
|
||||
packed = frame[cam_id]
|
||||
if packed.ndim == 3 and packed.shape[2] == 1:
|
||||
packed = packed[:, :, 0]
|
||||
|
||||
cam_meta = camera_frames.get(cam_id, {})
|
||||
packed_width = int(cam_meta.get("width", packed.shape[1]))
|
||||
height = int(cam_meta.get("height", packed.shape[0]))
|
||||
bayer = cam_meta.get("bayer_pattern", self.bayer_pattern)
|
||||
bit_depth = int(cam_meta.get("bit_depth", 10))
|
||||
|
||||
if bit_depth == 10:
|
||||
real_width = int((packed_width * 8) / 10)
|
||||
else:
|
||||
real_width = packed_width
|
||||
|
||||
rp = RawProcessorCore(
|
||||
sensor_width=real_width,
|
||||
sensor_height=height,
|
||||
bayer_pattern=bayer,
|
||||
)
|
||||
|
||||
raw16 = rp.unpack_raw10_packed(packed)
|
||||
|
||||
max_val = float((1 << bit_depth) - 1)
|
||||
single = np.clip(raw16.astype(np.float32) / max_val, 0.0, 1.0)[None, :, :]
|
||||
|
||||
arrays.append(single)
|
||||
channel_names.append(spec_name)
|
||||
|
||||
if len(arrays) < 2:
|
||||
raise RuntimeError("RAW_BRUTO requer RGB + pelo menos um canal espectral para inferência")
|
||||
|
||||
min_h = min(a.shape[1] for a in arrays)
|
||||
min_w = min(a.shape[2] for a in arrays)
|
||||
arrays = [a[:, :min_h, :min_w] for a in arrays]
|
||||
|
||||
raw_np = np.concatenate(arrays, axis=0)
|
||||
|
||||
if raw_np.shape[0] != channels_expected:
|
||||
raise RuntimeError(
|
||||
f"Tensor RAW_BRUTO montado com canais inesperados: {raw_np.shape[0]} | esperado={channels_expected} | got={channel_names}"
|
||||
)
|
||||
|
||||
return raw_np
|
||||
|
||||
# -------------------------------------------------
|
||||
# RGB ou MULTISPEC já pronto
|
||||
# -------------------------------------------------
|
||||
if frame_type == "RGB" or frame_type == "MULTISPEC":
|
||||
if not isinstance(frame, np.ndarray):
|
||||
raise RuntimeError(f"Frame {frame_type} esperado como ndarray")
|
||||
|
||||
if frame.ndim != 3:
|
||||
raise RuntimeError(f"Frame {frame_type} inválido: shape={frame.shape}")
|
||||
|
||||
if dtype_str == "uint8":
|
||||
raw_np = frame.astype(np.float32) / 255.0
|
||||
elif dtype_str == "float32":
|
||||
raw_np = frame.astype(np.float32)
|
||||
elif dtype_str == "uint16":
|
||||
raw_np = frame.astype(np.float32) / 65535.0
|
||||
else:
|
||||
raise RuntimeError(f"dtype {frame_type} não suportado: {dtype_str}")
|
||||
|
||||
if raw_np.shape[0] != channels_expected:
|
||||
raise RuntimeError(
|
||||
f"Frame {frame_type} com canais inesperados: {raw_np.shape[0]} | esperado={channels_expected}"
|
||||
)
|
||||
|
||||
return raw_np
|
||||
|
||||
|
||||
raise RuntimeError(f"frame_type não suportado para inferência: {frame_type}")
|
||||
|
||||
|
||||
def extract_camera_meta(self, meta_json: dict, cam_id: str) -> dict:
|
||||
cam_frames = meta_json.get("camera_frames", {}) or meta_json.get("stream_meta", {}).get("camera_frames", {})
|
||||
|
||||
cam = cam_frames.get(cam_id)
|
||||
if not cam:
|
||||
raise ValueError(f"Camera {cam_id} não encontrada no meta")
|
||||
|
||||
# Detecta RAW10 packed mono
|
||||
if int(cam.get("channels", 1)) == 1 and int(cam.get("bit_depth", 10)) == 10:
|
||||
packed_width = int(cam.get("width"))
|
||||
height = int(cam.get("height"))
|
||||
|
||||
# 🔥 converte packed → real
|
||||
real_width = int((packed_width * 8) / 10)
|
||||
|
||||
return {
|
||||
"camera_id": cam_id,
|
||||
"width": real_width,
|
||||
"height": height,
|
||||
"channels": 1,
|
||||
"bit_depth": 10,
|
||||
"shape": [height, packed_width], # packed shape
|
||||
"role": cam.get("role")
|
||||
}
|
||||
|
||||
# RGB
|
||||
else:
|
||||
width = int(cam.get("width"))
|
||||
height = int(cam.get("height"))
|
||||
channels = int(cam.get("channels", 3))
|
||||
|
||||
return {
|
||||
"camera_id": cam_id,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"channels": channels,
|
||||
"bit_depth": int(cam.get("bit_depth", 8)),
|
||||
"shape": [height, width, channels], # 🔥 AQUI está a correção
|
||||
"role": cam.get("role")
|
||||
}
|
||||
|
||||
# ============================================================
|
||||
# RAW10 PACKED
|
||||
# ============================================================
|
||||
|
||||
def packed_width_for_raw10(self, sensor_width: int = None) -> int:
|
||||
width = sensor_width if sensor_width is not None else self.sensor_width
|
||||
return math.ceil(width * 10 / 8)
|
||||
|
||||
def pack_raw10_packed(self, raw16: np.ndarray) -> np.ndarray:
|
||||
h, w = raw16.shape
|
||||
|
||||
if w % 4 != 0:
|
||||
raise ValueError(f"Width precisa ser múltiplo de 4 para pack otimizado. Veio {w}")
|
||||
|
||||
raw16 = np.clip(raw16, 0, 1023).astype(np.uint16)
|
||||
|
||||
p0 = raw16[:, 0::4]
|
||||
p1 = raw16[:, 1::4]
|
||||
p2 = raw16[:, 2::4]
|
||||
p3 = raw16[:, 3::4]
|
||||
|
||||
b0 = (p0 >> 2).astype(np.uint8)
|
||||
b1 = (p1 >> 2).astype(np.uint8)
|
||||
b2 = (p2 >> 2).astype(np.uint8)
|
||||
b3 = (p3 >> 2).astype(np.uint8)
|
||||
|
||||
b4 = (
|
||||
((p0 & 0x03) << 0) |
|
||||
((p1 & 0x03) << 2) |
|
||||
((p2 & 0x03) << 4) |
|
||||
((p3 & 0x03) << 6)
|
||||
).astype(np.uint8)
|
||||
|
||||
packed = np.empty((h, w // 4, 5), dtype=np.uint8)
|
||||
packed[:, :, 0] = b0
|
||||
packed[:, :, 1] = b1
|
||||
packed[:, :, 2] = b2
|
||||
packed[:, :, 3] = b3
|
||||
packed[:, :, 4] = b4
|
||||
|
||||
return packed.reshape(h, w // 4 * 5)
|
||||
|
||||
def load_raw10_packed_file(self, path: str, width: int, height: int) -> np.ndarray:
|
||||
packed_width = self.packed_width_for_raw10(width)
|
||||
|
||||
expected_size = height * packed_width
|
||||
actual_size = os.path.getsize(path)
|
||||
|
||||
if actual_size != expected_size:
|
||||
raise ValueError(
|
||||
f"Tamanho inválido RAW10: {actual_size}, esperado {expected_size} em {path}"
|
||||
)
|
||||
|
||||
packed = np.fromfile(path, dtype=np.uint8).reshape(height, packed_width)
|
||||
return self.unpack_raw10_packed(packed, sensor_width=width, sensor_height=height)
|
||||
|
||||
def save_raw10_packed_file(self, path: str, raw16: np.ndarray):
|
||||
packed = self.pack_raw10_packed(raw16)
|
||||
packed.tofile(path)
|
||||
|
||||
# ============================================================
|
||||
# RGB UINT8
|
||||
# ============================================================
|
||||
|
||||
def load_rgb_u8_file(self, path: str, shape) -> np.ndarray:
|
||||
arr = np.fromfile(path, dtype=np.uint8)
|
||||
|
||||
expected = np.prod(shape)
|
||||
if arr.size != expected:
|
||||
raise ValueError(
|
||||
f"Tamanho inválido RGB: {arr.size}, esperado {expected} em {path}"
|
||||
)
|
||||
|
||||
return arr.reshape(shape)
|
||||
|
||||
|
||||
def save_rgb_u8_file(self, path: str, arr: np.ndarray):
|
||||
arr.astype(np.uint8).tofile(path)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# DISPATCHER (O MAIS IMPORTANTE)
|
||||
# ============================================================
|
||||
|
||||
def load_native_bin(self, path: str, cam_meta: dict) -> np.ndarray:
|
||||
"""
|
||||
Decide automaticamente como carregar o .bin baseado no meta.
|
||||
"""
|
||||
|
||||
channels = int(cam_meta.get("channels", 1))
|
||||
bit_depth = int(cam_meta.get("bit_depth", 10))
|
||||
shape = cam_meta.get("shape")
|
||||
|
||||
if channels == 1 and bit_depth == 10:
|
||||
width = int(cam_meta["width"])
|
||||
height = int(cam_meta["height"])
|
||||
return self.load_raw10_packed_file(path, width, height)
|
||||
|
||||
elif channels == 3 and bit_depth == 8:
|
||||
return self.load_rgb_u8_file(path, shape)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Formato não suportado: channels={channels}, bit_depth={bit_depth}"
|
||||
)
|
||||
|
||||
|
||||
def save_native_bin(self, path: str, arr: np.ndarray, cam_meta: dict):
|
||||
"""
|
||||
Salva no formato correto baseado no meta.
|
||||
"""
|
||||
|
||||
channels = int(cam_meta.get("channels", 1))
|
||||
bit_depth = int(cam_meta.get("bit_depth", 10))
|
||||
|
||||
if channels == 1 and bit_depth == 10:
|
||||
self.save_raw10_packed_file(path, arr)
|
||||
|
||||
elif channels == 3 and bit_depth == 8:
|
||||
self.save_rgb_u8_file(path, arr)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Formato não suportado para salvar: channels={channels}, bit_depth={bit_depth}"
|
||||
)
|
||||
|
||||
|
|
@ -0,0 +1,125 @@
|
|||
import cv2
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class RawProcessorPreview:
|
||||
def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "GBRG"):
|
||||
self.sensor_width = sensor_width
|
||||
self.sensor_height = sensor_height
|
||||
self.bayer_pattern = bayer_pattern.upper()
|
||||
|
||||
def raw16_to_vis8(
|
||||
self, raw16: np.ndarray,
|
||||
black_level: Optional[int] = None,
|
||||
white_level: Optional[int] = None,
|
||||
gamma: float = 2.2,
|
||||
bit_depth: int = 10
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Conversão para visualização:
|
||||
- auto-level
|
||||
- gamma
|
||||
"""
|
||||
max_val = float((1 << bit_depth) - 1)
|
||||
|
||||
raw = raw16.astype(np.float32)
|
||||
|
||||
if black_level is None:
|
||||
black_level = float(raw.min())
|
||||
if white_level is None:
|
||||
white_level = float(raw.max())
|
||||
|
||||
if white_level <= black_level:
|
||||
norm = raw / max_val
|
||||
else:
|
||||
norm = (raw - black_level) / (white_level - black_level)
|
||||
|
||||
norm = np.clip(norm, 0.0, 1.0)
|
||||
|
||||
if gamma is not None and gamma > 0:
|
||||
norm = np.power(norm, 1.0 / gamma)
|
||||
|
||||
return (norm * 255.0).clip(0, 255).astype(np.uint8)
|
||||
|
||||
def _debayer_code(self):
|
||||
mapping = {
|
||||
# Troque de BayerGB para BayerGR para inverter R e B
|
||||
"GBRG": cv2.COLOR_BayerGR2BGR,
|
||||
"GRBG": cv2.COLOR_BayerGB2BGR,
|
||||
"RGGB": cv2.COLOR_BayerBG2BGR,
|
||||
"BGGR": cv2.COLOR_BayerRG2BGR,
|
||||
}
|
||||
if self.bayer_pattern not in mapping:
|
||||
raise ValueError(f"Padrão Bayer não suportado: {self.bayer_pattern}")
|
||||
return mapping[self.bayer_pattern]
|
||||
|
||||
def apply_preview_white_balance(self, bgr: np.ndarray, strength: float = 1.0) -> np.ndarray:
|
||||
"""
|
||||
Gray-world simples para deixar o preview mais agradável.
|
||||
Não usar no raw de treino.
|
||||
"""
|
||||
img = bgr.astype(np.float32)
|
||||
|
||||
mean_b = float(img[:, :, 0].mean())
|
||||
mean_g = float(img[:, :, 1].mean())
|
||||
mean_r = float(img[:, :, 2].mean())
|
||||
|
||||
mean_gray = (mean_b + mean_g + mean_r) / 3.0
|
||||
|
||||
eps = 1e-6
|
||||
gain_b = mean_gray / max(mean_b, eps)
|
||||
gain_g = mean_gray / max(mean_g, eps)
|
||||
gain_r = mean_gray / max(mean_r, eps)
|
||||
|
||||
# strength=1 aplica total, strength=0 não aplica
|
||||
gain_b = 1.0 + (gain_b - 1.0) * strength
|
||||
gain_g = 1.0 + (gain_g - 1.0) * strength
|
||||
gain_r = 1.0 + (gain_r - 1.0) * strength
|
||||
|
||||
img[:, :, 0] *= gain_b
|
||||
img[:, :, 1] *= gain_g
|
||||
img[:, :, 2] *= gain_r
|
||||
|
||||
return np.clip(img, 0, 255).astype(np.uint8)
|
||||
|
||||
def apply_preview_contrast(self, bgr: np.ndarray, alpha: float = 1.08, beta: float = 0.0) -> np.ndarray:
|
||||
"""
|
||||
Ajuste leve de contraste/brilho para preview.
|
||||
"""
|
||||
out = cv2.convertScaleAbs(bgr, alpha=alpha, beta=beta)
|
||||
return out
|
||||
|
||||
def raw16_to_preview_bgr(
|
||||
self,
|
||||
raw16: np.ndarray,
|
||||
gamma: float = 2.2,
|
||||
wb_strength: float = 0.8,
|
||||
apply_wb: bool = True,
|
||||
apply_contrast: bool = True,
|
||||
bit_depth: int = 10,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Pipeline de preview bonito:
|
||||
1. auto-level + gamma no mosaico
|
||||
2. demosaic
|
||||
3. white balance simples
|
||||
4. leve contraste final
|
||||
"""
|
||||
vis8 = self.raw16_to_vis8(raw16, gamma=gamma, bit_depth=bit_depth)
|
||||
bgr = cv2.cvtColor(vis8, self._debayer_code())
|
||||
|
||||
if apply_wb:
|
||||
bgr = self.apply_preview_white_balance(bgr, strength=wb_strength)
|
||||
|
||||
if apply_contrast:
|
||||
bgr = self.apply_preview_contrast(bgr, alpha=1.08, beta=0.0)
|
||||
|
||||
return bgr
|
||||
|
||||
def raw16_to_preview_jpg_bytes(self, raw16: np.ndarray, jpeg_quality: int = 95) -> bytes:
|
||||
bgr = self.raw16_to_preview_bgr(raw16)
|
||||
ok, enc = cv2.imencode(".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_quality)])
|
||||
if not ok:
|
||||
raise RuntimeError("Falha ao codificar preview JPG")
|
||||
return enc.tobytes()
|
||||
|
|
@ -0,0 +1,324 @@
|
|||
import json
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from numcodecs import Blosc
|
||||
import numpy as np
|
||||
|
||||
|
||||
class StreamReceiver:
|
||||
def __init__(self, host="0.0.0.0", port=6001):
|
||||
self.host = host
|
||||
self.port = port
|
||||
|
||||
self._server_sock = None
|
||||
self._client_sock = None
|
||||
self._thread = None
|
||||
self._running = False
|
||||
|
||||
self.last_frame = None
|
||||
self.last_meta = None
|
||||
self.last_receive_ts = None
|
||||
|
||||
self._codec = None
|
||||
self._codec_signature = None
|
||||
|
||||
@property
|
||||
def is_running(self):
|
||||
return self._running
|
||||
|
||||
# =========================================================
|
||||
# Lifecycle
|
||||
# =========================================================
|
||||
|
||||
def start(self):
|
||||
if self._running:
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._thread = threading.Thread(target=self._worker, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def stop(self):
|
||||
self._running = False
|
||||
|
||||
try:
|
||||
if self._client_sock:
|
||||
self._client_sock.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
if self._server_sock:
|
||||
self._server_sock.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self._client_sock = None
|
||||
self._server_sock = None
|
||||
|
||||
# =========================================================
|
||||
# Socket helpers
|
||||
# =========================================================
|
||||
|
||||
def _recv_exact(self, sock: socket.socket, n: int) -> bytes:
|
||||
chunks = []
|
||||
remaining = n
|
||||
|
||||
while remaining > 0:
|
||||
chunk = sock.recv(remaining)
|
||||
if not chunk:
|
||||
raise ConnectionError("Conexão encerrada durante recv")
|
||||
chunks.append(chunk)
|
||||
remaining -= len(chunk)
|
||||
|
||||
return b"".join(chunks)
|
||||
|
||||
# =========================================================
|
||||
# Worker principal
|
||||
# =========================================================
|
||||
|
||||
def _worker(self):
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
|
||||
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
server.bind((self.host, self.port))
|
||||
server.listen(1)
|
||||
server.settimeout(1.0)
|
||||
|
||||
self._server_sock = server
|
||||
print(f"[INFO] StreamReceiver ouvindo em {self.host}:{self.port}")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
client, addr = server.accept()
|
||||
except socket.timeout:
|
||||
continue
|
||||
|
||||
print(f"[INFO] StreamReceiver conectado por {addr}")
|
||||
self._client_sock = client
|
||||
|
||||
with client:
|
||||
while self._running:
|
||||
header_len = int.from_bytes(self._recv_exact(client, 4), "big")
|
||||
header_bytes = self._recv_exact(client, header_len)
|
||||
header = json.loads(header_bytes.decode("utf-8"))
|
||||
|
||||
payload_len = int.from_bytes(self._recv_exact(client, 4), "big")
|
||||
payload_comp = self._recv_exact(client, payload_len)
|
||||
|
||||
self._ensure_codec(header)
|
||||
if header.get("codec_family") == "none":
|
||||
payload = payload_comp
|
||||
else:
|
||||
payload = self._codec.decode(payload_comp)
|
||||
|
||||
expected = int(header["payload_size_raw"])
|
||||
if len(payload) != expected:
|
||||
raise ValueError(
|
||||
f"Tamanho descomprimido inválido: {len(payload)} != {expected}"
|
||||
)
|
||||
|
||||
payload_kind = header.get("payload_kind")
|
||||
if header.get("multi_payload", False) or payload_kind == "multi_array":
|
||||
frame = self._decode_multi_payload(payload, header)
|
||||
else:
|
||||
frame = self._decode_single_payload(payload, header)
|
||||
|
||||
self.last_frame = frame
|
||||
self.last_meta = header
|
||||
self.last_receive_ts = time.perf_counter()
|
||||
|
||||
print("[INFO] StreamReceiver cliente desconectado")
|
||||
self._client_sock = None
|
||||
|
||||
except Exception as e:
|
||||
print(f"[WARN] StreamReceiver encerrado com erro: {e}")
|
||||
|
||||
finally:
|
||||
self._running = False
|
||||
self._client_sock = None
|
||||
self._server_sock = None
|
||||
|
||||
# =========================================================
|
||||
# Decode helpers
|
||||
# =========================================================
|
||||
|
||||
def _numpy_dtype_from_string(self, dtype_str: str):
|
||||
mapping = {
|
||||
"uint8": np.uint8,
|
||||
"uint16": np.uint16,
|
||||
"float32": np.float32,
|
||||
}
|
||||
|
||||
if dtype_str not in mapping:
|
||||
raise RuntimeError(f"dtype não suportado: {dtype_str}")
|
||||
|
||||
return mapping[dtype_str]
|
||||
|
||||
def _numpy_dtype_from_header(self, header: dict):
|
||||
dtype_str = header.get("dtype") or header.get("output_dtype") or "uint8"
|
||||
|
||||
if dtype_str == "multi":
|
||||
# fallback conservador para protocolos mais antigos
|
||||
return np.uint16
|
||||
|
||||
return self._numpy_dtype_from_string(dtype_str)
|
||||
|
||||
def _reshape_from_shape(self, payload: bytes, dtype_str: str, shape):
|
||||
dtype = self._numpy_dtype_from_string(dtype_str)
|
||||
arr = np.frombuffer(payload, dtype=dtype)
|
||||
return arr.reshape(tuple(shape))
|
||||
|
||||
def _reshape_from_layout(self, payload: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
|
||||
dtype = self._numpy_dtype_from_string(dtype_str)
|
||||
arr = np.frombuffer(payload, dtype=dtype)
|
||||
|
||||
if layout == "HW":
|
||||
return arr.reshape(height, width)
|
||||
|
||||
if layout == "CHW":
|
||||
return arr.reshape(channels, height, width)
|
||||
|
||||
if layout == "HWC":
|
||||
return arr.reshape(height, width, channels)
|
||||
|
||||
raise RuntimeError(f"Layout não suportado: {layout}")
|
||||
|
||||
def _decode_single_payload(self, payload: bytes, header: dict):
|
||||
payload_parts = header.get("payload_parts", []) or []
|
||||
first_part = payload_parts[0] if payload_parts else {}
|
||||
|
||||
dtype_str = first_part.get("dtype") or header.get("dtype") or header.get("output_dtype") or "uint8"
|
||||
shape = first_part.get("shape")
|
||||
|
||||
if shape:
|
||||
return self._reshape_from_shape(payload, dtype_str, shape)
|
||||
|
||||
height = int(header.get("output_height", header.get("height")))
|
||||
width = int(header.get("output_width", header.get("width")))
|
||||
channels = int(header.get("output_channels", header.get("channels", 1)))
|
||||
layout = header.get("output_layout", "HWC")
|
||||
|
||||
return self._reshape_from_layout(
|
||||
payload=payload,
|
||||
dtype_str=dtype_str,
|
||||
layout=layout,
|
||||
width=width,
|
||||
height=height,
|
||||
channels=channels,
|
||||
)
|
||||
|
||||
def _decode_multi_payload(self, payload: bytes, header: dict):
|
||||
payload_parts = header.get("payload_parts", []) or []
|
||||
camera_frames = header.get("camera_frames", {}) or {}
|
||||
|
||||
frames = {}
|
||||
offset = 0
|
||||
|
||||
for part in payload_parts:
|
||||
cam_id = part.get("camera_id")
|
||||
if not cam_id:
|
||||
raise RuntimeError("payload_parts sem camera_id")
|
||||
|
||||
if offset + 4 > len(payload):
|
||||
raise RuntimeError("Payload multi truncado ao ler tamanho da parte")
|
||||
|
||||
part_size = int.from_bytes(payload[offset:offset + 4], "big")
|
||||
offset += 4
|
||||
|
||||
if offset + part_size > len(payload):
|
||||
raise RuntimeError(f"Payload multi truncado ao ler dados de {cam_id}")
|
||||
|
||||
part_bytes = payload[offset:offset + part_size]
|
||||
offset += part_size
|
||||
|
||||
cam_meta = camera_frames.get(cam_id, {})
|
||||
|
||||
dtype_str = part.get("dtype") or header.get("dtype") or header.get("output_dtype")
|
||||
shape = part.get("shape")
|
||||
|
||||
if dtype_str == "multi" or dtype_str is None:
|
||||
channels = int(part.get("channels", cam_meta.get("channels", 1)))
|
||||
dtype_str = "uint8" if channels > 1 else "uint16"
|
||||
|
||||
if shape:
|
||||
frame = self._reshape_from_shape(part_bytes, dtype_str, shape)
|
||||
else:
|
||||
width = int(part.get("width", cam_meta.get("width", 0)))
|
||||
height = int(part.get("height", cam_meta.get("height", 0)))
|
||||
channels = int(part.get("channels", cam_meta.get("channels", 1)))
|
||||
|
||||
layout = "HWC" if channels > 1 else "HW"
|
||||
|
||||
frame = self._reshape_from_layout(
|
||||
payload=part_bytes,
|
||||
dtype_str=dtype_str,
|
||||
layout=layout,
|
||||
width=width,
|
||||
height=height,
|
||||
channels=channels,
|
||||
)
|
||||
|
||||
frames[cam_id] = frame
|
||||
|
||||
if offset != len(payload):
|
||||
raise RuntimeError(
|
||||
f"Payload multi com bytes sobrando: consumidos={offset}, total={len(payload)}"
|
||||
)
|
||||
|
||||
return frames
|
||||
|
||||
# =========================================================
|
||||
# Codec
|
||||
# =========================================================
|
||||
|
||||
def _normalize_shuffle(self, shuffle_value):
|
||||
if isinstance(shuffle_value, int):
|
||||
return shuffle_value
|
||||
|
||||
mapping = {
|
||||
"NOSHUFFLE": Blosc.NOSHUFFLE,
|
||||
"SHUFFLE": Blosc.SHUFFLE,
|
||||
"BITSHUFFLE": Blosc.BITSHUFFLE,
|
||||
}
|
||||
|
||||
key = str(shuffle_value).upper()
|
||||
if key not in mapping:
|
||||
raise ValueError(f"shuffle inválido: {shuffle_value}")
|
||||
|
||||
return mapping[key]
|
||||
|
||||
def _build_codec_from_header(self, header: dict):
|
||||
family = header.get("codec_family")
|
||||
name = header.get("codec_name")
|
||||
params = dict(header.get("codec_params", {}))
|
||||
|
||||
if family != "numcodecs":
|
||||
raise ValueError(f"Família de codec não suportada: {family}")
|
||||
|
||||
if name == "blosc":
|
||||
params["shuffle"] = self._normalize_shuffle(params.get("shuffle", "SHUFFLE"))
|
||||
return Blosc(**params)
|
||||
|
||||
raise ValueError(f"Codec numcodecs não suportado: {name}")
|
||||
|
||||
def _get_codec_signature_from_header(self, header: dict):
|
||||
return (
|
||||
header.get("codec_family"),
|
||||
header.get("codec_name"),
|
||||
tuple(sorted(dict(header.get("codec_params", {})).items()))
|
||||
)
|
||||
|
||||
def _ensure_codec(self, header: dict):
|
||||
family = header.get("codec_family")
|
||||
|
||||
if family == "none":
|
||||
self._codec = None
|
||||
self._codec_signature = ("none", None, ())
|
||||
return
|
||||
|
||||
sig = self._get_codec_signature_from_header(header)
|
||||
if self._codec is None or self._codec_signature != sig:
|
||||
self._codec = self._build_codec_from_header(header)
|
||||
self._codec_signature = sig
|
||||
|
|
@ -6,7 +6,7 @@ import cv2
|
|||
import numpy as np
|
||||
import torch
|
||||
from gal5000.gal_service import Gal5000Camera
|
||||
from raw_segformer_service import RawSegformerService
|
||||
from Python.OAK.datasets.gal5000.raw_segformer_service import RawSegformerService
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Device: {device}")
|
||||
|
|
|
|||
|
|
@ -113,6 +113,17 @@ def desenhar_legenda_horizontal(colormap_rgb, classes, altura=30, largura_por_cl
|
|||
|
||||
return legenda
|
||||
|
||||
def _infer_ignore_id(ignore_rgb, default_id=255):
|
||||
import numpy as _np
|
||||
if isinstance(ignore_rgb, (list, tuple)):
|
||||
if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, _np.integer)):
|
||||
return int(ignore_rgb[0])
|
||||
if len(ignore_rgb) == 3:
|
||||
return default_id
|
||||
if isinstance(ignore_rgb, (int, _np.integer)):
|
||||
return int(ignore_rgb)
|
||||
return default_id
|
||||
|
||||
# ----------------------------
|
||||
# Helpers ROI
|
||||
# ----------------------------
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@ class CameraManager:
|
|||
self.initialized = False
|
||||
|
||||
self.camera_lock = RLock()
|
||||
self.frame_lock = Lock()
|
||||
|
||||
self.cameras_runtime = {}
|
||||
self._reconfigure_needed = False
|
||||
|
|
@ -45,43 +44,34 @@ class CameraManager:
|
|||
|
||||
"stop_event": threading.Event(),
|
||||
"thread": None,
|
||||
|
||||
"last_signature": None,
|
||||
"last_read_ts": None,
|
||||
"last_new_frame_ts": None,
|
||||
|
||||
"frame_lock": Lock(),
|
||||
}
|
||||
|
||||
# =========================================================
|
||||
# Seleção de câmeras necessárias
|
||||
# =========================================================
|
||||
def _make_frame_signature(self, frame: np.ndarray):
|
||||
if frame is None:
|
||||
return None
|
||||
|
||||
def _get_required_camera_ids(self):
|
||||
if hasattr(self.state, "get_required_camera_ids_for_frame_type"):
|
||||
return self.state.get_required_camera_ids_for_frame_type()
|
||||
h, w = frame.shape[:2]
|
||||
|
||||
frame_type = self.state.frame_type
|
||||
resolved_mode = self.state.resolve_capture_mode()
|
||||
# pega uma grade simples e barata
|
||||
step_y = max(1, h // 16)
|
||||
step_x = max(1, w // 16)
|
||||
|
||||
if frame_type == "RGB":
|
||||
cam = getattr(self.state, "rgb_camera_id", None)
|
||||
return [cam] if cam else []
|
||||
sample = frame[::step_y, ::step_x]
|
||||
|
||||
if frame_type == "MULTISPEC":
|
||||
ids = []
|
||||
for attr in ("rgb_camera_id", "re_camera_id", "nir_camera_id"):
|
||||
cam_id = getattr(self.state, attr, None)
|
||||
if cam_id:
|
||||
ids.append(cam_id)
|
||||
return ids
|
||||
|
||||
if frame_type == "RAW_BRUTO":
|
||||
if resolved_mode == "TRIPLE":
|
||||
return [cam.id for cam in self.state.cameras if cam.enabled]
|
||||
if resolved_mode == "SINGLE":
|
||||
for role in ("rgb", "re", "nir"):
|
||||
if hasattr(self.state, "get_active_camera_by_role"):
|
||||
cam = self.state.get_active_camera_by_role(role)
|
||||
if cam:
|
||||
return [cam.id]
|
||||
return []
|
||||
|
||||
return []
|
||||
# reduz para assinatura curtinha
|
||||
return (
|
||||
sample.shape,
|
||||
int(sample.mean()),
|
||||
int(sample.std()),
|
||||
int(sample[0, 0, 0]) if sample.ndim == 3 else int(sample[0, 0]),
|
||||
int(sample[-1, -1, 1]) if sample.ndim == 3 and sample.shape[2] > 1 else 0,
|
||||
)
|
||||
|
||||
# =========================================================
|
||||
# Inicialização
|
||||
|
|
@ -91,7 +81,7 @@ class CameraManager:
|
|||
with self.camera_lock:
|
||||
self.stop()
|
||||
|
||||
required_ids = set(self._get_required_camera_ids())
|
||||
required_ids = set(self.state.get_required_camera_ids_for_frame_type())
|
||||
if not required_ids:
|
||||
print("[WARN] Nenhuma câmera requerida para o frame_type/capture_mode atual")
|
||||
|
||||
|
|
@ -149,7 +139,7 @@ class CameraManager:
|
|||
picam2 = Picamera2(camera_num=cam.index)
|
||||
|
||||
config = picam2.create_video_configuration(
|
||||
main={"size": (640, 480), "format": "RGB888"},
|
||||
#main={"size": (640, 480), "format": "RGB888"},
|
||||
raw={"size": (cam.width, cam.height)},
|
||||
buffer_count=6
|
||||
)
|
||||
|
|
@ -305,7 +295,7 @@ class CameraManager:
|
|||
|
||||
raw = request.make_array("raw")
|
||||
|
||||
with self.frame_lock:
|
||||
with runtime["frame_lock"]:
|
||||
if (
|
||||
runtime["buffer"] is None or
|
||||
runtime["buffer"].shape != raw.shape or
|
||||
|
|
@ -329,11 +319,22 @@ class CameraManager:
|
|||
def _update_loop_opencv(self, runtime):
|
||||
cap = runtime["cap"]
|
||||
|
||||
target_fps = max(1.0, float(self.state.fps or 10))
|
||||
min_period = 1.0 / target_fps
|
||||
t_start = time.perf_counter()
|
||||
|
||||
ok, frame = cap.read()
|
||||
read_ts = time.perf_counter()
|
||||
|
||||
if not ok or frame is None:
|
||||
raise RuntimeError("Falha ao ler frame da câmera USB")
|
||||
|
||||
with self.frame_lock:
|
||||
signature = self._make_frame_signature(frame)
|
||||
|
||||
with runtime["frame_lock"]:
|
||||
runtime["last_read_ts"] = read_ts
|
||||
|
||||
if signature != runtime.get("last_signature"):
|
||||
if (
|
||||
runtime["buffer"] is None or
|
||||
runtime["buffer"].shape != frame.shape or
|
||||
|
|
@ -345,7 +346,14 @@ class CameraManager:
|
|||
|
||||
runtime["last_frame"] = runtime["buffer"]
|
||||
runtime["frame_id"] += 1
|
||||
runtime["frame_ts"] = time.perf_counter()
|
||||
runtime["frame_ts"] = read_ts
|
||||
runtime["last_new_frame_ts"] = read_ts
|
||||
runtime["last_signature"] = signature
|
||||
|
||||
dt = time.perf_counter() - t_start
|
||||
sleep_s = min_period - dt
|
||||
if sleep_s > 0:
|
||||
time.sleep(sleep_s)
|
||||
|
||||
# =========================================================
|
||||
# Leitura consolidada
|
||||
|
|
@ -354,12 +362,12 @@ class CameraManager:
|
|||
def capture_raw_frames(self):
|
||||
result = {}
|
||||
|
||||
with self.frame_lock:
|
||||
for cam_id, runtime in self.cameras_runtime.items():
|
||||
with runtime["frame_lock"]:
|
||||
if runtime["last_frame"] is None:
|
||||
continue
|
||||
|
||||
frame = runtime["last_frame"]
|
||||
frame = runtime["last_frame"].copy()
|
||||
h, w = frame.shape[:2]
|
||||
channels = frame.shape[2] if frame.ndim == 3 else 1
|
||||
|
||||
|
|
@ -369,7 +377,8 @@ class CameraManager:
|
|||
h,
|
||||
channels,
|
||||
runtime["frame_id"],
|
||||
runtime["frame_ts"]
|
||||
runtime["frame_ts"],
|
||||
runtime["last_read_ts"]
|
||||
)
|
||||
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ class FrameService:
|
|||
|
||||
return rp
|
||||
|
||||
def _capture_with_retry(self, required_sources, max_attempts=10, retry_delay_s=0.02):
|
||||
def _capture_with_retry(self, required_sources, previous_ids=None, max_attempts=10, retry_delay_s=0.02):
|
||||
last_frames = None
|
||||
|
||||
for _ in range(max_attempts):
|
||||
|
|
@ -55,10 +55,15 @@ class FrameService:
|
|||
ok = False
|
||||
break
|
||||
|
||||
frame, width, height, channels, frame_id, frame_ts = info
|
||||
frame, width, height, channels, frame_id, frame_ts, last_read_ts = info
|
||||
if frame is None or width <= 0 or height <= 0 or channels <= 0:
|
||||
ok = False
|
||||
break
|
||||
if previous_ids is not None:
|
||||
prev_id = previous_ids.get(cam_id)
|
||||
if prev_id is not None and frame_id <= prev_id:
|
||||
ok = False
|
||||
break
|
||||
|
||||
if ok:
|
||||
return frames
|
||||
|
|
@ -67,7 +72,7 @@ class FrameService:
|
|||
time.sleep(retry_delay_s)
|
||||
|
||||
raise RuntimeError(
|
||||
f"Capture retornou frames insuficientes após {max_attempts} tentativas. "
|
||||
f"Capture não atingiu frames válidos/novos após {max_attempts} tentativas... "
|
||||
f"required_sources={required_sources}, received_sources={list((last_frames or {}).keys())}"
|
||||
)
|
||||
|
||||
|
|
@ -170,6 +175,13 @@ class FrameService:
|
|||
has_csi_source = any(getattr(cam, "interface", "").upper() == "CSI" for cam in required_cams)
|
||||
|
||||
try:
|
||||
previous_ids = {}
|
||||
frames_before = self.camera.capture_raw_frames()
|
||||
for cam_id in required_sources:
|
||||
info = frames_before.get(cam_id)
|
||||
if info is not None:
|
||||
previous_ids[cam_id] = info[4] # frame_id
|
||||
|
||||
# Trigger só vale para as CSI RE/NIR.
|
||||
# A USB RGB não responde a esse trigger.
|
||||
if self.state.trigger_enabled and has_csi_source:
|
||||
|
|
@ -186,7 +198,7 @@ class FrameService:
|
|||
dt_settle = settle_end - settle_start
|
||||
|
||||
cap_start = time.perf_counter()
|
||||
raw_frames = self._capture_with_retry(required_sources)
|
||||
raw_frames = self._capture_with_retry(required_sources, previous_ids=previous_ids)
|
||||
cap_end = time.perf_counter()
|
||||
dt_capture = cap_end - cap_start
|
||||
|
||||
|
|
@ -208,12 +220,13 @@ class FrameService:
|
|||
|
||||
camera_frames_meta = {}
|
||||
for cam_id, info in raw_frames.items():
|
||||
frame, width, height, channels, cam_frame_id, cam_frame_ts = info
|
||||
frame, width, height, channels, cam_frame_id, cam_frame_ts, cam_last_read_ts = info
|
||||
cam = self._get_camera_spec(cam_id)
|
||||
|
||||
camera_frames_meta[cam_id] = {
|
||||
"camera_frame_id": int(cam_frame_id),
|
||||
"camera_frame_ts": cam_frame_ts,
|
||||
"camera_last_read_ts": cam_last_read_ts,
|
||||
"width": int(width),
|
||||
"height": int(height),
|
||||
"channels": int(channels),
|
||||
|
|
@ -295,7 +308,7 @@ class FrameService:
|
|||
|
||||
if len(sources) == 1:
|
||||
cam_id = sources[0]
|
||||
frame, width, height, channels, cam_frame_id, cam_frame_ts = raw_frames[cam_id]
|
||||
frame, width, height, channels, cam_frame_id, cam_frame_ts, cam_last_read_ts = raw_frames[cam_id]
|
||||
cam = self._get_camera_spec(cam_id)
|
||||
|
||||
if channels == 1:
|
||||
|
|
@ -333,7 +346,7 @@ class FrameService:
|
|||
sources_meta = []
|
||||
|
||||
for cam_id in sources:
|
||||
frame, width, height, channels, cam_frame_id, cam_frame_ts = raw_frames[cam_id]
|
||||
frame, width, height, channels, cam_frame_id, cam_frame_ts, cam_frame_last_read_ts = raw_frames[cam_id]
|
||||
cam = self._get_camera_spec(cam_id)
|
||||
frames_out[cam_id] = frame
|
||||
|
||||
|
|
@ -373,7 +386,7 @@ class FrameService:
|
|||
|
||||
def _process_rgb(self, raw_frames):
|
||||
cam_id = self.state.payload.sources[0]
|
||||
frame, width, height, channels, _, _ = raw_frames[cam_id]
|
||||
frame, width, height, channels, _, _, _ = raw_frames[cam_id]
|
||||
cam = self._get_camera_spec(cam_id)
|
||||
|
||||
if not self._is_usb_rgb_camera(cam):
|
||||
|
|
@ -429,7 +442,7 @@ class FrameService:
|
|||
raise RuntimeError("MULTISPEC requer uma câmera RGB ativa")
|
||||
|
||||
# RGB USB
|
||||
rgb_frame, rgb_w, rgb_h, rgb_channels, _, _ = raw_frames[rgb_cam.id]
|
||||
rgb_frame, rgb_w, rgb_h, rgb_channels, _, _, _ = raw_frames[rgb_cam.id]
|
||||
if rgb_channels != 3:
|
||||
raise RuntimeError(f"Frame RGB USB inválido: channels={rgb_channels}, shape={rgb_frame.shape}")
|
||||
|
||||
|
|
|
|||
|
|
@ -259,12 +259,23 @@ class ModuleServer:
|
|||
output_dtype = parse_output_dtype(msg.get("output_dtype", self.state.output_dtype))
|
||||
capture_mode = parse_capture_mode(msg.get("capture_mode", self.state.capture_mode))
|
||||
|
||||
old_frame_type = self.state.frame_type
|
||||
old_output_dtype = self.state.output_dtype
|
||||
old_capture_mode = self.state.capture_mode
|
||||
|
||||
self.state.frame_type = frame_type
|
||||
self.state.output_dtype = output_dtype
|
||||
self.state.capture_mode = capture_mode
|
||||
self.state.update_payload_spec()
|
||||
|
||||
if self.state.initialized:
|
||||
if (
|
||||
old_frame_type != self.state.frame_type or
|
||||
old_output_dtype != self.state.output_dtype or
|
||||
old_capture_mode != self.state.capture_mode
|
||||
):
|
||||
self._mark_reconfigure_needed()
|
||||
|
||||
self._restart_module_if_needed()
|
||||
return {
|
||||
"ok": True,
|
||||
|
|
@ -652,3 +663,7 @@ class ModuleServer:
|
|||
daemon=True
|
||||
)
|
||||
t.start()
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
|
||||
|
|
@ -205,56 +205,53 @@ class ModuleState:
|
|||
return None
|
||||
|
||||
def get_required_camera_ids_for_frame_type(self) -> List[str]:
|
||||
def enabled(role: str):
|
||||
return self.get_enabled_camera_by_role(role)
|
||||
|
||||
def get_all_enabled_camera_ids():
|
||||
return [cam.id for cam in self.cameras if cam.enabled]
|
||||
rgb_cam = enabled("rgb")
|
||||
re_cam = enabled("re")
|
||||
nir_cam = enabled("nir")
|
||||
|
||||
if self.frame_type == "RGB":
|
||||
cam = get_enabled_camera_by_role("rgb")
|
||||
return [cam.id] if cam else []
|
||||
|
||||
if self.frame_type == "MULTISPEC":
|
||||
ids = []
|
||||
re_cam = get_enabled_camera_by_role("re")
|
||||
nir_cam = get_enabled_camera_by_role("nir")
|
||||
rgb_cam = get_enabled_camera_by_role("rgb")
|
||||
|
||||
if re_cam:
|
||||
ids.append(re_cam.id)
|
||||
if nir_cam:
|
||||
ids.append(nir_cam.id)
|
||||
if rgb_cam:
|
||||
ids.append(rgb_cam.id)
|
||||
|
||||
return ids
|
||||
|
||||
if self.frame_type == "RAW_BRUTO":
|
||||
resolved = self.resolve_capture_mode()
|
||||
|
||||
if self.frame_type == "RGB":
|
||||
return [rgb_cam.id] if rgb_cam else []
|
||||
|
||||
if self.frame_type == "MULTISPEC":
|
||||
if resolved == "TRIPLE":
|
||||
if rgb_cam and re_cam and nir_cam:
|
||||
return [rgb_cam.id, re_cam.id, nir_cam.id]
|
||||
return []
|
||||
|
||||
if resolved == "DOUBLE":
|
||||
if rgb_cam and re_cam:
|
||||
return [rgb_cam.id, re_cam.id]
|
||||
if rgb_cam and nir_cam:
|
||||
return [rgb_cam.id, nir_cam.id]
|
||||
return []
|
||||
|
||||
# MULTISPEC não faz sentido com SINGLE
|
||||
return []
|
||||
|
||||
if self.frame_type == "RAW_BRUTO":
|
||||
if resolved == "TRIPLE":
|
||||
return [cam.id for cam in self.cameras if cam.enabled]
|
||||
|
||||
if resolved == "DOUBLE":
|
||||
ids = []
|
||||
rgb_cam = self.get_enabled_camera_by_role("rgb")
|
||||
re_cam = self.get_enabled_camera_by_role("re")
|
||||
nir_cam = self.get_enabled_camera_by_role("nir")
|
||||
|
||||
if rgb_cam:
|
||||
ids.append(rgb_cam.id)
|
||||
|
||||
if re_cam:
|
||||
ids.append(re_cam.id)
|
||||
elif nir_cam:
|
||||
ids.append(nir_cam.id)
|
||||
|
||||
return ids
|
||||
|
||||
if resolved == "SINGLE":
|
||||
for role in ("rgb", "re", "nir"):
|
||||
cam = self.get_enabled_camera_by_role(role)
|
||||
for cam in (rgb_cam, re_cam, nir_cam):
|
||||
if cam:
|
||||
return [cam.id]
|
||||
return []
|
||||
|
||||
return []
|
||||
|
||||
|
|
@ -356,27 +353,36 @@ class ModuleState:
|
|||
self.detected_mode = "NONE"
|
||||
|
||||
def resolve_capture_mode(self) -> str:
|
||||
enabled_re = self.get_enabled_camera_by_role("re") is not None
|
||||
enabled_nir = self.get_enabled_camera_by_role("nir") is not None
|
||||
enabled_rgb = self.get_enabled_camera_by_role("rgb") is not None
|
||||
connected_re = self.get_enabled_camera_by_role("re") is not None
|
||||
connected_nir = self.get_enabled_camera_by_role("nir") is not None
|
||||
connected_rgb = self.get_enabled_camera_by_role("rgb") is not None
|
||||
|
||||
if self.capture_mode == "AUTO":
|
||||
if enabled_rgb and enabled_re and enabled_nir and self.multi_camera_enabled:
|
||||
return "TRIPLE"
|
||||
if enabled_rgb and (enabled_re or enabled_nir):
|
||||
# MULTISPEC e RAW_BRUTO: preferir DOUBLE no automático
|
||||
if self.frame_type in ("MULTISPEC", "RAW_BRUTO"):
|
||||
if connected_rgb and (connected_re or connected_nir):
|
||||
return "DOUBLE"
|
||||
if enabled_rgb or enabled_re or enabled_nir:
|
||||
if self.frame_type == "RAW_BRUTO":
|
||||
if connected_rgb or connected_re or connected_nir:
|
||||
return "SINGLE"
|
||||
return "NONE"
|
||||
|
||||
if connected_rgb and connected_re and connected_nir and self.multi_camera_enabled:
|
||||
return "TRIPLE"
|
||||
if connected_rgb and (connected_re or connected_nir):
|
||||
return "DOUBLE"
|
||||
if connected_rgb or connected_re or connected_nir:
|
||||
return "SINGLE"
|
||||
return "NONE"
|
||||
|
||||
if self.capture_mode == "TRIPLE":
|
||||
return "TRIPLE" if (enabled_rgb and enabled_re and enabled_nir and self.multi_camera_enabled) else "NONE"
|
||||
return "TRIPLE" if (connected_rgb and connected_re and connected_nir and self.multi_camera_enabled) else "NONE"
|
||||
|
||||
if self.capture_mode == "DOUBLE":
|
||||
return "DOUBLE" if (enabled_rgb and (enabled_re or enabled_nir)) else "NONE"
|
||||
return "DOUBLE" if (connected_rgb and (connected_re or connected_nir)) else "NONE"
|
||||
|
||||
if self.capture_mode == "SINGLE":
|
||||
return "SINGLE" if (enabled_rgb or enabled_re or enabled_nir) else "NONE"
|
||||
return "SINGLE" if (connected_rgb or connected_re or connected_nir) else "NONE"
|
||||
|
||||
return "NONE"
|
||||
|
||||
|
|
@ -474,8 +480,8 @@ class ModuleState:
|
|||
layout="CHW",
|
||||
channels=3,
|
||||
channel_names=["R", "G", "B"],
|
||||
width=rgb_cam.width // 2,
|
||||
height=rgb_cam.height // 2,
|
||||
width=rgb_cam.width,
|
||||
height=rgb_cam.height,
|
||||
sources=[rgb_cam.id],
|
||||
complete=True,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ class StreamSender:
|
|||
self._frames_dropped = 0
|
||||
self._capture_errors = 0
|
||||
self._send_errors = 0
|
||||
self._last_queued_capture_signature = None
|
||||
self._last_sent_capture_signature = None
|
||||
|
||||
@property
|
||||
|
|
@ -66,6 +67,7 @@ class StreamSender:
|
|||
self.state.stream_host = host
|
||||
self.state.stream_port = port
|
||||
self.state.stream_fps = fps
|
||||
self._last_queued_capture_signature = None
|
||||
self._last_sent_capture_signature = None
|
||||
|
||||
self._thread_capture = threading.Thread(
|
||||
|
|
@ -125,6 +127,7 @@ class StreamSender:
|
|||
self.state.stream_host = None
|
||||
self.state.stream_port = None
|
||||
self.state.stream_fps = None
|
||||
self.state.stream_frame_id_sent = 0
|
||||
|
||||
def _clear_queue(self):
|
||||
while not self._queue.empty():
|
||||
|
|
@ -369,7 +372,7 @@ class StreamSender:
|
|||
continue
|
||||
|
||||
capture_signature = self._build_capture_signature(meta)
|
||||
if capture_signature is not None and capture_signature == self._last_sent_capture_signature:
|
||||
if capture_signature is not None and capture_signature == self._last_queued_capture_signature:
|
||||
time.sleep(0.001)
|
||||
continue
|
||||
|
||||
|
|
@ -412,7 +415,7 @@ class StreamSender:
|
|||
self._frames_dropped += 1
|
||||
|
||||
if queued:
|
||||
self._last_sent_capture_signature = capture_signature
|
||||
self._last_queued_capture_signature = capture_signature
|
||||
|
||||
if frame_interval > 0:
|
||||
next_deadline += frame_interval
|
||||
|
|
@ -451,7 +454,9 @@ class StreamSender:
|
|||
header["send_errors"] = self._send_errors
|
||||
|
||||
self._send_packet(sock, header, payload)
|
||||
self.state.stream_frame_id = header["frame_id"]
|
||||
self.state.stream_frame_id_sent = header["frame_id"]
|
||||
capture_signature = self._build_capture_signature(header)
|
||||
self._last_sent_capture_signature = capture_signature
|
||||
|
||||
except Exception as e:
|
||||
self._send_errors += 1
|
||||
|
|
|
|||
Loading…
Reference in New Issue