390 lines
12 KiB
Python
390 lines
12 KiB
Python
|
# -*- 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()
|