#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Normaliza/redimensiona PREVIEW + RAW + MASK (+MASK2), mantendo ESTRUTURA POR GRUPO. Entradas: MODELO/dataset/original/group//{previews,raws,masks,(masks2)} MODELO/dataset/augmented/group//{previews,raws,masks,(masks2)} Saídas (por resolução): MODELO/dataset//group//{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] pasta_base = os.path.join(MODELO, "dataset") labelmap_path = os.path.join(pasta_base, "labelmap.txt") 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_raw_gray_autodtype(path: str, src_hw: Tuple[int,int]) -> np.ndarray: """ Carrega RAW como (H,W) detectando uint8/uint16 pelo tamanho do arquivo. """ h, w = src_hw npx = h * w fsize = os.path.getsize(path) if fsize == npx: dtype = np.uint8 elif fsize == npx * 2: dtype = np.uint16 else: raise RuntimeError( f"Tamanho inesperado para RAW {path}: {fsize} bytes " f"(esperado {npx} (u8) ou {npx*2} (u16) para {w}x{h})." ) data = np.fromfile(path, dtype=dtype) if data.size != npx: raise RuntimeError(f"RAW {path}: size={data.size} != {npx} (HxW)") return data.reshape((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] mosaic = load_raw_gray_autodtype(raw_path, (src_h, src_w)) out_w, out_h = dim # dim = (W, H) raw4 = mosaic_to_raw4_resized_buf(mosaic, out_h, out_w) # float32, shape (4,H,W) em 0..1 # 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), } stats_path = os.path.join(pasta_base, "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)