ajustes hierarquia de arquivos para treinamento de IA
This commit is contained in:
parent
00142d2215
commit
7bcd259d9f
|
|
@ -1,327 +0,0 @@
|
|||
import os
|
||||
import time
|
||||
import json
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from raw_segformer_service import make_bgr_preview_from_raw
|
||||
from gal5000.gal_service import Gal5000Camera # ajuste o nome do módulo se estiver diferente
|
||||
|
||||
|
||||
# =========================
|
||||
# Helpers gerais
|
||||
# =========================
|
||||
|
||||
def clamp(v, lo, hi):
|
||||
return lo if v < lo else hi if v > hi else v
|
||||
|
||||
|
||||
def ts_name() -> str:
|
||||
"""Timestamp legível e único para nome de arquivo."""
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||||
|
||||
|
||||
def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
|
||||
"""
|
||||
Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
|
||||
Pensado pra deixar o preview bonitinho sem estourar tudo.
|
||||
"""
|
||||
x = np.asarray(x)
|
||||
if x.dtype != np.float32 and x.dtype != np.float64:
|
||||
x = x.astype(np.float32)
|
||||
|
||||
# Se o canal já está em 0..1, escala pra 0..255 antes de cortar
|
||||
if x.max() <= 1.5:
|
||||
x = x * 255.0
|
||||
|
||||
lo = np.percentile(x, p_lo)
|
||||
hi = np.percentile(x, p_hi)
|
||||
|
||||
if hi <= lo + 1e-3:
|
||||
y = x
|
||||
else:
|
||||
y = (x - lo) * (255.0 / (hi - lo))
|
||||
|
||||
return np.clip(y, 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def overlay_hud(
|
||||
img_bgr: np.ndarray,
|
||||
lines: list[str],
|
||||
base_h: int = 720,
|
||||
base_font_scale: float = 0.75,
|
||||
base_line_step: int = 28,
|
||||
):
|
||||
"""
|
||||
Escreve textos empilhados no canto superior esquerdo,
|
||||
ajustando o tamanho do texto de acordo com a altura da imagem.
|
||||
|
||||
base_h: altura de referência (ex: 720 ou a RAW_H original).
|
||||
"""
|
||||
h, w = img_bgr.shape[:2]
|
||||
|
||||
# Fator de escala com base na altura atual
|
||||
scale = h / float(base_h)
|
||||
|
||||
# Evita ficar microscópico em resoluções muito baixas
|
||||
scale = max(scale, 0.4)
|
||||
|
||||
font_scale = base_font_scale * scale
|
||||
line_step = int(base_line_step * scale)
|
||||
|
||||
# Espessuras proporcionais
|
||||
thick_outline = max(1, int(3 * scale))
|
||||
thick_text = max(1, int(2 * scale))
|
||||
|
||||
# Margem superior / esquerda também escaladas
|
||||
y = int(24 * scale)
|
||||
x = int(12 * scale)
|
||||
|
||||
for s in lines:
|
||||
# contorno preto
|
||||
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
|
||||
# texto branco
|
||||
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
|
||||
y += line_step
|
||||
|
||||
|
||||
def save_sample_raw4(
|
||||
base_dir: str,
|
||||
raw4: np.ndarray,
|
||||
preview_bgr: np.ndarray,
|
||||
meta: dict,
|
||||
):
|
||||
"""
|
||||
Salva:
|
||||
- RAW4 como .raw float32 (4,H,W)
|
||||
- preview RGB como .png
|
||||
- metadados como .json
|
||||
dentro de base_dir.
|
||||
"""
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
name = ts_name()
|
||||
|
||||
raw_path = os.path.join(base_dir, f"{name}.raw")
|
||||
png_path = os.path.join(base_dir, f"{name}.png")
|
||||
json_path = os.path.join(base_dir, f"{name}.json")
|
||||
|
||||
# RAW4
|
||||
#np.save(raw_path, raw4.astype(np.float32))
|
||||
raw4.astype(np.float32).tofile(raw_path)
|
||||
|
||||
# Preview
|
||||
cv2.imwrite(png_path, preview_bgr)
|
||||
|
||||
# Metadados
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return raw_path, png_path, json_path
|
||||
|
||||
|
||||
# =========================
|
||||
# MAIN
|
||||
# =========================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
||||
parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
||||
parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
|
||||
parser.add_argument("--dll_dir", default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls", help="Pasta onde está a VT_SDK64.dll (usada pelo Gal5000Camera).")
|
||||
parser.add_argument("--dll_name", default="VT_SDK64.dll", help="Nome da DLL da câmera.")
|
||||
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
|
||||
parser.add_argument("--no_ae", action="store_true", help="Desliga o AutoExposure do service (por padrão ele vem ligado).")
|
||||
parser.add_argument("--upscale", type=int, default=2, help="Fator de upscale visual do preview.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# ===== config =====
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
RAW_W = config["raw_size"][0]
|
||||
RAW_H = config["raw_size"][1]
|
||||
|
||||
# Define diretório de sessão:
|
||||
# dataset/cana_<estado>/<horario>/<YYYYMMDD>/
|
||||
session_dir = os.path.join(MODELO, args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
|
||||
os.makedirs(session_dir, exist_ok=True)
|
||||
|
||||
print("============================================")
|
||||
print("Coleta de dataset RAW4 - SegFormer B0")
|
||||
print(f"Cana : {args.cana}")
|
||||
print(f"Horário : {args.horario}")
|
||||
print(f"Saída : {session_dir}")
|
||||
print("============================================")
|
||||
|
||||
window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
|
||||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
||||
|
||||
auto_save = False
|
||||
last_auto_t = 0.0
|
||||
upscale = args.upscale
|
||||
|
||||
# Estatísticas simples de FPS
|
||||
t_fps = time.time()
|
||||
frames = 0
|
||||
fps = 0.0
|
||||
|
||||
last_msg = ""
|
||||
last_msg_t = 0.0
|
||||
|
||||
try:
|
||||
cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
|
||||
with cam:
|
||||
print("[CAM] Status inicial:", cam.get_status())
|
||||
|
||||
# opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
|
||||
cam.configure_fps(20)
|
||||
cam.start_streaming()
|
||||
|
||||
while True:
|
||||
t0 = time.time()
|
||||
raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
|
||||
t1 = time.time()
|
||||
|
||||
ae_dbg = dbg.get("ae", {}) or {}
|
||||
exp_raw = dbg.get("exp_raw", None)
|
||||
gain_a = dbg.get("gain_a", None)
|
||||
gain_d = dbg.get("gain_d", None)
|
||||
|
||||
apply_ir_comp = True
|
||||
ir_k_r = 0.4
|
||||
ir_k_g = 0.1
|
||||
ir_k_b = 0.5
|
||||
bgr = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=upscale > 0, preview_scale=upscale, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
||||
|
||||
# FPS
|
||||
frames += 1
|
||||
dt_fps = time.time() - t_fps
|
||||
if dt_fps >= 1.0:
|
||||
fps = frames / dt_fps
|
||||
frames = 0
|
||||
t_fps = time.time()
|
||||
|
||||
# AE info
|
||||
p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
|
||||
sat_disp = ae_dbg.get("sat", 0.0)
|
||||
hold = ae_dbg.get("hold", False)
|
||||
|
||||
ae_on = cam.is_auto_exposure_enabled()
|
||||
|
||||
# HUD principal
|
||||
lines = [
|
||||
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
|
||||
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
|
||||
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
|
||||
f"AEdbg: p95={p95_disp:.1f} sat={sat_disp:.3f} hold={hold}",
|
||||
"Keys: C/SPACE=save | A=auto-save | E=AE toggle | M=preview scale | Q/Esc=quit",
|
||||
]
|
||||
overlay_hud(bgr, lines, base_h=RAW_H)
|
||||
|
||||
# Mensagem rápida (ex: arquivo salvo)
|
||||
if last_msg and (time.time() - last_msg_t) < 2.0:
|
||||
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
|
||||
|
||||
cv2.imshow(window_name, bgr)
|
||||
|
||||
# Auto-save
|
||||
now = time.time()
|
||||
if auto_save and (now - last_auto_t) >= args.interval:
|
||||
meta = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"cana": args.cana,
|
||||
"horario": args.horario,
|
||||
"raw4_shape": list(raw4_base.shape),
|
||||
"out_h": RAW_H,
|
||||
"out_w": RAW_W,
|
||||
"ae_enabled": bool(ae_on),
|
||||
"exp_raw": int(exp_raw) if exp_raw is not None else None,
|
||||
"gain_a": int(gain_a) if gain_a is not None else None,
|
||||
"gain_d": int(gain_d) if gain_d is not None else None,
|
||||
"apply_ir_comp": apply_ir_comp,
|
||||
"ir_k_r": ir_k_r,
|
||||
"ir_k_g": ir_k_g,
|
||||
"ir_k_b": ir_k_b,
|
||||
"ae_dbg": {
|
||||
k: (float(v) if isinstance(v, (int, float, np.floating)) else v)
|
||||
for k, v in ae_dbg.items()
|
||||
},
|
||||
"note": "autosave",
|
||||
}
|
||||
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
||||
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
||||
last_msg = f"SALVO (auto): {os.path.basename(raw_path)}"
|
||||
last_msg_t = now
|
||||
last_auto_t = now
|
||||
|
||||
# Teclado
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
if k in (ord("q"), ord("Q"), 27): # Q ou ESC
|
||||
break
|
||||
|
||||
elif k in (ord("a"), ord("A")):
|
||||
auto_save = not auto_save
|
||||
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("e"), ord("E")):
|
||||
cam.enable_auto_exposure(not ae_on)
|
||||
last_msg = f"AE -> {'ON' if cam.is_auto_exposure_enabled() else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("m"), ord("M")):
|
||||
upscale = 0 if upscale else args.upscale
|
||||
last_msg = f"Preview UPSCALE -> {upscale}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord("c"), ord("C"), 32): # C ou SPACE
|
||||
meta = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"cana": args.cana,
|
||||
"horario": args.horario,
|
||||
"raw4_shape": list(raw4_base.shape),
|
||||
"out_h": RAW_H,
|
||||
"out_w": RAW_W,
|
||||
"ae_enabled": bool(ae_on),
|
||||
"exp_raw": int(exp_raw) if exp_raw is not None else None,
|
||||
"gain_a": int(gain_a) if gain_a is not None else None,
|
||||
"gain_d": int(gain_d) if gain_d is not None else None,
|
||||
"apply_ir_comp": apply_ir_comp,
|
||||
"ir_k_r": ir_k_r,
|
||||
"ir_k_g": ir_k_g,
|
||||
"ir_k_b": ir_k_b,
|
||||
"ae_dbg": {
|
||||
k2: (float(v2) if isinstance(v2, (int, float, np.floating)) else v2)
|
||||
for k2, v2 in ae_dbg.items()
|
||||
},
|
||||
"note": "manual",
|
||||
}
|
||||
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False)
|
||||
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
||||
last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
# você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)
|
||||
|
||||
# Só pra não ficar rodando a 1000 FPS na UI
|
||||
# mas sem travar muito a captura
|
||||
dt_loop = time.time() - t0
|
||||
if dt_loop < 0.001:
|
||||
time.sleep(0.001)
|
||||
|
||||
finally:
|
||||
cv2.destroyAllWindows()
|
||||
print("Fim da captura.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -7,7 +7,7 @@ import argparse
|
|||
import numpy as np
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config.json", "r") as f:
|
||||
with open("config_oak.json", "r") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
|||
# ====================== Configurações base ======================
|
||||
|
||||
def carregar_config_e_paths():
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ import albumentations as A
|
|||
import argparse
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera", ".")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ import numpy as np
|
|||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
MODEL_NAME = config["model_name"]
|
||||
|
|
@ -351,7 +351,7 @@ def main(args):
|
|||
|
||||
# Se quiser, você pode nomear os canais explicitamente
|
||||
# dependendo da convenção do raw4:
|
||||
channel_names = ["R", "G", "IR", "B"]
|
||||
channel_names = ["R", "G", "B"]
|
||||
|
||||
stats = {
|
||||
"channels": channel_names[:len(mean_list)],
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ import random
|
|||
import argparse
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
with open("config_oak.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
|
|
|||
|
|
@ -258,7 +258,7 @@ def run_one_epoch(model: nn.Module,
|
|||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", default="config.json")
|
||||
parser.add_argument("--config", default="config_oak.json")
|
||||
parser.add_argument("--epochs", type=int, default=120)
|
||||
parser.add_argument("--batch", type=int, default=1) # <<< default seguro pra 8GB
|
||||
parser.add_argument("--lr", type=float, default=6e-5)
|
||||
|
|
|
|||
|
|
@ -595,7 +595,7 @@ def apply_warmup(optimizer, base_lr: float, epoch: int, warmup_epochs: int):
|
|||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
|
||||
ap.add_argument("--config", default="config.json")
|
||||
ap.add_argument("--config", default="config_oak.json")
|
||||
ap.add_argument("--epochs", type=int, default=120)
|
||||
ap.add_argument("--batch", type=int, default=4)
|
||||
ap.add_argument("--num_workers", type=int, default=4)
|
||||
|
|
|
|||
|
|
@ -215,7 +215,7 @@ def main():
|
|||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# Lê config do projeto (mesmo padrão do fastscnn)
|
||||
with open("config.json", "r") as f:
|
||||
with open("config_oak.json", "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
|
|
|
|||
|
|
@ -457,7 +457,7 @@ def main():
|
|||
device = find_device()
|
||||
print("Device:", device)
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
with open("config_oak.json", "r") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
MODEL_NAME = config["model_name"] # ex: "segformer_b3"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
{
|
||||
"camera": "oak-d",
|
||||
"modelo": "segformer_b0",
|
||||
"model_name": "nav_big",
|
||||
"model_name": "nav_mit",
|
||||
"dual_head": false,
|
||||
"main_class_name": "navegavel",
|
||||
"es_classes": "",
|
||||
|
|
@ -13,5 +13,5 @@
|
|||
"shaves": 3,
|
||||
"channels": 3,
|
||||
"use_ndvi": false,
|
||||
"backbone": "nvidia/segformer-b0-finetuned-ade-512-512"
|
||||
"backbone": "nvidia/mit-b0"
|
||||
}
|
||||
|
|
@ -0,0 +1,422 @@
|
|||
import os
|
||||
import time
|
||||
import json
|
||||
import math
|
||||
import ctypes as C
|
||||
from ctypes import wintypes as W
|
||||
from datetime import datetime
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
# =========================
|
||||
# CONFIG
|
||||
# =========================
|
||||
SDK_DIR = os.path.join(os.path.dirname(__file__), "dlls")
|
||||
DLL_NAME = "VT_SDK64.dll"
|
||||
|
||||
# Onde salvar o dataset
|
||||
OUT_ROOT = os.path.join(os.path.dirname(__file__), "dataset")
|
||||
SESSION_DIR = os.path.join(OUT_ROOT, datetime.now().strftime("%Y%m%d"))
|
||||
os.makedirs(SESSION_DIR, exist_ok=True)
|
||||
|
||||
# Camera scan/open
|
||||
DEVICE_UDEF = 0
|
||||
DEVICE_INDEX = 0
|
||||
DATA_RAW = 0
|
||||
|
||||
# RAW geometry (se mudar no futuro, ajuste)
|
||||
RAW_W = 2592
|
||||
RAW_H = 2056
|
||||
|
||||
TIMEOUT_MS = 2000
|
||||
WINDOW_NAME = "GAL5000 Dataset Capture (C/SPACE=save | A=auto-save | E=AE toggle | Q=quit)"
|
||||
|
||||
# Preview
|
||||
UPSCALE = 2
|
||||
|
||||
# Auto-save
|
||||
CAPTURE_INTERVAL_S = 1.0
|
||||
|
||||
# Param IDs (VT_Param.h)
|
||||
BUF_SIZE = 256
|
||||
PARAM_ID_SENSOR_EXPOSURETIMERAW = 0x00003010
|
||||
PARAM_ID_SENSOR_GAINANALOGRAW = 0x00003020
|
||||
PARAM_ID_SENSOR_GAINDIGITRAW = 0x0000302A
|
||||
|
||||
# PARAM_VALUETYPE
|
||||
VALUE_INT = 0
|
||||
VALUE_FLOAT = 1
|
||||
VALUE_STR = 2
|
||||
|
||||
# Exposure/Gain limits (ajuste depois conforme o sensor aceitar)
|
||||
EXP_MIN = 1
|
||||
EXP_MAX = 20000
|
||||
|
||||
GAIN_A_MIN, GAIN_A_MAX = 0, 255
|
||||
GAIN_D_MIN, GAIN_D_MAX = 0, 255
|
||||
|
||||
# =========================
|
||||
# Helpers
|
||||
# =========================
|
||||
def ck(ret: int, name: str):
|
||||
if ret != 0:
|
||||
raise RuntimeError(f"{name} falhou, ret={ret}")
|
||||
|
||||
def ts_name() -> str:
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||||
|
||||
def clamp(v, lo, hi):
|
||||
return lo if v < lo else hi if v > hi else v
|
||||
|
||||
def norm8(x, p_lo=2, p_hi=98):
|
||||
lo = np.percentile(x, p_lo)
|
||||
hi = np.percentile(x, p_hi)
|
||||
if hi <= lo + 1:
|
||||
return x.astype(np.uint8)
|
||||
y = (x.astype(np.float32) - lo) * (255.0 / (hi - lo))
|
||||
return np.clip(y, 0, 255).astype(np.uint8)
|
||||
|
||||
def make_rgb_preview(raw: np.ndarray, upscale=2) -> np.ndarray:
|
||||
# pattern:
|
||||
# R G
|
||||
# IR B
|
||||
R = raw[0::2, 0::2]
|
||||
G = raw[0::2, 1::2]
|
||||
B = raw[1::2, 1::2]
|
||||
|
||||
Rn, Gn, Bn = norm8(R), norm8(G), norm8(B)
|
||||
bgr = np.dstack([Bn, Gn, Rn]) # OpenCV usa BGR
|
||||
if upscale and upscale != 1:
|
||||
bgr = cv2.resize(bgr, (bgr.shape[1]*upscale, bgr.shape[0]*upscale), interpolation=cv2.INTER_NEAREST)
|
||||
return bgr
|
||||
|
||||
def measure_raw_g_metrics(raw: np.ndarray):
|
||||
"""
|
||||
Mede brilho no canal G cru usando uma ROI na base (mais parecido com chão).
|
||||
Retorna p90/p95 e fração saturada.
|
||||
"""
|
||||
G = raw[0::2, 1::2] # H/2 x W/2
|
||||
h2, w2 = G.shape
|
||||
|
||||
# ROI: base da imagem, cortando laterais
|
||||
y0, y1 = int(h2 * 0.55), int(h2 * 0.95)
|
||||
x0, x1 = int(w2 * 0.15), int(w2 * 0.85)
|
||||
roi = G[y0:y1, x0:x1]
|
||||
|
||||
p90 = float(np.percentile(roi, 90))
|
||||
p95 = float(np.percentile(roi, 95))
|
||||
sat = float(np.mean(roi >= 250))
|
||||
return p90, p95, sat
|
||||
|
||||
class RobustAE:
|
||||
"""
|
||||
Controle soft de exposure (sem depender do GET da camera):
|
||||
- mede p95 do canal G cru em ROI
|
||||
- usa EMA + deadband (pra não ficar "descendo até 16" como você viu)
|
||||
- passo multiplicativo em log, com limite de passo
|
||||
"""
|
||||
def __init__(self,
|
||||
exp_min=EXP_MIN, exp_max=EXP_MAX,
|
||||
target_p95=140.0,
|
||||
deadband=6.0,
|
||||
k=0.12,
|
||||
max_step=0.10,
|
||||
ema_alpha=0.20,
|
||||
sat_limit=0.01):
|
||||
self.exp_min = exp_min
|
||||
self.exp_max = exp_max
|
||||
self.target = target_p95
|
||||
self.deadband = deadband
|
||||
self.k = k
|
||||
self.max_step = max_step
|
||||
self.ema_alpha = ema_alpha
|
||||
self.sat_limit = sat_limit
|
||||
self.p95_ema = None
|
||||
|
||||
def step(self, raw, exp_raw):
|
||||
p90, p95, sat = measure_raw_g_metrics(raw)
|
||||
|
||||
# EMA do p95 (estabiliza)
|
||||
if self.p95_ema is None:
|
||||
self.p95_ema = p95
|
||||
else:
|
||||
self.p95_ema = (1 - self.ema_alpha) * self.p95_ema + self.ema_alpha * p95
|
||||
|
||||
e = self.target - self.p95_ema # erro em nível de pixel
|
||||
|
||||
# deadband: segura a mão perto do alvo
|
||||
if abs(e) <= self.deadband and sat <= self.sat_limit:
|
||||
return exp_raw, {"p90": p90, "p95": p95, "p95_ema": self.p95_ema, "sat": sat, "hold": True}
|
||||
|
||||
# saturou: garante redução
|
||||
if sat > self.sat_limit:
|
||||
step = -min(self.max_step, 0.12)
|
||||
else:
|
||||
ratio = (self.target + 1e-6) / (self.p95_ema + 1e-6)
|
||||
step = self.k * math.log(ratio)
|
||||
step = max(-self.max_step, min(self.max_step, step))
|
||||
|
||||
new_exp = int(round(exp_raw * math.exp(step)))
|
||||
new_exp = max(self.exp_min, min(self.exp_max, new_exp))
|
||||
|
||||
return new_exp, {"p90": p90, "p95": p95, "p95_ema": self.p95_ema, "sat": sat, "step": step, "hold": False}
|
||||
|
||||
# =========================
|
||||
# STRUCTS + Param API
|
||||
# =========================
|
||||
class VT_FRAMEINFO(C.Structure):
|
||||
_fields_ = [
|
||||
("lFrameID", W.DWORD),
|
||||
("lBufSize", W.DWORD),
|
||||
("lWidth", W.DWORD),
|
||||
("lHeight", W.DWORD),
|
||||
("lPixBits", C.c_ubyte),
|
||||
("_pad0", C.c_ubyte * 3),
|
||||
("pBufPtr", C.POINTER(C.c_ubyte)),
|
||||
("lFrameStatus", W.DWORD),
|
||||
("lPixType", W.DWORD),
|
||||
("lTimeStamp", W.DWORD),
|
||||
("_reserve", W.DWORD * 8),
|
||||
]
|
||||
|
||||
class VT_DEVPARAM(C.Structure):
|
||||
_fields_ = [
|
||||
("bUseName", W.BOOL),
|
||||
("lParamByID", W.DWORD),
|
||||
("lParamByName", C.c_char * BUF_SIZE),
|
||||
]
|
||||
|
||||
def devparam_by_id(pid: int) -> VT_DEVPARAM:
|
||||
p = VT_DEVPARAM()
|
||||
p.bUseName = False
|
||||
p.lParamByID = pid
|
||||
p.lParamByName = b""
|
||||
return p
|
||||
|
||||
# =========================
|
||||
# DLL LOAD + prototypes
|
||||
# =========================
|
||||
os.add_dll_directory(SDK_DIR)
|
||||
dll = C.WinDLL(os.path.join(SDK_DIR, DLL_NAME))
|
||||
print("DLL carregada OK:", dll)
|
||||
|
||||
dll.VT_DeviceScan.argtypes = [C.POINTER(C.c_ubyte), C.c_int]
|
||||
dll.VT_DeviceScan.restype = C.c_int
|
||||
|
||||
dll.VT_DeviceOpen.argtypes = [C.c_void_p, C.POINTER(W.HANDLE), C.c_int, C.c_int]
|
||||
dll.VT_DeviceOpen.restype = C.c_int
|
||||
|
||||
dll.VT_SingleFrameCapture.argtypes = [W.HANDLE, C.POINTER(VT_FRAMEINFO), C.c_int, C.c_int, W.BOOL]
|
||||
dll.VT_SingleFrameCapture.restype = C.c_int
|
||||
|
||||
dll.VT_DeviceClose.argtypes = [C.POINTER(W.HANDLE)]
|
||||
dll.VT_DeviceClose.restype = C.c_int
|
||||
|
||||
dll.VT_ParamGetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
|
||||
dll.VT_ParamGetValue.restype = C.c_int
|
||||
|
||||
dll.VT_ParamSetValue.argtypes = [W.HANDLE, VT_DEVPARAM, C.c_void_p, C.c_int]
|
||||
dll.VT_ParamSetValue.restype = C.c_int
|
||||
|
||||
def param_set_int(h: W.HANDLE, pid: int, value: int):
|
||||
p = devparam_by_id(pid)
|
||||
v = C.c_int(int(value))
|
||||
ret = dll.VT_ParamSetValue(h, p, C.byref(v), VALUE_INT)
|
||||
ck(ret, f"VT_ParamSetValue({hex(pid)})")
|
||||
|
||||
def capture_raw8(h: W.HANDLE) -> np.ndarray:
|
||||
fi = VT_FRAMEINFO()
|
||||
ret = dll.VT_SingleFrameCapture(h, C.byref(fi), DATA_RAW, TIMEOUT_MS, True)
|
||||
ck(ret, "VT_SingleFrameCapture")
|
||||
|
||||
w, hh = int(fi.lWidth), int(fi.lHeight)
|
||||
buf = C.string_at(fi.pBufPtr, fi.lBufSize)
|
||||
arr = np.frombuffer(buf, dtype=np.uint8)
|
||||
|
||||
needed = w * hh
|
||||
if arr.size < needed:
|
||||
arr = np.pad(arr, (0, needed - arr.size), mode="constant", constant_values=0)
|
||||
arr = arr[:needed].reshape(hh, w)
|
||||
return arr
|
||||
|
||||
def overlay_hud(img_bgr, lines):
|
||||
y = 28
|
||||
for s in lines:
|
||||
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 3, cv2.LINE_AA)
|
||||
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255,255,255), 2, cv2.LINE_AA)
|
||||
y += 28
|
||||
|
||||
def save_sample(raw: np.ndarray, bgr_preview: np.ndarray, meta: dict):
|
||||
name = ts_name()
|
||||
raw_path = os.path.join(SESSION_DIR, f"{name}.raw")
|
||||
png_path = os.path.join(SESSION_DIR, f"{name}.png")
|
||||
json_path = os.path.join(SESSION_DIR, f"{name}.json")
|
||||
|
||||
raw.tofile(raw_path)
|
||||
cv2.imwrite(png_path, bgr_preview)
|
||||
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return raw_path, png_path, json_path
|
||||
|
||||
def main():
|
||||
# scan
|
||||
n = C.c_ubyte(0)
|
||||
ck(dll.VT_DeviceScan(C.byref(n), DEVICE_UDEF), "VT_DeviceScan")
|
||||
if n.value == 0:
|
||||
raise RuntimeError("Nenhuma câmera encontrada.")
|
||||
|
||||
# open
|
||||
idx = C.c_ubyte(0)
|
||||
h = W.HANDLE()
|
||||
ck(dll.VT_DeviceOpen(C.byref(idx), C.byref(h), DEVICE_INDEX, DEVICE_UDEF), "VT_DeviceOpen")
|
||||
print("DeviceOpen OK, handle=", h.value)
|
||||
print("Saving to:", SESSION_DIR)
|
||||
|
||||
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
|
||||
|
||||
# Estado local (não dependemos de GET)
|
||||
exp_raw = 1500
|
||||
gain_a = 0
|
||||
gain_d = 0
|
||||
|
||||
# Aplica estado inicial
|
||||
try:
|
||||
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
|
||||
param_set_int(h, PARAM_ID_SENSOR_GAINANALOGRAW, gain_a)
|
||||
param_set_int(h, PARAM_ID_SENSOR_GAINDIGITRAW, gain_d)
|
||||
except Exception as e:
|
||||
print("[WARN] Falhou set inicial:", e)
|
||||
|
||||
ae = RobustAE(target_p95=140.0, deadband=6.0, k=0.12, max_step=0.10, ema_alpha=0.20, sat_limit=0.01)
|
||||
ae_on = True
|
||||
auto_save = False
|
||||
last_auto_t = 0.0
|
||||
|
||||
# FPS
|
||||
t0 = time.time()
|
||||
frames = 0
|
||||
fps = 0.0
|
||||
|
||||
last_msg = ""
|
||||
last_msg_t = 0.0
|
||||
|
||||
try:
|
||||
while True:
|
||||
raw = capture_raw8(h)
|
||||
|
||||
# soft AE
|
||||
ae_dbg = {}
|
||||
if ae_on:
|
||||
new_exp, ae_dbg = ae.step(raw, exp_raw)
|
||||
if new_exp != exp_raw:
|
||||
exp_raw = new_exp
|
||||
try:
|
||||
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
|
||||
except Exception as e:
|
||||
# se set falhar, desliga AE pra não ficar insistindo
|
||||
print("[ERR] set exposure:", e)
|
||||
ae_on = False
|
||||
|
||||
# preview RGB bonitão
|
||||
rgb_clean = make_rgb_preview(raw, upscale=UPSCALE)
|
||||
bgr = rgb_clean.copy()
|
||||
|
||||
# FPS
|
||||
frames += 1
|
||||
dt = time.time() - t0
|
||||
if dt >= 1.0:
|
||||
fps = frames / dt
|
||||
frames = 0
|
||||
t0 = time.time()
|
||||
|
||||
# HUD
|
||||
lines = [
|
||||
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Interval: {CAPTURE_INTERVAL_S:.1f}s",
|
||||
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
|
||||
f"AEdbg: p95={ae_dbg.get('p95_ema', ae_dbg.get('p95', 0)):.1f} sat={ae_dbg.get('sat', 0):.3f} hold={ae_dbg.get('hold', False)}",
|
||||
"Keys: C/SPACE=save | A=toggle autosave | E=toggle AE | +/- exp | Q/ESC quit",
|
||||
]
|
||||
overlay_hud(bgr, lines)
|
||||
|
||||
# msg pós-save
|
||||
if last_msg and (time.time() - last_msg_t) < 2.0:
|
||||
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2, cv2.LINE_AA)
|
||||
|
||||
cv2.imshow(WINDOW_NAME, bgr)
|
||||
|
||||
# autosave
|
||||
now = time.time()
|
||||
if auto_save and (now - last_auto_t) >= CAPTURE_INTERVAL_S:
|
||||
meta = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"raw_w": RAW_W, "raw_h": RAW_H,
|
||||
"exp_raw": int(exp_raw),
|
||||
"gain_a": int(gain_a),
|
||||
"gain_d": int(gain_d),
|
||||
"ae_on": bool(ae_on),
|
||||
"note": "autosave",
|
||||
}
|
||||
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
|
||||
last_msg = f"SAVED: {os.path.basename(raw_path)}"
|
||||
last_msg_t = now
|
||||
last_auto_t = now
|
||||
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
if k in (ord('q'), ord('Q'), 27):
|
||||
break
|
||||
|
||||
elif k in (ord('a'), ord('A')):
|
||||
auto_save = not auto_save
|
||||
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord('e'), ord('E')):
|
||||
ae_on = not ae_on
|
||||
last_msg = f"AE -> {'ON' if ae_on else 'OFF'}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord('c'), ord('C'), 32): # C ou SPACE
|
||||
meta = {
|
||||
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
||||
"raw_w": RAW_W, "raw_h": RAW_H,
|
||||
"exp_raw": int(exp_raw),
|
||||
"gain_a": int(gain_a),
|
||||
"gain_d": int(gain_d),
|
||||
"ae_on": bool(ae_on),
|
||||
"note": "manual",
|
||||
}
|
||||
raw_path, png_path, json_path = save_sample(raw, rgb_clean, meta)
|
||||
last_msg = f"SAVED: {os.path.basename(raw_path)}"
|
||||
last_msg_t = time.time()
|
||||
|
||||
elif k in (ord('+'), ord('=')):
|
||||
exp_raw = clamp(exp_raw + 200, EXP_MIN, EXP_MAX)
|
||||
try:
|
||||
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
|
||||
except Exception as e:
|
||||
print("[ERR] manual exp +:", e)
|
||||
|
||||
elif k in (ord('-'), ord('_')):
|
||||
exp_raw = clamp(exp_raw - 200, EXP_MIN, EXP_MAX)
|
||||
try:
|
||||
param_set_int(h, PARAM_ID_SENSOR_EXPOSURETIMERAW, exp_raw)
|
||||
except Exception as e:
|
||||
print("[ERR] manual exp -:", e)
|
||||
|
||||
finally:
|
||||
try:
|
||||
ret = dll.VT_DeviceClose(C.byref(h))
|
||||
if ret != 0:
|
||||
print("VT_DeviceClose retornou:", ret)
|
||||
except Exception as e:
|
||||
print("Erro ao fechar:", e)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
print("Fim.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -22,9 +22,9 @@ RAW_WIDTH = RESOLUCAO[0] # largura real do raw (ex: 1296)
|
|||
RAW_HEIGHT = RESOLUCAO[1] # altura real do raw (ex: 1028)
|
||||
RAW_EXTS = [".raw"] # se tiver outro, adiciona aqui
|
||||
|
||||
PASTA_FINAL_PREVIEWS = os.path.join(MODELO, "dataset", "original", "previews")
|
||||
PASTA_FINAL_MASKS = os.path.join(MODELO, "dataset", "original", "masks")
|
||||
PASTA_FINAL_RAWS = os.path.join(MODELO, "dataset", "original", "raws")
|
||||
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
|
||||
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
|
||||
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
|
||||
|
||||
# Cor da classe será obtida via labelmap
|
||||
# COR_CLASSE_RGB = (128, 0, 0)
|
||||
|
|
@ -200,7 +200,7 @@ def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None:
|
|||
return None
|
||||
|
||||
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False):
|
||||
source = os.path.join(MODELO, "dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join(MODELO, "dataset", "original")
|
||||
source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original")
|
||||
pasta_new_previews = os.path.join(source, "group", grupo, "previews")
|
||||
pasta_new_masks = os.path.join(source, "group", grupo, "masks")
|
||||
pasta_new_raws = os.path.join(source, "group", grupo, "raws")
|
||||
|
|
@ -296,7 +296,7 @@ if __name__ == "__main__":
|
|||
args = ap.parse_args()
|
||||
|
||||
# Caminho padrão pro labelmap, caso não seja informado
|
||||
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
|
||||
labelmap_path = os.path.join("dataset", "labelmap.txt")
|
||||
|
||||
if not os.path.exists(labelmap_path):
|
||||
raise SystemExit(f"Labelmap não encontrado em: {labelmap_path}")
|
||||
|
|
@ -12,15 +12,15 @@ with open("config.json", "r") as f:
|
|||
MODELO = config["camera"]
|
||||
|
||||
# Raiz das brutas agrupadas
|
||||
PASTA_BRUTAS_GROUP_ROOT = os.path.join(MODELO, "dataset", "brutas")
|
||||
PASTA_BRUTAS_GROUP_ROOT = os.path.join("dataset", "brutas")
|
||||
|
||||
# Onde as máscaras novas (rotuladas externamente) são colocadas
|
||||
PASTA_NEW_MASKS = os.path.join(MODELO, "dataset", "new_masks")
|
||||
PASTA_NEW_MASKS = os.path.join("dataset", "new_masks")
|
||||
|
||||
# Destino final do dataset consolidado
|
||||
PASTA_FINAL_PREVIEWS = os.path.join(MODELO, "dataset", "original", "previews")
|
||||
PASTA_FINAL_RAWS = os.path.join(MODELO, "dataset", "original", "raws")
|
||||
PASTA_FINAL_MASKS = os.path.join(MODELO, "dataset", "original", "masks")
|
||||
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
|
||||
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
|
||||
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
|
||||
|
||||
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
||||
EXT_MASKS = (".png",)
|
||||
|
|
@ -12,14 +12,14 @@ with open("config.json", "r") as f:
|
|||
MODELO = config["camera"]
|
||||
|
||||
# Raiz das brutas (todas as canas/horários/grupos)
|
||||
PASTA_BRUTAS_ROOT = os.path.join(MODELO, "dataset", "brutas")
|
||||
PASTA_BRUTAS_ROOT = os.path.join("dataset", "brutas")
|
||||
|
||||
# Onde você coloca os previews selecionados (tudo misturado)
|
||||
PASTA_SELECTED_PREVIEWS = os.path.join(MODELO, "dataset", "selected_previews")
|
||||
PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews")
|
||||
|
||||
# Destino final, organizado por grupo:
|
||||
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
|
||||
PASTA_ORIGINAL_GROUP_ROOT = os.path.join(MODELO, "dataset", "original", "group")
|
||||
PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group")
|
||||
|
||||
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
||||
EXT_RAWS = (".raw",)
|
||||
|
|
@ -309,8 +309,8 @@ def build_cli():
|
|||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
originals_dir = os.path.join(MODELO, "dataset", "original")
|
||||
labelmap_dir = os.path.join(MODELO, "dataset", "labelmap.txt")
|
||||
originals_dir = os.path.join("dataset", "original")
|
||||
labelmap_dir = os.path.join("dataset", "labelmap.txt")
|
||||
processar(
|
||||
originals_dir=originals_dir,
|
||||
labelmap_path=labelmap_dir,
|
||||
|
|
@ -6,15 +6,15 @@ Augmenta imagens/máscaras *por grupo* usando RAW4 (float32, 0..1, R G IR B).
|
|||
Novo fluxo (GAL5000):
|
||||
|
||||
Entrada:
|
||||
MODELO/dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
|
||||
MODELO/dataset/original/group/<grupo>/masks (.png/.jpg)
|
||||
(opcional) MODELO/dataset/original/group/<grupo>/masks2
|
||||
dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
|
||||
dataset/original/group/<grupo>/masks (.png/.jpg)
|
||||
(opcional) dataset/original/group/<grupo>/masks2
|
||||
|
||||
Saída:
|
||||
MODELO/dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
|
||||
MODELO/dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
|
||||
MODELO/dataset/augmented/group/<grupo>/masks
|
||||
MODELO/dataset/augmented/group/<grupo>/masks2
|
||||
dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
|
||||
dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
|
||||
dataset/augmented/group/<grupo>/masks
|
||||
dataset/augmented/group/<grupo>/masks2
|
||||
|
||||
Regras:
|
||||
- GEOMETRIA (flip, shift, rotate) aplicada em RAW + masks (+ masks2).
|
||||
|
|
@ -44,7 +44,7 @@ RESOLUCAO = config.get("raw_size")
|
|||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
||||
# Pastas base
|
||||
DATASET_BASE = os.path.join(MODELO, "dataset")
|
||||
DATASET_BASE = os.path.join("dataset")
|
||||
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
|
||||
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
|
||||
|
||||
|
|
@ -4,11 +4,11 @@
|
|||
Normaliza/redimensiona PREVIEW + RAW + MASK (+MASK2), mantendo ESTRUTURA POR GRUPO.
|
||||
|
||||
Entradas:
|
||||
MODELO/dataset/original/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
MODELO/dataset/augmented/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
dataset/original/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
dataset/augmented/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
|
||||
Saídas (por resolução):
|
||||
MODELO/dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
|
||||
Conversão de máscara:
|
||||
- Lê máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids
|
||||
|
|
@ -37,9 +37,10 @@ USE_MASKS2 = config["dual_head"]
|
|||
RESOLUCAO = tuple(config["resolucao"]) # [W,H]
|
||||
MODEL_NAME = config["model_name"]
|
||||
CHANNELS = int(config.get("channels", 4))
|
||||
pasta_base = os.path.join(MODELO, "dataset")
|
||||
pasta_base = os.path.join("dataset")
|
||||
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
|
||||
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
|
||||
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
|
||||
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag)
|
||||
|
||||
RESOLUCOES = {f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1])}
|
||||
FONTES = ["original", "augmented"]
|
||||
|
|
@ -5,10 +5,10 @@ Split estratificado por GRUPO com **val/test só do ORIGINAL** e
|
|||
garantia de NÃO VAZAMENTO entre splits, agora para estrutura RAW:
|
||||
|
||||
Lê de:
|
||||
MODELO/dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
|
||||
Escreve em:
|
||||
MODELO/dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}
|
||||
dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}
|
||||
|
||||
Definições:
|
||||
- "Família" = todas as variações da MESMA base original:
|
||||
|
|
@ -33,8 +33,8 @@ USE_MASKS2 = config.get("dual_head", False)
|
|||
RESOLUCAO = tuple(config.get("resolucao"))
|
||||
|
||||
# Pastas (ajustadas para PREVIEWS/RAWS)
|
||||
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
|
||||
pasta_destino = os.path.join(MODELO, "dataset", "split")
|
||||
pasta_origem = os.path.join("dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
|
||||
pasta_destino = os.path.join("dataset", "split")
|
||||
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXT = ".png" # máscaras normalizadas em PNG
|
||||
|
|
@ -376,7 +376,6 @@ def main():
|
|||
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
|
||||
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
|
||||
|
||||
ap.add_argument("--modelo", type=str, default=None, help="Sobrescreve MODELO do config.json.")
|
||||
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")
|
||||
|
||||
ap.add_argument("--cap-train-families", type=str, default="",
|
||||
|
|
@ -384,7 +383,6 @@ def main():
|
|||
|
||||
args = ap.parse_args()
|
||||
|
||||
modelo = args.modelo or MODELO
|
||||
if args.resolucao:
|
||||
try:
|
||||
w, h = args.resolucao.lower().split("x")
|
||||
|
|
@ -406,8 +404,8 @@ def main():
|
|||
caps_map = parse_cap_map(args.cap_train_families)
|
||||
|
||||
global pasta_origem, pasta_destino
|
||||
pasta_origem = os.path.join(modelo, "dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
|
||||
pasta_destino = os.path.join(modelo, "dataset", "split")
|
||||
pasta_origem = os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
|
||||
pasta_destino = os.path.join("dataset", "split")
|
||||
|
||||
soma = args.train + args.val + args.test
|
||||
if soma <= 0:
|
||||
|
|
@ -358,10 +358,10 @@ def main():
|
|||
|
||||
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
|
||||
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
|
||||
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, experiment_tag)
|
||||
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, experiment_tag)
|
||||
|
||||
#save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
|
||||
dataset_path = os.path.join(MODELO, "dataset")
|
||||
#save_path = os.path.join("backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
|
||||
dataset_path = os.path.join("dataset")
|
||||
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
|
||||
|
|
@ -378,7 +378,7 @@ def main():
|
|||
# Caminho padrão: dentro do dataset, nome do arquivo de stats
|
||||
# (ajusta aqui pro nome que você realmente usou: norm_stats.json, por ex.)
|
||||
#norm_stats_path = os.path.join(save_path, "norm_stats.json")
|
||||
norm_stats_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME, stats_source_tag, "norm_stats.json")
|
||||
norm_stats_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag, "norm_stats.json")
|
||||
if args.norm_stats is not None:
|
||||
norm_stats_path = args.norm_stats
|
||||
|
||||
|
|
@ -547,7 +547,6 @@ def main():
|
|||
strict=False,
|
||||
)
|
||||
patch_segformer_input_channels(model, in_ch=CHANNELS)
|
||||
|
||||
elif FUSION_MODE == "dual_branch":
|
||||
model = build_dual_branch_segformer_model(
|
||||
num_classes=num_classes,
|
||||
|
|
@ -243,7 +243,7 @@ def main():
|
|||
print(f"[cfg] resolucao nominal RAW: {W}x{H}")
|
||||
print(f"[cfg] channels={CHANNELS} use_ndvi={USE_NDVI}")
|
||||
|
||||
dataset_path = os.path.join(MODELO, "dataset")
|
||||
dataset_path = os.path.join("dataset")
|
||||
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
||||
|
||||
# ==========================
|
||||
|
|
@ -254,7 +254,7 @@ def main():
|
|||
|
||||
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
|
||||
|
||||
norm_stats_path = os.path.join(MODELO, "backup", modelo_folder, MODEL_NAME, experiment_tag, "norm_stats.json")
|
||||
norm_stats_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag, "norm_stats.json")
|
||||
if args.norm_stats is not None:
|
||||
norm_stats_path = args.norm_stats
|
||||
if os.path.isfile(norm_stats_path):
|
||||
|
|
@ -307,7 +307,7 @@ def main():
|
|||
if args.ckpt is not None:
|
||||
ckpt_path = args.ckpt
|
||||
else:
|
||||
save_path = os.path.join(MODELO, "backup", modelo_folder, MODEL_NAME, experiment_tag)
|
||||
save_path = os.path.join("backup", modelo_folder, MODEL_NAME, experiment_tag)
|
||||
ckpt_path = os.path.join(save_path, "best_miou.pt")
|
||||
|
||||
if not os.path.isfile(ckpt_path):
|
||||
|
|
@ -1,19 +1,19 @@
|
|||
{
|
||||
"camera": "gal5000",
|
||||
"modelo": "segformer_b1",
|
||||
"model_name": "pulv_new",
|
||||
"model_name": "pulv_mit",
|
||||
"dual_head": false,
|
||||
"main_class_name": "cana",
|
||||
"es_classes": "",
|
||||
"model_to_use": "geral",
|
||||
"raw_size": [1296, 1028],
|
||||
"resolucao": [1008, 800],
|
||||
"resolucao": [1024, 800],
|
||||
"roi_inicio": 0.0,
|
||||
"roi_tamanho": 1.0,
|
||||
"shaves": 3,
|
||||
"channels": 4,
|
||||
"use_ndvi": false,
|
||||
"backbone": "nvidia/segformer-b1-finetuned-ade-512-512",
|
||||
"backbone": "nvidia/mit-b1",
|
||||
"fusion_mode": "stacked",
|
||||
"stats_source_tag": "stacked_raw4"
|
||||
}
|
||||
|
|
@ -0,0 +1,390 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Regenera previews a partir dos RAWs antigos da GAL5000,
|
||||
aplicando compensação de IR no preview.
|
||||
|
||||
Estrutura esperada:
|
||||
dataset/brutas/group/{GRUPO}/
|
||||
masks/
|
||||
previews/
|
||||
raws/
|
||||
metas/
|
||||
|
||||
Uso exemplo:
|
||||
python regen_previews_from_raws.py ^
|
||||
--root dataset/brutas/group ^
|
||||
--raw-h 1028 ^
|
||||
--raw-w 1296 ^
|
||||
--layout rgirb ^
|
||||
--ir-k-r 0.8 ^
|
||||
--ir-k-g 0.4 ^
|
||||
--ir-k-b 0.9
|
||||
|
||||
Observações:
|
||||
- Para .raw, informe --raw-h e --raw-w
|
||||
- Para .npy/.npz, o shape é lido automaticamente
|
||||
- O script recria os previews em cada pasta previews do grupo
|
||||
- Por padrão sobrescreve os previews existentes
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
cv2 = None
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# Preview
|
||||
# ---------------------------------------------------------
|
||||
def make_bgr_preview_from_raw(
|
||||
raw_np: np.ndarray,
|
||||
rgirb: bool,
|
||||
preview_fast: bool,
|
||||
preview_scale: int = 2,
|
||||
apply_ir_comp: bool = True,
|
||||
ir_k_r: float = 0.40,
|
||||
ir_k_g: float = 0.10,
|
||||
ir_k_b: float = 0.50,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Gera preview BGR (OpenCV) a partir do tensor raw_np (C,H,W) float32 em 0..1.
|
||||
|
||||
Casos esperados:
|
||||
- rgirb=False:
|
||||
raw_np = [R, G, B] ou [R, G, B, NDVI]
|
||||
- rgirb=True:
|
||||
raw_np = [R, G, IR, B] ou [R, G, IR, B, NDVI]
|
||||
"""
|
||||
assert raw_np.ndim == 3, "raw_np deve ser (C,H,W)"
|
||||
C = raw_np.shape[0]
|
||||
if C not in (3, 4, 5):
|
||||
raise RuntimeError(f"Esperado C=3, 4 ou 5, veio {C}")
|
||||
|
||||
raw_np = raw_np.astype(np.float32, copy=False)
|
||||
|
||||
r = raw_np[0]
|
||||
g = raw_np[1]
|
||||
b = raw_np[3] if rgirb else raw_np[2]
|
||||
|
||||
if preview_fast:
|
||||
if preview_scale > 1:
|
||||
r = r[::preview_scale, ::preview_scale]
|
||||
g = g[::preview_scale, ::preview_scale]
|
||||
b = b[::preview_scale, ::preview_scale]
|
||||
|
||||
if rgirb and apply_ir_comp and C >= 4:
|
||||
ir = raw_np[2]
|
||||
if preview_scale > 1:
|
||||
ir = ir[::preview_scale, ::preview_scale]
|
||||
|
||||
r = np.clip(r - ir_k_r * ir, 0.0, 1.0)
|
||||
g = np.clip(g - ir_k_g * ir, 0.0, 1.0)
|
||||
b = np.clip(b - ir_k_b * ir, 0.0, 1.0)
|
||||
|
||||
bgr = np.stack([b, g, r], axis=0)
|
||||
bgr = np.power(np.clip(bgr, 0.0, 1.0), 1 / 1.8)
|
||||
bgr8 = (bgr * 255.0).clip(0, 255).astype(np.uint8)
|
||||
return np.transpose(bgr8, (1, 2, 0)).copy()
|
||||
|
||||
if rgirb and apply_ir_comp and C >= 4:
|
||||
ir = raw_np[2]
|
||||
r = np.clip(r - ir_k_r * ir, 0.0, 1.0)
|
||||
g = np.clip(g - ir_k_g * ir, 0.0, 1.0)
|
||||
b = np.clip(b - ir_k_b * ir, 0.0, 1.0)
|
||||
|
||||
def stretch_channel(x: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray:
|
||||
lo = np.percentile(x, p_low)
|
||||
hi = np.percentile(x, p_high)
|
||||
|
||||
if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
|
||||
return np.clip(x, 0.0, 1.0)
|
||||
|
||||
x = (x - lo) / (hi - lo)
|
||||
return np.clip(x, 0.0, 1.0)
|
||||
|
||||
r = stretch_channel(r)
|
||||
g = stretch_channel(g)
|
||||
b = stretch_channel(b)
|
||||
|
||||
bgr = np.stack([b, g, r], axis=0).astype(np.float32)
|
||||
bgr = np.power(np.clip(bgr, 0.0, 1.0), 1 / 2.0)
|
||||
|
||||
bgr8 = (bgr * 255.0).clip(0, 255).astype(np.uint8)
|
||||
return np.transpose(bgr8, (1, 2, 0)).copy()
|
||||
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# IO helpers
|
||||
# ---------------------------------------------------------
|
||||
def save_bgr_image(path: Path, bgr: np.ndarray) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if cv2 is not None:
|
||||
ok = cv2.imwrite(str(path), bgr)
|
||||
if not ok:
|
||||
raise RuntimeError(f"Falha ao salvar imagem: {path}")
|
||||
else:
|
||||
rgb = bgr[..., ::-1]
|
||||
Image.fromarray(rgb).save(path)
|
||||
|
||||
|
||||
def to_chw_float01(arr: np.ndarray, layout: str) -> np.ndarray:
|
||||
"""
|
||||
Converte entrada para (C,H,W) float32 0..1.
|
||||
|
||||
layout:
|
||||
- rgirb
|
||||
- rgbir
|
||||
- rgb
|
||||
"""
|
||||
arr = np.asarray(arr)
|
||||
|
||||
if arr.ndim != 3:
|
||||
raise RuntimeError(f"Esperava array 3D, veio shape={arr.shape}")
|
||||
|
||||
# HWC -> CHW
|
||||
if arr.shape[-1] in (3, 4, 5) and arr.shape[0] not in (3, 4, 5):
|
||||
arr = np.transpose(arr, (2, 0, 1))
|
||||
|
||||
if arr.shape[0] not in (3, 4, 5):
|
||||
raise RuntimeError(f"Não consegui interpretar canais em shape={arr.shape}")
|
||||
|
||||
arr = arr.astype(np.float32, copy=False)
|
||||
|
||||
# Normalização para 0..1
|
||||
if arr.dtype == np.uint8:
|
||||
arr = arr / 255.0
|
||||
elif arr.dtype == np.uint16:
|
||||
arr = arr / 65535.0
|
||||
else:
|
||||
# Se já vier float mas fora de 0..1, tenta ajustar
|
||||
maxv = float(np.nanmax(arr)) if arr.size else 1.0
|
||||
if maxv > 1.0:
|
||||
arr = arr / maxv
|
||||
|
||||
arr = np.clip(arr, 0.0, 1.0)
|
||||
|
||||
# Reorganiza para o contrato do preview
|
||||
# Queremos:
|
||||
# rgirb=True -> [R,G,IR,B]
|
||||
# rgirb=False -> [R,G,B]
|
||||
if layout == "rgirb":
|
||||
# já assume [R,G,IR,B] ou [R,G,IR,B,NDVI]
|
||||
return arr
|
||||
elif layout == "rgbir":
|
||||
# [R,G,B,IR] -> [R,G,IR,B]
|
||||
if arr.shape[0] < 4:
|
||||
raise RuntimeError("layout=rgbir exige pelo menos 4 canais")
|
||||
if arr.shape[0] == 4:
|
||||
arr = arr[[0, 1, 3, 2], :, :]
|
||||
else:
|
||||
# [R,G,B,IR,NDVI] -> [R,G,IR,B,NDVI]
|
||||
arr = arr[[0, 1, 3, 2, 4], :, :]
|
||||
return arr
|
||||
elif layout == "rgb":
|
||||
return arr[:3]
|
||||
else:
|
||||
raise ValueError(f"layout inválido: {layout}")
|
||||
|
||||
|
||||
def load_raw_file(raw_path: Path, raw_hw: Optional[Tuple[int, int]], layout: str) -> np.ndarray:
|
||||
"""
|
||||
Retorna (C,H,W) float32 em 0..1.
|
||||
|
||||
Suporta:
|
||||
- .npy
|
||||
- .npz
|
||||
- .raw
|
||||
|
||||
Para .raw:
|
||||
- mosaico uint8 HxW em padrão 2x2 R,G / IR,B
|
||||
- ou RAW4 float32 (4,H,W) salvo em [R,G,IR,B]
|
||||
"""
|
||||
ext = raw_path.suffix.lower()
|
||||
|
||||
if ext == ".npy":
|
||||
arr = np.load(raw_path)
|
||||
return to_chw_float01(arr, layout)
|
||||
|
||||
if ext == ".npz":
|
||||
z = np.load(raw_path)
|
||||
key = list(z.keys())[0]
|
||||
arr = z[key]
|
||||
return to_chw_float01(arr, layout)
|
||||
|
||||
if ext == ".raw":
|
||||
if raw_hw is None:
|
||||
raise RuntimeError(f"{raw_path.name}: para .raw informe --raw-h e --raw-w")
|
||||
|
||||
H, W = raw_hw
|
||||
size_bytes = raw_path.stat().st_size
|
||||
mosa_bytes = H * W
|
||||
raw4_bytes = 4 * H * W * 4 # float32
|
||||
|
||||
if size_bytes == mosa_bytes:
|
||||
# mosaico uint8 cru
|
||||
arr = np.fromfile(raw_path, dtype=np.uint8).reshape(H, W)
|
||||
|
||||
if (H % 2) != 0 or (W % 2) != 0:
|
||||
raise RuntimeError(f"{raw_path.name}: H e W precisam ser pares para mosaico 2x2")
|
||||
|
||||
r_sub = arr[0::2, 0::2]
|
||||
g_sub = arr[0::2, 1::2]
|
||||
ir_sub = arr[1::2, 0::2]
|
||||
b_sub = arr[1::2, 1::2]
|
||||
|
||||
if cv2 is not None:
|
||||
r = cv2.resize(r_sub, (W, H), interpolation=cv2.INTER_LINEAR)
|
||||
g = cv2.resize(g_sub, (W, H), interpolation=cv2.INTER_LINEAR)
|
||||
ir = cv2.resize(ir_sub, (W, H), interpolation=cv2.INTER_LINEAR)
|
||||
b = cv2.resize(b_sub, (W, H), interpolation=cv2.INTER_LINEAR)
|
||||
else:
|
||||
r = np.array(Image.fromarray(r_sub).resize((W, H), resample=Image.BILINEAR))
|
||||
g = np.array(Image.fromarray(g_sub).resize((W, H), resample=Image.BILINEAR))
|
||||
ir = np.array(Image.fromarray(ir_sub).resize((W, H), resample=Image.BILINEAR))
|
||||
b = np.array(Image.fromarray(b_sub).resize((W, H), resample=Image.BILINEAR))
|
||||
|
||||
chw = np.stack([r, g, ir, b], axis=0).astype(np.float32) / 255.0
|
||||
return chw
|
||||
|
||||
if size_bytes == raw4_bytes:
|
||||
# RAW4 float32 salvo como (4,H,W) em [R,G,IR,B]
|
||||
arr = np.fromfile(raw_path, dtype=np.float32).reshape(4, H, W)
|
||||
arr = np.clip(arr, 0.0, 1.0)
|
||||
return to_chw_float01(arr, layout="rgirb")
|
||||
|
||||
raise RuntimeError(
|
||||
f"{raw_path.name}: tamanho inesperado {size_bytes} bytes "
|
||||
f"(esperado mosaico={mosa_bytes} ou raw4 float32={raw4_bytes})"
|
||||
)
|
||||
|
||||
raise RuntimeError(f"Extensão não suportada: {raw_path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# Processamento
|
||||
# ---------------------------------------------------------
|
||||
def process_group(
|
||||
group_dir: Path,
|
||||
raw_hw: Optional[Tuple[int, int]],
|
||||
layout: str,
|
||||
preview_fast: bool,
|
||||
preview_scale: int,
|
||||
apply_ir_comp: bool,
|
||||
ir_k_r: float,
|
||||
ir_k_g: float,
|
||||
ir_k_b: float,
|
||||
overwrite: bool,
|
||||
) -> tuple[int, int]:
|
||||
raws_dir = group_dir / "raws"
|
||||
previews_dir = group_dir / "previews"
|
||||
|
||||
if not raws_dir.is_dir():
|
||||
return 0, 0
|
||||
|
||||
previews_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
raw_files = []
|
||||
for ext in ("*.raw", "*.npy", "*.npz"):
|
||||
raw_files.extend(sorted(raws_dir.glob(ext)))
|
||||
|
||||
done = 0
|
||||
failed = 0
|
||||
|
||||
for raw_path in raw_files:
|
||||
out_path = previews_dir / f"{raw_path.stem}.jpg"
|
||||
|
||||
if out_path.exists() and not overwrite:
|
||||
continue
|
||||
|
||||
try:
|
||||
raw_np = load_raw_file(raw_path, raw_hw=raw_hw, layout=layout)
|
||||
|
||||
bgr = make_bgr_preview_from_raw(
|
||||
raw_np=raw_np,
|
||||
rgirb=(layout == "rgirb" or layout == "rgbir"),
|
||||
preview_fast=preview_fast,
|
||||
preview_scale=preview_scale,
|
||||
apply_ir_comp=apply_ir_comp,
|
||||
ir_k_r=ir_k_r,
|
||||
ir_k_g=ir_k_g,
|
||||
ir_k_b=ir_k_b,
|
||||
)
|
||||
|
||||
save_bgr_image(out_path, bgr)
|
||||
done += 1
|
||||
print(f"[OK] {group_dir.name}/{raw_path.name} -> {out_path.name}")
|
||||
|
||||
except Exception as e:
|
||||
failed += 1
|
||||
print(f"[ERRO] {group_dir.name}/{raw_path.name}: {e}")
|
||||
|
||||
return done, failed
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Regenera previews a partir dos RAWs antigos.")
|
||||
parser.add_argument("--root", type=str, required=True, help="Pasta group, ex: dataset/brutas/group")
|
||||
parser.add_argument("--raw-h", type=int, default=None, help="Altura do RAW para arquivos .raw")
|
||||
parser.add_argument("--raw-w", type=int, default=None, help="Largura do RAW para arquivos .raw")
|
||||
parser.add_argument("--layout", type=str, default="rgirb", choices=["rgirb", "rgbir", "rgb"], help="Layout dos canais dos RAWs")
|
||||
parser.add_argument("--preview-fast", action="store_true", help="Usa modo rápido")
|
||||
parser.add_argument("--preview-scale", type=int, default=2, help="Escala no preview_fast")
|
||||
parser.add_argument("--no-ir-comp", action="store_true", help="Desliga compensação de IR")
|
||||
parser.add_argument("--ir-k-r", type=float, default=0.8)
|
||||
parser.add_argument("--ir-k-g", type=float, default=0.4)
|
||||
parser.add_argument("--ir-k-b", type=float, default=0.9)
|
||||
parser.add_argument("--no-overwrite", action="store_true", help="Não sobrescreve previews existentes")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
root = Path(args.root)
|
||||
if not root.is_dir():
|
||||
raise RuntimeError(f"Pasta root não encontrada: {root}")
|
||||
|
||||
raw_hw = None
|
||||
if args.raw_h is not None and args.raw_w is not None:
|
||||
raw_hw = (args.raw_h, args.raw_w)
|
||||
|
||||
total_done = 0
|
||||
total_failed = 0
|
||||
|
||||
group_dirs = [p for p in sorted(root.iterdir()) if p.is_dir()]
|
||||
if not group_dirs:
|
||||
raise RuntimeError(f"Nenhum grupo encontrado em: {root}")
|
||||
|
||||
for group_dir in group_dirs:
|
||||
done, failed = process_group(
|
||||
group_dir=group_dir,
|
||||
raw_hw=raw_hw,
|
||||
layout=args.layout,
|
||||
preview_fast=args.preview_fast,
|
||||
preview_scale=args.preview_scale,
|
||||
apply_ir_comp=not args.no_ir_comp,
|
||||
ir_k_r=args.ir_k_r,
|
||||
ir_k_g=args.ir_k_g,
|
||||
ir_k_b=args.ir_k_b,
|
||||
overwrite=not args.no_overwrite,
|
||||
)
|
||||
total_done += done
|
||||
total_failed += failed
|
||||
|
||||
print("\n============================================")
|
||||
print("Regeneração concluída")
|
||||
print(f"Previews gerados : {total_done}")
|
||||
print(f"Falhas : {total_failed}")
|
||||
print("============================================")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,669 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
raw_triplet_sorter.py
|
||||
-------------------------------------------------
|
||||
Ferramenta para classificar manualmente amostras RAW do dataset
|
||||
(arquivos .raw + .png + .json) em pastas de labels,
|
||||
usando atalhos de teclado.
|
||||
|
||||
Fluxo:
|
||||
- Entrada: uma ou mais pastas de sessão (cada uma contendo N arquivos):
|
||||
- NOME.raw
|
||||
- NOME.png
|
||||
- NOME.json
|
||||
- O script mostra o preview (PNG) e você usa teclas 1..9/0 para enviar
|
||||
o TRIPLO (raw+preview+json) para uma pasta de saída organizada por label.
|
||||
|
||||
Estrutura de saída:
|
||||
out_root/
|
||||
<label>/
|
||||
raws/
|
||||
previews/
|
||||
metas/
|
||||
masks/
|
||||
|
||||
Teclas:
|
||||
- 1..9, 0 -> envia para a label correspondente
|
||||
- Espaço / n / seta direita -> pular (skip)
|
||||
- p / seta esquerda -> voltar
|
||||
- b -> undo (desfaz última ação)
|
||||
- q / Esc -> sair
|
||||
|
||||
Depedências:
|
||||
- Python 3.8+
|
||||
- Pillow (PIL): pip install pillow
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Dict, Optional
|
||||
|
||||
import tkinter as tk
|
||||
from tkinter import messagebox, filedialog
|
||||
|
||||
try:
|
||||
from PIL import Image, ImageTk
|
||||
except Exception as e:
|
||||
print("ERROR: Pillow (PIL) is required. Install with: pip install pillow", file=sys.stderr)
|
||||
raise
|
||||
|
||||
# Sufixos/padrões esperados:
|
||||
# - preview: NOME.png
|
||||
# - raw: NOME.raw
|
||||
# - meta: NOME.json
|
||||
PREVIEW_EXT = ".png"
|
||||
RAW_EXT = ".raw"
|
||||
META_EXT = ".json"
|
||||
|
||||
|
||||
def find_triplets_in_folder(folder: Path) -> List[Tuple[Path, Path, Path, str]]:
|
||||
"""
|
||||
Encontra tripletas (raw, preview, json) em uma pasta:
|
||||
NOME.png
|
||||
NOME.raw
|
||||
NOME.json
|
||||
|
||||
Retorna lista de tuplas (preview_path, raw_path, json_path, base_name).
|
||||
"""
|
||||
if not folder.is_dir():
|
||||
return []
|
||||
|
||||
previews: Dict[str, Path] = {}
|
||||
raws: Dict[str, Path] = {}
|
||||
metas: Dict[str, Path] = {}
|
||||
|
||||
for p in folder.rglob("*"):
|
||||
if not p.is_file():
|
||||
continue
|
||||
stem = p.stem
|
||||
suffix = p.suffix.lower()
|
||||
|
||||
# Preview
|
||||
if suffix == PREVIEW_EXT:
|
||||
base = stem
|
||||
previews[base] = p
|
||||
continue
|
||||
|
||||
# RAW
|
||||
if suffix == RAW_EXT:
|
||||
base = stem
|
||||
raws[base] = p
|
||||
continue
|
||||
|
||||
# JSON
|
||||
if suffix == META_EXT:
|
||||
base = stem
|
||||
metas[base] = p
|
||||
continue
|
||||
|
||||
triplets = []
|
||||
for base, prev_path in previews.items():
|
||||
raw_path = raws.get(base)
|
||||
meta_path = metas.get(base)
|
||||
if raw_path is not None and meta_path is not None:
|
||||
triplets.append((prev_path, raw_path, meta_path, base))
|
||||
|
||||
# Ordena por base (se for número, ordena numérico)
|
||||
def sort_key(t):
|
||||
b = t[3]
|
||||
return (0, int(b)) if b.isdigit() else (1, b)
|
||||
|
||||
triplets.sort(key=sort_key)
|
||||
return triplets
|
||||
|
||||
|
||||
def collect_all_triplets(folders: List[Path]) -> List[Tuple[Path, Path, Path, str]]:
|
||||
all_tr = []
|
||||
for f in folders:
|
||||
all_tr.extend(find_triplets_in_folder(f))
|
||||
# já vem ordenado por pasta, mas garantimos ordenação global
|
||||
def sort_key(t):
|
||||
b = t[3]
|
||||
return (0, int(b)) if b.isdigit() else (1, b)
|
||||
all_tr.sort(key=sort_key)
|
||||
return all_tr
|
||||
|
||||
|
||||
class ActionLogger:
|
||||
"""
|
||||
Loga ações em:
|
||||
- CSV: sorting_log.csv
|
||||
- JSONL: sorting_log.jsonl
|
||||
Serve também para suportar "resume" (não reprocessar o que já foi classificado/pulado).
|
||||
"""
|
||||
|
||||
def __init__(self, out_root: Path):
|
||||
self.out_root = out_root
|
||||
self.log_csv = out_root / "sorting_log.csv"
|
||||
self.log_json = out_root / "sorting_log.jsonl"
|
||||
self.seen = set()
|
||||
|
||||
if self.log_csv.exists():
|
||||
try:
|
||||
with self.log_csv.open("r", newline="", encoding="utf-8") as f:
|
||||
reader = csv.DictReader(f)
|
||||
for row in reader:
|
||||
key = row.get("preview_path", "")
|
||||
if key:
|
||||
self.seen.add(key)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def already_logged(self, preview_path: Path) -> bool:
|
||||
return str(preview_path) in self.seen
|
||||
|
||||
def log(
|
||||
self,
|
||||
action: str,
|
||||
preview_path: Path,
|
||||
raw_path: Path,
|
||||
meta_path: Path,
|
||||
label: Optional[str] = None,
|
||||
dest_preview: Optional[Path] = None,
|
||||
dest_raw: Optional[Path] = None,
|
||||
dest_meta: Optional[Path] = None,
|
||||
):
|
||||
self.out_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# CSV
|
||||
new_file = not self.log_csv.exists()
|
||||
with self.log_csv.open("a", newline="", encoding="utf-8") as f:
|
||||
writer = csv.writer(f)
|
||||
if new_file:
|
||||
writer.writerow([
|
||||
"timestamp",
|
||||
"action",
|
||||
"label",
|
||||
"preview_path",
|
||||
"raw_path",
|
||||
"meta_path",
|
||||
"dest_preview",
|
||||
"dest_raw",
|
||||
"dest_meta",
|
||||
])
|
||||
writer.writerow([
|
||||
datetime.now().isoformat(timespec="seconds"),
|
||||
action,
|
||||
label or "",
|
||||
str(preview_path),
|
||||
str(raw_path),
|
||||
str(meta_path),
|
||||
str(dest_preview or ""),
|
||||
str(dest_raw or ""),
|
||||
str(dest_meta or ""),
|
||||
])
|
||||
|
||||
# JSONL
|
||||
with self.log_json.open("a", encoding="utf-8") as f:
|
||||
rec = {
|
||||
"ts": datetime.now().isoformat(timespec="seconds"),
|
||||
"action": action,
|
||||
"label": label,
|
||||
"preview_path": str(preview_path),
|
||||
"raw_path": str(raw_path),
|
||||
"meta_path": str(meta_path),
|
||||
"dest_preview": str(dest_preview) if dest_preview else None,
|
||||
"dest_raw": str(dest_raw) if dest_raw else None,
|
||||
"dest_meta": str(dest_meta) if dest_meta else None,
|
||||
}
|
||||
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
|
||||
self.seen.add(str(preview_path))
|
||||
|
||||
|
||||
class RawTripletSorterApp:
|
||||
def __init__(
|
||||
self,
|
||||
triplets: List[Tuple[Path, Path, Path, str]],
|
||||
labels: List[str],
|
||||
out_root: Path,
|
||||
move: bool,
|
||||
resume: bool,
|
||||
display_height: int = 512,
|
||||
):
|
||||
self.all_triplets = triplets
|
||||
self.labels = labels
|
||||
self.out_root = out_root
|
||||
self.move = move
|
||||
self.display_height = display_height
|
||||
self.logger = ActionLogger(out_root)
|
||||
|
||||
if resume:
|
||||
self.all_triplets = [
|
||||
t for t in self.all_triplets if not self.logger.already_logged(t[0])
|
||||
]
|
||||
|
||||
# Estado de sessão
|
||||
self.idx = 0
|
||||
self.history: List[dict] = [] # para undo
|
||||
|
||||
# Prepara pastas destino
|
||||
for label in self.labels:
|
||||
(out_root / label / "raws").mkdir(parents=True, exist_ok=True)
|
||||
(out_root / label / "previews").mkdir(parents=True, exist_ok=True)
|
||||
(out_root / label / "metas").mkdir(parents=True, exist_ok=True)
|
||||
(out_root / label / "masks").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# UI
|
||||
self.root = tk.Tk()
|
||||
self.root.title("Agrobotics RAW Triplet Sorter (preview | raw | json)")
|
||||
self.root.geometry("1100x750")
|
||||
self.root.bind("<Key>", self.on_key)
|
||||
|
||||
# Topo: info + legenda
|
||||
self.top_frame = tk.Frame(self.root)
|
||||
self.top_frame.pack(side=tk.TOP, fill=tk.X)
|
||||
|
||||
self.info_label = tk.Label(self.top_frame, text="", font=("Segoe UI", 11))
|
||||
self.info_label.pack(side=tk.LEFT, padx=10, pady=6)
|
||||
|
||||
self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10))
|
||||
self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6)
|
||||
|
||||
# Área da imagem
|
||||
self.img_frame = tk.Frame(self.root)
|
||||
self.img_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
|
||||
|
||||
self.preview_label = tk.Label(self.img_frame)
|
||||
self.preview_label.pack(side=tk.LEFT, expand=True, padx=6, pady=6)
|
||||
|
||||
# Status (última ação)
|
||||
self.status_var = tk.StringVar(value="Pronto.")
|
||||
self.status_label = tk.Label(self.root, textvariable=self.status_var, font=("Segoe UI", 10), anchor="w")
|
||||
self.status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=4)
|
||||
|
||||
# Rodapé: ajuda
|
||||
self.footer = tk.Label(
|
||||
self.root,
|
||||
text="1..9/0 = labels | espaço/n/→ = próxima (skip) | p/← = anterior | b = undo | q/Esc = sair",
|
||||
font=("Segoe UI", 10),
|
||||
)
|
||||
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
|
||||
|
||||
# Primeira renderização
|
||||
self.render()
|
||||
|
||||
def build_legend_text(self) -> str:
|
||||
parts = []
|
||||
for i, label in enumerate(self.labels, start=1):
|
||||
key = i if i <= 9 else 0 # 0 = 10ª label
|
||||
parts.append(f"[{key}] {label}")
|
||||
return " | ".join(parts)
|
||||
|
||||
def pil_load_preview(self, preview_path: Path) -> ImageTk.PhotoImage:
|
||||
img = Image.open(preview_path).convert("RGB")
|
||||
# escala pra altura desejada mantendo proporção
|
||||
h_target = self.display_height
|
||||
w, h = img.size
|
||||
new_w = int(w * (h_target / h))
|
||||
img = img.resize((new_w, h_target), Image.BILINEAR)
|
||||
return ImageTk.PhotoImage(img)
|
||||
|
||||
def render(self):
|
||||
if not self.all_triplets:
|
||||
messagebox.showinfo("Fim", "Não há amostras para exibir (talvez tudo já foi classificado?).")
|
||||
self.root.destroy()
|
||||
return
|
||||
|
||||
self.idx = max(0, min(self.idx, len(self.all_triplets) - 1))
|
||||
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx]
|
||||
|
||||
try:
|
||||
tk_img = self.pil_load_preview(preview_path)
|
||||
self.preview_tk = tk_img
|
||||
self.preview_label.configure(image=self.preview_tk)
|
||||
except Exception as e:
|
||||
self.info_label.configure(text=f"Erro ao abrir preview: {e}")
|
||||
return
|
||||
|
||||
self.info_label.configure(
|
||||
text=f"{self.idx+1}/{len(self.all_triplets)} | base='{base}' | PREVIEW: {preview_path.name}"
|
||||
)
|
||||
|
||||
def do_copy_or_move(self, src: Path, dst: Path):
|
||||
if self.move:
|
||||
shutil.move(str(src), str(dst))
|
||||
else:
|
||||
shutil.copy2(str(src), str(dst))
|
||||
|
||||
def unique_path(self, p: Path) -> Path:
|
||||
"""Se o caminho já existe, gera um novo com sufixo __k."""
|
||||
if not p.exists():
|
||||
return p
|
||||
stem, ext = p.stem, p.suffix
|
||||
k = 1
|
||||
while True:
|
||||
cand = p.with_name(f"{stem}__{k}{ext}")
|
||||
if not cand.exists():
|
||||
return cand
|
||||
k += 1
|
||||
|
||||
def send_to_label(self, label_index: int):
|
||||
if label_index < 0 or label_index >= len(self.labels):
|
||||
return
|
||||
label = self.labels[label_index]
|
||||
|
||||
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx]
|
||||
|
||||
dst_preview = self.out_root / label / "previews" / preview_path.name
|
||||
dst_raw = self.out_root / label / "raws" / raw_path.name
|
||||
dst_meta = self.out_root / label / "metas" / meta_path.name
|
||||
|
||||
dst_preview = self.unique_path(dst_preview)
|
||||
dst_raw = self.unique_path(dst_raw)
|
||||
dst_meta = self.unique_path(dst_meta)
|
||||
|
||||
try:
|
||||
self.do_copy_or_move(preview_path, dst_preview)
|
||||
self.do_copy_or_move(raw_path, dst_raw)
|
||||
self.do_copy_or_move(meta_path, dst_meta)
|
||||
|
||||
self.logger.log(
|
||||
"assign",
|
||||
preview_path=preview_path,
|
||||
raw_path=raw_path,
|
||||
meta_path=meta_path,
|
||||
label=label,
|
||||
dest_preview=dst_preview,
|
||||
dest_raw=dst_raw,
|
||||
dest_meta=dst_meta,
|
||||
)
|
||||
|
||||
self.history.append({
|
||||
"action": "assign",
|
||||
"label": label,
|
||||
"preview_src": preview_path,
|
||||
"raw_src": raw_path,
|
||||
"meta_src": meta_path,
|
||||
"preview_dst": dst_preview,
|
||||
"raw_dst": dst_raw,
|
||||
"meta_dst": dst_meta,
|
||||
"moved": self.move,
|
||||
"index": self.idx,
|
||||
})
|
||||
|
||||
self.status_var.set(
|
||||
f"{'Movido' if self.move else 'Copiado'} → '{label}': {preview_path.name} (+raw+json)"
|
||||
)
|
||||
|
||||
self.idx += 1
|
||||
if self.idx >= len(self.all_triplets):
|
||||
messagebox.showinfo("Concluído", "Você chegou ao final da fila!")
|
||||
self.root.destroy()
|
||||
return
|
||||
self.render()
|
||||
|
||||
except Exception as e:
|
||||
messagebox.showerror("Erro", f"Falha ao copiar/mover: {e}")
|
||||
self.status_var.set(f"ERRO: {e}")
|
||||
|
||||
def undo(self):
|
||||
if not self.history:
|
||||
return
|
||||
last = self.history.pop()
|
||||
if last["action"] != "assign":
|
||||
return
|
||||
try:
|
||||
if last["moved"]:
|
||||
# mover de volta pro source
|
||||
shutil.move(str(last["preview_dst"]), str(last["preview_src"]))
|
||||
shutil.move(str(last["raw_dst"]), str(last["raw_src"]))
|
||||
shutil.move(str(last["meta_dst"]), str(last["meta_src"]))
|
||||
else:
|
||||
# apagar os arquivos de destino
|
||||
for p in [last["preview_dst"], last["raw_dst"], last["meta_dst"]]:
|
||||
if p.exists():
|
||||
p.unlink()
|
||||
|
||||
self.logger.log(
|
||||
"undo",
|
||||
preview_path=last["preview_src"],
|
||||
raw_path=last["raw_src"],
|
||||
meta_path=last["meta_src"],
|
||||
label=last["label"],
|
||||
dest_preview=last["preview_dst"],
|
||||
dest_raw=last["raw_dst"],
|
||||
dest_meta=last["meta_dst"],
|
||||
)
|
||||
|
||||
self.idx = max(0, min(last.get("index", self.idx), len(self.all_triplets) - 1))
|
||||
self.status_var.set(f"Desfeito ← '{last['label']}': {last['preview_src'].name}")
|
||||
self.render()
|
||||
except Exception as e:
|
||||
messagebox.showerror("Erro", f"Falha no undo: {e}")
|
||||
self.status_var.set(f"ERRO: {e}")
|
||||
|
||||
def on_key(self, event):
|
||||
ch = event.keysym.lower()
|
||||
|
||||
# 1..9 / 0 -> labels
|
||||
if ch in [str(i) for i in range(1, 10)] or ch == "0":
|
||||
label_idx = 9 if ch == "0" else int(ch) - 1
|
||||
self.send_to_label(label_idx)
|
||||
return
|
||||
|
||||
# espaço / n / seta direita = skip / próxima
|
||||
if ch in ("space", "s", "n", "right"):
|
||||
preview_path, raw_path, meta_path, base = self.all_triplets[self.idx]
|
||||
self.logger.log("skip", preview_path, raw_path, meta_path, label=None)
|
||||
self.status_var.set(f"Pulada → {preview_path.name}")
|
||||
self.idx += 1
|
||||
if self.idx >= len(self.all_triplets):
|
||||
messagebox.showinfo("Concluído", "Você chegou ao final da fila!")
|
||||
self.root.destroy()
|
||||
return
|
||||
self.render()
|
||||
return
|
||||
|
||||
# p / seta esquerda = voltar
|
||||
if ch in ("p", "left"):
|
||||
self.idx = max(0, self.idx - 1)
|
||||
self.status_var.set("Voltou uma imagem.")
|
||||
self.render()
|
||||
return
|
||||
|
||||
# b = undo
|
||||
if ch in ("b",):
|
||||
self.undo()
|
||||
return
|
||||
|
||||
# q / Esc = sair
|
||||
if ch in ("q", "escape"):
|
||||
self.root.destroy()
|
||||
return
|
||||
|
||||
def run(self):
|
||||
self.root.mainloop()
|
||||
|
||||
|
||||
class SetupWindow:
|
||||
"""
|
||||
Janela de setup (igual ao sorter anterior, mas voltada para RAW triplets).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.root = tk.Tk()
|
||||
self.root.title("Configurar - RAW Triplet Sorter")
|
||||
self.root.geometry("720x520")
|
||||
|
||||
# Pastas de entrada
|
||||
frm_in = tk.LabelFrame(self.root, text="Pastas de entrada (sessões de captura RAW)")
|
||||
frm_in.pack(fill=tk.BOTH, expand=False, padx=10, pady=8)
|
||||
|
||||
self.inputs_listbox = tk.Listbox(frm_in, height=6)
|
||||
self.inputs_listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=6, pady=6)
|
||||
|
||||
btns_in = tk.Frame(frm_in)
|
||||
btns_in.pack(side=tk.RIGHT, fill=tk.Y, padx=6, pady=6)
|
||||
tk.Button(btns_in, text="Adicionar pasta...", command=self.add_input).pack(fill=tk.X, pady=2)
|
||||
tk.Button(btns_in, text="Remover selecionada", command=self.remove_selected).pack(fill=tk.X, pady=2)
|
||||
tk.Button(btns_in, text="Limpar lista", command=self.clear_inputs).pack(fill=tk.X, pady=2)
|
||||
|
||||
# Pasta de saída
|
||||
frm_out = tk.LabelFrame(self.root, text="Pasta de saída (raiz do dataset rotulado)")
|
||||
frm_out.pack(fill=tk.X, expand=False, padx=10, pady=8)
|
||||
self.out_root_var = tk.StringVar(value="")
|
||||
tk.Entry(frm_out, textvariable=self.out_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6)
|
||||
tk.Button(frm_out, text="Escolher...", command=self.choose_out_root).pack(side=tk.RIGHT, padx=6, pady=6)
|
||||
|
||||
# Labels
|
||||
frm_labels = tk.LabelFrame(self.root, text="Labels (classes) separadas por vírgula")
|
||||
frm_labels.pack(fill=tk.X, expand=False, padx=10, pady=8)
|
||||
self.labels_var = tk.StringVar(value="chao,chao_cana,chao_erva,chao_cana_erva,cana,cana_erva,erva")
|
||||
tk.Entry(frm_labels, textvariable=self.labels_var).pack(fill=tk.X, padx=6, pady=6)
|
||||
|
||||
# Opções
|
||||
frm_opts = tk.LabelFrame(self.root, text="Opções")
|
||||
frm_opts.pack(fill=tk.X, expand=False, padx=10, pady=8)
|
||||
self.move_var = tk.BooleanVar(value=False)
|
||||
self.resume_var = tk.BooleanVar(value=True)
|
||||
self.height_var = tk.IntVar(value=512)
|
||||
|
||||
tk.Checkbutton(frm_opts, text="Mover em vez de copiar", variable=self.move_var).pack(anchor="w", padx=6)
|
||||
tk.Checkbutton(frm_opts, text="Retomar (pular itens já logados)", variable=self.resume_var).pack(anchor="w", padx=6)
|
||||
|
||||
frm_height = tk.Frame(frm_opts)
|
||||
frm_height.pack(fill=tk.X, padx=6, pady=6)
|
||||
tk.Label(frm_height, text="Altura de exibição (px):").pack(side=tk.LEFT)
|
||||
tk.Entry(frm_height, textvariable=self.height_var, width=6).pack(side=tk.LEFT, padx=6)
|
||||
|
||||
# Botão iniciar
|
||||
tk.Button(self.root, text="Iniciar classificação", command=self.start).pack(pady=10)
|
||||
|
||||
# Hint
|
||||
tk.Label(
|
||||
self.root,
|
||||
text="Teclas: 1..9 (0=10ª), espaço/n/→=pular, p/←=anterior, b=undo, q/Esc=sair",
|
||||
fg="#555",
|
||||
).pack(pady=4)
|
||||
|
||||
self.result = None
|
||||
|
||||
def add_input(self):
|
||||
p = filedialog.askdirectory(title="Selecione a pasta de sessão (contendo .raw, .png, .json)")
|
||||
if p:
|
||||
self.inputs_listbox.insert(tk.END, p)
|
||||
|
||||
def remove_selected(self):
|
||||
sel = list(self.inputs_listbox.curselection())
|
||||
for idx in reversed(sel):
|
||||
self.inputs_listbox.delete(idx)
|
||||
|
||||
def clear_inputs(self):
|
||||
self.inputs_listbox.delete(0, tk.END)
|
||||
|
||||
def choose_out_root(self):
|
||||
p = filedialog.askdirectory(title="Selecione a pasta raiz de saída")
|
||||
if p:
|
||||
self.out_root_var.set(p)
|
||||
|
||||
def start(self):
|
||||
inputs = [self.inputs_listbox.get(i) for i in range(self.inputs_listbox.size())]
|
||||
out_root = self.out_root_var.get().strip()
|
||||
labels_raw = self.labels_var.get().strip()
|
||||
labels = [s.strip() for s in labels_raw.split(",") if s.strip()]
|
||||
|
||||
if not inputs:
|
||||
messagebox.showwarning("Faltando entradas", "Adicione pelo menos uma pasta de entrada.")
|
||||
return
|
||||
if not out_root:
|
||||
messagebox.showwarning("Faltando saída", "Escolha a pasta raiz de saída.")
|
||||
return
|
||||
if not labels:
|
||||
messagebox.showwarning("Faltando labels", "Informe pelo menos uma label (classe).")
|
||||
return
|
||||
|
||||
self.result = (
|
||||
[Path(p) for p in inputs],
|
||||
labels,
|
||||
Path(out_root),
|
||||
self.move_var.get(),
|
||||
self.resume_var.get(),
|
||||
self.height_var.get(),
|
||||
)
|
||||
self.root.destroy()
|
||||
|
||||
def run(self):
|
||||
self.root.mainloop()
|
||||
return self.result
|
||||
|
||||
|
||||
def run_with_gui_setup():
|
||||
setup = SetupWindow()
|
||||
res = setup.run()
|
||||
if not res:
|
||||
return
|
||||
input_folders, labels, out_root, move, resume, height = res
|
||||
|
||||
triplets = collect_all_triplets(input_folders)
|
||||
if not triplets:
|
||||
messagebox.showinfo(
|
||||
"Sem amostras",
|
||||
"Nenhuma tripleta encontrada.\nCertifique-se de que existam arquivos .raw, .png, .json.",
|
||||
)
|
||||
return
|
||||
|
||||
app = RawTripletSorterApp(
|
||||
triplets=triplets,
|
||||
labels=labels,
|
||||
out_root=out_root,
|
||||
move=move,
|
||||
resume=resume,
|
||||
display_height=height,
|
||||
)
|
||||
app.run()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Sortear manualmente amostras (raw + preview + json) em pastas de labels, usando hotkeys."
|
||||
)
|
||||
parser.add_argument("--inputs", nargs="+", help="Pastas de sessão contendo .raw, .png e .json")
|
||||
parser.add_argument("--labels", nargs="+", help="Labels (classes) mapeadas para teclas 1..9/0")
|
||||
parser.add_argument("--out-root", help="Pasta raiz de saída")
|
||||
parser.add_argument("--move", action="store_true", help="Mover em vez de copiar")
|
||||
parser.add_argument("--resume", action="store_true", help="Pular itens já presentes no sorting_log.csv")
|
||||
parser.add_argument("--display-height", type=int, default=512, help="Altura de exibição do preview (px)")
|
||||
parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir a GUI de setup (usar apenas args CLI)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Se faltou algo e GUI está permitida, abre o setup.
|
||||
if (not args.inputs or not args.labels or not args.out_root) and not args.no_gui_setup:
|
||||
run_with_gui_setup()
|
||||
return
|
||||
|
||||
# Modo apenas CLI
|
||||
if not args.inputs or not args.labels or not args.out_root:
|
||||
print("ERRO: É preciso informar --inputs, --labels e --out-root ou usar a GUI (sem --no-gui-setup).")
|
||||
sys.exit(1)
|
||||
|
||||
input_folders = [Path(p) for p in args.inputs]
|
||||
labels = args.labels
|
||||
out_root = Path(args.out_root)
|
||||
|
||||
triplets = collect_all_triplets(input_folders)
|
||||
if not triplets:
|
||||
print("Nenhuma tripleta encontrada nas pastas informadas.")
|
||||
sys.exit(1)
|
||||
|
||||
app = RawTripletSorterApp(
|
||||
triplets=triplets,
|
||||
labels=labels,
|
||||
out_root=out_root,
|
||||
move=args.move,
|
||||
resume=args.resume,
|
||||
display_height=args.display_height,
|
||||
)
|
||||
app.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,318 @@
|
|||
import json
|
||||
import os
|
||||
import cv2
|
||||
import csv
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from utils import carregar_labelmap_completo # importa da sua utils.py
|
||||
|
||||
# ===================================================
|
||||
# ⚙️ Configurações base vindas do config.json
|
||||
# ===================================================
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
RESOLUCAO = config["raw_size"]
|
||||
|
||||
RAW_WIDTH = RESOLUCAO[0] # largura real do raw (ex: 1296)
|
||||
RAW_HEIGHT = RESOLUCAO[1] # altura real do raw (ex: 1028)
|
||||
RAW_EXTS = [".raw"] # se tiver outro, adiciona aqui
|
||||
|
||||
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
|
||||
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
|
||||
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
|
||||
|
||||
# Cor da classe será obtida via labelmap
|
||||
# COR_CLASSE_RGB = (128, 0, 0)
|
||||
|
||||
COPIAR_IMAGENS = True
|
||||
|
||||
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
|
||||
MASK_EXT_OUT = ".png" # saída das máscaras sempre PNG
|
||||
FORCAR_SOBRESCRITA_NEW_MASK = False # sobrescrever máscara em new_masks se já existir
|
||||
VALIDAR_DIM_MASK_EXISTENTE = True # se existir, validar dimensões
|
||||
|
||||
MANIFESTO = "manifest.csv"
|
||||
|
||||
# ===================================================
|
||||
|
||||
def garantir_pasta(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
def criar_mask_solida(dim_h, dim_w, cor_rgb):
|
||||
r, g, b = cor_rgb
|
||||
mask_bgr = np.zeros((dim_h, dim_w, 3), dtype=np.uint8)
|
||||
mask_bgr[:] = (b, g, r) # OpenCV usa BGR
|
||||
return mask_bgr
|
||||
|
||||
def caminho_mask_new_para_img(pasta_new_masks, caminho_img):
|
||||
base = os.path.splitext(os.path.basename(caminho_img))[0]
|
||||
return os.path.join(pasta_new_masks, base + MASK_EXT_OUT)
|
||||
|
||||
def ler_dim(caminho_img):
|
||||
ext = os.path.splitext(caminho_img)[1].lower()
|
||||
|
||||
# Caso seja RAW: usamos dimensões fixas
|
||||
if ext in RAW_EXTS:
|
||||
h, w = RAW_HEIGHT, RAW_WIDTH
|
||||
return (h, w), None
|
||||
|
||||
# Caso normal: JPG/PNG/etc
|
||||
img = cv2.imread(caminho_img, cv2.IMREAD_COLOR)
|
||||
if img is None:
|
||||
raise RuntimeError(f"Erro ao abrir: {caminho_img}")
|
||||
h, w = img.shape[:2]
|
||||
return (h, w), img
|
||||
|
||||
def salvar_mask_solidaria(pasta_new_masks, caminho_img, cor_rgb, forcar=False, validar_dim=True):
|
||||
(h, w), _ = ler_dim(caminho_img)
|
||||
caminho_mask = caminho_mask_new_para_img(pasta_new_masks, caminho_img)
|
||||
|
||||
if os.path.exists(caminho_mask) and not forcar:
|
||||
if validar_dim:
|
||||
m = cv2.imread(caminho_mask, cv2.IMREAD_COLOR)
|
||||
if m is None or m.shape[0] != h or m.shape[1] != w:
|
||||
print(f"[ALERTA] Máscara existente com dimensão diferente, recriando: {caminho_mask}")
|
||||
else:
|
||||
print(f"[OK] Já existe (mantida): {os.path.basename(caminho_mask)}")
|
||||
return caminho_mask
|
||||
else:
|
||||
print(f"[OK] Já existe (mantida): {os.path.basename(caminho_mask)}")
|
||||
return caminho_mask
|
||||
|
||||
mask_bgr = criar_mask_solida(h, w, cor_rgb)
|
||||
garantir_pasta(os.path.dirname(caminho_mask))
|
||||
cv2.imwrite(caminho_mask, mask_bgr)
|
||||
print(f"[CRIADA] {os.path.basename(caminho_mask)} ({w}x{h})")
|
||||
return caminho_mask
|
||||
|
||||
# ========= NOVO: suporte a cópia tripla (img + mask + raw) =========
|
||||
|
||||
def _base_esta_ocupado(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
|
||||
"""
|
||||
Verifica se algum dos arquivos (img/mask/raw) com esse base já existe.
|
||||
"""
|
||||
if os.path.exists(os.path.join(dir_img, base + img_ext)):
|
||||
return True
|
||||
if os.path.exists(os.path.join(dir_mask, base + mask_ext)):
|
||||
return True
|
||||
if raw_ext and dir_raw:
|
||||
if os.path.exists(os.path.join(dir_raw, base + raw_ext)):
|
||||
return True
|
||||
return False
|
||||
|
||||
def gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
|
||||
"""
|
||||
Gera um 'base' comum livre para img/mask/raw.
|
||||
"""
|
||||
if not _base_esta_ocupado(base, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
|
||||
return base
|
||||
|
||||
i = 1
|
||||
while True:
|
||||
candidate = f"{base}_{i:03d}"
|
||||
if not _base_esta_ocupado(candidate, img_ext, mask_ext, raw_ext, dir_img, dir_mask, dir_raw):
|
||||
return candidate
|
||||
i += 1
|
||||
|
||||
def copiar_triplo(caminho_img_src, caminho_mask_src, caminho_raw_src, dest_img_dir, dest_mask_dir, dest_raw_dir):
|
||||
"""
|
||||
Copia imagem, máscara e (se existir) o RAW correspondente
|
||||
mantendo o mesmo base (com renome em caso de conflito).
|
||||
|
||||
Retorna (dst_img_path, dst_mask_path, dst_raw_path_ou_None).
|
||||
"""
|
||||
garantir_pasta(dest_img_dir)
|
||||
garantir_pasta(dest_mask_dir)
|
||||
if dest_raw_dir:
|
||||
garantir_pasta(dest_raw_dir)
|
||||
|
||||
base = os.path.splitext(os.path.basename(caminho_img_src))[0]
|
||||
|
||||
img_ext = os.path.splitext(caminho_img_src)[1].lower()
|
||||
mask_ext = ".png"
|
||||
raw_ext = os.path.splitext(caminho_raw_src)[1].lower() if caminho_raw_src else None
|
||||
|
||||
new_base = gerar_base_disponivel(base, img_ext, mask_ext, raw_ext, dest_img_dir, dest_mask_dir, dest_raw_dir)
|
||||
|
||||
dst_img_path = os.path.join(dest_img_dir, new_base + img_ext)
|
||||
dst_mask_path = os.path.join(dest_mask_dir, new_base + mask_ext)
|
||||
dst_raw_path = os.path.join(dest_raw_dir, new_base + raw_ext) if (caminho_raw_src and dest_raw_dir and raw_ext) else None
|
||||
|
||||
shutil.copy2(caminho_img_src, dst_img_path)
|
||||
shutil.copy2(caminho_mask_src, dst_mask_path)
|
||||
if caminho_raw_src and dst_raw_path:
|
||||
shutil.copy2(caminho_raw_src, dst_raw_path)
|
||||
print(f"[COPIADO] {os.path.basename(dst_img_path)} | {os.path.basename(dst_mask_path)} | {os.path.basename(dst_raw_path)}")
|
||||
else:
|
||||
print(f"[COPIADO] {os.path.basename(dst_img_path)} | {os.path.basename(dst_mask_path)} (sem RAW)")
|
||||
|
||||
return dst_img_path, dst_mask_path, dst_raw_path
|
||||
|
||||
# ===================================================
|
||||
|
||||
def obter_cor_da_classe(nome_classe: str, caminho_labelmap: str):
|
||||
"""
|
||||
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()
|
||||
target_id = None
|
||||
for cid, cname in id_para_nome.items():
|
||||
if cname.strip().lower() == alvo:
|
||||
target_id = cid
|
||||
break
|
||||
|
||||
if target_id is None:
|
||||
disponiveis = ", ".join(sorted(id_para_nome.values()))
|
||||
raise RuntimeError(
|
||||
f"Classe '{nome_classe}' não encontrada no labelmap.\n"
|
||||
f"Classes disponíveis: {disponiveis}"
|
||||
)
|
||||
|
||||
try:
|
||||
cor_rgb = colormap_rgb[target_id]
|
||||
except IndexError:
|
||||
raise RuntimeError(
|
||||
f"ID {target_id} da classe '{nome_classe}' não possui cor correspondente em colormap_rgb."
|
||||
)
|
||||
|
||||
print(f"[LABELMAP] Classe '{nome_classe}' -> ID={target_id} -> cor RGB={cor_rgb}")
|
||||
return cor_rgb
|
||||
|
||||
def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None:
|
||||
"""
|
||||
Dado o nome do preview (ex: 20260122_091833_277.png),
|
||||
tenta achar o RAW correspondente em PASTA_NEW_RAWS:
|
||||
20260122_091833_277.raw (ou outras extensões de RAW_EXTS)
|
||||
"""
|
||||
base = os.path.splitext(nome_img)[0]
|
||||
for ext in RAW_EXTS:
|
||||
candidate = os.path.join(pasta_new_raws, base + ext)
|
||||
if os.path.exists(candidate):
|
||||
return candidate
|
||||
return None
|
||||
|
||||
def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False):
|
||||
source = os.path.join("dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join("dataset", "original")
|
||||
pasta_new_previews = os.path.join(source, "group", grupo, "previews")
|
||||
pasta_new_masks = os.path.join(source, "group", grupo, "masks")
|
||||
pasta_new_raws = os.path.join(source, "group", grupo, "raws")
|
||||
|
||||
garantir_pasta(pasta_new_previews)
|
||||
garantir_pasta(pasta_new_masks)
|
||||
garantir_pasta(pasta_new_raws)
|
||||
garantir_pasta(PASTA_FINAL_PREVIEWS)
|
||||
garantir_pasta(PASTA_FINAL_MASKS)
|
||||
garantir_pasta(PASTA_FINAL_RAWS)
|
||||
|
||||
registros = []
|
||||
total, criadas_mask, copiados, puladas, erros = 0, 0, 0, 0, 0
|
||||
|
||||
for nome in os.listdir(pasta_new_previews):
|
||||
if not nome.lower().endswith(EXT_IMAGENS):
|
||||
continue
|
||||
total += 1
|
||||
caminho_img = os.path.join(pasta_new_previews, nome)
|
||||
try:
|
||||
# 1) criar máscara sólida em new_masks
|
||||
antes = os.path.exists(caminho_mask_new_para_img(pasta_new_masks, caminho_img))
|
||||
caminho_mask_new = salvar_mask_solidaria(
|
||||
pasta_new_masks,
|
||||
caminho_img,
|
||||
cor_classe_rgb,
|
||||
forcar=FORCAR_SOBRESCRITA_NEW_MASK,
|
||||
validar_dim=VALIDAR_DIM_MASK_EXISTENTE
|
||||
)
|
||||
if caminho_mask_new and not antes:
|
||||
criadas_mask += 1
|
||||
|
||||
# 2) copiar imagem + máscara + RAW para as pastas finais (com renome se necessário)
|
||||
if fazer_copia_final:
|
||||
raw_src = localizar_raw_correspondente(pasta_new_raws, nome)
|
||||
if raw_src is None:
|
||||
print(f"[AVISO] RAW não encontrado para {nome} em {pasta_new_raws}, copiando só preview+mask.")
|
||||
dst_img, dst_mask, dst_raw = copiar_triplo(
|
||||
caminho_img,
|
||||
caminho_mask_new,
|
||||
raw_src,
|
||||
PASTA_FINAL_PREVIEWS,
|
||||
PASTA_FINAL_MASKS,
|
||||
PASTA_FINAL_RAWS,
|
||||
)
|
||||
copiados += 1
|
||||
registros.append([
|
||||
caminho_img,
|
||||
caminho_mask_new,
|
||||
raw_src or "",
|
||||
dst_img,
|
||||
dst_mask,
|
||||
dst_raw or "",
|
||||
])
|
||||
else:
|
||||
puladas += 1
|
||||
|
||||
except Exception as e:
|
||||
erros += 1
|
||||
print(f"[ERRO] {nome}: {e}")
|
||||
|
||||
if manifesto_csv and registros:
|
||||
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
|
||||
w = csv.writer(f)
|
||||
w.writerow([
|
||||
"src_image",
|
||||
"src_mask",
|
||||
"src_raw",
|
||||
"dst_image",
|
||||
"dst_mask",
|
||||
"dst_raw",
|
||||
])
|
||||
w.writerows(registros)
|
||||
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
|
||||
|
||||
print(f"\nResumo: total_imgs={total} | masks_criadas={criadas_mask} | copiados={copiados} | puladas={puladas} | erros={erros}")
|
||||
|
||||
def build_cli():
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Gera máscaras sólidas para novas imagens de UMA classe (via labelmap) e copia para dataset final com dedup."
|
||||
)
|
||||
ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
||||
ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
||||
ap.add_argument("--classe", required=True, help="Nome da classe (como está no labelmap.txt). Ex: chao, cana, erva")
|
||||
ap.add_argument("--grupo", required=True, help="Nome da pasta grupo. Ex: chao, cana, erva, chao_cana, chao_cana_erva")
|
||||
ap.add_argument("--no-copy", action="store_true", help="Não copia para as pastas finais (só cria masks em new_masks).")
|
||||
ap.add_argument("--from-originals", action="store_true", help="Faz o procedimento na pasta em originais")
|
||||
ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a gerar (ou vazio para não gerar).")
|
||||
return ap
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
|
||||
# Caminho padrão pro labelmap, caso não seja informado
|
||||
labelmap_path = os.path.join("dataset", "labelmap.txt")
|
||||
|
||||
if not os.path.exists(labelmap_path):
|
||||
raise SystemExit(f"Labelmap não encontrado em: {labelmap_path}")
|
||||
|
||||
# Descobre a cor da classe no labelmap
|
||||
cor_rgb = obter_cor_da_classe(args.classe, labelmap_path)
|
||||
|
||||
fazer_copia = not args.no_copy
|
||||
manifesto_csv = args.manifest if args.manifest else None
|
||||
|
||||
processar_novas_imagens(
|
||||
cana=args.cana,
|
||||
horario=args.horario,
|
||||
grupo=args.grupo,
|
||||
cor_classe_rgb=cor_rgb,
|
||||
fazer_copia_final=fazer_copia,
|
||||
manifesto_csv=manifesto_csv,
|
||||
orignais=args.from_originals
|
||||
)
|
||||
|
|
@ -0,0 +1,205 @@
|
|||
import json
|
||||
import os
|
||||
import shutil
|
||||
import argparse
|
||||
import csv
|
||||
|
||||
# ================= CONFIG =================
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
|
||||
# Raiz das brutas agrupadas
|
||||
PASTA_BRUTAS_GROUP_ROOT = os.path.join("dataset", "brutas")
|
||||
|
||||
# Onde as máscaras novas (rotuladas externamente) são colocadas
|
||||
PASTA_NEW_MASKS = os.path.join("dataset", "new_masks")
|
||||
|
||||
# Destino final do dataset consolidado
|
||||
PASTA_FINAL_PREVIEWS = os.path.join("dataset", "original", "previews")
|
||||
PASTA_FINAL_RAWS = os.path.join("dataset", "original", "raws")
|
||||
PASTA_FINAL_MASKS = os.path.join("dataset", "original", "masks")
|
||||
|
||||
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
||||
EXT_MASKS = (".png",)
|
||||
EXT_RAWS = (".raw",)
|
||||
|
||||
# =========================================
|
||||
|
||||
def garantir_pasta(p: str):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
def indexar_brutas_group(root_group: str):
|
||||
"""
|
||||
Varrre dataset/brutas/group recursivamente e monta dois dicionários:
|
||||
base -> caminho_preview
|
||||
base -> caminho_raw
|
||||
|
||||
Considera que a estrutura é:
|
||||
brutas/group/<grupo>/previews/*.png|jpg
|
||||
brutas/group/<grupo>/raws/*.raw
|
||||
"""
|
||||
previews_map = {}
|
||||
raws_map = {}
|
||||
|
||||
if not os.path.isdir(root_group):
|
||||
print(f"[AVISO] Pasta de grupos não existe: {root_group}")
|
||||
return previews_map, raws_map
|
||||
|
||||
for dirpath, dirnames, filenames in os.walk(root_group):
|
||||
base_dir = os.path.basename(dirpath).lower()
|
||||
|
||||
if base_dir == "previews":
|
||||
for nome in filenames:
|
||||
if not nome.lower().endswith(EXT_PREVIEWS):
|
||||
continue
|
||||
base = os.path.splitext(nome)[0]
|
||||
# Se já existir, mantemos o primeiro encontrado (pode logar conflito se quiser)
|
||||
if base not in previews_map:
|
||||
previews_map[base] = os.path.join(dirpath, nome)
|
||||
|
||||
elif base_dir == "raws":
|
||||
for nome in filenames:
|
||||
if not nome.lower().endswith(EXT_RAWS):
|
||||
continue
|
||||
base = os.path.splitext(nome)[0]
|
||||
if base not in raws_map:
|
||||
raws_map[base] = os.path.join(dirpath, nome)
|
||||
|
||||
print(f"[INDEX] Previews indexados: {len(previews_map)}")
|
||||
print(f"[INDEX] Raws indexados : {len(raws_map)}")
|
||||
return previews_map, raws_map
|
||||
|
||||
def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
|
||||
"""
|
||||
Devolve um caminho disponível em dest_dir com base + ext.
|
||||
Se já existir, adiciona _001, _002, ...
|
||||
"""
|
||||
p = os.path.join(dest_dir, base + ext)
|
||||
if not os.path.exists(p):
|
||||
return p
|
||||
|
||||
i = 1
|
||||
while True:
|
||||
p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}")
|
||||
if not os.path.exists(p):
|
||||
return p
|
||||
i += 1
|
||||
|
||||
def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
|
||||
garantir_pasta(PASTA_FINAL_PREVIEWS)
|
||||
garantir_pasta(PASTA_FINAL_RAWS)
|
||||
garantir_pasta(PASTA_FINAL_MASKS)
|
||||
|
||||
if not os.path.isdir(PASTA_NEW_MASKS):
|
||||
raise SystemExit(f"[ERRO] Pasta new_masks não existe: {PASTA_NEW_MASKS}")
|
||||
|
||||
# 1) Indexa tudo que existe em brutas/group
|
||||
previews_map, raws_map = indexar_brutas_group(PASTA_BRUTAS_GROUP_ROOT)
|
||||
|
||||
registros = []
|
||||
total, copiados, ignorados = 0, 0, 0
|
||||
|
||||
for nome in os.listdir(PASTA_NEW_MASKS):
|
||||
if not nome.lower().endswith(EXT_MASKS):
|
||||
continue
|
||||
|
||||
total += 1
|
||||
base = os.path.splitext(nome)[0]
|
||||
caminho_mask_src = os.path.join(PASTA_NEW_MASKS, nome)
|
||||
|
||||
preview_src = previews_map.get(base)
|
||||
raw_src = raws_map.get(base)
|
||||
|
||||
if preview_src is None or raw_src is None:
|
||||
ignorados += 1
|
||||
motivo = []
|
||||
if preview_src is None:
|
||||
motivo.append("preview")
|
||||
if raw_src is None:
|
||||
motivo.append("raw")
|
||||
print(f"[SKIP] {base} -> faltando: {', '.join(motivo)}")
|
||||
# NÃO move a máscara, ela fica em new_masks pra você analisar depois
|
||||
continue
|
||||
|
||||
prev_ext = os.path.splitext(preview_src)[1].lower()
|
||||
mask_ext = os.path.splitext(nome)[1].lower()
|
||||
raw_ext = os.path.splitext(raw_src)[1].lower()
|
||||
|
||||
# Gera um base final único usando o preview como referência
|
||||
dst_preview = nome_disponivel(PASTA_FINAL_PREVIEWS, base, prev_ext)
|
||||
new_base = os.path.splitext(os.path.basename(dst_preview))[0]
|
||||
|
||||
dst_raw = os.path.join(PASTA_FINAL_RAWS, new_base + raw_ext)
|
||||
dst_mask = os.path.join(PASTA_FINAL_MASKS, new_base + mask_ext)
|
||||
|
||||
# Copia preview + raw, move ou copia a máscara
|
||||
shutil.copy2(preview_src, dst_preview)
|
||||
shutil.copy2(raw_src, dst_raw)
|
||||
if mover:
|
||||
shutil.move(caminho_mask_src, dst_mask)
|
||||
else:
|
||||
shutil.copy2(caminho_mask_src, dst_mask)
|
||||
|
||||
copiados += 1
|
||||
registros.append([
|
||||
preview_src,
|
||||
raw_src,
|
||||
caminho_mask_src,
|
||||
dst_preview,
|
||||
dst_raw,
|
||||
dst_mask,
|
||||
])
|
||||
|
||||
print(f"[OK] {new_base}")
|
||||
|
||||
if manifesto_csv and registros:
|
||||
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
|
||||
w = csv.writer(f)
|
||||
w.writerow([
|
||||
"src_preview",
|
||||
"src_raw",
|
||||
"src_mask",
|
||||
"dst_preview",
|
||||
"dst_raw",
|
||||
"dst_mask",
|
||||
])
|
||||
w.writerows(registros)
|
||||
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
|
||||
|
||||
print(f"\nResumo: total_masks={total} | ingestas={copiados} | ignoradas={ignorados}")
|
||||
|
||||
# ================= CLI =================
|
||||
|
||||
def build_cli():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Ingere máscaras novas (dataset/new_masks) procurando previews e raws "
|
||||
"nas subpastas de dataset/brutas/group/** e copiando tudo para dataset/original."
|
||||
)
|
||||
)
|
||||
ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
||||
ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
||||
ap.add_argument("--copy", action="store_true", help="Copia as máscaras em vez de mover (por padrão, move e esvazia new_masks conforme ingere).")
|
||||
ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).")
|
||||
return ap
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
|
||||
manifesto = args.manifest if args.manifest.strip() else None
|
||||
mover = not args.copy
|
||||
|
||||
PASTA_BRUTAS_GROUP_ROOT = os.path.join(PASTA_BRUTAS_GROUP_ROOT, f"cana_{args.cana}", args.horario, "group")
|
||||
|
||||
print(f"[INFO] MODELO : {MODELO}")
|
||||
print(f"[INFO] Brutas (group root): {PASTA_BRUTAS_GROUP_ROOT}")
|
||||
print(f"[INFO] New masks : {PASTA_NEW_MASKS}")
|
||||
print(f"[INFO] Mover máscaras? : {mover}")
|
||||
print(f"[INFO] Manifesto : {manifesto or '(nenhum)'}")
|
||||
print()
|
||||
|
||||
processar_new_masks(mover=mover, manifesto_csv=manifesto)
|
||||
|
|
@ -0,0 +1,328 @@
|
|||
import json
|
||||
import os
|
||||
import shutil
|
||||
import argparse
|
||||
import csv
|
||||
|
||||
# ================= CONFIG =================
|
||||
|
||||
with open("config.json", "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
|
||||
# Raiz das brutas (todas as canas/horários/grupos)
|
||||
PASTA_BRUTAS_ROOT = os.path.join("dataset", "brutas")
|
||||
|
||||
# Onde você coloca os previews selecionados (tudo misturado)
|
||||
PASTA_SELECTED_PREVIEWS = os.path.join("dataset", "selected_previews")
|
||||
|
||||
# Destino final, organizado por grupo:
|
||||
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
|
||||
PASTA_ORIGINAL_GROUP_ROOT = os.path.join("dataset", "original", "group")
|
||||
|
||||
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
||||
EXT_RAWS = (".raw",)
|
||||
EXT_MASKS = (".png",)
|
||||
# Ajusta se suas metas tiverem outra extensão
|
||||
EXT_METAS = (".json", ".yml", ".yaml", ".txt", ".csv")
|
||||
|
||||
|
||||
# ================= HELPERS =================
|
||||
|
||||
def garantir_pasta(p: str):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
|
||||
def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
|
||||
"""
|
||||
Devolve um caminho disponível em dest_dir com base + ext.
|
||||
Se já existir, adiciona _001, _002, ...
|
||||
"""
|
||||
p = os.path.join(dest_dir, base + ext)
|
||||
if not os.path.exists(p):
|
||||
return p
|
||||
|
||||
i = 1
|
||||
while True:
|
||||
p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}")
|
||||
if not os.path.exists(p):
|
||||
return p
|
||||
i += 1
|
||||
|
||||
|
||||
def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None:
|
||||
"""
|
||||
Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'.
|
||||
Retorna o caminho completo ou None.
|
||||
"""
|
||||
if not os.path.isdir(pasta):
|
||||
return None
|
||||
|
||||
for nome in os.listdir(pasta):
|
||||
nome_lower = nome.lower()
|
||||
root, ext = os.path.splitext(nome_lower)
|
||||
if root == base.lower() and ext in exts:
|
||||
return os.path.join(pasta, nome)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def indexar_brutas_por_raw(root_brutas: str):
|
||||
"""
|
||||
Varre dataset/brutas recursivamente, olhando pastas 'raws' e montando índice:
|
||||
|
||||
base -> {
|
||||
"raw": caminho_raw,
|
||||
"cana": "cana_alta" | "cana_baixa" | ... (se conseguir inferir),
|
||||
"horario": "meio_dia" | "cedo" | ... (se conseguir inferir),
|
||||
"grupo": "<GRUPO>",
|
||||
"group_dir": caminho_da_pasta_do_grupo
|
||||
}
|
||||
|
||||
Estrutura esperada (relativa a root_brutas):
|
||||
cana_alta/meio_dia/group/<GRUPO>/raws/*.raw
|
||||
"""
|
||||
index = {}
|
||||
|
||||
if not os.path.isdir(root_brutas):
|
||||
print(f"[AVISO] Pasta de brutas não existe: {root_brutas}")
|
||||
return index
|
||||
|
||||
for dirpath, dirnames, filenames in os.walk(root_brutas):
|
||||
base_dir = os.path.basename(dirpath).lower()
|
||||
|
||||
if base_dir != "raws":
|
||||
continue
|
||||
|
||||
# dirpath = .../cana_x/horario/group/<GRUPO>/raws
|
||||
group_dir = os.path.dirname(dirpath) # .../cana_x/horario/group/<GRUPO>
|
||||
|
||||
rel = os.path.relpath(dirpath, root_brutas)
|
||||
parts = rel.split(os.sep)
|
||||
|
||||
# Defaults
|
||||
cana = None
|
||||
horario = None
|
||||
grupo = None
|
||||
|
||||
if len(parts) >= 5:
|
||||
# [0] = cana_<...>
|
||||
# [1] = horario
|
||||
# [2] = "group"
|
||||
# [3] = <GRUPO>
|
||||
cana = parts[0]
|
||||
horario = parts[1]
|
||||
# parts[2] deve ser "group"
|
||||
grupo = parts[3]
|
||||
else:
|
||||
# fallback bem genérico
|
||||
if "group" in parts:
|
||||
i = parts.index("group")
|
||||
if i + 1 < len(parts):
|
||||
grupo = parts[i + 1]
|
||||
|
||||
for nome in filenames:
|
||||
if not nome.lower().endswith(EXT_RAWS):
|
||||
continue
|
||||
|
||||
base = os.path.splitext(nome)[0]
|
||||
|
||||
if base in index:
|
||||
# Conflito (mesmo base em dois lugares) -> loga e mantém o primeiro
|
||||
print(f"[CONFLITO] base repetida em raws: {base}")
|
||||
continue
|
||||
|
||||
raw_path = os.path.join(dirpath, nome)
|
||||
|
||||
index[base] = {
|
||||
"raw": raw_path,
|
||||
"cana": cana,
|
||||
"horario": horario,
|
||||
"grupo": grupo,
|
||||
"group_dir": group_dir,
|
||||
}
|
||||
|
||||
print(f"[INDEX] Entradas indexadas por RAW: {len(index)}")
|
||||
return index
|
||||
|
||||
|
||||
def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | None = None):
|
||||
if not os.path.isdir(PASTA_SELECTED_PREVIEWS):
|
||||
raise SystemExit(f"[ERRO] Pasta selected_previews não existe: {PASTA_SELECTED_PREVIEWS}")
|
||||
|
||||
# Indexa todas as brutas a partir dos RAWs
|
||||
index_raws = indexar_brutas_por_raw(PASTA_BRUTAS_ROOT)
|
||||
|
||||
registros = []
|
||||
total, movidos, ignorados = 0, 0, 0
|
||||
|
||||
for nome in os.listdir(PASTA_SELECTED_PREVIEWS):
|
||||
caminho_preview_sel = os.path.join(PASTA_SELECTED_PREVIEWS, nome)
|
||||
|
||||
if not os.path.isfile(caminho_preview_sel):
|
||||
continue
|
||||
|
||||
if not nome.lower().endswith(EXT_PREVIEWS):
|
||||
continue
|
||||
|
||||
total += 1
|
||||
|
||||
base = os.path.splitext(nome)[0]
|
||||
|
||||
info = index_raws.get(base)
|
||||
if info is None:
|
||||
ignorados += 1
|
||||
print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via RAW)")
|
||||
continue
|
||||
|
||||
grupo = info.get("grupo") or "unknown"
|
||||
cana = info.get("cana") or "unknown"
|
||||
horario = info.get("horario") or "unknown"
|
||||
group_dir = info["group_dir"]
|
||||
raw_src = info["raw"]
|
||||
|
||||
# Pastas irmãs em brutas
|
||||
metas_src_dir = os.path.join(group_dir, "metas")
|
||||
masks_src_dir = os.path.join(group_dir, "masks")
|
||||
|
||||
meta_src = buscar_por_base(metas_src_dir, base, EXT_METAS)
|
||||
mask_src = buscar_por_base(masks_src_dir, base, EXT_MASKS)
|
||||
|
||||
# Destino: dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
|
||||
dest_group_root = os.path.join(PASTA_ORIGINAL_GROUP_ROOT, grupo)
|
||||
dest_prev_dir = os.path.join(dest_group_root, "previews")
|
||||
dest_raw_dir = os.path.join(dest_group_root, "raws")
|
||||
dest_meta_dir = os.path.join(dest_group_root, "metas")
|
||||
dest_mask_dir = os.path.join(dest_group_root, "masks")
|
||||
|
||||
garantir_pasta(dest_prev_dir)
|
||||
garantir_pasta(dest_raw_dir)
|
||||
garantir_pasta(dest_meta_dir)
|
||||
garantir_pasta(dest_mask_dir)
|
||||
|
||||
# Define extensões
|
||||
prev_ext_sel = os.path.splitext(nome)[1].lower()
|
||||
raw_ext = os.path.splitext(raw_src)[1].lower()
|
||||
|
||||
# Gera nome final único com base no preview
|
||||
dst_preview = nome_disponivel(dest_prev_dir, base, prev_ext_sel)
|
||||
new_base = os.path.splitext(os.path.basename(dst_preview))[0]
|
||||
|
||||
dst_raw = os.path.join(dest_raw_dir, new_base + raw_ext)
|
||||
dst_meta = None
|
||||
dst_mask = None
|
||||
|
||||
if meta_src is not None:
|
||||
meta_ext = os.path.splitext(meta_src)[1].lower()
|
||||
dst_meta = os.path.join(dest_meta_dir, new_base + meta_ext)
|
||||
|
||||
if mask_src is not None:
|
||||
mask_ext = os.path.splitext(mask_src)[1].lower()
|
||||
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
|
||||
|
||||
# Copia/move preview selecionado
|
||||
if copy_only:
|
||||
shutil.copy2(caminho_preview_sel, dst_preview)
|
||||
else:
|
||||
shutil.move(caminho_preview_sel, dst_preview)
|
||||
|
||||
# Copia RAW e metas/masks (se existirem)
|
||||
if copy_only:
|
||||
shutil.copy2(raw_src, dst_raw)
|
||||
else:
|
||||
shutil.move(raw_src, dst_raw)
|
||||
|
||||
if meta_src is not None:
|
||||
if copy_only:
|
||||
shutil.copy2(meta_src, dst_meta)
|
||||
else:
|
||||
shutil.move(meta_src, dst_meta)
|
||||
|
||||
if mask_src is not None:
|
||||
if copy_only:
|
||||
shutil.copy2(mask_src, dst_mask)
|
||||
else:
|
||||
shutil.move(mask_src, dst_mask)
|
||||
|
||||
movidos += 1
|
||||
|
||||
registros.append([
|
||||
base,
|
||||
grupo,
|
||||
cana,
|
||||
horario,
|
||||
caminho_preview_sel,
|
||||
raw_src,
|
||||
meta_src or "",
|
||||
mask_src or "",
|
||||
dst_preview,
|
||||
dst_raw,
|
||||
dst_meta or "",
|
||||
dst_mask or "",
|
||||
])
|
||||
|
||||
print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario}")
|
||||
|
||||
if manifesto_csv and registros:
|
||||
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
|
||||
w = csv.writer(f)
|
||||
w.writerow([
|
||||
"base",
|
||||
"grupo",
|
||||
"cana",
|
||||
"horario",
|
||||
"src_preview_selected",
|
||||
"src_raw",
|
||||
"src_meta",
|
||||
"src_mask",
|
||||
"dst_preview",
|
||||
"dst_raw",
|
||||
"dst_meta",
|
||||
"dst_mask",
|
||||
])
|
||||
w.writerows(registros)
|
||||
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
|
||||
|
||||
print(f"\nResumo: total_selected={total} | organizadas={movidos} | ignoradas={ignorados}")
|
||||
|
||||
|
||||
# ================= CLI =================
|
||||
|
||||
def build_cli():
|
||||
ap = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Organiza previews selecionadas (dataset/selected_previews) "
|
||||
"descobrindo cana/horário/grupo em dataset/brutas e copiando/movendo "
|
||||
"preview + raw + meta (+ mask se existir) para dataset/original/group/{GRUPO}."
|
||||
)
|
||||
)
|
||||
ap.add_argument(
|
||||
"--copy",
|
||||
action="store_true",
|
||||
help="Copia os previews em vez de mover (por padrão, move e esvazia selected_previews conforme organiza)."
|
||||
)
|
||||
ap.add_argument(
|
||||
"--manifest",
|
||||
default="",
|
||||
help="Caminho para CSV de manifesto (opcional)."
|
||||
)
|
||||
return ap
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
|
||||
manifesto = args.manifest if args.manifest.strip() else None
|
||||
copy_only = args.copy
|
||||
|
||||
print(f"[INFO] MODELO : {MODELO}")
|
||||
print(f"[INFO] Brutas root : {PASTA_BRUTAS_ROOT}")
|
||||
print(f"[INFO] Selected previews : {PASTA_SELECTED_PREVIEWS}")
|
||||
print(f"[INFO] Original group root : {PASTA_ORIGINAL_GROUP_ROOT}")
|
||||
print(f"[INFO] Copy only? : {copy_only}")
|
||||
print(f"[INFO] Manifesto : {manifesto or '(nenhum)'}")
|
||||
print()
|
||||
|
||||
organizar_selected_previews(copy_only=copy_only, manifesto_csv=manifesto)
|
||||
|
|
@ -0,0 +1,322 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
Versão OTIMIZADA — Agrupamento ultra-rápido por classes presentes na máscara.
|
||||
Substitui extrair_ids_presentes por uma versão 50× mais rápida.
|
||||
|
||||
Entrada:
|
||||
originals/
|
||||
previews/
|
||||
raws/
|
||||
masks/
|
||||
(opcional) masks2/
|
||||
|
||||
Saída:
|
||||
originals/group/<grupo>/previews
|
||||
/raws
|
||||
/masks
|
||||
/masks2
|
||||
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import csv
|
||||
import json
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config["camera"]
|
||||
USE_MASKS2 = config["dual_head"]
|
||||
|
||||
EXT_PREVIEWS = (".jpg", ".jpeg", ".png")
|
||||
EXT_MASKS = (".png", ".jpg", ".jpeg")
|
||||
EXT_MASKS2 = (".png", ".jpg", ".jpeg")
|
||||
EXT_RAW = (".raw",)
|
||||
|
||||
MANIFESTO_DEFAULT = "manifest.csv"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 🔥 EXTRAÇÃO DE CLASSES — VERSÃO ULTRA OTIMIZADA
|
||||
# ============================================================
|
||||
|
||||
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
|
||||
"""
|
||||
Extração ultra-rápida de IDs presentes na máscara.
|
||||
|
||||
Estratégia:
|
||||
1) Máscara 1 canal → np.unique (instantâneo).
|
||||
2) Máscara 3 canais:
|
||||
- se assume_rgb=True: labelmap está em RGB,
|
||||
convertemos a imagem BGR -> RGB para casar com cor_para_id.
|
||||
- se assume_rgb=False: labelmap está em BGR,
|
||||
mantemos a imagem em BGR.
|
||||
3) Amostragem em grid + converter_mask_rgb_para_ids.
|
||||
4) Fallback full-scan se necessário.
|
||||
"""
|
||||
m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
|
||||
if m is None:
|
||||
raise RuntimeError(f"Falha ao abrir máscara: {mask_path}")
|
||||
|
||||
# -------------------------------------------------------
|
||||
# CASO 1: máscara indexada (1 canal) — instantâneo
|
||||
# -------------------------------------------------------
|
||||
if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1):
|
||||
ids = np.unique(m)
|
||||
return set(int(v) for v in ids)
|
||||
|
||||
# -------------------------------------------------------
|
||||
# CASO 2: máscara RGB
|
||||
# -------------------------------------------------------
|
||||
if assume_rgb:
|
||||
# Labelmap em RGB, OpenCV em BGR -> converte
|
||||
img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB)
|
||||
else:
|
||||
# Labelmap em BGR, OpenCV já em BGR -> usa direto
|
||||
img = m
|
||||
|
||||
# Agora cor_para_id e img estão no MESMO espaço de cor
|
||||
mapa_rgb = cor_para_id
|
||||
max_classes = len(cor_para_id)
|
||||
|
||||
# ---------- AMOSTRAGEM ----------
|
||||
step = 8 # pode virar 4 se quiser mais precisão, 16 se quiser mais velocidade
|
||||
amostra = img[::step, ::step]
|
||||
amostra_ids = converter_mask_rgb_para_ids(amostra, mapa_rgb, ignore_id=255)
|
||||
ids = set(int(x) for x in np.unique(amostra_ids) if x != 255)
|
||||
|
||||
if len(ids) >= max_classes:
|
||||
return ids
|
||||
|
||||
# ---------- FULL-SCAN OTIMIZADO (último caso) ----------
|
||||
full_ids = converter_mask_rgb_para_ids(img, mapa_rgb, ignore_id=255)
|
||||
ids = set(int(x) for x in np.unique(full_ids) if x != 255)
|
||||
return ids
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 🔧 Helpers
|
||||
# ============================================================
|
||||
|
||||
def garantir_pasta(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
def nome_disponivel(dest_dir, base_name, ext):
|
||||
cand = os.path.join(dest_dir, base_name + ext)
|
||||
if not os.path.exists(cand):
|
||||
return cand
|
||||
i = 1
|
||||
while True:
|
||||
cand = os.path.join(dest_dir, f"{base_name}_{i:03d}{ext}")
|
||||
if not os.path.exists(cand):
|
||||
return cand
|
||||
i += 1
|
||||
|
||||
def mapear_por_base_priorizando_png(pasta, exts):
|
||||
if not pasta or not os.path.isdir(pasta):
|
||||
return {}
|
||||
mapa = {}
|
||||
for nome in os.listdir(pasta):
|
||||
lower = nome.lower()
|
||||
if not lower.endswith(exts):
|
||||
continue
|
||||
base, ext = os.path.splitext(nome)
|
||||
cam = os.path.join(pasta, nome)
|
||||
if base not in mapa:
|
||||
mapa[base] = cam
|
||||
else:
|
||||
atual_ext = os.path.splitext(mapa[base])[1].lower()
|
||||
if atual_ext != ".png" and ext.lower() == ".png":
|
||||
mapa[base] = cam
|
||||
return mapa
|
||||
|
||||
def inferir_ignore_id(ignore_rgb, cor_para_id):
|
||||
if ignore_rgb is None:
|
||||
return None
|
||||
if isinstance(ignore_rgb, (list, tuple)) and len(ignore_rgb) == 3:
|
||||
return cor_para_id.get(tuple(ignore_rgb))
|
||||
if isinstance(ignore_rgb, int):
|
||||
return ignore_rgb
|
||||
return None
|
||||
|
||||
def montar_nome_grupo(ids_presentes, id_para_nome):
|
||||
"""
|
||||
Constrói o nome do grupo respeitando a ordem natural dos IDs do labelmap.
|
||||
Ex: {0,1,2} -> chao_cana_obstaculo
|
||||
"""
|
||||
if not ids_presentes:
|
||||
return "sem_classe"
|
||||
nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)]
|
||||
return "_".join(nomes)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 📦 PIPELINE PRINCIPAL
|
||||
# ============================================================
|
||||
|
||||
def copiar_ou_mover_tripla(preview_src, raw_src, mask_src,
|
||||
dest_prev_dir, dest_raw_dir, dest_mask_dir,
|
||||
mover=False,
|
||||
mask2_src=None, dest_mask2_dir=None):
|
||||
|
||||
garantir_pasta(dest_prev_dir)
|
||||
garantir_pasta(dest_raw_dir)
|
||||
garantir_pasta(dest_mask_dir)
|
||||
|
||||
base_prev = os.path.splitext(os.path.basename(preview_src))[0]
|
||||
prev_ext = os.path.splitext(preview_src)[1].lower()
|
||||
raw_ext = os.path.splitext(raw_src)[1].lower()
|
||||
mask_ext = os.path.splitext(mask_src)[1].lower()
|
||||
|
||||
dst_prev = nome_disponivel(dest_prev_dir, base_prev, prev_ext)
|
||||
new_base = os.path.splitext(os.path.basename(dst_prev))[0]
|
||||
|
||||
dst_raw = os.path.join(dest_raw_dir, new_base + raw_ext)
|
||||
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
|
||||
|
||||
if os.path.exists(dst_raw) or os.path.exists(dst_mask):
|
||||
dst_mask = nome_disponivel(dest_mask_dir, new_base, mask_ext)
|
||||
new_base = os.path.splitext(os.path.basename(dst_mask))[0]
|
||||
dst_prev = nome_disponivel(dest_prev_dir, new_base, prev_ext)
|
||||
dst_raw = nome_disponivel(dest_raw_dir, new_base, raw_ext)
|
||||
|
||||
if mover:
|
||||
shutil.move(preview_src, dst_prev)
|
||||
shutil.move(raw_src, dst_raw)
|
||||
shutil.move(mask_src, dst_mask)
|
||||
if mask2_src:
|
||||
shutil.move(mask2_src, os.path.join(dest_mask2_dir, new_base + os.path.splitext(mask2_src)[1]))
|
||||
else:
|
||||
shutil.copy2(preview_src, dst_prev)
|
||||
shutil.copy2(raw_src, dst_raw)
|
||||
shutil.copy2(mask_src, dst_mask)
|
||||
if mask2_src:
|
||||
garantir_pasta(dest_mask2_dir)
|
||||
shutil.copy2(mask2_src, os.path.join(dest_mask2_dir, new_base + os.path.splitext(mask2_src)[1]))
|
||||
|
||||
return dst_prev, dst_raw, dst_mask, None
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 🚀 PROCESSAMENTO
|
||||
# ============================================================
|
||||
|
||||
def processar(originals_dir, labelmap_path, mover=False,
|
||||
manifesto=MANIFESTO_DEFAULT, validar_dim=True,
|
||||
estrito=False, labelmap_bgr=False):
|
||||
|
||||
originals_dir = os.path.abspath(originals_dir)
|
||||
previews_dir = os.path.join(originals_dir, "previews")
|
||||
raws_dir = os.path.join(originals_dir, "raws")
|
||||
masks_dir = os.path.join(originals_dir, "masks")
|
||||
masks2_dir = os.path.join(originals_dir, "masks2")
|
||||
group_dir = os.path.join(originals_dir, "group")
|
||||
|
||||
if not os.path.isdir(previews_dir) or not os.path.isdir(raws_dir) or not os.path.isdir(masks_dir):
|
||||
raise RuntimeError("Estrutura inválida em originals/")
|
||||
|
||||
usar_masks2 = USE_MASKS2 and os.path.isdir(masks2_dir)
|
||||
|
||||
cor_para_id, colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
ignore_id = inferir_ignore_id(ignore_rgb, cor_para_id)
|
||||
|
||||
garantir_pasta(group_dir)
|
||||
|
||||
mapa_previews = mapear_por_base_priorizando_png(previews_dir, EXT_PREVIEWS)
|
||||
mapa_masks = mapear_por_base_priorizando_png(masks_dir, EXT_MASKS)
|
||||
mapa_raws = {os.path.splitext(n)[0]: os.path.join(raws_dir, n)
|
||||
for n in os.listdir(raws_dir) if n.lower().endswith(EXT_RAW)}
|
||||
mapa_masks2 = mapear_por_base_priorizando_png(masks2_dir, EXT_MASKS2) if usar_masks2 else {}
|
||||
|
||||
registros = []
|
||||
totais = {"total_masks":0, "processados":0, "pulados":0}
|
||||
|
||||
for base, mask_path in sorted(mapa_masks.items()):
|
||||
totais["total_masks"] += 1
|
||||
|
||||
prev_path = mapa_previews.get(base)
|
||||
raw_path = mapa_raws.get(base)
|
||||
|
||||
if not prev_path or not raw_path:
|
||||
totais["pulados"] += 1
|
||||
continue
|
||||
|
||||
# ------------------------------
|
||||
# EXTRAÇÃO ULTRA RÁPIDA
|
||||
# ------------------------------
|
||||
ids_presentes = extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=not labelmap_bgr)
|
||||
|
||||
if ignore_id in ids_presentes:
|
||||
ids_presentes.discard(ignore_id)
|
||||
|
||||
grupo = montar_nome_grupo(ids_presentes, id_para_nome)
|
||||
|
||||
# ------------------------------
|
||||
# CRIA PASTAS
|
||||
# ------------------------------
|
||||
dest_prev_dir = os.path.join(group_dir, grupo, "previews")
|
||||
dest_raw_dir = os.path.join(group_dir, grupo, "raws")
|
||||
dest_mask_dir = os.path.join(group_dir, grupo, "masks")
|
||||
dest_mask2_dir= os.path.join(group_dir, grupo, "masks2") if usar_masks2 else None
|
||||
|
||||
dst_prev, dst_raw, dst_mask, _ = copiar_ou_mover_tripla(
|
||||
prev_path, raw_path, mask_path,
|
||||
dest_prev_dir, dest_raw_dir, dest_mask_dir,
|
||||
mover=mover,
|
||||
mask2_src=mapa_masks2.get(base),
|
||||
dest_mask2_dir=dest_mask2_dir
|
||||
)
|
||||
|
||||
registros.append([grupo, prev_path, raw_path, mask_path, dst_prev, dst_raw, dst_mask])
|
||||
totais["processados"] += 1
|
||||
|
||||
# ------------------------------
|
||||
# MANIFESTO
|
||||
# ------------------------------
|
||||
if manifesto:
|
||||
with open(manifesto, "w", newline="", encoding="utf-8") as f:
|
||||
w = csv.writer(f)
|
||||
w.writerow(["grupo","src_preview","src_raw","src_mask","dst_preview","dst_raw","dst_mask"])
|
||||
for r in registros:
|
||||
w.writerow(r)
|
||||
|
||||
print("\nFinalizado!")
|
||||
print("Total masks:", totais["total_masks"])
|
||||
print("Processados:", totais["processados"])
|
||||
print("Pulados :", totais["pulados"])
|
||||
|
||||
|
||||
# ============================================================
|
||||
# CLI
|
||||
# ============================================================
|
||||
|
||||
def build_cli():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--move", action="store_true", help="Mover em vez de copiar")
|
||||
ap.add_argument("--manifest", default=MANIFESTO_DEFAULT)
|
||||
ap.add_argument("--no-validate", action="store_true")
|
||||
ap.add_argument("--strict", action="store_true")
|
||||
ap.add_argument("--labels-bgr", action="store_true")
|
||||
return ap
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = build_cli()
|
||||
args = ap.parse_args()
|
||||
originals_dir = os.path.join("dataset", "original")
|
||||
labelmap_dir = os.path.join("dataset", "labelmap.txt")
|
||||
processar(
|
||||
originals_dir=originals_dir,
|
||||
labelmap_path=labelmap_dir,
|
||||
mover=args.move,
|
||||
manifesto=args.manifest,
|
||||
validar_dim=not args.no_validate,
|
||||
estrito=args.strict,
|
||||
labelmap_bgr=args.labels_bgr
|
||||
)
|
||||
|
|
@ -0,0 +1,708 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Augmenta imagens/máscaras *por grupo* usando RAW4 (float32, 0..1, R G IR B).
|
||||
|
||||
Novo fluxo (GAL5000):
|
||||
|
||||
Entrada:
|
||||
dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
|
||||
dataset/original/group/<grupo>/masks (.png/.jpg)
|
||||
(opcional) dataset/original/group/<grupo>/masks2
|
||||
|
||||
Saída:
|
||||
dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
|
||||
dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
|
||||
dataset/augmented/group/<grupo>/masks
|
||||
dataset/augmented/group/<grupo>/masks2
|
||||
|
||||
Regras:
|
||||
- GEOMETRIA (flip, shift, rotate) aplicada em RAW + masks (+ masks2).
|
||||
- BLUR / MOTION BLUR aplicado no RAW, canal por canal, com o mesmo kernel.
|
||||
- Preview é SEMPRE derivado do RAW já transformado (geom + blur).
|
||||
- Máscaras não sofrem blur, apenas geom.
|
||||
|
||||
Uso:
|
||||
python _5_augmentation_raw.py --copies 5
|
||||
python _5_augmentation_raw.py --copies 5 --groups chao,chao_erva,cana
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import cv2
|
||||
from PIL import Image
|
||||
import albumentations as A
|
||||
from albumentations import ReplayCompose
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera", ".")
|
||||
RESOLUCAO = config.get("raw_size")
|
||||
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
|
||||
# Pastas base
|
||||
DATASET_BASE = os.path.join("dataset")
|
||||
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
|
||||
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
|
||||
|
||||
# Fallback (modo antigo, sem grupos) – mantido por compat
|
||||
ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images")
|
||||
ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks")
|
||||
ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2")
|
||||
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
|
||||
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
|
||||
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
|
||||
|
||||
# Extensões aceitas
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXTS = (".png", ".jpg", ".jpeg")
|
||||
MSK2_EXTS = (".png", ".jpg", ".jpeg")
|
||||
RAW_EXTS = (".raw",)
|
||||
|
||||
|
||||
def garantir_dir(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
|
||||
# ===============================
|
||||
# Pipelines de augmentations
|
||||
# ===============================
|
||||
|
||||
# 1) Geométricas: aplicam em preview “temporário”, mask, mask2 e RAW.
|
||||
# RAW é passado como 'raw' em additional_targets para receber só geom.
|
||||
train_geo = ReplayCompose([
|
||||
A.HorizontalFlip(p=0.5),
|
||||
A.ShiftScaleRotate(
|
||||
shift_limit=0.01,
|
||||
scale_limit=0.10,
|
||||
rotate_limit=5,
|
||||
border_mode=cv2.BORDER_REFLECT_101,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
p=0.30
|
||||
),
|
||||
], additional_targets={
|
||||
'mask2': 'mask',
|
||||
'raw': 'mask', # trata raw como “mask” para receber apenas geom
|
||||
})
|
||||
|
||||
|
||||
# ===============================
|
||||
# Funções de blur físico no RAW
|
||||
# ===============================
|
||||
|
||||
def motion_blur_kernel(ksize=5, angle=0.0):
|
||||
"""
|
||||
Gera um kernel de motion blur 2D com tamanho ksize e ângulo em graus.
|
||||
Usado tanto para RAW quanto para preview.
|
||||
"""
|
||||
ksize = int(ksize)
|
||||
if ksize < 3:
|
||||
ksize = 3
|
||||
if ksize % 2 == 0:
|
||||
ksize += 1
|
||||
|
||||
kernel = np.zeros((ksize, ksize), dtype=np.float32)
|
||||
# linha horizontal no centro
|
||||
kernel[ksize // 2, :] = 1.0
|
||||
|
||||
# rotaciona a linha para o ângulo desejado
|
||||
center = (ksize / 2.0 - 0.5, ksize / 2.0 - 0.5)
|
||||
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
||||
kernel = cv2.warpAffine(kernel, M, (ksize, ksize))
|
||||
|
||||
s = kernel.sum()
|
||||
if s > 0:
|
||||
kernel /= s
|
||||
return kernel
|
||||
|
||||
|
||||
def apply_coherent_blur_raw(raw, p_motion=0.20, p_gauss=0.20):
|
||||
"""
|
||||
Aplica blur/motion blur de forma COERENTE no RAW, canal por canal.
|
||||
|
||||
- Mesmo kernel em TODOS os canais.
|
||||
- Máscaras não passam aqui.
|
||||
- Se nada for sorteado, retorna o RAW original.
|
||||
"""
|
||||
if raw is None:
|
||||
return None
|
||||
|
||||
r = np.random.rand()
|
||||
if r >= (p_motion + p_gauss):
|
||||
# sem blur nesta amostra
|
||||
return raw
|
||||
|
||||
# escolhe tamanho de kernel
|
||||
k_choices = [3, 5, 7]
|
||||
ksize = int(np.random.choice(k_choices))
|
||||
|
||||
if r < p_motion:
|
||||
# Motion blur com ângulo pequeno (vibração / movimento leve)
|
||||
angle = float(np.random.uniform(-20.0, 20.0))
|
||||
kernel = motion_blur_kernel(ksize=ksize, angle=angle)
|
||||
else:
|
||||
# Gaussian blur (leve desfoque de foco)
|
||||
if ksize % 2 == 0:
|
||||
ksize += 1
|
||||
g1d = cv2.getGaussianKernel(ksize, 0)
|
||||
kernel = g1d @ g1d.T # vira kernel 2D
|
||||
|
||||
# Aplica kernel canal a canal
|
||||
if raw.ndim == 2:
|
||||
raw_blur = cv2.filter2D(raw, ddepth=-1, kernel=kernel)
|
||||
else:
|
||||
raw_blur = np.empty_like(raw)
|
||||
h, w, c = raw.shape
|
||||
for ch in range(c):
|
||||
raw_blur[..., ch] = cv2.filter2D(
|
||||
raw[..., ch],
|
||||
ddepth=-1,
|
||||
kernel=kernel,
|
||||
borderType=cv2.BORDER_REFLECT_101,
|
||||
)
|
||||
|
||||
# Garante que continua em [0,1] para float
|
||||
if np.issubdtype(raw_blur.dtype, np.floating):
|
||||
raw_blur = np.clip(raw_blur, 0.0, 1.0)
|
||||
|
||||
return raw_blur
|
||||
|
||||
|
||||
# ===============================
|
||||
# IO de imagens / RAW
|
||||
# ===============================
|
||||
|
||||
def raw_to_preview_rgb(raw_np: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
raw_np: (H,W,C) float32 em 0..1, C=4 ou 5.
|
||||
|
||||
Layout esperado (após _load_raw):
|
||||
- C=4 -> [R,G,B,IR]
|
||||
- C=5 -> [R,G,B,IR,NDVI] (se um dia tiver)
|
||||
|
||||
Usamos sempre R,G,B para o preview e devolvemos (H,W,3) uint8.
|
||||
"""
|
||||
assert raw_np.ndim == 3, "raw_np deve ser (H,W,C)"
|
||||
|
||||
H, W, C = raw_np.shape
|
||||
if C not in (4, 5):
|
||||
raise RuntimeError(f"Esperado C=4 ou 5, veio {C} em raw_to_preview_rgb")
|
||||
|
||||
# RAW já vem como [R,G,B,IR] do _load_raw
|
||||
r = raw_np[..., 0].astype(np.float32)
|
||||
g = raw_np[..., 1].astype(np.float32)
|
||||
b = raw_np[..., 2].astype(np.float32)
|
||||
|
||||
# Monta RGB em (H,W,3) para trabalhar
|
||||
rgb = np.stack([r, g, b], axis=-1) # (H,W,3)
|
||||
|
||||
# Auto white balance simples (referenciado no G)
|
||||
means = rgb.reshape(-1, 3).mean(axis=0) # (3,)
|
||||
ref = means[1] # canal G
|
||||
gains = ref / (means + 1e-6)
|
||||
gains = np.clip(gains, 0.5, 2.0)
|
||||
rgb_wb = rgb * gains # broadcasting em (3,)
|
||||
|
||||
# Força do WB
|
||||
alpha = 1.0 # teste 0.4..0.8 se quiser misturar com original
|
||||
if alpha >= 0.999:
|
||||
rgb = rgb_wb
|
||||
else:
|
||||
rgb_orig = rgb.copy()
|
||||
rgb = (1 - alpha) * rgb_orig + alpha * rgb_wb
|
||||
|
||||
rgb = np.clip(rgb, 0.0, 1.0)
|
||||
|
||||
# Gamma mais suave
|
||||
rgb = np.power(rgb, 1/1.9, out=rgb)
|
||||
|
||||
# Converte pra uint8 OpenCV-friendly
|
||||
rgb8 = (rgb * 255.0).clip(0, 255).astype(np.uint8) # (H,W,3)
|
||||
|
||||
return rgb8
|
||||
|
||||
|
||||
def load_rgb(path):
|
||||
im = cv2.imread(path, cv2.IMREAD_COLOR)
|
||||
if im is None:
|
||||
raise FileNotFoundError(path)
|
||||
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
||||
|
||||
|
||||
def save_rgb(path, arr_rgb):
|
||||
Image.fromarray(arr_rgb).save(path)
|
||||
|
||||
|
||||
def _load_raw(path: str, raw_hw) -> np.ndarray:
|
||||
ext = os.path.splitext(path)[1].lower()
|
||||
if ext == ".npz":
|
||||
z = np.load(path)
|
||||
key = list(z.keys())[0]
|
||||
arr = z[key]
|
||||
elif ext == ".npy":
|
||||
arr = np.load(path)
|
||||
elif ext == ".raw":
|
||||
if raw_hw is None:
|
||||
raise RuntimeError(
|
||||
f"Para ler .raw preciso de raw_hw=(H,W). "
|
||||
f"Passe raw_hw ao criar o dataset."
|
||||
)
|
||||
H, W = raw_hw
|
||||
size_bytes = os.path.getsize(path)
|
||||
mosa_bytes = H * W # mosaico uint8
|
||||
raw4_bytes = 4 * H * W * 4 # 4 canais * H * W * 4 bytes (float32)
|
||||
if size_bytes == mosa_bytes:
|
||||
# MODO ANTIGO: mosaico uint8
|
||||
arr = np.fromfile(path, dtype=np.uint8)
|
||||
raw2d = arr.reshape(H, W)
|
||||
if (H % 2) != 0 or (W % 2) != 0:
|
||||
raise RuntimeError(f"raw_hw deve ser par em H e W p/ mosaico 2x2 (R,G,IR,B), veio H={H}, W={W}")
|
||||
H2, W2 = H // 2, W // 2
|
||||
r_sub = raw2d[0::2, 0::2]
|
||||
g_sub = raw2d[0::2, 1::2]
|
||||
ir_sub = raw2d[1::2, 0::2]
|
||||
b_sub = raw2d[1::2, 1::2]
|
||||
from PIL import Image
|
||||
def upsample(ch_2d: np.ndarray) -> np.ndarray:
|
||||
im = Image.fromarray(ch_2d) # uint8
|
||||
im = im.resize((W, H), resample=Image.BILINEAR)
|
||||
return np.array(im)
|
||||
r_full = upsample(r_sub)
|
||||
g_full = upsample(g_sub)
|
||||
ir_full = upsample(ir_sub)
|
||||
b_full = upsample(b_sub)
|
||||
# (H, W, 4) uint8 [R,G,IR,B]
|
||||
arr = np.stack([r_full, g_full, ir_full, b_full], axis=-1).astype(np.uint8)
|
||||
# REORGANIZA PARA CONTRATO INTERNO [R,G,B,IR]
|
||||
arr = arr[..., [0, 1, 3, 2]]
|
||||
elif size_bytes == raw4_bytes:
|
||||
# NOVO MODO: RAW4 float32 salvo pelo normalize
|
||||
arr_f32 = np.fromfile(path, dtype=np.float32)
|
||||
raw4 = arr_f32.reshape(4, H, W) # (C,H,W) [R,G,IR,B]
|
||||
arr = np.transpose(raw4, (1, 2, 0)) # (H,W,4) [R,G,IR,B]
|
||||
# REORGANIZA PARA CONTRATO INTERNO: [R,G,B,IR]
|
||||
arr = arr[..., [0, 1, 3, 2]]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Tamanho inesperado em {path}: {size_bytes} bytes "
|
||||
f"(esperado {mosa_bytes} ou {raw4_bytes})"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"Extensão não suportada para RAW: {ext} ({path})")
|
||||
# Até aqui queremos (H,W,4) base
|
||||
if arr.ndim != 3:
|
||||
raise RuntimeError(f"RAW precisa ser 3D, veio {arr.shape} em {path}")
|
||||
if arr.shape[0] == 4 and arr.shape[-1] != 4:
|
||||
arr = np.transpose(arr, (1, 2, 0)) # (H,W,4)
|
||||
if arr.shape[-1] != 4:
|
||||
raise RuntimeError(f"Esperava 4 canais base, veio shape={arr.shape} em {path}")
|
||||
return arr # (H,W,4)
|
||||
|
||||
def _scale_to_float01(raw: np.ndarray, max_value = None) -> np.ndarray:
|
||||
if raw.dtype == np.uint16:
|
||||
mv = float(max_value) if max_value is not None else 65535.0
|
||||
elif raw.dtype == np.uint8:
|
||||
mv = float(max_value) if max_value is not None else 255.0
|
||||
else:
|
||||
# float já vem normalizado ou não; se passar max_value, aplica
|
||||
mv = float(max_value) if max_value is not None else None
|
||||
raw_f = raw.astype(np.float32)
|
||||
if mv is not None and mv > 0:
|
||||
raw_f = raw_f / mv
|
||||
# clip pra não explodir
|
||||
raw_f = np.clip(raw_f, 0.0, 1.0)
|
||||
return raw_f
|
||||
|
||||
|
||||
def save_raw_any(path, arr):
|
||||
arr = np.asarray(arr)
|
||||
|
||||
if arr.ndim == 3 and arr.shape[-1] == 4:
|
||||
# arr está no contrato interno [R,G,B,IR]
|
||||
arr_cam = internal_rgbir_to_camera_raw4(arr) # [R,G,IR,B]
|
||||
arr_cam = np.transpose(arr_cam, (2, 0, 1)) # (4,H,W)
|
||||
arr_cam.astype(np.float32).tofile(path)
|
||||
else:
|
||||
arr.tofile(path)
|
||||
|
||||
|
||||
# ===============================
|
||||
# Helpers de diretório / mapeamento
|
||||
# ===============================
|
||||
|
||||
def list_groups(root):
|
||||
"""
|
||||
Lista grupos válidos para o NOVO fluxo:
|
||||
- group/<g>/raws
|
||||
- group/<g>/masks
|
||||
(previews/images não são mais obrigatórios)
|
||||
"""
|
||||
if not os.path.isdir(root):
|
||||
return []
|
||||
grupos = []
|
||||
for name in sorted(os.listdir(root)):
|
||||
gdir = os.path.join(root, name)
|
||||
if not os.path.isdir(gdir):
|
||||
continue
|
||||
has_masks = os.path.isdir(os.path.join(gdir, "masks"))
|
||||
has_raws = os.path.isdir(os.path.join(gdir, "raws"))
|
||||
if has_masks and has_raws:
|
||||
grupos.append(name)
|
||||
return grupos
|
||||
|
||||
|
||||
def map_by_base_priorizando_png(msk_dir, exts):
|
||||
"""Mapeia arquivos por base (prioriza .png)."""
|
||||
by_base = {}
|
||||
if not os.path.isdir(msk_dir):
|
||||
return by_base
|
||||
for fname in os.listdir(msk_dir):
|
||||
f_lower = fname.lower()
|
||||
if not f_lower.endswith(exts):
|
||||
continue
|
||||
base, ext = os.path.splitext(fname)
|
||||
cand = os.path.join(msk_dir, fname)
|
||||
if base not in by_base:
|
||||
by_base[base] = cand
|
||||
else:
|
||||
cur_ext = os.path.splitext(by_base[base])[1].lower()
|
||||
if cur_ext != ".png" and ext.lower() == ".png":
|
||||
by_base[base] = cand
|
||||
return by_base
|
||||
|
||||
|
||||
def map_raws_by_base(raw_dir):
|
||||
by_base = {}
|
||||
if not os.path.isdir(raw_dir):
|
||||
return by_base
|
||||
for fname in os.listdir(raw_dir):
|
||||
f_lower = fname.lower()
|
||||
if not f_lower.endswith(RAW_EXTS):
|
||||
continue
|
||||
base, _ = os.path.splitext(fname)
|
||||
by_base[base] = os.path.join(raw_dir, fname)
|
||||
return by_base
|
||||
|
||||
|
||||
def ensure_aug_dirs(group_name=None, use_masks2=False, use_raw=False, is_preview_mode=True):
|
||||
"""
|
||||
Cria diretórios de saída para o grupo ou modo antigo.
|
||||
|
||||
is_preview_mode:
|
||||
True -> saídas em 'previews' + 'raws'
|
||||
False -> saídas em 'images' + (sem raws)
|
||||
"""
|
||||
if group_name:
|
||||
base = os.path.join(AUG_GROUP_ROOT, group_name)
|
||||
if is_preview_mode:
|
||||
img_out = os.path.join(base, "previews")
|
||||
else:
|
||||
img_out = os.path.join(base, "images")
|
||||
msk_out = os.path.join(base, "masks")
|
||||
msk2_out = os.path.join(base, "masks2") if use_masks2 else None
|
||||
raw_out = os.path.join(base, "raws") if use_raw else None
|
||||
else:
|
||||
# modo legacy
|
||||
img_out = AUG_OLD_IMG
|
||||
msk_out = AUG_OLD_MSK
|
||||
msk2_out = AUG_OLD_MSK2 if use_masks2 else None
|
||||
raw_out = None
|
||||
|
||||
garantir_dir(img_out)
|
||||
garantir_dir(msk_out)
|
||||
if use_masks2 and msk2_out:
|
||||
garantir_dir(msk2_out)
|
||||
if use_raw and raw_out:
|
||||
garantir_dir(raw_out)
|
||||
return img_out, msk_out, msk2_out, raw_out
|
||||
|
||||
|
||||
# ===============================
|
||||
# Núcleo: augmentação de um sample
|
||||
# ===============================
|
||||
|
||||
def augment_sample(
|
||||
raw_path,
|
||||
msk_path,
|
||||
img_out_dir,
|
||||
msk_out_dir,
|
||||
copies,
|
||||
msk2_path=None,
|
||||
msk2_out_dir=None,
|
||||
raw_out_dir=None,
|
||||
aug_suffix="aug",
|
||||
):
|
||||
"""
|
||||
Faz a augmentação a partir de UM RAW4 + máscara (e opcionalmente máscara2).
|
||||
|
||||
Passos:
|
||||
1) Carrega máscara (para saber HxW).
|
||||
2) Carrega RAW4 float32 usando HxW da máscara.
|
||||
3) Gera preview RGB a partir do RAW4 (para servir de "image" no Albumentations).
|
||||
4) Aplica GEOMETRIA (mesma em preview, mask, mask2 e raw).
|
||||
5) Aplica BLUR/MOTION BLUR no RAW geométrico, canal a canal, com o mesmo kernel.
|
||||
6) Gera preview final a partir do RAW blurado.
|
||||
7) Salva:
|
||||
- preview_n em img_out_dir
|
||||
- mask_n em msk_out_dir
|
||||
- raw_n em raw_out_dir
|
||||
- mask2_n em msk2_out_dir (se houver)
|
||||
"""
|
||||
base_raw, raw_ext = os.path.splitext(os.path.basename(raw_path))
|
||||
base = base_raw
|
||||
|
||||
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
|
||||
if base_msk != base:
|
||||
print(f"[WARN] RAW '{base}' e máscara '{base_msk}' com bases diferentes, mas seguindo mesmo assim.")
|
||||
|
||||
msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
|
||||
|
||||
# 1) Máscaras
|
||||
msk = load_rgb(msk_path)
|
||||
msk2 = load_rgb(msk2_path) if msk2_path else None
|
||||
|
||||
# 2) RAW4
|
||||
raw = _load_raw(raw_path, raw_hw=(RESOLUCAO[1], RESOLUCAO[0])) # (H,W,4) uint8
|
||||
raw = _scale_to_float01(raw) # float32 0..1
|
||||
|
||||
# 3) Preview inicial derivado do RAW (antes de geom/blur, só referência)
|
||||
img = raw_to_preview_rgb(raw)
|
||||
|
||||
gen = 0
|
||||
for i in range(copies):
|
||||
# 4) GEOMETRIA: mesma transformação em img, mask, mask2 e raw
|
||||
if msk2 is not None and msk2_out_dir:
|
||||
aug_geo = train_geo(image=img, mask=msk, mask2=msk2, raw=raw)
|
||||
else:
|
||||
aug_geo = train_geo(image=img, mask=msk, raw=raw)
|
||||
|
||||
msk_g = aug_geo["mask"]
|
||||
raw_g = aug_geo.get("raw", None)
|
||||
msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None
|
||||
|
||||
# 5) BLUR FÍSICO: aplica (ou não) o MESMO blur no RAW geométrico
|
||||
raw_b = apply_coherent_blur_raw(raw_g)
|
||||
|
||||
# 6) Preview final derivado do RAW blurado
|
||||
img_b = raw_to_preview_rgb(raw_b)
|
||||
|
||||
# 7) Salva cópias
|
||||
out_img = os.path.join(img_out_dir, f"{base}_{aug_suffix}_{i:02d}.png")
|
||||
out_msk = os.path.join(msk_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk_ext}")
|
||||
save_rgb(out_img, img_b)
|
||||
save_rgb(out_msk, msk_g)
|
||||
|
||||
if msk2_g is not None and msk2_out_dir:
|
||||
out_msk2 = os.path.join(msk2_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk2_ext}")
|
||||
save_rgb(out_msk2, msk2_g)
|
||||
|
||||
if raw_b is not None and raw_out_dir:
|
||||
out_raw = os.path.join(raw_out_dir, f"{base}_{aug_suffix}_{i:02d}{raw_ext}")
|
||||
save_raw_any(out_raw, raw_b)
|
||||
|
||||
gen += 1
|
||||
|
||||
return gen
|
||||
|
||||
|
||||
# ===============================
|
||||
# Processamento por grupo
|
||||
# ===============================
|
||||
|
||||
def process_group(group_name, copies, limit=None, seed=42, aug_suffix="aug"):
|
||||
"""
|
||||
Processa um grupo único usando APENAS:
|
||||
- group/<g>/raws
|
||||
- group/<g>/masks (+ masks2 opcional)
|
||||
|
||||
Previews de entrada são ignoradas.
|
||||
As de saída são geradas do RAW4 (geom + blur).
|
||||
"""
|
||||
gdir = os.path.join(ORIG_GROUP_ROOT, group_name)
|
||||
|
||||
raw_dir = os.path.join(gdir, "raws")
|
||||
msk_dir = os.path.join(gdir, "masks")
|
||||
msk2_dir = os.path.join(gdir, "masks2")
|
||||
|
||||
if not (os.path.isdir(raw_dir) and os.path.isdir(msk_dir)):
|
||||
print(f"[WARN] Grupo '{group_name}' inválido (precisa de 'raws' e 'masks'). Pulando.")
|
||||
return 0
|
||||
|
||||
use_raw = True
|
||||
use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
|
||||
|
||||
raw_map = map_raws_by_base(raw_dir)
|
||||
items = sorted(raw_map.items())
|
||||
|
||||
if limit is not None and limit > 0:
|
||||
rng = np.random.default_rng(seed)
|
||||
if limit < len(items):
|
||||
idx = rng.choice(len(items), size=limit, replace=False)
|
||||
idx = sorted(idx.tolist())
|
||||
items = [items[i] for i in idx]
|
||||
|
||||
msk_map = map_by_base_priorizando_png(msk_dir, MSK_EXTS)
|
||||
msk2_map = map_by_base_priorizando_png(msk2_dir, MSK2_EXTS) if use_masks2 else {}
|
||||
|
||||
img_out_dir, msk_out_dir, msk2_out_dir, raw_out_dir = ensure_aug_dirs(
|
||||
group_name,
|
||||
use_masks2=use_masks2,
|
||||
use_raw=use_raw,
|
||||
is_preview_mode=True
|
||||
)
|
||||
|
||||
count = 0
|
||||
for base, raw_file in items:
|
||||
msk_file = msk_map.get(base)
|
||||
if not msk_file:
|
||||
print(f"[WARN] [{group_name}] Máscara não encontrada para RAW {base}, pulando.")
|
||||
continue
|
||||
|
||||
msk2_file = msk2_map.get(base) if use_masks2 else None
|
||||
if use_masks2 and not msk2_file:
|
||||
print(f"[WARN] [{group_name}] mask2 não encontrada para {base}, gerando só RAW+preview+mask.")
|
||||
|
||||
try:
|
||||
count += augment_sample(
|
||||
raw_path=raw_file,
|
||||
msk_path=msk_file,
|
||||
img_out_dir=img_out_dir,
|
||||
msk_out_dir=msk_out_dir,
|
||||
copies=copies,
|
||||
msk2_path=msk2_file,
|
||||
msk2_out_dir=msk2_out_dir,
|
||||
raw_out_dir=raw_out_dir,
|
||||
aug_suffix=aug_suffix,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[ERRO] [{group_name}] {base}: {e}")
|
||||
print(f"[OK] Grupo '{group_name}' → {count} amostras geradas.")
|
||||
return count
|
||||
|
||||
|
||||
# ===============================
|
||||
# Fallback legacy (sem RAW)
|
||||
# ===============================
|
||||
|
||||
def process_legacy(copies):
|
||||
"""
|
||||
Fallback: modo sem grupos (original/images e original/masks).
|
||||
Mantido por compat, mas não usa RAW.
|
||||
"""
|
||||
if not (os.path.isdir(ORIG_OLD_IMG) and os.path.isdir(ORIG_OLD_MSK)):
|
||||
print("[WARN] Modo legacy não encontrado. Nada a fazer.")
|
||||
return 0
|
||||
|
||||
imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||||
msk_map = map_by_base_priorizando_png(ORIG_OLD_MSK, MSK_EXTS)
|
||||
use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
|
||||
msk2_map = map_by_base_priorizando_png(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}
|
||||
|
||||
img_out_dir, msk_out_dir, msk2_out_dir, _ = ensure_aug_dirs(
|
||||
group_name=None,
|
||||
use_masks2=use_masks2,
|
||||
use_raw=False,
|
||||
is_preview_mode=False
|
||||
)
|
||||
|
||||
count = 0
|
||||
for img_file in sorted(imgs):
|
||||
base, img_ext = os.path.splitext(img_file)
|
||||
msk_file = msk_map.get(base)
|
||||
if not msk_file:
|
||||
print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
|
||||
continue
|
||||
msk2_file = msk2_map.get(base) if use_masks2 else None
|
||||
if use_masks2 and not msk2_file:
|
||||
print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.")
|
||||
|
||||
try:
|
||||
img_path = os.path.join(ORIG_OLD_IMG, img_file)
|
||||
img = load_rgb(img_path)
|
||||
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_file))
|
||||
msk = load_rgb(msk_file)
|
||||
msk2 = load_rgb(msk2_file) if msk2_file else None
|
||||
|
||||
for i in range(copies):
|
||||
if msk2 is not None and msk2_out_dir:
|
||||
aug_geo = train_geo(image=img, mask=msk, mask2=msk2)
|
||||
else:
|
||||
aug_geo = train_geo(image=img, mask=msk)
|
||||
|
||||
img_g = aug_geo["image"]
|
||||
msk_g = aug_geo["mask"]
|
||||
msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None
|
||||
|
||||
out_img = os.path.join(img_out_dir, f"{base}_aug_{i:02d}{img_ext}")
|
||||
out_msk = os.path.join(msk_out_dir, f"{base_msk}_aug_{i:02d}{msk_ext}")
|
||||
save_rgb(out_img, img_g)
|
||||
save_rgb(out_msk, msk_g)
|
||||
|
||||
if msk2_g is not None and msk2_out_dir:
|
||||
msk2_ext = os.path.splitext(msk2_file)[1]
|
||||
out_msk2 = os.path.join(msk2_out_dir, f"{base_msk}_aug_{i:02d}{msk2_ext}")
|
||||
save_rgb(out_msk2, msk2_g)
|
||||
|
||||
count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERRO] (legacy) {img_file}: {e}")
|
||||
print(f"[OK] Legacy → {count} amostras geradas.")
|
||||
return count
|
||||
|
||||
def internal_rgbir_to_camera_raw4(arr: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Converte do contrato interno (H,W,4) [R,G,B,IR]
|
||||
para o contrato bruto da câmera (H,W,4) [R,G,IR,B].
|
||||
"""
|
||||
if arr is None:
|
||||
return None
|
||||
if arr.ndim != 3 or arr.shape[-1] != 4:
|
||||
raise RuntimeError(f"Esperado array (H,W,4), veio {arr.shape}")
|
||||
return arr[..., [0, 1, 3, 2]]
|
||||
|
||||
|
||||
# ===============================
|
||||
# main
|
||||
# ===============================
|
||||
|
||||
def main(copies=5, groups_csv=None, limit=None, seed=42, suffix="aug"):
|
||||
total = 0
|
||||
if os.path.isdir(ORIG_GROUP_ROOT):
|
||||
grupos = list_groups(ORIG_GROUP_ROOT)
|
||||
if groups_csv:
|
||||
want = {g.strip() for g in groups_csv.split(",") if g.strip()}
|
||||
grupos = [g for g in grupos if g in want]
|
||||
if not grupos:
|
||||
print("[WARN] Nenhum grupo válido encontrado após filtro.")
|
||||
if not grupos:
|
||||
print("[WARN] Nenhum grupo encontrado em original/group. Tentando modo legacy...")
|
||||
total += process_legacy(copies)
|
||||
else:
|
||||
print(f"Grupos encontrados: {', '.join(grupos)}")
|
||||
for g in grupos:
|
||||
total += process_group(g, copies, limit=limit, seed=seed, aug_suffix=suffix)
|
||||
else:
|
||||
total += process_legacy(copies)
|
||||
|
||||
print(f"\nAugmentation completed! Total: {total} amostras geradas.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser(description="Augmentação por grupos usando RAW4 (preview gerado do RAW).")
|
||||
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem (default=5).")
|
||||
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula (ex: chao,erva_cana).")
|
||||
ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de imagens originais do grupo a augmentar.")
|
||||
ap.add_argument("--seed", type=int, default=42, help="Seed para seleção reproduzível quando usar --limit.")
|
||||
ap.add_argument("--suffix", type=str, default="aug", help="Sufixo usado no nome dos arquivos gerados.")
|
||||
args = ap.parse_args()
|
||||
main(
|
||||
copies=args.copies,
|
||||
groups_csv=args.groups,
|
||||
limit=args.limit,
|
||||
seed=args.seed,
|
||||
suffix=args.suffix,
|
||||
)
|
||||
|
|
@ -0,0 +1,325 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Normaliza/redimensiona PREVIEW + RAW + MASK (+MASK2), mantendo ESTRUTURA POR GRUPO.
|
||||
|
||||
Entradas:
|
||||
dataset/original/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
dataset/augmented/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
|
||||
Saídas (por resolução):
|
||||
dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
|
||||
|
||||
Conversão de máscara:
|
||||
- Lê máscara RGB e converte para IDs via utils.converter_mask_rgb_para_ids
|
||||
- ignore_id conforme labelmap (default 255)
|
||||
|
||||
RAW:
|
||||
- Detecta dtype (uint8/uint16) pelo tamanho do arquivo
|
||||
- Carrega como (H,W), redimensiona, salva em .raw
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import cv2
|
||||
import numpy as np
|
||||
from typing import Dict, List, Tuple
|
||||
from gal5000.gal_service import mosaic_to_raw4_resized_buf
|
||||
from raw_segformer_service import _infer_ignore_id
|
||||
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
|
||||
|
||||
# ===== config =====
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
USE_MASKS2 = config["dual_head"]
|
||||
RESOLUCAO = tuple(config["resolucao"]) # [W,H]
|
||||
MODEL_NAME = config["model_name"]
|
||||
CHANNELS = int(config.get("channels", 4))
|
||||
pasta_base = os.path.join("dataset")
|
||||
labelmap_path = os.path.join(pasta_base, "labelmap.txt")
|
||||
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
|
||||
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag)
|
||||
|
||||
RESOLUCOES = {f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1])}
|
||||
FONTES = ["original", "augmented"]
|
||||
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXTS = (".png", ".jpg", ".jpeg")
|
||||
MSK2_EXTS = (".png", ".jpg", ".jpeg")
|
||||
RAW_EXTS = (".raw",)
|
||||
|
||||
# === Acumuladores globais para mean/std dos canais RAW4 ===
|
||||
GLOBAL_SUM = None # soma por canal
|
||||
GLOBAL_SUMSQ = None # soma dos quadrados por canal
|
||||
GLOBAL_PIXELS = 0 # n de pixels por canal (H*W por imagem)
|
||||
|
||||
def garantir_dir(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
def list_groups_raw(root) -> List[str]:
|
||||
"""
|
||||
Lista grupos válidos no modo RAW:
|
||||
tem masks e previews (raws opcional, mas esperado).
|
||||
"""
|
||||
if not os.path.isdir(root):
|
||||
return []
|
||||
grupos = []
|
||||
for name in sorted(os.listdir(root)):
|
||||
gdir = os.path.join(root, name)
|
||||
if not os.path.isdir(gdir):
|
||||
continue
|
||||
if os.path.isdir(os.path.join(gdir, "masks")) and os.path.isdir(os.path.join(gdir, "previews")):
|
||||
grupos.append(name)
|
||||
return grupos
|
||||
|
||||
def map_by_base_priorizando_png(dir_path: str, exts: Tuple[str, ...]) -> Dict[str, str]:
|
||||
by_base = {}
|
||||
if not os.path.isdir(dir_path):
|
||||
return by_base
|
||||
for fname in os.listdir(dir_path):
|
||||
low = fname.lower()
|
||||
if not low.endswith(exts):
|
||||
continue
|
||||
base, ext = os.path.splitext(fname)
|
||||
cand = os.path.join(dir_path, fname)
|
||||
if base not in by_base:
|
||||
by_base[base] = cand
|
||||
else:
|
||||
cur_ext = os.path.splitext(by_base[base])[1].lower()
|
||||
if cur_ext != ".png" and ext.lower() == ".png":
|
||||
by_base[base] = cand
|
||||
return by_base
|
||||
|
||||
def map_raws_by_base(raw_dir: str) -> Dict[str, str]:
|
||||
by_base = {}
|
||||
if not os.path.isdir(raw_dir):
|
||||
return by_base
|
||||
for fname in os.listdir(raw_dir):
|
||||
if fname.lower().endswith(RAW_EXTS):
|
||||
base, _ = os.path.splitext(fname)
|
||||
by_base[base] = os.path.join(raw_dir, fname)
|
||||
return by_base
|
||||
|
||||
def load_raw4_float(path: str, src_hw: Tuple[int,int]) -> np.ndarray:
|
||||
h, w = src_hw
|
||||
npx = h * w
|
||||
fsize = os.path.getsize(path)
|
||||
|
||||
num_floats = fsize // 4
|
||||
if num_floats != 4 * npx:
|
||||
raise RuntimeError(
|
||||
f"RAW {path}: esperado 4 canais float32, "
|
||||
f"mas num_floats={num_floats}, H*W={npx}"
|
||||
)
|
||||
|
||||
data = np.fromfile(path, dtype=np.float32)
|
||||
return data.reshape(4, h, w) # (C,H,W)
|
||||
|
||||
def save_raw(path: str, arr: np.ndarray):
|
||||
np.asarray(arr).tofile(path)
|
||||
|
||||
def normalize_mask_ids(mask_path: str, cor_para_id, ignore_id: int, dim: Tuple[int,int]) -> np.ndarray:
|
||||
msk_bgr = cv2.imread(mask_path, cv2.IMREAD_COLOR)
|
||||
if msk_bgr is None:
|
||||
raise RuntimeError(f"Erro ao ler máscara: {mask_path}")
|
||||
msk_rgb = cv2.cvtColor(msk_bgr, cv2.COLOR_BGR2RGB)
|
||||
ids = converter_mask_rgb_para_ids(msk_rgb, cor_para_id, ignore_id)
|
||||
ids_res = cv2.resize(ids, dim, interpolation=cv2.INTER_NEAREST)
|
||||
return ids_res
|
||||
|
||||
def normalize_mask2(mask2_path: str, dim: Tuple[int,int]) -> np.ndarray:
|
||||
m2 = cv2.imread(mask2_path, cv2.IMREAD_UNCHANGED)
|
||||
if m2 is None:
|
||||
raise RuntimeError(f"Erro ao ler máscara2: {mask2_path}")
|
||||
if len(m2.shape) == 3:
|
||||
m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY)
|
||||
else:
|
||||
m2g = m2
|
||||
_, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY)
|
||||
m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST)
|
||||
return m2res
|
||||
|
||||
def normalize_group_raw(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: int, groups_except: str = "") -> int:
|
||||
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
|
||||
grupos = list_groups_raw(fonte_root)
|
||||
if not grupos:
|
||||
return 0
|
||||
|
||||
not_want = {g.strip() for g in groups_except.split(",") if g.strip()}
|
||||
total = 0
|
||||
|
||||
for nome_res, dim in RESOLUCOES.items():
|
||||
out_root = os.path.join(pasta_base, nome_res, "group")
|
||||
for grupo in grupos:
|
||||
if grupo in not_want:
|
||||
print(f"[WARN] Grupo desconsiderado: {grupo}")
|
||||
continue
|
||||
|
||||
in_prev = os.path.join(fonte_root, grupo, "previews")
|
||||
in_raw = os.path.join(fonte_root, grupo, "raws")
|
||||
in_msk = os.path.join(fonte_root, grupo, "masks")
|
||||
in_msk2 = os.path.join(fonte_root, grupo, "masks2")
|
||||
|
||||
if not (os.path.isdir(in_prev) and os.path.isdir(in_msk)):
|
||||
print(f"[WARN] Grupo inválido (sem previews/masks): {grupo}")
|
||||
continue
|
||||
|
||||
usar_raw = os.path.isdir(in_raw)
|
||||
usar_msk2 = USE_MASKS2 and os.path.isdir(in_msk2)
|
||||
|
||||
out_prev = os.path.join(out_root, grupo, "previews")
|
||||
out_raw = os.path.join(out_root, grupo, "raws") if usar_raw else None
|
||||
out_msk = os.path.join(out_root, grupo, "masks")
|
||||
out_msk2 = os.path.join(out_root, grupo, "masks2") if usar_msk2 else None
|
||||
|
||||
garantir_dir(out_prev)
|
||||
garantir_dir(out_msk)
|
||||
if usar_raw and out_raw:
|
||||
garantir_dir(out_raw)
|
||||
if usar_msk2 and out_msk2:
|
||||
garantir_dir(out_msk2)
|
||||
|
||||
prev_files = [f for f in os.listdir(in_prev) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||||
msk_map = map_by_base_priorizando_png(in_msk, MSK_EXTS)
|
||||
raw_map = map_raws_by_base(in_raw) if usar_raw else {}
|
||||
msk2_map = map_by_base_priorizando_png(in_msk2, MSK2_EXTS) if usar_msk2 else {}
|
||||
|
||||
n = len(prev_files)
|
||||
for i, fname in enumerate(sorted(prev_files), 1):
|
||||
base, ext = os.path.splitext(fname)
|
||||
prev_path = os.path.join(in_prev, fname)
|
||||
msk_path = msk_map.get(base)
|
||||
raw_path = raw_map.get(base) if usar_raw else None
|
||||
msk2_path = msk2_map.get(base) if usar_msk2 else None
|
||||
|
||||
if not msk_path:
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] Sem máscara p/ {fname}, pulando.")
|
||||
continue
|
||||
|
||||
# --- preview ---
|
||||
prev_bgr = cv2.imread(prev_path, cv2.IMREAD_COLOR)
|
||||
if prev_bgr is None:
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] Falha ao ler preview: {prev_path}")
|
||||
continue
|
||||
prev_res = cv2.resize(prev_bgr, dim, interpolation=cv2.INTER_AREA)
|
||||
|
||||
# nomes saída com prefixo (igual o normalize atual)
|
||||
out_name_prev = f"{fonte_nome}_{fname}"
|
||||
out_name_base = os.path.splitext(out_name_prev)[0] # pra raw/masks
|
||||
|
||||
cv2.imwrite(os.path.join(out_prev, out_name_prev), prev_res)
|
||||
|
||||
# --- mask ids ---
|
||||
ids_res = normalize_mask_ids(msk_path, cor_para_id, ignore_id, dim)
|
||||
cv2.imwrite(os.path.join(out_msk, out_name_base + ".png"), ids_res)
|
||||
|
||||
# --- mask2 ---
|
||||
if usar_msk2 and out_msk2:
|
||||
if msk2_path:
|
||||
m2res = normalize_mask2(msk2_path, dim)
|
||||
cv2.imwrite(os.path.join(out_msk2, out_name_base + ".png"), m2res)
|
||||
else:
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname}")
|
||||
|
||||
# --- raw ---
|
||||
if usar_raw and out_raw:
|
||||
if raw_path:
|
||||
src_h, src_w = prev_bgr.shape[:2]
|
||||
|
||||
raw4 = load_raw4_float(raw_path, (src_h, src_w)) # (4, Hsrc, Wsrc)
|
||||
out_w, out_h = dim
|
||||
|
||||
if (src_w, src_h) != (out_w, out_h):
|
||||
# redimensiona cada canal
|
||||
chans = []
|
||||
for k in range(raw4.shape[0]):
|
||||
ch = raw4[k]
|
||||
ch_res = cv2.resize(ch, (out_w, out_h), interpolation=cv2.INTER_AREA)
|
||||
chans.append(ch_res.astype(np.float32))
|
||||
raw4 = np.stack(chans, axis=0) # (4, out_h, out_w)
|
||||
|
||||
# Atualiza acumuladores de stats
|
||||
# raw4: (C,H,W) -> (C,N)
|
||||
c, hh, ww = raw4.shape
|
||||
if GLOBAL_SUM is None:
|
||||
GLOBAL_SUM = np.zeros(c, dtype=np.float64)
|
||||
GLOBAL_SUMSQ = np.zeros(c, dtype=np.float64)
|
||||
|
||||
flat = raw4.reshape(c, -1).astype(np.float64)
|
||||
GLOBAL_SUM += flat.sum(axis=1)
|
||||
GLOBAL_SUMSQ += (flat ** 2).sum(axis=1)
|
||||
GLOBAL_PIXELS += hh * ww # por canal é o mesmo H*W
|
||||
|
||||
# Salva como float32 "linearzão" (4 * H * W floats)
|
||||
save_raw(os.path.join(out_raw, out_name_base + ".raw"), raw4.astype(np.float32))
|
||||
else:
|
||||
print(f"[WARN] [{fonte_nome} | {grupo}] Sem RAW p/ {fname} (seguindo só preview+mask).")
|
||||
|
||||
total += 1
|
||||
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n} → {fname}")
|
||||
|
||||
return total
|
||||
|
||||
def main(args):
|
||||
cor_para_id, _colormap_rgb, _id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||||
ignore_id = _infer_ignore_id(ignore_rgb, default_id=255)
|
||||
|
||||
total_geral = 0
|
||||
|
||||
# ORIGINAL
|
||||
orig_group = os.path.join(pasta_base, "original", "group")
|
||||
if os.path.isdir(orig_group):
|
||||
total_geral += normalize_group_raw(orig_group, "original", cor_para_id, ignore_id, groups_except=args.groups_except)
|
||||
else:
|
||||
print("[WARN] Não achei original/group (modo RAW).")
|
||||
|
||||
# AUGMENTED
|
||||
aug_group = os.path.join(pasta_base, "augmented", "group")
|
||||
if os.path.isdir(aug_group):
|
||||
total_geral += normalize_group_raw(aug_group, "augmented", cor_para_id, ignore_id, groups_except=args.groups_except)
|
||||
else:
|
||||
print("[WARN] Não achei augmented/group (modo RAW).")
|
||||
|
||||
print(f"\n✅ Concluído! Total normalizados: {total_geral}")
|
||||
|
||||
# === calcula mean/std globais e salva em JSON ===
|
||||
global GLOBAL_SUM, GLOBAL_SUMSQ, GLOBAL_PIXELS
|
||||
if GLOBAL_SUM is not None and GLOBAL_PIXELS > 0:
|
||||
# média e variância por canal
|
||||
mean = (GLOBAL_SUM / GLOBAL_PIXELS)
|
||||
var = (GLOBAL_SUMSQ / GLOBAL_PIXELS) - mean**2
|
||||
std = np.sqrt(np.maximum(var, 1e-6))
|
||||
|
||||
# Converte para list pra salvar em JSON
|
||||
mean_list = mean.tolist()
|
||||
std_list = std.tolist()
|
||||
|
||||
# Se quiser, você pode nomear os canais explicitamente
|
||||
# dependendo da convenção do raw4:
|
||||
channel_names = ["R", "G", "IR", "B"]
|
||||
|
||||
stats = {
|
||||
"channels": channel_names[:len(mean_list)],
|
||||
"mean": mean_list,
|
||||
"std": std_list,
|
||||
"pixels_per_channel": int(GLOBAL_PIXELS),
|
||||
}
|
||||
|
||||
garantir_dir(save_path)
|
||||
stats_path = os.path.join(save_path, "norm_stats.json")
|
||||
with open(stats_path, "w", encoding="utf-8") as f:
|
||||
json.dump(stats, f, indent=2, ensure_ascii=False)
|
||||
|
||||
print(f"📁 Stats salvos em: {stats_path}")
|
||||
print(f" mean: {mean_list}")
|
||||
print(f" std : {std_list}")
|
||||
else:
|
||||
print("⚠️ Nenhum RAW processado, não há stats para salvar.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
ap = argparse.ArgumentParser(description="Normalize por grupos (RAW: previews/raws/masks)")
|
||||
ap.add_argument("--groups-except", type=str, default="", help="Grupos para não usar, separados por vírgula.")
|
||||
args = ap.parse_args()
|
||||
main(args)
|
||||
|
|
@ -0,0 +1,451 @@
|
|||
#!/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
|
||||
import shutil
|
||||
import random
|
||||
import argparse
|
||||
|
||||
# ⚙️ Configurações
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
MODELO = config.get("camera")
|
||||
USE_MASKS2 = config.get("dual_head", False)
|
||||
RESOLUCAO = tuple(config.get("resolucao"))
|
||||
|
||||
# Pastas (ajustadas para PREVIEWS/RAWS)
|
||||
pasta_origem = os.path.join("dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
|
||||
pasta_destino = os.path.join("dataset", "split")
|
||||
|
||||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||||
MSK_EXT = ".png" # máscaras normalizadas em PNG
|
||||
RAW_EXT = ".raw"
|
||||
|
||||
# Regex para identificar famílias
|
||||
RE_ORIGINAL_PREFIX = re.compile(r'^original_(.+)$', re.IGNORECASE)
|
||||
RE_AUGMENTED_FAMILY = re.compile(r'^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$', re.IGNORECASE)
|
||||
RE_AUG_SUFFIX = re.compile(r'_aug[a-zA-Z0-9]*_\d+$', re.IGNORECASE)
|
||||
|
||||
|
||||
def garantir(p):
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
|
||||
def lista_grupos(root):
|
||||
"""Lista grupos que têm previews + masks (raws opcional)."""
|
||||
if not os.path.isdir(root):
|
||||
return []
|
||||
out = []
|
||||
for g in sorted(os.listdir(root)):
|
||||
gdir = os.path.join(root, g)
|
||||
if not os.path.isdir(gdir):
|
||||
continue
|
||||
if os.path.isdir(os.path.join(gdir, "previews")) and os.path.isdir(os.path.join(gdir, "masks")):
|
||||
out.append(g)
|
||||
return out
|
||||
|
||||
|
||||
def listar_previews(prev_dir):
|
||||
if not os.path.isdir(prev_dir):
|
||||
return []
|
||||
fs = []
|
||||
for f in os.listdir(prev_dir):
|
||||
ext = os.path.splitext(f.lower())[1]
|
||||
if ext in IMG_EXTS:
|
||||
fs.append(f)
|
||||
return sorted(fs)
|
||||
|
||||
|
||||
def mask_from_image_name(img_name):
|
||||
base, _ = os.path.splitext(img_name)
|
||||
return base + MSK_EXT
|
||||
|
||||
|
||||
def mask2_from_image_name(img_name):
|
||||
base, _ = os.path.splitext(img_name)
|
||||
return base + MSK_EXT # masks2 também normalizadas em PNG
|
||||
|
||||
|
||||
def raw_from_image_name(img_name):
|
||||
base, _ = os.path.splitext(img_name)
|
||||
return base + RAW_EXT
|
||||
|
||||
|
||||
def classify_source_and_family(filename_no_ext):
|
||||
"""
|
||||
Retorna (source, family_key)
|
||||
source ∈ {"original", "augmented", "unknown"}
|
||||
family_key = base associada ao original (sem prefixo/sufixos), ex: "foo_001"
|
||||
"""
|
||||
m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
|
||||
if m:
|
||||
return "original", m.group(1)
|
||||
|
||||
m = RE_AUGMENTED_FAMILY.match(filename_no_ext)
|
||||
if m:
|
||||
return "augmented", m.group(1)
|
||||
|
||||
# legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
|
||||
if RE_AUG_SUFFIX.search(filename_no_ext):
|
||||
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
|
||||
return "augmented", fam
|
||||
|
||||
return "unknown", filename_no_ext
|
||||
|
||||
|
||||
def build_family_index(prev_dir, msk_dir):
|
||||
"""
|
||||
Constroi índice de famílias a partir de previews/masks.
|
||||
Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]}
|
||||
(strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe)
|
||||
"""
|
||||
familias = {}
|
||||
imgs = listar_previews(prev_dir)
|
||||
for img_name in imgs:
|
||||
base_no_ext, ext = os.path.splitext(img_name)
|
||||
mask_name = mask_from_image_name(img_name)
|
||||
if not os.path.exists(os.path.join(msk_dir, mask_name)):
|
||||
continue # garante pareamento preview/mask
|
||||
|
||||
source, fam = classify_source_and_family(base_no_ext)
|
||||
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
|
||||
d["all"].append(img_name)
|
||||
if source == "original":
|
||||
d["original"] = img_name
|
||||
elif source == "augmented":
|
||||
d["augmented"].append(img_name)
|
||||
else:
|
||||
# trata como original desconhecido para não perder dado
|
||||
if d["original"] is None:
|
||||
d["original"] = img_name
|
||||
else:
|
||||
d["augmented"].append(img_name)
|
||||
return familias
|
||||
|
||||
|
||||
def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
||||
"""
|
||||
Mesmo esquema do script original: calcula quantas FAMÍLIAS vão para train/val/test.
|
||||
"""
|
||||
n_train = int(round(n * p_train))
|
||||
n_val = int(round(n * p_val))
|
||||
n_test = n - n_train - n_val
|
||||
|
||||
if n_test < 0:
|
||||
excesso = -n_test
|
||||
take_train = min(excesso, max(0, n_train))
|
||||
n_train -= take_train
|
||||
excesso -= take_train
|
||||
if excesso > 0:
|
||||
take_val = min(excesso, max(0, n_val))
|
||||
n_val -= take_val
|
||||
excesso -= take_val
|
||||
n_test = 0
|
||||
|
||||
min_sum = min_train + min_val + min_test
|
||||
if n >= min_sum:
|
||||
n_train = max(n_train, min_train)
|
||||
n_val = max(n_val, min_val)
|
||||
n_test = max(n_test, min_test)
|
||||
|
||||
total = n_train + n_val + n_test
|
||||
while total > n:
|
||||
if n_test > min_test:
|
||||
n_test -= 1
|
||||
elif n_val > min_val:
|
||||
n_val -= 1
|
||||
elif n_train > min_train:
|
||||
n_train -= 1
|
||||
else:
|
||||
break
|
||||
total = n_train + n_val + n_test
|
||||
while total < n:
|
||||
if n_train - min_train <= n_val - min_val:
|
||||
n_train += 1
|
||||
else:
|
||||
n_val += 1
|
||||
total = n_train + n_val + n_test
|
||||
else:
|
||||
n_train = min(n, max(1, min_train))
|
||||
resto = n - n_train
|
||||
n_val = max(0, min(resto, min_val))
|
||||
n_test = max(0, resto - n_val)
|
||||
|
||||
# ajuste final
|
||||
diff = n - (n_train + n_val + n_test)
|
||||
if diff != 0:
|
||||
if diff > 0:
|
||||
# adiciona em train, depois val
|
||||
take = min(diff, n - n_train)
|
||||
n_train += take
|
||||
diff -= take
|
||||
if diff > 0:
|
||||
n_val += diff
|
||||
else:
|
||||
diff = -diff
|
||||
# tira de test, depois val
|
||||
take = min(diff, n_test)
|
||||
n_test -= take
|
||||
diff -= take
|
||||
if diff > 0:
|
||||
n_val -= diff
|
||||
|
||||
return n_train, n_val, n_test
|
||||
|
||||
|
||||
def copiar(
|
||||
nomes,
|
||||
src_prev_dir,
|
||||
src_msk_dir,
|
||||
dst_prev_dir,
|
||||
dst_msk_dir,
|
||||
src_msk2_dir=None,
|
||||
dst_msk2_dir=None,
|
||||
src_raw_dir=None,
|
||||
dst_raw_dir=None,
|
||||
):
|
||||
"""
|
||||
Copia preview + mask (+ mask2) (+ raw) para o split.
|
||||
"""
|
||||
garantir(dst_prev_dir)
|
||||
garantir(dst_msk_dir)
|
||||
|
||||
use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir))
|
||||
use_raw = bool(src_raw_dir and dst_raw_dir and os.path.isdir(src_raw_dir))
|
||||
|
||||
if use_msk2:
|
||||
garantir(dst_msk2_dir)
|
||||
if use_raw:
|
||||
garantir(dst_raw_dir)
|
||||
|
||||
moved = 0
|
||||
for nome in nomes:
|
||||
mask_name = mask_from_image_name(nome)
|
||||
src_prev = os.path.join(src_prev_dir, nome)
|
||||
src_msk = os.path.join(src_msk_dir, mask_name)
|
||||
|
||||
if not (os.path.exists(src_prev) and os.path.exists(src_msk)):
|
||||
continue
|
||||
|
||||
shutil.copy2(src_prev, os.path.join(dst_prev_dir, nome))
|
||||
shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name))
|
||||
|
||||
if use_msk2:
|
||||
m2_name = mask2_from_image_name(nome)
|
||||
src_m2 = os.path.join(src_msk2_dir, m2_name)
|
||||
if os.path.exists(src_m2):
|
||||
shutil.copy2(src_m2, os.path.join(dst_msk2_dir, m2_name))
|
||||
|
||||
if use_raw:
|
||||
raw_name = raw_from_image_name(nome)
|
||||
src_raw = os.path.join(src_raw_dir, raw_name)
|
||||
if os.path.exists(src_raw):
|
||||
shutil.copy2(src_raw, os.path.join(dst_raw_dir, raw_name))
|
||||
|
||||
moved += 1
|
||||
return moved
|
||||
|
||||
|
||||
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
|
||||
src_prev_dir = os.path.join(pasta_origem, group_name, "previews")
|
||||
src_msk_dir = os.path.join(pasta_origem, group_name, "masks")
|
||||
src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2")
|
||||
src_raw_dir = os.path.join(pasta_origem, group_name, "raws")
|
||||
|
||||
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
|
||||
use_raw = os.path.isdir(src_raw_dir)
|
||||
|
||||
familias = build_family_index(src_prev_dir, src_msk_dir)
|
||||
|
||||
# apenas famílias que têm ORIGINAL para participar de val/test
|
||||
familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
|
||||
total_familias = len(familias_originais)
|
||||
if total_familias == 0:
|
||||
print(f"[{group_name}] 0 famílias com original, pulando.")
|
||||
return {"train": 0, "val": 0, "test": 0, "familias": 0}
|
||||
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(familias_originais)
|
||||
|
||||
n_tr, n_va, n_te = allocate_counts(
|
||||
total_familias, p_train, p_val, p_test,
|
||||
mins["train"], mins["val"], mins["test"]
|
||||
)
|
||||
|
||||
fam_train = set(familias_originais[:n_tr])
|
||||
fam_val = set(familias_originais[n_tr:n_tr+n_va])
|
||||
fam_test = set(familias_originais[n_tr+n_va: n_tr+n_va+n_te])
|
||||
|
||||
# CAP por grupo (apenas no TRAIN)
|
||||
if caps_map and group_name in caps_map:
|
||||
cap = caps_map[group_name]
|
||||
if len(fam_train) > cap:
|
||||
fam_list = list(fam_train)
|
||||
rng.shuffle(fam_list)
|
||||
kept = set(fam_list[:cap])
|
||||
dropped = set(fam_list[cap:])
|
||||
fam_train = kept
|
||||
print(f"[{group_name}] cap-train-families={cap} → mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN")
|
||||
|
||||
nomes_train, nomes_val, nomes_test = [], [], []
|
||||
|
||||
for fam, d in familias.items():
|
||||
if fam in fam_train:
|
||||
if d["original"]:
|
||||
nomes_train.append(d["original"])
|
||||
if d["augmented"]:
|
||||
nomes_train.extend(d["augmented"])
|
||||
elif fam in fam_val:
|
||||
if d["original"]:
|
||||
nomes_val.append(d["original"])
|
||||
elif fam in fam_test:
|
||||
if d["original"]:
|
||||
nomes_test.append(d["original"])
|
||||
else:
|
||||
pass
|
||||
|
||||
# dest dirs
|
||||
dest_train_prev = os.path.join(pasta_destino, "train", "group", group_name, "previews")
|
||||
dest_train_msk = os.path.join(pasta_destino, "train", "group", group_name, "masks")
|
||||
dest_val_prev = os.path.join(pasta_destino, "val", "group", group_name, "previews")
|
||||
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
|
||||
dest_test_prev = os.path.join(pasta_destino, "test", "group", group_name, "previews")
|
||||
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks")
|
||||
|
||||
dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None
|
||||
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
|
||||
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
|
||||
|
||||
dest_train_raw = os.path.join(pasta_destino, "train", "group", group_name, "raws") if use_raw else None
|
||||
dest_val_raw = os.path.join(pasta_destino, "val", "group", group_name, "raws") if use_raw else None
|
||||
dest_test_raw = os.path.join(pasta_destino, "test", "group", group_name, "raws") if use_raw else None
|
||||
|
||||
m_train = copiar(
|
||||
nomes_train,
|
||||
src_prev_dir, src_msk_dir,
|
||||
dest_train_prev, dest_train_msk,
|
||||
src_msk2_dir, dest_train_msk2,
|
||||
src_raw_dir, dest_train_raw
|
||||
)
|
||||
m_val = copiar(
|
||||
nomes_val,
|
||||
src_prev_dir, src_msk_dir,
|
||||
dest_val_prev, dest_val_msk,
|
||||
src_msk2_dir, dest_val_msk2,
|
||||
src_raw_dir, dest_val_raw
|
||||
)
|
||||
m_test = copiar(
|
||||
nomes_test,
|
||||
src_prev_dir, src_msk_dir,
|
||||
dest_test_prev, dest_test_msk,
|
||||
src_msk2_dir, dest_test_msk2,
|
||||
src_raw_dir, dest_test_raw
|
||||
)
|
||||
|
||||
print(f"[{group_name}] famílias={total_familias} → train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}")
|
||||
return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (RAW: previews/raws/masks).")
|
||||
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).")
|
||||
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação (default=0.29).")
|
||||
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste (default=0.01).")
|
||||
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento (default=42).")
|
||||
|
||||
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
|
||||
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
|
||||
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
|
||||
|
||||
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")
|
||||
|
||||
ap.add_argument("--cap-train-families", type=str, default="",
|
||||
help="Mapa 'grupo:cap,...' p/ limitar número de FAMÍLIAS no TRAIN. Ex.: 'chao:350'")
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.resolucao:
|
||||
try:
|
||||
w, h = args.resolucao.lower().split("x")
|
||||
resolucao = (int(w), int(h))
|
||||
except Exception:
|
||||
resolucao = RESOLUCAO
|
||||
else:
|
||||
resolucao = RESOLUCAO
|
||||
|
||||
def parse_cap_map(s):
|
||||
caps = {}
|
||||
if not s:
|
||||
return caps
|
||||
for item in s.split(","):
|
||||
k, v = item.strip().split(":")
|
||||
caps[k.strip()] = int(v)
|
||||
return caps
|
||||
|
||||
caps_map = parse_cap_map(args.cap_train_families)
|
||||
|
||||
global pasta_origem, pasta_destino
|
||||
pasta_origem = os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
|
||||
pasta_destino = os.path.join("dataset", "split")
|
||||
|
||||
soma = args.train + args.val + args.test
|
||||
if soma <= 0:
|
||||
raise ValueError("Soma de proporções deve ser > 0.")
|
||||
p_train = args.train / soma
|
||||
p_val = args.val / soma
|
||||
p_test = args.test / soma
|
||||
|
||||
mins = {
|
||||
"train": max(0, args.min_train),
|
||||
"val": max(0, args.min_val),
|
||||
"test": max(0, args.min_test),
|
||||
}
|
||||
|
||||
garantir(pasta_destino)
|
||||
|
||||
grupos = lista_grupos(pasta_origem)
|
||||
if not grupos:
|
||||
print(f"[WARN] Nenhum grupo encontrado em: {pasta_origem}")
|
||||
return
|
||||
|
||||
random.seed(args.seed)
|
||||
|
||||
total_global = {"train": 0, "val": 0, "test": 0, "familias": 0}
|
||||
print(f"Grupos: {', '.join(grupos)}")
|
||||
print(f"Proporções normalizadas: train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
|
||||
print(f"Mínimos por grupo (famílias): train={mins['train']} val={mins['val']} test={mins['test']}")
|
||||
|
||||
for g in grupos:
|
||||
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map)
|
||||
for k in total_global.keys():
|
||||
total_global[k] += res.get(k, 0)
|
||||
|
||||
print("\nResumo global (imagens copiadas):")
|
||||
print(f" train: {total_global['train']}")
|
||||
print(f" val: {total_global['val']}")
|
||||
print(f" test: {total_global['test']}")
|
||||
print(f" famílias (total): {total_global['familias']}")
|
||||
print("\n✅ Split RAW sem vazamento concluído!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,779 @@
|
|||
#python _8_train_segformer_b3_raw.py --config config.json --epochs 20 --batch 4 --lr 3e-5 --wd 0.01 --num_workers 4 --amp --amp_val --grad_accum 2 --class_weights auto
|
||||
|
||||
|
||||
import os
|
||||
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128")
|
||||
|
||||
import json
|
||||
import time
|
||||
import argparse
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch.amp import autocast, GradScaler
|
||||
import torch.nn.functional as F
|
||||
|
||||
from raw_segformer_service import (RawSegDataset, normalize_raw, patch_segformer_input_channels, build_raw_segformer_model, build_dual_branch_segformer_model)
|
||||
|
||||
|
||||
# ----------------------------
|
||||
# Utils
|
||||
# ----------------------------
|
||||
def set_seed(seed: int = 42):
|
||||
import random
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def update_confusion_matrix(
|
||||
cm: torch.Tensor,
|
||||
preds: torch.Tensor,
|
||||
labels: torch.Tensor,
|
||||
num_classes: int,
|
||||
ignore_index: int = 255,
|
||||
):
|
||||
preds = preds.view(-1)
|
||||
labels = labels.view(-1)
|
||||
|
||||
valid = labels != ignore_index
|
||||
preds = preds[valid]
|
||||
labels = labels[valid]
|
||||
|
||||
idx = labels * num_classes + preds
|
||||
bins = torch.bincount(idx, minlength=num_classes * num_classes)
|
||||
cm += bins.view(num_classes, num_classes)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_iou_from_cm(
|
||||
cm: torch.Tensor, eps: float = 1e-6
|
||||
) -> Tuple[float, List[float]]:
|
||||
cm = cm.float()
|
||||
tp = torch.diag(cm)
|
||||
fp = cm.sum(0) - tp
|
||||
fn = cm.sum(1) - tp
|
||||
denom = tp + fp + fn + eps
|
||||
iou = (tp / denom).cpu().tolist()
|
||||
miou = float(np.mean(iou))
|
||||
return miou, iou
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float:
|
||||
cm = cm.float()
|
||||
return (torch.diag(cm).sum() / (cm.sum() + eps)).item()
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
path: str,
|
||||
model: nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
scaler: Optional[GradScaler],
|
||||
epoch: int,
|
||||
best_miou: float,
|
||||
best_main_iou: float,
|
||||
extra: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
ckpt = {
|
||||
"epoch": epoch,
|
||||
"model": model.state_dict(),
|
||||
"optimizer": optimizer.state_dict(),
|
||||
"best_miou": best_miou,
|
||||
"best_main_iou": best_main_iou,
|
||||
}
|
||||
if scaler is not None:
|
||||
ckpt["scaler"] = scaler.state_dict()
|
||||
if extra:
|
||||
ckpt["extra"] = extra
|
||||
torch.save(ckpt, path)
|
||||
|
||||
|
||||
def load_checkpoint(
|
||||
path: str,
|
||||
model: nn.Module,
|
||||
optimizer: Optional[torch.optim.Optimizer] = None,
|
||||
scaler: Optional[GradScaler] = None,
|
||||
map_location: str = "cpu",
|
||||
) -> Dict[str, Any]:
|
||||
ckpt = torch.load(path, map_location=map_location, weights_only=False)
|
||||
model.load_state_dict(ckpt["model"], strict=True)
|
||||
if optimizer is not None and "optimizer" in ckpt:
|
||||
optimizer.load_state_dict(ckpt["optimizer"])
|
||||
if scaler is not None and "scaler" in ckpt:
|
||||
scaler.load_state_dict(ckpt["scaler"])
|
||||
return ckpt
|
||||
|
||||
|
||||
def default_collate(batch):
|
||||
imgs, masks = [], []
|
||||
for item in batch:
|
||||
img = item["image"] if isinstance(item, dict) else item[0]
|
||||
mask = item["mask"] if isinstance(item, dict) else item[1]
|
||||
|
||||
if isinstance(img, np.ndarray):
|
||||
img = torch.from_numpy(img)
|
||||
if isinstance(mask, np.ndarray):
|
||||
mask = torch.from_numpy(mask)
|
||||
|
||||
imgs.append(img.float())
|
||||
masks.append(mask.long())
|
||||
return torch.stack(imgs, 0), torch.stack(masks, 0)
|
||||
|
||||
|
||||
# ----------------------------
|
||||
# Train/Val loop
|
||||
# ----------------------------
|
||||
def dice_loss(
|
||||
logits: torch.Tensor, # (B, C, H, W)
|
||||
target: torch.Tensor, # (B, H, W) com ids
|
||||
num_classes: int,
|
||||
ignore_index: int = -100,
|
||||
smooth: float = 1.0
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Dice loss multi-classe (macro) com suporte a ignore_index.
|
||||
Retorna um escalar.
|
||||
"""
|
||||
# Softmax -> probabilidades
|
||||
probs = torch.softmax(logits, dim=1) # (B,C,H,W)
|
||||
|
||||
# Máscara de pixels válidos
|
||||
valid = (target != ignore_index) # (B,H,W)
|
||||
|
||||
# Se por algum motivo não tem pixel válido no batch
|
||||
if valid.sum() == 0:
|
||||
return logits.new_tensor(0.0)
|
||||
|
||||
# One-hot do target (clamp só pra evitar index inválido)
|
||||
target_clamped = target.clone()
|
||||
target_clamped[~valid] = 0 # coloca qualquer classe nos ignore, pois vamos mascarar depois
|
||||
|
||||
# GARANTIR tipo inteiro pra one_hot
|
||||
target_clamped = target_clamped.long()
|
||||
|
||||
target_1h = F.one_hot(target_clamped, num_classes=num_classes) # (B,H,W,C)
|
||||
target_1h = target_1h.permute(0, 3, 1, 2).float() # (B,C,H,W)
|
||||
|
||||
# Aplica máscara de validade em probs e target
|
||||
valid_f = valid.unsqueeze(1).float() # (B,1,H,W)
|
||||
probs = probs * valid_f
|
||||
target_1h = target_1h * valid_f
|
||||
|
||||
# Dice por classe (macro)
|
||||
dims = (0, 2, 3) # soma em batch e espaço
|
||||
inter = (probs * target_1h).sum(dims)
|
||||
den = probs.sum(dims) + target_1h.sum(dims)
|
||||
|
||||
dice_per_class = (2.0 * inter + smooth) / (den + smooth) # (C,)
|
||||
loss = 1.0 - dice_per_class.mean()
|
||||
return loss
|
||||
|
||||
def run_one_epoch(
|
||||
model: nn.Module,
|
||||
loader: DataLoader,
|
||||
optimizer: Optional[torch.optim.Optimizer],
|
||||
device: torch.device,
|
||||
num_classes: int,
|
||||
ignore_index: int,
|
||||
criterion: nn.Module,
|
||||
amp: bool,
|
||||
scaler: Optional[GradScaler],
|
||||
train: bool,
|
||||
grad_accum: int = 1,
|
||||
normalizer=None
|
||||
) -> Dict[str, Any]:
|
||||
|
||||
model.train(train)
|
||||
total_loss = 0.0
|
||||
cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)
|
||||
|
||||
t0 = time.time()
|
||||
n_batches = 0
|
||||
|
||||
with torch.set_grad_enabled(train):
|
||||
if train and optimizer is not None:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
for step, (imgs, masks) in enumerate(loader):
|
||||
# imgs: (B,C,H,W), C=4 ou 5
|
||||
imgs = imgs.to(device, non_blocking=True)
|
||||
masks = masks.to(device, non_blocking=True) # (B,H,W)
|
||||
|
||||
if normalizer is not None:
|
||||
imgs = normalizer(imgs)
|
||||
else:
|
||||
imgs = normalize_raw(imgs)
|
||||
|
||||
with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
|
||||
out = model(pixel_values=imgs)
|
||||
logits = out.logits
|
||||
|
||||
if logits.shape[-2:] != masks.shape[-2:]:
|
||||
logits = torch.nn.functional.interpolate(
|
||||
logits,
|
||||
size=masks.shape[-2:],
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
ce = criterion(logits, masks) # CrossEntropyLoss com weight/ignore_index
|
||||
dice = dice_loss(logits=logits, target=masks, num_classes=num_classes, ignore_index=ignore_index, smooth=1.0)
|
||||
loss = 0.7 * ce + 0.3 * dice
|
||||
#loss = criterion(logits, masks)
|
||||
|
||||
if train and grad_accum > 1:
|
||||
loss = loss / grad_accum
|
||||
|
||||
if train and optimizer is not None:
|
||||
if amp and scaler is not None and device.type == "cuda":
|
||||
scaler.scale(loss).backward()
|
||||
if ((step + 1) % grad_accum) == 0:
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
else:
|
||||
loss.backward()
|
||||
if ((step + 1) % grad_accum) == 0:
|
||||
optimizer.step()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
total_loss += float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0)
|
||||
n_batches += 1
|
||||
|
||||
preds = torch.argmax(logits, dim=1)
|
||||
update_confusion_matrix(
|
||||
cm,
|
||||
preds,
|
||||
masks,
|
||||
num_classes=num_classes,
|
||||
ignore_index=ignore_index,
|
||||
)
|
||||
|
||||
dt = time.time() - t0
|
||||
avg_loss = total_loss / max(n_batches, 1)
|
||||
miou, iou_per_class = compute_iou_from_cm(cm)
|
||||
acc = compute_pixel_acc_from_cm(cm)
|
||||
|
||||
return {
|
||||
"loss": avg_loss,
|
||||
"miou": miou,
|
||||
"iou_per_class": iou_per_class,
|
||||
"acc": acc,
|
||||
"time_s": dt,
|
||||
}
|
||||
|
||||
|
||||
def estimate_class_weights(
|
||||
ds: Dataset,
|
||||
num_classes: int,
|
||||
ignore_index: int = 255,
|
||||
max_samples: int = 800,
|
||||
) -> torch.Tensor:
|
||||
n = min(len(ds), max_samples)
|
||||
idxs = np.random.choice(len(ds), size=n, replace=False)
|
||||
counts = np.zeros(num_classes, dtype=np.float64)
|
||||
|
||||
for i in idxs:
|
||||
item = ds[i]
|
||||
mask = item["mask"] if isinstance(item, dict) else item[1]
|
||||
m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
|
||||
|
||||
m = m.reshape(-1)
|
||||
m = m[m != ignore_index]
|
||||
if m.size == 0:
|
||||
continue
|
||||
|
||||
counts += np.bincount(m, minlength=num_classes)[:num_classes]
|
||||
|
||||
freq = counts / (counts.sum() + 1e-12)
|
||||
freq = np.clip(freq, 1e-12, 1.0)
|
||||
weights = 1.0 / np.log(1.02 + freq)
|
||||
weights = weights / weights.mean()
|
||||
return torch.tensor(weights, dtype=torch.float32)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", default="config.json")
|
||||
parser.add_argument("--epochs", type=int, default=120)
|
||||
parser.add_argument("--batch", type=int, default=2)
|
||||
parser.add_argument("--lr", type=float, default=6e-5)
|
||||
parser.add_argument("--wd", type=float, default=0.01)
|
||||
parser.add_argument("--num_workers", type=int, default=4)
|
||||
|
||||
parser.add_argument("--save_every", type=int, default=10)
|
||||
parser.add_argument("--resume", action="store_true")
|
||||
parser.add_argument("--resume_ckpt", type=str, default=None)
|
||||
|
||||
parser.add_argument("--amp", action="store_true")
|
||||
parser.add_argument("--amp_val", action="store_true")
|
||||
parser.add_argument("--grad_accum", type=int, default=1)
|
||||
parser.add_argument("--grad_ckpt", action="store_true")
|
||||
|
||||
parser.add_argument("--ignore_index", type=int, default=255)
|
||||
parser.add_argument("--class_weights", type=str, default="auto")
|
||||
parser.add_argument("--main_class", type=str, default=None)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
|
||||
# RAW specifics
|
||||
parser.add_argument("--raw_max", type=float, default=None, help="Ex: 1023, 4095, 16383, 65535. Se None: auto por dtype.")
|
||||
parser.add_argument("--resize_h", type=int, default=None)
|
||||
parser.add_argument("--resize_w", type=int, default=None)
|
||||
|
||||
parser.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal (ex: norm_stats.json).")
|
||||
parser.add_argument("--fusion_mode", type=str, default=None)
|
||||
|
||||
args = parser.parse_args()
|
||||
set_seed(args.seed)
|
||||
|
||||
with open(args.config, "r") as f:
|
||||
config = json.load(f)
|
||||
|
||||
MODELO = config["camera"]
|
||||
MODEL_NAME = config["model_name"]
|
||||
RESOLUCAO = config["resolucao"]
|
||||
W, H = RESOLUCAO[0], RESOLUCAO[1]
|
||||
MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower()
|
||||
|
||||
# novos parâmetros do config
|
||||
USE_NDVI = bool(config.get("use_ndvi", False))
|
||||
CHANNELS = int(config.get("channels", 4))
|
||||
BACKBONE = config.get("backbone")
|
||||
|
||||
if args.main_class is not None:
|
||||
MAIN_CLASS_NAME = args.main_class.lower()
|
||||
|
||||
FUSION_MODE = config.get("fusion_mode", "stacked")
|
||||
if args.fusion_mode is not None:
|
||||
FUSION_MODE = args.fusion_mode
|
||||
|
||||
if FUSION_MODE == "dual_branch" and CHANNELS not in (4, 5):
|
||||
raise ValueError("dual_branch requer channels=4 ou 5")
|
||||
|
||||
stats_source_tag = config.get("stats_source_tag", "stacked_raw4")
|
||||
experiment_tag = f"{FUSION_MODE}_raw{CHANNELS}"
|
||||
save_path = os.path.join("backup", config["modelo"], MODEL_NAME, experiment_tag)
|
||||
|
||||
#save_path = os.path.join("backup", config["modelo"], MODEL_NAME, f"raw{CHANNELS}")
|
||||
dataset_path = os.path.join("dataset")
|
||||
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Device: {device}")
|
||||
print(f"Config: channels={CHANNELS} use_ndvi={USE_NDVI}")
|
||||
|
||||
# ==========================
|
||||
# Normalizador (fixo ou dinâmico)
|
||||
# ==========================
|
||||
norm_stats = None
|
||||
normalizer = None
|
||||
|
||||
# Caminho padrão: dentro do dataset, nome do arquivo de stats
|
||||
# (ajusta aqui pro nome que você realmente usou: norm_stats.json, por ex.)
|
||||
#norm_stats_path = os.path.join(save_path, "norm_stats.json")
|
||||
norm_stats_path = os.path.join("backup", config["modelo"], MODEL_NAME, stats_source_tag, "norm_stats.json")
|
||||
if args.norm_stats is not None:
|
||||
norm_stats_path = args.norm_stats
|
||||
|
||||
if norm_stats_path is not None and os.path.exists(norm_stats_path):
|
||||
with open(norm_stats_path, "r", encoding="utf-8") as f:
|
||||
norm_stats = json.load(f)
|
||||
|
||||
stats_channels = norm_stats.get("channels", [])
|
||||
stats_mean = norm_stats.get("mean", [])
|
||||
stats_std = norm_stats.get("std", [])
|
||||
|
||||
print(f"[NORM] usando stats fixos de: {norm_stats_path}")
|
||||
print(f"[NORM] channels={stats_channels}")
|
||||
print(f"[NORM] mean={stats_mean}")
|
||||
print(f"[NORM] std ={stats_std}")
|
||||
|
||||
# JSON veio da normalização RAW4: ["R","G","IR","B"] em 0..1
|
||||
# Nosso tensor de treino já vem em contrato interno:
|
||||
# 4 canais: [R,G,B,IR]
|
||||
# 5 canais: [R,G,B,IR,NDVI01]
|
||||
# Então reordenamos os stats para [R,G,B,IR] e,
|
||||
# se tiver NDVI, deixamos o 5º canal sem normalizar (mean=0, std=1).
|
||||
|
||||
# Garante que temos pelo menos R,G,IR,B
|
||||
idx_by_name = {name: i for i, name in enumerate(stats_channels)}
|
||||
required = ["R", "G", "IR", "B"]
|
||||
if not all(ch in idx_by_name for ch in required):
|
||||
print("[NORM] AVISO: norm_stats não contém todos os canais R,G,IR,B. Mantendo normalize_raw dinâmico.")
|
||||
normalizer = None
|
||||
else:
|
||||
# ordem interna bruta do JSON: R,G,IR,B
|
||||
mean_arr = np.array(stats_mean, dtype=np.float32)
|
||||
std_arr = np.array(stats_std, dtype=np.float32)
|
||||
|
||||
# reordena para o contrato interno do modelo: [R,G,B,IR]
|
||||
desired_order = ["R", "G", "B", "IR"]
|
||||
mean_rgbi = []
|
||||
std_rgbi = []
|
||||
for ch_name in desired_order:
|
||||
i = idx_by_name[ch_name]
|
||||
mean_rgbi.append(mean_arr[i])
|
||||
std_rgbi.append(std_arr[i])
|
||||
mean_rgbi = np.array(mean_rgbi, dtype=np.float32)
|
||||
std_rgbi = np.array(std_rgbi, dtype=np.float32)
|
||||
|
||||
if CHANNELS == 3 and not USE_NDVI:
|
||||
mean3 = mean_rgbi[:3]
|
||||
std3 = std_rgbi[:3]
|
||||
|
||||
mean_t = torch.tensor(mean3, dtype=torch.float32, device=device).view(1, 3, 1, 1)
|
||||
std_t = torch.tensor(std3, dtype=torch.float32, device=device).view(1, 3, 1, 1)
|
||||
|
||||
def normalizer(x: torch.Tensor) -> torch.Tensor:
|
||||
return (x - mean_t) / std_t
|
||||
|
||||
print("[NORM] Normalização fixa por canal ativada para [R,G,B].")
|
||||
|
||||
elif CHANNELS == 4 and not USE_NDVI:
|
||||
# [R,G,B,IR]
|
||||
mean_t = torch.tensor(mean_rgbi, dtype=torch.float32, device=device).view(1, 4, 1, 1)
|
||||
std_t = torch.tensor(std_rgbi, dtype=torch.float32, device=device).view(1, 4, 1, 1)
|
||||
|
||||
def normalizer(x: torch.Tensor) -> torch.Tensor:
|
||||
# x: (B,4,H,W) em 0..1
|
||||
return (x - mean_t) / std_t
|
||||
|
||||
print("[NORM] Normalização fixa por canal ativada para [R,G,B,IR].")
|
||||
|
||||
elif CHANNELS == 5 and USE_NDVI:
|
||||
# [R,G,B,IR,NDVI01]
|
||||
# NDVI já está em 0..1 e não queremos mexer, então:
|
||||
# mean_ndvi = 0, std_ndvi = 1 -> x_ndvi sai intacto.
|
||||
mean5 = np.concatenate([mean_rgbi, np.array([0.0], dtype=np.float32)], axis=0)
|
||||
std5 = np.concatenate([std_rgbi, np.array([1.0], dtype=np.float32)], axis=0)
|
||||
|
||||
mean_t = torch.tensor(mean5, dtype=torch.float32, device=device).view(1, 5, 1, 1)
|
||||
std_t = torch.tensor(std5, dtype=torch.float32, device=device).view(1, 5, 1, 1)
|
||||
|
||||
def normalizer(x: torch.Tensor) -> torch.Tensor:
|
||||
# x: (B,5,H,W) em 0..1, com NDVI no último canal
|
||||
return (x - mean_t) / std_t
|
||||
|
||||
print("[NORM] Normalização fixa por canal ativada para [R,G,B,IR] (NDVI permanece inalterado).")
|
||||
|
||||
else:
|
||||
print("[NORM] AVISO: norm_stats JSON é para 4 canais RAW (R,G,IR,B), "
|
||||
f"mas config está channels={CHANNELS}, use_ndvi={USE_NDVI}. "
|
||||
"Mantendo normalize_raw dinâmico.")
|
||||
normalizer = None
|
||||
else:
|
||||
if norm_stats_path:
|
||||
print(f"[NORM] Caminho de norm_stats não encontrado: {norm_stats_path}. Usando normalize_raw dinâmico.")
|
||||
else:
|
||||
print("[NORM] norm_stats não informado. Usando normalize_raw dinâmico.")
|
||||
|
||||
resize_hw = None
|
||||
if args.resize_h is not None and args.resize_w is not None:
|
||||
resize_hw = (args.resize_h, args.resize_w)
|
||||
else:
|
||||
# Se quiser, pode forçar pra RESOLUCAO do config (H,W)
|
||||
# resize_hw = (RESOLUCAO[1], RESOLUCAO[0])
|
||||
pass
|
||||
|
||||
# Datasets (RAW com 4 ou 5 canais)
|
||||
ds_train = RawSegDataset(
|
||||
os.path.join(dataset_path, "split", "train"),
|
||||
labelmap_path=labelmap_path,
|
||||
max_value=args.raw_max,
|
||||
resize_hw=resize_hw,
|
||||
raw_hw=(H, W),
|
||||
use_ndvi=USE_NDVI,
|
||||
channels=CHANNELS,
|
||||
)
|
||||
ds_val = RawSegDataset(
|
||||
os.path.join(dataset_path, "split", "val"),
|
||||
labelmap_path=labelmap_path,
|
||||
max_value=args.raw_max,
|
||||
resize_hw=resize_hw,
|
||||
raw_hw=(H, W),
|
||||
use_ndvi=USE_NDVI,
|
||||
channels=CHANNELS,
|
||||
)
|
||||
|
||||
# classes
|
||||
if ds_train.classes is None:
|
||||
raise RuntimeError("Crie labelmap.txt (um nome por linha) para definirmos num_classes com segurança.")
|
||||
|
||||
num_classes = len(ds_train.classes)
|
||||
class_id_by_name = {
|
||||
name.lower(): cid for cid, name in ds_train.classes.items()
|
||||
}
|
||||
class_name_by_id = {cid: name for cid, name in ds_train.classes.items()}
|
||||
|
||||
main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None)
|
||||
if main_class_id is None:
|
||||
print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.")
|
||||
else:
|
||||
print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}")
|
||||
|
||||
dl_train = DataLoader(
|
||||
ds_train,
|
||||
batch_size=args.batch,
|
||||
shuffle=True,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=True,
|
||||
collate_fn=default_collate,
|
||||
drop_last=True,
|
||||
)
|
||||
dl_val = DataLoader(
|
||||
ds_val,
|
||||
batch_size=1,
|
||||
shuffle=False,
|
||||
num_workers=max(2, args.num_workers // 2),
|
||||
pin_memory=True,
|
||||
collate_fn=default_collate,
|
||||
drop_last=False,
|
||||
)
|
||||
|
||||
if FUSION_MODE == "stacked":
|
||||
model = build_raw_segformer_model(
|
||||
num_classes=num_classes,
|
||||
channels=CHANNELS,
|
||||
backbone=BACKBONE,
|
||||
device=device,
|
||||
ckpt_path=None,
|
||||
strict=False,
|
||||
)
|
||||
patch_segformer_input_channels(model, in_ch=CHANNELS)
|
||||
elif FUSION_MODE == "dual_branch":
|
||||
model = build_dual_branch_segformer_model(
|
||||
num_classes=num_classes,
|
||||
channels=CHANNELS,
|
||||
backbone=BACKBONE,
|
||||
device=device,
|
||||
ckpt_path=None,
|
||||
strict=False,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"fusion_mode inválido: {FUSION_MODE}")
|
||||
|
||||
if args.grad_ckpt:
|
||||
try:
|
||||
model.gradient_checkpointing_enable()
|
||||
print("[OK] gradient checkpointing enabled")
|
||||
except Exception as e:
|
||||
print(f"[WARN] grad_ckpt não suportado ({type(e).__name__}: {e}). Seguindo sem.")
|
||||
|
||||
model.to(device)
|
||||
|
||||
# Loss weights
|
||||
if args.class_weights.lower() == "none":
|
||||
weights = None
|
||||
elif args.class_weights.lower() == "auto":
|
||||
w = estimate_class_weights(ds_train, num_classes=num_classes, ignore_index=args.ignore_index)
|
||||
weights = w.to(device)
|
||||
print("Class weights (auto):", w.cpu().numpy().round(3).tolist())
|
||||
else:
|
||||
parts = [float(x) for x in args.class_weights.split(",")]
|
||||
if len(parts) != num_classes:
|
||||
raise ValueError(f"class_weights manual precisa ter {num_classes} valores, recebeu {len(parts)}.")
|
||||
weights = torch.tensor(parts, dtype=torch.float32, device=device)
|
||||
|
||||
criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=args.ignore_index)
|
||||
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd)
|
||||
|
||||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=6, threshold=1e-4)
|
||||
|
||||
scaler = GradScaler(enabled=args.amp and device.type == "cuda")
|
||||
|
||||
last_ckpt_path = os.path.join(save_path, "last.pt")
|
||||
best_miou_path = os.path.join(save_path, "best_miou.pt")
|
||||
best_main_path = os.path.join(save_path, "best_main.pt")
|
||||
|
||||
start_epoch = 1
|
||||
best_miou = -1.0
|
||||
best_main_iou = -1.0
|
||||
|
||||
resume_path = args.resume_ckpt if args.resume_ckpt is not None else last_ckpt_path
|
||||
|
||||
if args.resume and os.path.exists(resume_path):
|
||||
ckpt = load_checkpoint(
|
||||
resume_path,
|
||||
model,
|
||||
optimizer,
|
||||
scaler=scaler,
|
||||
map_location="cpu",
|
||||
)
|
||||
start_epoch = int(ckpt["epoch"]) + 1
|
||||
best_miou = float(ckpt.get("best_miou", -1.0))
|
||||
best_main_iou = float(ckpt.get("best_main_iou", -1.0))
|
||||
print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f} {resume_path}")
|
||||
|
||||
def pretty_iou(iou_list):
|
||||
return " | ".join(
|
||||
[
|
||||
f"{class_name_by_id.get(cid, cid)}:{v:.3f}"
|
||||
for cid, v in enumerate(iou_list)
|
||||
]
|
||||
)
|
||||
|
||||
# Debug opcional de layout de canais/máscara
|
||||
if False:
|
||||
sample = ds_train[0]
|
||||
img = sample["image"] # tensor (C,H,W)
|
||||
mask = sample["mask"] # tensor (H,W)
|
||||
import cv2
|
||||
|
||||
img_np = img.cpu().numpy() # (C,H,W)
|
||||
mask_np = mask.cpu().numpy()
|
||||
|
||||
use_ndvi_dbg = getattr(ds_train, "use_ndvi", False)
|
||||
channels_dbg = getattr(ds_train, "channels", img_np.shape[0])
|
||||
|
||||
if use_ndvi_dbg:
|
||||
if channels_dbg == 4:
|
||||
# [R,G,B,NDVI]
|
||||
r = (img_np[0] * 255).astype("uint8")
|
||||
g = (img_np[1] * 255).astype("uint8")
|
||||
b = (img_np[2] * 255).astype("uint8")
|
||||
ndvi = (img_np[3] * 255).astype("uint8")
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_r.png"), r)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_g.png"), g)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_b.png"), b)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_ndvi.png"), ndvi)
|
||||
elif channels_dbg == 5:
|
||||
# [R,G,B,IR,NDVI]
|
||||
r = (img_np[0] * 255).astype("uint8")
|
||||
g = (img_np[1] * 255).astype("uint8")
|
||||
b = (img_np[2] * 255).astype("uint8")
|
||||
ir = (img_np[3] * 255).astype("uint8")
|
||||
ndvi = (img_np[4] * 255).astype("uint8")
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_r.png"), r)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_g.png"), g)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_b.png"), b)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_ir.png"), ir)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_ndvi.png"), ndvi)
|
||||
else:
|
||||
# layout clássico [R,G,IR,B] (C deve ser 4)
|
||||
r = (img_np[0] * 255).astype("uint8")
|
||||
g = (img_np[1] * 255).astype("uint8")
|
||||
ir = (img_np[2] * 255).astype("uint8")
|
||||
b = (img_np[3] * 255).astype("uint8")
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_r.png"), r)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_g.png"), g)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_ir.png"), ir)
|
||||
cv2.imwrite(os.path.join(save_path, "dbg_b.png"), b)
|
||||
|
||||
cv2.imwrite(
|
||||
os.path.join(save_path, "dbg_mask.png"), mask_np.astype("uint8")
|
||||
)
|
||||
|
||||
for epoch in range(start_epoch, args.epochs + 1):
|
||||
lr_now = optimizer.param_groups[0]["lr"]
|
||||
print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====")
|
||||
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
tr = run_one_epoch(
|
||||
model=model,
|
||||
loader=dl_train,
|
||||
optimizer=optimizer,
|
||||
device=device,
|
||||
num_classes=num_classes,
|
||||
ignore_index=args.ignore_index,
|
||||
criterion=criterion,
|
||||
amp=args.amp,
|
||||
scaler=scaler,
|
||||
train=True,
|
||||
grad_accum=max(1, args.grad_accum),
|
||||
normalizer=normalizer,
|
||||
)
|
||||
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
va = run_one_epoch(
|
||||
model=model,
|
||||
loader=dl_val,
|
||||
optimizer=None,
|
||||
device=device,
|
||||
num_classes=num_classes,
|
||||
ignore_index=args.ignore_index,
|
||||
criterion=criterion,
|
||||
amp=args.amp_val,
|
||||
scaler=None,
|
||||
train=False,
|
||||
grad_accum=1,
|
||||
normalizer=normalizer,
|
||||
)
|
||||
|
||||
scheduler.step(va["loss"])
|
||||
main_iou = (
|
||||
va["miou"]
|
||||
if main_class_id is None
|
||||
else va["iou_per_class"][main_class_id]
|
||||
)
|
||||
|
||||
print(f"TRAIN: loss={tr['loss']:.4f} acc={tr['acc']:.4f} miou={tr['miou']:.4f} (t={tr['time_s']:.1f}s)")
|
||||
print(f"VAL : loss={va['loss']:.4f} acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={float(main_iou):.4f} (t={va['time_s']:.1f}s)")
|
||||
print("IoU per class:", pretty_iou(va["iou_per_class"]))
|
||||
|
||||
save_checkpoint(
|
||||
last_ckpt_path,
|
||||
model,
|
||||
optimizer,
|
||||
scaler=scaler,
|
||||
epoch=epoch,
|
||||
best_miou=best_miou,
|
||||
best_main_iou=best_main_iou,
|
||||
extra={
|
||||
"val_loss": va["loss"],
|
||||
"val_miou": va["miou"],
|
||||
"val_main_iou": float(main_iou),
|
||||
},
|
||||
)
|
||||
|
||||
if args.save_every > 0 and (epoch % args.save_every == 0):
|
||||
save_checkpoint(
|
||||
os.path.join(save_path, f"epoch_{epoch:04d}.pt"),
|
||||
model,
|
||||
optimizer,
|
||||
scaler=scaler,
|
||||
epoch=epoch,
|
||||
best_miou=best_miou,
|
||||
best_main_iou=best_main_iou,
|
||||
)
|
||||
|
||||
if va["miou"] > best_miou:
|
||||
best_miou = va["miou"]
|
||||
save_checkpoint(
|
||||
best_miou_path,
|
||||
model,
|
||||
optimizer,
|
||||
scaler=scaler,
|
||||
epoch=epoch,
|
||||
best_miou=best_miou,
|
||||
best_main_iou=best_main_iou,
|
||||
)
|
||||
print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}")
|
||||
|
||||
if float(main_iou) > best_main_iou:
|
||||
best_main_iou = float(main_iou)
|
||||
save_checkpoint(
|
||||
best_main_path,
|
||||
model,
|
||||
optimizer,
|
||||
scaler=scaler,
|
||||
epoch=epoch,
|
||||
best_miou=best_miou,
|
||||
best_main_iou=best_main_iou,
|
||||
)
|
||||
print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,535 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Teste/visualização do SegFormer RAW (4 ou 5 canais) treinado no _8_train_segformer_b3_raw.py
|
||||
|
||||
Modos:
|
||||
|
||||
1) Dataset (com GT):
|
||||
- Usa RawSegDataset (raws + masks) a partir de split (train/val/test)
|
||||
- Mostra: [preview RAW (RGB)] | [mask GT] | [overlay predito]
|
||||
|
||||
2) Pasta sem máscara (--test_folder sem masks):
|
||||
- Faz inferência 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()
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
{
|
||||
"camera": "multispec_module",
|
||||
"modelo": "segformer_b1",
|
||||
"model_name": "pulv_mit",
|
||||
"dual_head": false,
|
||||
"main_class_name": "cana",
|
||||
"es_classes": "",
|
||||
"model_to_use": "geral",
|
||||
"raw_size": [1296, 1028],
|
||||
"resolucao": [1024, 800],
|
||||
"roi_inicio": 0.0,
|
||||
"roi_tamanho": 1.0,
|
||||
"shaves": 3,
|
||||
"channels": 5,
|
||||
"use_ndvi": false,
|
||||
"backbone": "nvidia/mit-b1",
|
||||
"fusion_mode": "stacked",
|
||||
"stats_source_tag": "stacked_raw5"
|
||||
}
|
||||
|
|
@ -1,146 +0,0 @@
|
|||
import socket
|
||||
import json
|
||||
import numpy as np
|
||||
import base64
|
||||
import time
|
||||
|
||||
|
||||
class MultiSpectralService:
|
||||
def __init__(self, host="192.168.105.6", port=5000, timeout=5):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout = timeout
|
||||
self.sock = None
|
||||
self.file = None
|
||||
|
||||
def __enter__(self):
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
self.disconnect()
|
||||
|
||||
def connect(self):
|
||||
if self.sock is not None:
|
||||
return
|
||||
|
||||
self.sock = socket.create_connection((self.host, self.port), timeout=self.timeout)
|
||||
self.sock.settimeout(self.timeout)
|
||||
self.file = self.sock.makefile("r", encoding="utf-8")
|
||||
|
||||
def disconnect(self):
|
||||
try:
|
||||
if self.file:
|
||||
self.file.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
try:
|
||||
if self.sock:
|
||||
self.sock.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
self.file = None
|
||||
self.sock = None
|
||||
|
||||
def _send_command(self, payload: dict) -> dict:
|
||||
if self.sock is None:
|
||||
self.connect()
|
||||
|
||||
data = (json.dumps(payload) + "\n").encode("utf-8")
|
||||
self.sock.sendall(data)
|
||||
|
||||
line = self.file.readline()
|
||||
if not line:
|
||||
self.disconnect()
|
||||
raise RuntimeError("Conexão encerrada pelo servidor")
|
||||
|
||||
return json.loads(line.strip())
|
||||
|
||||
def ping(self):
|
||||
return self._send_command({"cmd": "ping"})
|
||||
|
||||
def get_status(self):
|
||||
return self._send_command({"cmd": "get_status"})
|
||||
|
||||
def get_config(self):
|
||||
return self._send_command({"cmd": "get_config"})
|
||||
|
||||
def begin(self):
|
||||
return self._send_command({"cmd": "begin"})
|
||||
|
||||
def stop(self):
|
||||
return self._send_command({"cmd": "stop"})
|
||||
|
||||
def set_fps(self, fps: int):
|
||||
return self._send_command({"cmd": "set_fps", "value": fps})
|
||||
|
||||
def set_jpeg_quality(self, quality: int):
|
||||
return self._send_command({"cmd": "set_jpeg_quality", "value": quality})
|
||||
|
||||
def set_resolution(self, width: int, height: int):
|
||||
return self._send_command({
|
||||
"cmd": "set_resolution",
|
||||
"width": width,
|
||||
"height": height
|
||||
})
|
||||
|
||||
def capture_jpg_disk(self):
|
||||
return self._send_command({"cmd": "capture_jpg_disk"})
|
||||
|
||||
def capture_jpg_bytes(self):
|
||||
return self._send_command({"cmd": "capture_jpg_bytes"})
|
||||
|
||||
def capture_jpg_base64(self):
|
||||
resp = self._send_command({"cmd": "capture_jpg_base64"})
|
||||
if not resp.get("ok"):
|
||||
raise RuntimeError(resp.get("error", "Falha ao capturar JPG"))
|
||||
return base64.b64decode(resp["data"])
|
||||
|
||||
def capture_frame_array(self):
|
||||
t0 = time.perf_counter()
|
||||
|
||||
resp = self._send_command({"cmd": "capture_frame"})
|
||||
if not resp.get("ok"):
|
||||
raise RuntimeError(resp.get("error", "Falha ao capturar frame"))
|
||||
|
||||
raw = base64.b64decode(resp["data"])
|
||||
width = resp["width"]
|
||||
height = resp["height"]
|
||||
channels = resp["channels"]
|
||||
arr = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, channels)
|
||||
|
||||
t1 = time.perf_counter()
|
||||
|
||||
meta = {
|
||||
"frame_type": resp.get("frame_type"),
|
||||
"width": width,
|
||||
"height": height,
|
||||
"channels": channels,
|
||||
"dtype": resp.get("dtype"),
|
||||
"size": resp.get("size"),
|
||||
|
||||
# 👇 TELEMETRIA PI
|
||||
"ts_pi": resp.get("ts_pi"),
|
||||
"dt_trigger": resp.get("dt_trigger"),
|
||||
"dt_settle": resp.get("dt_settle"),
|
||||
"dt_capture": resp.get("dt_capture"),
|
||||
"dt_total_pi": resp.get("dt_total_pi"),
|
||||
|
||||
# 👇 TELEMETRIA PC
|
||||
"dt_total_pc": t1 - t0,
|
||||
}
|
||||
|
||||
return arr, meta
|
||||
|
||||
def start_stream(self, host: str, port: int, fps: float):
|
||||
return self._send_command({
|
||||
"cmd": "start_stream",
|
||||
"host": host,
|
||||
"port": port,
|
||||
"fps": fps
|
||||
})
|
||||
|
||||
def stop_stream(self):
|
||||
return self._send_command({"cmd": "stop_stream"})
|
||||
|
|
@ -1,139 +0,0 @@
|
|||
import json
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import lz4.frame
|
||||
import zstandard as zstd
|
||||
from numcodecs import Blosc
|
||||
import numpy as np
|
||||
|
||||
|
||||
class StreamReceiver:
|
||||
def __init__(self, host="0.0.0.0", port=6001):
|
||||
self.host = host
|
||||
self.port = port
|
||||
|
||||
self._server_sock = None
|
||||
self._client_sock = None
|
||||
self._thread = None
|
||||
self._running = False
|
||||
|
||||
self.last_frame = None
|
||||
self.last_meta = None
|
||||
self.last_receive_ts = None
|
||||
|
||||
self._zstd_d = zstd.ZstdDecompressor()
|
||||
self._codec = Blosc(cname="lz4", clevel=1, shuffle=Blosc.SHUFFLE)
|
||||
|
||||
@property
|
||||
def is_running(self):
|
||||
return self._running
|
||||
|
||||
def start(self):
|
||||
if self._running:
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._thread = threading.Thread(target=self._worker, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def stop(self):
|
||||
self._running = False
|
||||
|
||||
try:
|
||||
if self._client_sock:
|
||||
self._client_sock.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
try:
|
||||
if self._server_sock:
|
||||
self._server_sock.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
self._client_sock = None
|
||||
self._server_sock = None
|
||||
|
||||
def _recv_exact(self, sock: socket.socket, n: int) -> bytes:
|
||||
chunks = []
|
||||
remaining = n
|
||||
|
||||
while remaining > 0:
|
||||
chunk = sock.recv(remaining)
|
||||
if not chunk:
|
||||
raise ConnectionError("Conexão encerrada durante recv")
|
||||
chunks.append(chunk)
|
||||
remaining -= len(chunk)
|
||||
|
||||
return b"".join(chunks)
|
||||
|
||||
def _worker(self):
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
|
||||
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
server.bind((self.host, self.port))
|
||||
server.listen(1)
|
||||
server.settimeout(1.0)
|
||||
|
||||
self._server_sock = server
|
||||
print(f"[INFO] StreamReceiver ouvindo em {self.host}:{self.port}")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
client, addr = server.accept()
|
||||
except socket.timeout:
|
||||
continue
|
||||
|
||||
print(f"[INFO] StreamReceiver conectado por {addr}")
|
||||
self._client_sock = client
|
||||
|
||||
with client:
|
||||
while self._running:
|
||||
header_len = int.from_bytes(self._recv_exact(client, 4), "big")
|
||||
header_bytes = self._recv_exact(client, header_len)
|
||||
header = json.loads(header_bytes.decode("utf-8"))
|
||||
|
||||
payload_len = int.from_bytes(self._recv_exact(client, 4), "big")
|
||||
payload_comp = self._recv_exact(client, payload_len)
|
||||
codec = header["codec"]
|
||||
|
||||
if codec == "lz4":
|
||||
payload = lz4.frame.decompress(payload_comp)
|
||||
elif codec == "zstd":
|
||||
payload = self._zstd_d.decompress(payload_comp)
|
||||
elif codec == "numcodecs":
|
||||
payload = self._codec.decode(payload_comp)
|
||||
else:
|
||||
raise ValueError(f"Codec não suportado: {codec}")
|
||||
|
||||
expected = header["payload_size_raw"]
|
||||
if len(payload) != expected:
|
||||
raise ValueError(
|
||||
f"Tamanho descomprimido inválido: {len(payload)} != {expected}"
|
||||
)
|
||||
|
||||
height = header["height"]
|
||||
width = header["width"]
|
||||
channels = header["channels"]
|
||||
dtype = np.uint8 if header["dtype"] == "uint8" else None
|
||||
|
||||
if dtype is None:
|
||||
raise RuntimeError(f"dtype não suportado: {header['dtype']}")
|
||||
|
||||
frame = np.frombuffer(payload, dtype=dtype).reshape(height, width, channels)
|
||||
|
||||
self.last_frame = frame
|
||||
self.last_meta = header
|
||||
self.last_receive_ts = time.perf_counter()
|
||||
|
||||
print("[INFO] StreamReceiver cliente desconectado")
|
||||
self._client_sock = None
|
||||
|
||||
except Exception as e:
|
||||
print(f"[WARN] StreamReceiver encerrado com erro: {e}")
|
||||
|
||||
finally:
|
||||
self._running = False
|
||||
self._client_sock = None
|
||||
self._server_sock = None
|
||||
|
|
@ -1,19 +0,0 @@
|
|||
import time
|
||||
from multispectral_service import MultiSpectralService
|
||||
|
||||
svc = MultiSpectralService(host="192.168.105.6", port=5000)
|
||||
svc.connect()
|
||||
|
||||
print("BEGIN:", svc.begin())
|
||||
|
||||
tempos_pi = []
|
||||
tempos_pc = []
|
||||
|
||||
for i in range(10):
|
||||
frame, meta = svc.capture_frame_array()
|
||||
tempos_pi.append(meta["dt_total_pi"])
|
||||
tempos_pc.append(meta["dt_total_pc"])
|
||||
print(i, meta["dt_total_pi"], meta["dt_total_pc"])
|
||||
|
||||
print("STOP:", svc.stop())
|
||||
svc.disconnect()
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
import time
|
||||
from multispectral_service import MultiSpectralService
|
||||
|
||||
svc = MultiSpectralService(host="192.168.105.6", port=5000)
|
||||
|
||||
svc.connect()
|
||||
|
||||
print("PING:", svc.ping())
|
||||
print("STATUS:", svc.get_status())
|
||||
print("SET FPS:", svc.set_fps(15))
|
||||
print("SET JPG:", svc.set_jpeg_quality(85))
|
||||
print("SET RES:", svc.set_resolution(1280, 720))
|
||||
print("BEGIN:", svc.begin())
|
||||
jpg = None
|
||||
for i in range(1, 6): # Começa em 1 e vai até 5
|
||||
t0 = time.time()
|
||||
jpg = svc.capture_jpg_base64()
|
||||
status = "OK" if jpg is not None else "Falha"
|
||||
tempo = time.time() - t0
|
||||
print(f"CAPTURE {i} JPG: {status}, Tempo: {tempo:.4f}")
|
||||
print("STOP:", svc.stop())
|
||||
print("CONFIG:", svc.get_config())
|
||||
print("STATUS FINAL:", svc.get_status())
|
||||
|
||||
svc.disconnect()
|
||||
|
||||
if jpg is not None:
|
||||
with open("capture.jpg", "wb") as f:
|
||||
f.write(jpg)
|
||||
|
|
@ -1,55 +0,0 @@
|
|||
import time
|
||||
from multispectral_service import MultiSpectralService
|
||||
from stream_receiver import StreamReceiver
|
||||
|
||||
|
||||
STREAM_PORT = 6001
|
||||
PI_HOST = "192.168.105.6"
|
||||
PC_HOST = "192.168.105.5"
|
||||
|
||||
|
||||
def main():
|
||||
receiver = StreamReceiver(host="0.0.0.0", port=STREAM_PORT)
|
||||
receiver.start()
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
svc = MultiSpectralService(host=PI_HOST, port=5000, timeout=10)
|
||||
svc.connect()
|
||||
|
||||
try:
|
||||
_fps = 15
|
||||
print("SET RES:", svc.set_resolution(640, 480))
|
||||
print("SET FPS:", svc.set_fps(_fps))
|
||||
print("BEGIN:", svc.begin())
|
||||
|
||||
print("START STREAM:", svc.start_stream(PC_HOST, STREAM_PORT, fps=_fps))
|
||||
|
||||
t0 = time.perf_counter()
|
||||
last_frame_id = -1
|
||||
|
||||
while time.perf_counter() - t0 < 5:
|
||||
meta = receiver.last_meta
|
||||
frame = receiver.last_frame
|
||||
|
||||
if meta is not None and meta["frame_id"] != last_frame_id:
|
||||
last_frame_id = meta["frame_id"]
|
||||
print(meta)
|
||||
#print(
|
||||
# f"frame_id={meta['frame_id']} "
|
||||
# f"shape={frame.shape if frame is not None else None} "
|
||||
# f"dt_total_pi={meta.get('dt_total_pi'):.4f}"
|
||||
#)
|
||||
|
||||
time.sleep(0.02)
|
||||
|
||||
print("STOP STREAM:", svc.stop_stream())
|
||||
print("STOP:", svc.stop())
|
||||
|
||||
finally:
|
||||
svc.disconnect()
|
||||
receiver.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue