agrobot_base/Python/OAK/datasets/gal5000/_6_normalize.py

#!/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 sys
import cv2
import numpy as np
from typing import Dict, List, Tuple
from gal5000.gal_service import mosaic_to_raw4_resized_buf
from raw_segformer_service import _infer_ignore_id

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids

# ===== config =====
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)