# -*- 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()