1622 lines
52 KiB
Python
1622 lines
52 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
_5_augmentation_raw_oak.py
|
|
|
|
Augmentation multiespectral para OAK-FCC-3 usando os .bin RAW das câmeras,
|
|
máscaras e metadados do dataset.
|
|
|
|
Objetivo:
|
|
- Ler amostras em dataset/original/group/<grupo>/
|
|
- Carregar RGB/RE/NIR a partir dos .bin RAW10 packed
|
|
- Aplicar augmentation geométrico sincronizado em todos os canais + máscara
|
|
- Aplicar augmentation radiométrico/físico nos RAWs
|
|
- Salvar nova amostra em dataset/augmented/group/<grupo>/
|
|
- Gerar preview RGB novo para inspeção visual
|
|
|
|
Estrutura esperada de entrada:
|
|
|
|
dataset/original/group/<grupo>/
|
|
bins/
|
|
masks/
|
|
metas/
|
|
previews/ # opcional
|
|
|
|
Estrutura de saída:
|
|
|
|
dataset/augmented/group/<grupo>/
|
|
bins/
|
|
masks/
|
|
metas/
|
|
previews/
|
|
|
|
Exemplos:
|
|
|
|
python _5_augmentation_raw_oak.py --copies 3 --clear-dst
|
|
|
|
python _5_augmentation_raw_oak.py ^
|
|
--copies 3 ^
|
|
--groups chao,chao_cana,chao_erva,chao_cana_erva ^
|
|
--clear-dst
|
|
|
|
python _5_augmentation_raw_oak.py ^
|
|
--group-copies chao:1,chao_cana:3,chao_erva:3,chao_cana_erva:5 ^
|
|
--clear-dst
|
|
|
|
python _5_augmentation_raw_oak.py --dry-run
|
|
|
|
Observações importantes:
|
|
- Este script NÃO augmenta preview PNG como fonte principal.
|
|
- O preview é somente consequência visual do RGB RAW augmentado.
|
|
- A geometria é sempre igual em RGB/RE/NIR/mask.
|
|
- Radiometria usa mesma base global + pequenas variações por banda/câmera.
|
|
- O script tenta ser tolerante a nomes de arquivo, mas espera que base da amostra
|
|
esteja preservada nos nomes dos bins/masks/metas.
|
|
"""
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import shutil
|
|
from dataclasses import dataclass, asdict
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple, Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
try:
|
|
from core.raw_processor_core import RawProcessorCore
|
|
_HAS_RAW_PROCESSOR_CORE = True
|
|
except Exception:
|
|
RawProcessorCore = None
|
|
_HAS_RAW_PROCESSOR_CORE = False
|
|
|
|
_CORE_PREVIEW_CACHE = {}
|
|
|
|
|
|
# ============================================================
|
|
# Configuração base
|
|
# ============================================================
|
|
|
|
IMG_EXTS = (".png", ".jpg", ".jpeg")
|
|
META_EXTS = (".json",)
|
|
BIN_EXTS = (".bin",)
|
|
|
|
ROLE_ALIASES = {
|
|
"rgb": ["rgb", "color", "cor", "cam_a", "cama", "CAM_A"],
|
|
"re": ["re", "rededge", "red_edge", "red-edge", "cam_b", "camb", "CAM_B"],
|
|
"nir": ["nir", "ir", "infra", "infrared", "cam_c", "camc", "CAM_C"],
|
|
}
|
|
|
|
CAM_KEYS = ["CAM_A", "CAM_B", "CAM_C", "cam_a", "cam_b", "cam_c", "cam0", "cam1", "cam2"]
|
|
|
|
DEFAULT_BIT_DEPTH = 10
|
|
DEFAULT_BAYER_PATTERN = "BGGR"
|
|
DEFAULT_RAW_FORMAT = "RAW10_PACKED"
|
|
|
|
|
|
# ============================================================
|
|
# Data classes
|
|
# ============================================================
|
|
|
|
@dataclass
|
|
class CameraBin:
|
|
role: str
|
|
cam_key: str
|
|
path: Path
|
|
width: int
|
|
height: int
|
|
bit_depth: int = DEFAULT_BIT_DEPTH
|
|
raw_format: str = DEFAULT_RAW_FORMAT
|
|
bayer_pattern: str = DEFAULT_BAYER_PATTERN
|
|
data: Optional[np.ndarray] = None
|
|
|
|
|
|
@dataclass
|
|
class AugParams:
|
|
seed: int
|
|
|
|
# Geometria, igual para tudo
|
|
flip_h: bool
|
|
shift_x_frac: float
|
|
shift_y_frac: float
|
|
scale: float
|
|
rotate_deg: float
|
|
perspective: bool
|
|
perspective_strength: float
|
|
|
|
# Radiometria base
|
|
exposure_mult: float
|
|
gamma: float
|
|
|
|
# Ganhos por papel espectral
|
|
rgb_channel_gain: Tuple[float, float, float]
|
|
re_gain: float
|
|
nir_gain: float
|
|
|
|
# Efeitos físicos leves
|
|
shadow_enabled: bool
|
|
shadow_strength: float
|
|
shadow_angle_deg: float
|
|
highlight_enabled: bool
|
|
highlight_strength: float
|
|
noise_enabled: bool
|
|
noise_sigma_dn: float
|
|
blur_enabled: bool
|
|
blur_kernel: int
|
|
|
|
|
|
# ============================================================
|
|
# Utilitários gerais
|
|
# ============================================================
|
|
|
|
def read_config_model() -> str:
|
|
"""Tenta ler config.json para manter compatibilidade com seus scripts atuais."""
|
|
cfg = Path("config.json")
|
|
if not cfg.exists():
|
|
return "."
|
|
|
|
try:
|
|
with cfg.open("r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
return str(config.get("camera", "."))
|
|
except Exception:
|
|
return "."
|
|
|
|
|
|
def read_config_value(key: str, default=None):
|
|
cfg = Path("config.json")
|
|
if not cfg.exists():
|
|
return default
|
|
|
|
try:
|
|
with cfg.open("r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
return config.get(key, default)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def ensure_dir(path: Path) -> None:
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def clear_dir(path: Path) -> None:
|
|
if path.exists():
|
|
shutil.rmtree(path)
|
|
ensure_dir(path)
|
|
|
|
|
|
def normalizar_base(stem: str) -> str:
|
|
"""Remove sufixos comuns para casar preview/mask/meta/bin."""
|
|
suffixes = [
|
|
"_rgb", "_RGB", "_Rgb", "_color", "_COLOR",
|
|
"_re", "_RE", "_rededge", "_red_edge", "_nir", "_NIR",
|
|
"_cam_a", "_cam_b", "_cam_c", "_CAM_A", "_CAM_B", "_CAM_C",
|
|
"_cam0", "_cam1", "_cam2", "_CAM0", "_CAM1", "_CAM2",
|
|
"_image", "_img", "_frame", "_preview", "_previews",
|
|
"_mask", "_masks", "_seg", "_SEG", "_segment", "_segmentacao",
|
|
"_meta", "_metadata",
|
|
]
|
|
|
|
out = stem
|
|
changed = True
|
|
while changed:
|
|
changed = False
|
|
for sfx in suffixes:
|
|
if out.endswith(sfx):
|
|
out = out[: -len(sfx)]
|
|
changed = True
|
|
break
|
|
return out
|
|
|
|
|
|
def list_groups(src_root: Path) -> List[str]:
|
|
if not src_root.exists():
|
|
return []
|
|
|
|
groups = []
|
|
for p in sorted(src_root.iterdir()):
|
|
if not p.is_dir():
|
|
continue
|
|
if (p / "bins").is_dir() and (p / "masks").is_dir() and (p / "metas").is_dir():
|
|
groups.append(p.name)
|
|
return groups
|
|
|
|
|
|
def map_files_by_base(folder: Path, exts: Tuple[str, ...]) -> Dict[str, Path]:
|
|
by_base: Dict[str, Path] = {}
|
|
if not folder.is_dir():
|
|
return by_base
|
|
|
|
priority = {
|
|
".json": 0,
|
|
".png": 0,
|
|
".bin": 0,
|
|
".jpg": 1,
|
|
".jpeg": 2,
|
|
}
|
|
|
|
for p in sorted(folder.iterdir()):
|
|
if not p.is_file():
|
|
continue
|
|
ext = p.suffix.lower()
|
|
if ext not in exts:
|
|
continue
|
|
|
|
base = normalizar_base(p.stem)
|
|
if base not in by_base:
|
|
by_base[base] = p
|
|
else:
|
|
cur_ext = by_base[base].suffix.lower()
|
|
if priority.get(ext, 99) < priority.get(cur_ext, 99):
|
|
by_base[base] = p
|
|
|
|
return by_base
|
|
|
|
|
|
def load_json(path: Path) -> Dict[str, Any]:
|
|
with path.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def save_json(path: Path, data: Dict[str, Any]) -> None:
|
|
ensure_dir(path.parent)
|
|
with path.open("w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def load_mask(path: Path) -> np.ndarray:
|
|
im = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
|
|
if im is None:
|
|
raise FileNotFoundError(path)
|
|
|
|
if im.ndim == 2:
|
|
return im
|
|
|
|
if im.shape[2] == 4:
|
|
im = cv2.cvtColor(im, cv2.COLOR_BGRA2RGBA)
|
|
else:
|
|
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
|
return im
|
|
|
|
|
|
def save_mask(path: Path, mask: np.ndarray) -> None:
|
|
ensure_dir(path.parent)
|
|
if mask.ndim == 2:
|
|
Image.fromarray(mask).save(path)
|
|
else:
|
|
Image.fromarray(mask).save(path)
|
|
|
|
|
|
# ============================================================
|
|
# RAW10 pack/unpack
|
|
# ============================================================
|
|
|
|
def unpack_raw10_packed(buf: bytes, width: int, height: int) -> np.ndarray:
|
|
"""
|
|
Desempacota MIPI RAW10 packed no padrão mais comum usado por câmeras/DepthAI:
|
|
|
|
Para cada grupo de 5 bytes:
|
|
byte0 = bits [9:2] do pixel 0
|
|
byte1 = bits [9:2] do pixel 1
|
|
byte2 = bits [9:2] do pixel 2
|
|
byte3 = bits [9:2] do pixel 3
|
|
byte4 = bits [1:0] dos 4 pixels, empilhados em pares de bits
|
|
|
|
Retorna uint16 com valores 0..1023.
|
|
|
|
Importante:
|
|
- A versão anterior tratava byte0..byte3 como bits baixos e byte4 como bits altos.
|
|
Isso embaralha o RAW e gera exatamente aquele padrão de ruído/colorido sem imagem.
|
|
"""
|
|
expected_pixels = width * height
|
|
arr = np.frombuffer(buf, dtype=np.uint8)
|
|
|
|
expected_bytes = (expected_pixels * 10 + 7) // 8
|
|
if arr.size < expected_bytes:
|
|
raise ValueError(
|
|
f"RAW10 menor que esperado: bytes={arr.size}, esperado>={expected_bytes}, "
|
|
f"shape={width}x{height}"
|
|
)
|
|
|
|
arr = arr[:expected_bytes]
|
|
groups = arr.size // 5
|
|
main = arr[: groups * 5].reshape(-1, 5).astype(np.uint16)
|
|
|
|
p0 = (main[:, 0] << 2) | ((main[:, 4] >> 0) & 0x03)
|
|
p1 = (main[:, 1] << 2) | ((main[:, 4] >> 2) & 0x03)
|
|
p2 = (main[:, 2] << 2) | ((main[:, 4] >> 4) & 0x03)
|
|
p3 = (main[:, 3] << 2) | ((main[:, 4] >> 6) & 0x03)
|
|
|
|
out = np.empty(groups * 4, dtype=np.uint16)
|
|
out[0::4] = p0
|
|
out[1::4] = p1
|
|
out[2::4] = p2
|
|
out[3::4] = p3
|
|
|
|
if out.size < expected_pixels:
|
|
raise ValueError(f"RAW10 gerou pixels insuficientes: {out.size} < {expected_pixels}")
|
|
|
|
return out[:expected_pixels].reshape(height, width)
|
|
|
|
|
|
def pack_raw10_packed(raw: np.ndarray) -> bytes:
|
|
"""
|
|
Empacota uint16 0..1023 para MIPI RAW10 packed no mesmo padrão do unpack:
|
|
|
|
byte0..byte3 = bits altos [9:2]
|
|
byte4 = bits baixos [1:0] de p0,p1,p2,p3
|
|
"""
|
|
flat = np.asarray(raw, dtype=np.uint16).reshape(-1)
|
|
flat = np.clip(flat, 0, 1023).astype(np.uint16)
|
|
|
|
pad = (-flat.size) % 4
|
|
if pad:
|
|
flat = np.pad(flat, (0, pad), mode="edge")
|
|
|
|
p = flat.reshape(-1, 4)
|
|
out = np.empty((p.shape[0], 5), dtype=np.uint8)
|
|
|
|
out[:, 0] = ((p[:, 0] >> 2) & 0xFF).astype(np.uint8)
|
|
out[:, 1] = ((p[:, 1] >> 2) & 0xFF).astype(np.uint8)
|
|
out[:, 2] = ((p[:, 2] >> 2) & 0xFF).astype(np.uint8)
|
|
out[:, 3] = ((p[:, 3] >> 2) & 0xFF).astype(np.uint8)
|
|
out[:, 4] = (
|
|
((p[:, 0] & 0x03) << 0)
|
|
| ((p[:, 1] & 0x03) << 2)
|
|
| ((p[:, 2] & 0x03) << 4)
|
|
| ((p[:, 3] & 0x03) << 6)
|
|
).astype(np.uint8)
|
|
|
|
return out.reshape(-1).tobytes()
|
|
|
|
|
|
def pack_raw10_packed_array(raw: np.ndarray) -> np.ndarray:
|
|
"""
|
|
RAW16 0..1023 -> ndarray uint8 RAW10 packed shape=(height, packed_width).
|
|
Esse é o formato esperado pelo RawProcessorCore.decode_stream_cameras.
|
|
"""
|
|
h, w = raw.shape[:2]
|
|
packed = np.frombuffer(pack_raw10_packed(raw), dtype=np.uint8)
|
|
packed_w = int(math.ceil(int(w) * 10 / 8))
|
|
return packed.reshape(int(h), packed_w)
|
|
|
|
|
|
def read_raw_bin(path: Path, width: int, height: int, raw_format: str) -> np.ndarray:
|
|
raw_format_u = (raw_format or "").upper()
|
|
buf = path.read_bytes()
|
|
|
|
if "RAW10" in raw_format_u or "PACKED" in raw_format_u:
|
|
return unpack_raw10_packed(buf, width, height)
|
|
|
|
# Fallback: RAW16 little-endian com valores possivelmente 0..1023.
|
|
arr = np.frombuffer(buf, dtype=np.uint16)
|
|
expected = width * height
|
|
if arr.size < expected:
|
|
raise ValueError(f"RAW16 menor que esperado em {path}: {arr.size} < {expected}")
|
|
return arr[:expected].reshape(height, width)
|
|
|
|
|
|
def write_raw_bin(path: Path, raw: np.ndarray, raw_format: str) -> None:
|
|
ensure_dir(path.parent)
|
|
raw_format_u = (raw_format or "").upper()
|
|
|
|
raw10 = np.clip(np.rint(raw), 0, 1023).astype(np.uint16)
|
|
|
|
if "RAW10" in raw_format_u or "PACKED" in raw_format_u:
|
|
path.write_bytes(pack_raw10_packed(raw10))
|
|
else:
|
|
path.write_bytes(raw10.astype(np.uint16).tobytes())
|
|
|
|
|
|
# ============================================================
|
|
# Metadados e descoberta de bins
|
|
# ============================================================
|
|
|
|
def _text_has_any(text: str, needles: List[str]) -> bool:
|
|
t = text.lower()
|
|
return any(n.lower() in t for n in needles)
|
|
|
|
|
|
def role_from_text(text: str) -> Optional[str]:
|
|
for role, aliases in ROLE_ALIASES.items():
|
|
if _text_has_any(text, aliases):
|
|
return role
|
|
return None
|
|
|
|
|
|
def normalize_role(role: Optional[str], cam_key: Optional[str] = None) -> Optional[str]:
|
|
if role:
|
|
r = str(role).lower().strip()
|
|
if r in ("rgb", "color", "cor"):
|
|
return "rgb"
|
|
if r in ("re", "rededge", "red_edge", "red-edge"):
|
|
return "re"
|
|
if r in ("nir", "ir", "infrared", "infra"):
|
|
return "nir"
|
|
|
|
if cam_key:
|
|
ck = str(cam_key).upper()
|
|
# Convenção atual mais provável: CAM_A=RGB, CAM_B=RE, CAM_C=NIR.
|
|
if ck == "CAM_A":
|
|
return "rgb"
|
|
if ck == "CAM_B":
|
|
return "re"
|
|
if ck == "CAM_C":
|
|
return "nir"
|
|
|
|
return None
|
|
|
|
|
|
def extract_camera_info(meta: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Tenta extrair camera_info de vários formatos possíveis.
|
|
Retorna dict cam_key -> info.
|
|
"""
|
|
candidates = []
|
|
|
|
for key in ["camera_info", "cameras", "camera_infos", "payload_sources_info"]:
|
|
if isinstance(meta.get(key), dict):
|
|
candidates.append(meta[key])
|
|
|
|
# Algumas versões guardam dentro de meta["capture"] ou meta["meta"]
|
|
for parent_key in ["capture", "meta", "frame_meta"]:
|
|
parent = meta.get(parent_key)
|
|
if isinstance(parent, dict):
|
|
for key in ["camera_info", "cameras", "camera_infos"]:
|
|
if isinstance(parent.get(key), dict):
|
|
candidates.append(parent[key])
|
|
|
|
if candidates:
|
|
cam_info = candidates[0]
|
|
out = {}
|
|
for k, v in cam_info.items():
|
|
if isinstance(v, dict):
|
|
out[str(k)] = v
|
|
return out
|
|
|
|
# Fallback mínimo caso o meta já tenha campos diretos.
|
|
out = {}
|
|
for ck in CAM_KEYS:
|
|
if isinstance(meta.get(ck), dict):
|
|
out[ck] = meta[ck]
|
|
return out
|
|
|
|
|
|
def get_int_any(d: Dict[str, Any], keys: List[str], default: Optional[int] = None) -> Optional[int]:
|
|
for k in keys:
|
|
if k in d and d[k] is not None:
|
|
try:
|
|
return int(d[k])
|
|
except Exception:
|
|
pass
|
|
return default
|
|
|
|
|
|
def get_str_any(d: Dict[str, Any], keys: List[str], default: str = "") -> str:
|
|
for k in keys:
|
|
if k in d and d[k] is not None:
|
|
return str(d[k])
|
|
return default
|
|
|
|
|
|
def find_bin_for_camera(bins_dir: Path, base: str, role: str, cam_key: str) -> Optional[Path]:
|
|
"""
|
|
Localiza o bin de uma câmera por base + role/cam_key.
|
|
É tolerante a nomes tipo:
|
|
base_CAM_A.bin
|
|
base_rgb.bin
|
|
base_cam0.bin
|
|
base_arquivo_CAM_A.bin
|
|
"""
|
|
if not bins_dir.is_dir():
|
|
return None
|
|
|
|
role_aliases = ROLE_ALIASES.get(role, [role])
|
|
cam_key_variants = {cam_key, cam_key.upper(), cam_key.lower()}
|
|
|
|
candidates = []
|
|
for p in sorted(bins_dir.glob("*.bin")):
|
|
name = p.stem
|
|
name_l = name.lower()
|
|
base_l = base.lower()
|
|
if not name_l.startswith(base_l):
|
|
continue
|
|
|
|
score = 0
|
|
if name_l == base_l:
|
|
score += 1
|
|
if any(v.lower() in name_l for v in cam_key_variants):
|
|
score += 10
|
|
if any(a.lower() in name_l for a in role_aliases):
|
|
score += 12
|
|
|
|
# Compatibilidade com cam0/cam1/cam2 caso necessário.
|
|
if role == "rgb" and "cam0" in name_l:
|
|
score += 3
|
|
if role == "re" and "cam1" in name_l:
|
|
score += 3
|
|
if role == "nir" and "cam2" in name_l:
|
|
score += 3
|
|
|
|
if score > 0:
|
|
candidates.append((score, p))
|
|
|
|
if not candidates:
|
|
return None
|
|
|
|
candidates.sort(key=lambda x: x[0], reverse=True)
|
|
return candidates[0][1]
|
|
|
|
|
|
def discover_sample_cameras(meta: Dict[str, Any], bins_dir: Path, base: str) -> Dict[str, CameraBin]:
|
|
cam_info = extract_camera_info(meta)
|
|
out: Dict[str, CameraBin] = {}
|
|
|
|
# Primeiro caminho: usar camera_info do meta.
|
|
for cam_key, info in cam_info.items():
|
|
role = normalize_role(info.get("role") or info.get("papel") or info.get("type"), cam_key)
|
|
if role not in ("rgb", "re", "nir"):
|
|
role = role_from_text(cam_key) or role_from_text(json.dumps(info, ensure_ascii=False))
|
|
|
|
if role not in ("rgb", "re", "nir"):
|
|
continue
|
|
|
|
width = get_int_any(info, ["width", "w", "sensor_width", "cols", "shape_w"])
|
|
height = get_int_any(info, ["height", "h", "sensor_height", "rows", "shape_h"])
|
|
|
|
if width is None or height is None:
|
|
# Fallback para campos globais do meta.
|
|
width = get_int_any(meta, ["width", "sensor_width", "raw_width", "w"])
|
|
height = get_int_any(meta, ["height", "sensor_height", "raw_height", "h"])
|
|
|
|
if width is None or height is None:
|
|
continue
|
|
|
|
bit_depth = get_int_any(info, ["bit_depth", "bits", "raw_bit_depth"], DEFAULT_BIT_DEPTH)
|
|
raw_format = get_str_any(info, ["raw_format", "format", "encoding"], DEFAULT_RAW_FORMAT)
|
|
bayer_pattern = get_str_any(info, ["bayer_pattern", "bayer", "pattern"], DEFAULT_BAYER_PATTERN)
|
|
|
|
bin_path = find_bin_for_camera(bins_dir, base, role, cam_key)
|
|
if bin_path is None:
|
|
continue
|
|
|
|
out[role] = CameraBin(
|
|
role=role,
|
|
cam_key=str(cam_key),
|
|
path=bin_path,
|
|
width=int(width),
|
|
height=int(height),
|
|
bit_depth=int(bit_depth or DEFAULT_BIT_DEPTH),
|
|
raw_format=raw_format or DEFAULT_RAW_FORMAT,
|
|
bayer_pattern=bayer_pattern or DEFAULT_BAYER_PATTERN,
|
|
)
|
|
|
|
# Segundo caminho: se o meta não ajudou, tenta descobrir por nome.
|
|
if len(out) < 3:
|
|
global_w = get_int_any(meta, ["width", "sensor_width", "raw_width", "w"])
|
|
global_h = get_int_any(meta, ["height", "sensor_height", "raw_height", "h"])
|
|
|
|
for role, cam_key in [("rgb", "CAM_A"), ("re", "CAM_B"), ("nir", "CAM_C")]:
|
|
if role in out:
|
|
continue
|
|
if global_w is None or global_h is None:
|
|
continue
|
|
bin_path = find_bin_for_camera(bins_dir, base, role, cam_key)
|
|
if bin_path is None:
|
|
continue
|
|
out[role] = CameraBin(
|
|
role=role,
|
|
cam_key=cam_key,
|
|
path=bin_path,
|
|
width=int(global_w),
|
|
height=int(global_h),
|
|
bit_depth=DEFAULT_BIT_DEPTH,
|
|
raw_format=DEFAULT_RAW_FORMAT,
|
|
bayer_pattern=DEFAULT_BAYER_PATTERN,
|
|
)
|
|
|
|
return out
|
|
|
|
|
|
# ============================================================
|
|
# Preview RGB simples a partir de RAW Bayer
|
|
# ============================================================
|
|
|
|
def _normalize_preview_channel(x: np.ndarray, lo_p: float = 1.0, hi_p: float = 99.5, gamma: float = 1.0 / 2.2) -> np.ndarray:
|
|
"""
|
|
Normalização visual robusta para preview, não científica.
|
|
Transforma um canal RAW/float em uint8 bonito para inspeção humana.
|
|
"""
|
|
y = x.astype(np.float32)
|
|
lo = float(np.percentile(y, lo_p))
|
|
hi = float(np.percentile(y, hi_p))
|
|
if hi <= lo:
|
|
hi = lo + 1.0
|
|
y = np.clip((y - lo) / (hi - lo), 0.0, 1.0)
|
|
if gamma and abs(gamma - 1.0) > 1e-6:
|
|
y = np.power(y, gamma)
|
|
return np.clip(y * 255.0, 0, 255).astype(np.uint8)
|
|
|
|
|
|
def _get_bayer_cv2_code_like_core(bayer_pattern: str | None, algorithm: str = "ea") -> int:
|
|
"""
|
|
Replica o mapeamento usado pelo RawProcessorCore._get_bayer_cv2_code.
|
|
|
|
Observação importante:
|
|
O mapeamento OpenCV parece invertido à primeira vista, mas é exatamente
|
|
o contrato usado no core atual para gerar o RGB do treinamento/inferência.
|
|
"""
|
|
p = str(bayer_pattern or DEFAULT_BAYER_PATTERN or "RGGB").upper()
|
|
algo = str(algorithm or "ea").lower()
|
|
|
|
if algo in ("bilinear", "linear", "fast", "normal"):
|
|
code_map = {
|
|
"BGGR": cv2.COLOR_BayerRG2RGB,
|
|
"RGGB": cv2.COLOR_BayerBG2RGB,
|
|
"GRBG": cv2.COLOR_BayerGR2RGB,
|
|
"GBRG": cv2.COLOR_BayerGB2RGB,
|
|
}
|
|
elif algo in ("ea", "edge_aware", "edge-aware"):
|
|
code_map = {
|
|
"BGGR": cv2.COLOR_BayerRG2RGB_EA,
|
|
"RGGB": cv2.COLOR_BayerBG2RGB_EA,
|
|
"GRBG": cv2.COLOR_BayerGR2RGB_EA,
|
|
"GBRG": cv2.COLOR_BayerGB2RGB_EA,
|
|
}
|
|
else:
|
|
raise ValueError(f"demosaic_algorithm inválido: {algorithm}")
|
|
|
|
if p not in code_map:
|
|
raise ValueError(f"Padrão Bayer não suportado para demosaic: {p}")
|
|
|
|
return code_map[p]
|
|
|
|
|
|
def demosaic_raw10_to_rgb01(raw: np.ndarray, bayer_pattern: str = DEFAULT_BAYER_PATTERN, algorithm: str = "ea") -> np.ndarray:
|
|
"""
|
|
RAW Bayer 0..1023 -> RGB HWC float32 0..1 usando o mesmo mapeamento do core.
|
|
"""
|
|
raw_u16 = np.clip(raw, 0, 1023).astype(np.uint16)
|
|
cv_code = _get_bayer_cv2_code_like_core(bayer_pattern, algorithm=algorithm)
|
|
rgb16 = cv2.cvtColor(raw_u16, cv_code)
|
|
rgb = rgb16.astype(np.float32) / 1023.0
|
|
return np.clip(rgb, 0.0, 1.0).astype(np.float32, copy=False)
|
|
|
|
|
|
def remosaic_rgb01_to_bayer_raw10(rgb: np.ndarray, bayer_pattern: str = DEFAULT_BAYER_PATTERN) -> np.ndarray:
|
|
"""
|
|
RGB HWC float32 0..1 -> mosaico Bayer RAW10 uint16.
|
|
|
|
Isso é necessário porque NÃO podemos aplicar warp diretamente no mosaico Bayer.
|
|
O fluxo correto para augmentar RGB RAW é:
|
|
raw Bayer -> demosaic RGB -> augmentation -> remosaic Bayer -> pack RAW10.
|
|
|
|
Não é uma reconstrução óptica perfeita, mas preserva o contrato RAW10 Bayer para
|
|
o normalize/RawProcessorCore e evita o xadrez colorido causado por interpolar
|
|
diretamente o mosaico.
|
|
"""
|
|
rgb = np.clip(rgb.astype(np.float32), 0.0, 1.0)
|
|
h, w = rgb.shape[:2]
|
|
raw = np.empty((h, w), dtype=np.float32)
|
|
|
|
p = str(bayer_pattern or DEFAULT_BAYER_PATTERN or "RGGB").upper()
|
|
|
|
if p == "RGGB":
|
|
raw[0::2, 0::2] = rgb[0::2, 0::2, 0] # R
|
|
raw[0::2, 1::2] = rgb[0::2, 1::2, 1] # G
|
|
raw[1::2, 0::2] = rgb[1::2, 0::2, 1] # G
|
|
raw[1::2, 1::2] = rgb[1::2, 1::2, 2] # B
|
|
elif p == "BGGR":
|
|
raw[0::2, 0::2] = rgb[0::2, 0::2, 2] # B
|
|
raw[0::2, 1::2] = rgb[0::2, 1::2, 1] # G
|
|
raw[1::2, 0::2] = rgb[1::2, 0::2, 1] # G
|
|
raw[1::2, 1::2] = rgb[1::2, 1::2, 0] # R
|
|
elif p == "GRBG":
|
|
raw[0::2, 0::2] = rgb[0::2, 0::2, 1] # G
|
|
raw[0::2, 1::2] = rgb[0::2, 1::2, 0] # R
|
|
raw[1::2, 0::2] = rgb[1::2, 0::2, 2] # B
|
|
raw[1::2, 1::2] = rgb[1::2, 1::2, 1] # G
|
|
elif p == "GBRG":
|
|
raw[0::2, 0::2] = rgb[0::2, 0::2, 1] # G
|
|
raw[0::2, 1::2] = rgb[0::2, 1::2, 2] # B
|
|
raw[1::2, 0::2] = rgb[1::2, 0::2, 0] # R
|
|
raw[1::2, 1::2] = rgb[1::2, 1::2, 1] # G
|
|
else:
|
|
raise ValueError(f"Padrão Bayer não suportado para remosaic: {p}")
|
|
|
|
return np.clip(np.rint(raw * 1023.0), 0, 1023).astype(np.uint16)
|
|
|
|
|
|
def rgb01_to_preview_rgb8(rgb: np.ndarray) -> np.ndarray:
|
|
rgb = np.clip(rgb.astype(np.float32), 0.0, 1.0)
|
|
rgb8 = np.zeros(rgb.shape, dtype=np.uint8)
|
|
for c in range(3):
|
|
rgb8[:, :, c] = _normalize_preview_channel(rgb[:, :, c], lo_p=1.0, hi_p=99.5, gamma=1.0 / 2.2)
|
|
return rgb8
|
|
|
|
|
|
def debayer_raw10_to_rgb8(raw: np.ndarray, bayer_pattern: str = DEFAULT_BAYER_PATTERN) -> np.ndarray:
|
|
"""
|
|
Preview visual a partir do RAW Bayer, usando o mesmo mapeamento Bayer do core.
|
|
"""
|
|
rgb01 = demosaic_raw10_to_rgb01(raw, bayer_pattern=bayer_pattern, algorithm="ea")
|
|
return rgb01_to_preview_rgb8(rgb01)
|
|
|
|
|
|
def _resolve_module_params_path(meta: Dict[str, Any]) -> Optional[str]:
|
|
candidates = []
|
|
|
|
for key in ("camera_params_json", "module_params_json"):
|
|
v = meta.get(key)
|
|
if v:
|
|
candidates.append(str(v))
|
|
|
|
stream_meta = meta.get("stream_meta") if isinstance(meta.get("stream_meta"), dict) else None
|
|
if stream_meta:
|
|
for key in ("camera_params_json", "module_params_json"):
|
|
v = stream_meta.get(key)
|
|
if v:
|
|
candidates.append(str(v))
|
|
|
|
cfg_module = read_config_value("module_params_json", None)
|
|
if cfg_module:
|
|
candidates.append(str(cfg_module))
|
|
|
|
candidates.append("calibration/module_params.json")
|
|
|
|
for c in candidates:
|
|
p = Path(c)
|
|
if p.is_file():
|
|
return str(p)
|
|
|
|
return None
|
|
|
|
|
|
def _build_processing_meta_for_core(meta: Dict[str, Any], cameras: Dict[str, CameraBin]) -> Dict[str, Any]:
|
|
"""
|
|
Monta o meta que o RawProcessorCore espera no build_infer_tensor_from_stream.
|
|
O validador faz isso a partir de stream_meta + camera_info. Aqui criamos/atualizamos
|
|
esse pacote usando os bins augmentados.
|
|
"""
|
|
stream_meta = copy.deepcopy(meta.get("stream_meta") if isinstance(meta.get("stream_meta"), dict) else {})
|
|
|
|
camera_info = stream_meta.get("camera_info")
|
|
if not isinstance(camera_info, dict):
|
|
camera_info = copy.deepcopy(meta.get("camera_info") if isinstance(meta.get("camera_info"), dict) else {})
|
|
|
|
camera_frames = stream_meta.get("camera_frames")
|
|
if not isinstance(camera_frames, dict):
|
|
camera_frames = copy.deepcopy(meta.get("camera_frames") if isinstance(meta.get("camera_frames"), dict) else {})
|
|
|
|
for role, cam in cameras.items():
|
|
info = dict(camera_info.get(cam.cam_key, {}) or {})
|
|
info.update({
|
|
"role": role,
|
|
"width": int(cam.width),
|
|
"height": int(cam.height),
|
|
"bit_depth": int(cam.bit_depth),
|
|
"raw_format": cam.raw_format or DEFAULT_RAW_FORMAT,
|
|
"packed": True,
|
|
"channels": 1,
|
|
"bayer_pattern": cam.bayer_pattern or DEFAULT_BAYER_PATTERN,
|
|
})
|
|
camera_info[cam.cam_key] = info
|
|
camera_frames[cam.cam_key] = dict(info)
|
|
|
|
stream_meta["frame_type"] = "RAW_BRUTO"
|
|
stream_meta["camera_info"] = camera_info
|
|
stream_meta["camera_frames"] = camera_frames
|
|
stream_meta["payload_sources"] = [cameras[r].cam_key for r in ("rgb", "re", "nir") if r in cameras]
|
|
|
|
# A normalização radiométrica do core busca frame_controls no próprio meta.
|
|
# Preservamos os controles originais se existirem.
|
|
for key in ("frame_controls", "actual_camera_controls", "startup_camera_controls"):
|
|
if key not in stream_meta and isinstance(meta.get(key), dict):
|
|
stream_meta[key] = copy.deepcopy(meta[key])
|
|
|
|
return stream_meta
|
|
|
|
|
|
def get_preview_core_cached(sensor_width: int, sensor_height: int, bayer: str, calib_path: str):
|
|
"""
|
|
Reutiliza o RawProcessorCore entre amostras para evitar recarregar flat-field
|
|
e recriar caches de gain/remap toda hora.
|
|
"""
|
|
if not _HAS_RAW_PROCESSOR_CORE:
|
|
return None
|
|
|
|
key = (
|
|
int(sensor_width),
|
|
int(sensor_height),
|
|
str(bayer).upper(),
|
|
str(Path(calib_path).resolve()) if calib_path else "",
|
|
)
|
|
|
|
core = _CORE_PREVIEW_CACHE.get(key)
|
|
if core is not None:
|
|
return core
|
|
|
|
core = RawProcessorCore(
|
|
sensor_width=int(sensor_width),
|
|
sensor_height=int(sensor_height),
|
|
bayer_pattern=str(bayer).upper(),
|
|
calibration_json_path=calib_path,
|
|
)
|
|
|
|
_CORE_PREVIEW_CACHE[key] = core
|
|
return core
|
|
|
|
|
|
def build_core_preview_from_augmented_raws(meta: Dict[str, Any], cameras: Dict[str, CameraBin]) -> Optional[np.ndarray]:
|
|
"""
|
|
Gera o preview salvo passando os bins augmentados pelo RawProcessorCore,
|
|
igual ao check_saved_files faz para montar o MULTISPEC final.
|
|
|
|
Saída: RGB uint8 do tensor final [R,G,B], pronto para salvar com PIL.
|
|
"""
|
|
if not _HAS_RAW_PROCESSOR_CORE:
|
|
return None
|
|
|
|
required = ("rgb", "re", "nir")
|
|
if any(r not in cameras or cameras[r].data is None for r in required):
|
|
return None
|
|
|
|
sensor_width = int(meta.get("sensor_width") or cameras["rgb"].width)
|
|
sensor_height = int(meta.get("sensor_height") or cameras["rgb"].height)
|
|
bayer = str(meta.get("bayer_pattern") or cameras["rgb"].bayer_pattern or DEFAULT_BAYER_PATTERN)
|
|
calib_path = _resolve_module_params_path(meta)
|
|
|
|
if not calib_path:
|
|
return None
|
|
|
|
frame = {}
|
|
for role in required:
|
|
cam = cameras[role]
|
|
frame[cam.cam_key] = pack_raw10_packed_array(cam.data)
|
|
|
|
processing_meta = _build_processing_meta_for_core(meta, cameras)
|
|
|
|
core = get_preview_core_cached(
|
|
sensor_width=sensor_width,
|
|
sensor_height=sensor_height,
|
|
bayer=bayer,
|
|
calib_path=calib_path,
|
|
)
|
|
|
|
if core is None:
|
|
return None
|
|
|
|
tensor = core.build_infer_tensor_from_stream(frame, processing_meta, 5)
|
|
if tensor is None or not isinstance(tensor, np.ndarray) or tensor.ndim != 3 or tensor.shape[0] < 3:
|
|
return None
|
|
|
|
rgb_hwc = np.transpose(tensor[:3].astype(np.float32), (1, 2, 0))
|
|
rgb8 = np.clip(rgb_hwc * 255.0, 0, 255).astype(np.uint8)
|
|
return rgb8
|
|
|
|
|
|
def build_beauty_preview_from_augmented_raws(cameras: Dict[str, CameraBin]) -> np.ndarray:
|
|
"""
|
|
Fallback visual a partir dos RAWs augmentados, sem passar pelo core.
|
|
Preferimos build_core_preview_from_augmented_raws sempre que possível.
|
|
"""
|
|
rgb_cam = cameras.get("rgb")
|
|
if rgb_cam is None or rgb_cam.data is None:
|
|
raise RuntimeError("Sem RAW RGB para gerar preview.")
|
|
|
|
return debayer_raw10_to_rgb8(rgb_cam.data, rgb_cam.bayer_pattern)
|
|
|
|
|
|
def save_preview_from_augmented_raws(path: Path, cameras: Dict[str, CameraBin], meta: Optional[Dict[str, Any]] = None) -> str:
|
|
ensure_dir(path.parent)
|
|
|
|
method = "fallback_cam_a_debayer_preview"
|
|
rgb8 = None
|
|
|
|
if meta is not None:
|
|
try:
|
|
rgb8 = build_core_preview_from_augmented_raws(meta, cameras)
|
|
if rgb8 is not None:
|
|
method = "raw_processor_core_multispec_rgb_final"
|
|
except Exception as e:
|
|
print(f"[WARN] Preview via RawProcessorCore falhou em {path.name}: {e}. Usando fallback CAM_A.")
|
|
rgb8 = None
|
|
|
|
if rgb8 is None:
|
|
rgb8 = build_beauty_preview_from_augmented_raws(cameras)
|
|
|
|
Image.fromarray(rgb8).save(path)
|
|
return method
|
|
|
|
|
|
# ============================================================
|
|
# Augmentation
|
|
# ============================================================
|
|
|
|
def sample_aug_params(rng: random.Random, seed_value: int) -> AugParams:
|
|
# Geometria conservadora.
|
|
perspective = rng.random() < 0.15
|
|
blur_enabled = rng.random() < 0.15
|
|
|
|
return AugParams(
|
|
seed=seed_value,
|
|
flip_h=rng.random() < 0.50,
|
|
shift_x_frac=rng.uniform(-0.015, 0.015),
|
|
shift_y_frac=rng.uniform(-0.015, 0.015),
|
|
scale=rng.uniform(0.92, 1.08),
|
|
rotate_deg=rng.uniform(-4.0, 4.0),
|
|
perspective=perspective,
|
|
perspective_strength=rng.uniform(0.002, 0.012) if perspective else 0.0,
|
|
exposure_mult=rng.uniform(0.75, 1.30),
|
|
gamma=rng.uniform(0.92, 1.08),
|
|
rgb_channel_gain=(
|
|
rng.uniform(0.92, 1.08),
|
|
rng.uniform(0.92, 1.08),
|
|
rng.uniform(0.92, 1.08),
|
|
),
|
|
re_gain=rng.uniform(0.85, 1.20),
|
|
nir_gain=rng.uniform(0.85, 1.20),
|
|
shadow_enabled=rng.random() < 0.25,
|
|
shadow_strength=rng.uniform(0.12, 0.35),
|
|
shadow_angle_deg=rng.uniform(0.0, 180.0),
|
|
highlight_enabled=rng.random() < 0.15,
|
|
highlight_strength=rng.uniform(0.05, 0.18),
|
|
noise_enabled=rng.random() < 0.30,
|
|
noise_sigma_dn=rng.uniform(1.5, 6.0),
|
|
blur_enabled=blur_enabled,
|
|
blur_kernel=3 if blur_enabled else 1,
|
|
)
|
|
|
|
|
|
def build_affine_matrix(width: int, height: int, p: AugParams) -> np.ndarray:
|
|
cx = width * 0.5
|
|
cy = height * 0.5
|
|
M = cv2.getRotationMatrix2D((cx, cy), p.rotate_deg, p.scale)
|
|
M[0, 2] += p.shift_x_frac * width
|
|
M[1, 2] += p.shift_y_frac * height
|
|
return M
|
|
|
|
|
|
def warp_image(img: np.ndarray, M: np.ndarray, out_size: Tuple[int, int], is_mask: bool) -> np.ndarray:
|
|
interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_LINEAR
|
|
return cv2.warpAffine(
|
|
img,
|
|
M,
|
|
out_size,
|
|
flags=interp,
|
|
borderMode=cv2.BORDER_REFLECT_101,
|
|
)
|
|
|
|
|
|
def build_perspective_matrix(width: int, height: int, p: AugParams, rng: random.Random) -> Optional[np.ndarray]:
|
|
if not p.perspective:
|
|
return None
|
|
|
|
s = p.perspective_strength
|
|
dx = width * s
|
|
dy = height * s
|
|
|
|
src = np.float32([
|
|
[0, 0],
|
|
[width - 1, 0],
|
|
[width - 1, height - 1],
|
|
[0, height - 1],
|
|
])
|
|
|
|
dst = src + np.float32([
|
|
[rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
|
|
[rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
|
|
[rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
|
|
[rng.uniform(-dx, dx), rng.uniform(-dy, dy)],
|
|
])
|
|
|
|
return cv2.getPerspectiveTransform(src, dst)
|
|
|
|
|
|
def warp_perspective(img: np.ndarray, H: np.ndarray, out_size: Tuple[int, int], is_mask: bool) -> np.ndarray:
|
|
interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_LINEAR
|
|
return cv2.warpPerspective(
|
|
img,
|
|
H,
|
|
out_size,
|
|
flags=interp,
|
|
borderMode=cv2.BORDER_REFLECT_101,
|
|
)
|
|
|
|
|
|
def apply_geometry_raw(raw: np.ndarray, p: AugParams, rng: random.Random) -> np.ndarray:
|
|
h, w = raw.shape[:2]
|
|
|
|
out = raw
|
|
if p.flip_h:
|
|
out = cv2.flip(out, 1)
|
|
|
|
M = build_affine_matrix(w, h, p)
|
|
out = warp_image(out, M, (w, h), is_mask=False)
|
|
|
|
H = build_perspective_matrix(w, h, p, rng)
|
|
if H is not None:
|
|
out = warp_perspective(out, H, (w, h), is_mask=False)
|
|
|
|
return out
|
|
|
|
|
|
def apply_geometry_mask(mask: np.ndarray, p: AugParams, rng: random.Random) -> np.ndarray:
|
|
h, w = mask.shape[:2]
|
|
|
|
out = mask
|
|
if p.flip_h:
|
|
out = cv2.flip(out, 1)
|
|
|
|
M = build_affine_matrix(w, h, p)
|
|
out = warp_image(out, M, (w, h), is_mask=True)
|
|
|
|
H = build_perspective_matrix(w, h, p, rng)
|
|
if H is not None:
|
|
out = warp_perspective(out, H, (w, h), is_mask=True)
|
|
|
|
return out
|
|
|
|
|
|
def make_spatial_light_map(shape: Tuple[int, int], p: AugParams) -> np.ndarray:
|
|
h, w = shape
|
|
yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
|
|
xx = (xx / max(w - 1, 1)) - 0.5
|
|
yy = (yy / max(h - 1, 1)) - 0.5
|
|
|
|
light = np.ones((h, w), dtype=np.float32)
|
|
|
|
if p.shadow_enabled:
|
|
theta = math.radians(p.shadow_angle_deg)
|
|
direction = math.cos(theta) * xx + math.sin(theta) * yy
|
|
direction = (direction - direction.min()) / max(direction.max() - direction.min(), 1e-6)
|
|
shadow = 1.0 - p.shadow_strength * direction
|
|
light *= shadow
|
|
|
|
if p.highlight_enabled:
|
|
# Mancha larga e suave, simulando região de sol/reflexo.
|
|
cx = np.random.uniform(-0.25, 0.25)
|
|
cy = np.random.uniform(-0.25, 0.25)
|
|
sigma = np.random.uniform(0.20, 0.42)
|
|
d2 = (xx - cx) ** 2 + (yy - cy) ** 2
|
|
blob = np.exp(-d2 / (2 * sigma * sigma))
|
|
light *= 1.0 + p.highlight_strength * blob
|
|
|
|
return np.clip(light, 0.50, 1.45).astype(np.float32)
|
|
|
|
|
|
def apply_gamma_raw(raw: np.ndarray, gamma: float) -> np.ndarray:
|
|
if abs(gamma - 1.0) < 1e-3:
|
|
return raw
|
|
x = np.clip(raw.astype(np.float32) / 1023.0, 0, 1)
|
|
x = np.power(x, gamma)
|
|
return x * 1023.0
|
|
|
|
|
|
def apply_bayer_channel_gains(raw: np.ndarray, gains_rgb: Tuple[float, float, float], pattern: str) -> np.ndarray:
|
|
"""
|
|
Aplica ganhos aproximados por canal em mosaico Bayer.
|
|
Para preview/treino, isso simula variação de balanço/cor no sensor RGB.
|
|
"""
|
|
out = raw.astype(np.float32).copy()
|
|
r_gain, g_gain, b_gain = gains_rgb
|
|
pat = (pattern or DEFAULT_BAYER_PATTERN).upper()
|
|
|
|
if pat == "RGGB":
|
|
out[0::2, 0::2] *= r_gain
|
|
out[0::2, 1::2] *= g_gain
|
|
out[1::2, 0::2] *= g_gain
|
|
out[1::2, 1::2] *= b_gain
|
|
elif pat == "BGGR":
|
|
out[0::2, 0::2] *= b_gain
|
|
out[0::2, 1::2] *= g_gain
|
|
out[1::2, 0::2] *= g_gain
|
|
out[1::2, 1::2] *= r_gain
|
|
elif pat == "GRBG":
|
|
out[0::2, 0::2] *= g_gain
|
|
out[0::2, 1::2] *= r_gain
|
|
out[1::2, 0::2] *= b_gain
|
|
out[1::2, 1::2] *= g_gain
|
|
elif pat == "GBRG":
|
|
out[0::2, 0::2] *= g_gain
|
|
out[0::2, 1::2] *= b_gain
|
|
out[1::2, 0::2] *= r_gain
|
|
out[1::2, 1::2] *= g_gain
|
|
else:
|
|
out *= float(np.mean(gains_rgb))
|
|
|
|
return out
|
|
|
|
|
|
def apply_radiometry_rgb01(rgb: np.ndarray, p: AugParams, rng_np: np.random.Generator, light_map_cache: Dict[Tuple[int, int], np.ndarray]) -> np.ndarray:
|
|
"""
|
|
Radiometria para RGB já demosaicado.
|
|
Evita mexer no mosaico Bayer diretamente.
|
|
"""
|
|
h, w = rgb.shape[:2]
|
|
x = np.clip(rgb.astype(np.float32), 0.0, 1.0)
|
|
|
|
x *= np.float32(p.exposure_mult)
|
|
|
|
gains = np.array(p.rgb_channel_gain, dtype=np.float32).reshape(1, 1, 3)
|
|
x *= gains
|
|
|
|
key = (h, w)
|
|
if key not in light_map_cache:
|
|
light_map_cache[key] = make_spatial_light_map((h, w), p)
|
|
x *= light_map_cache[key][:, :, None]
|
|
|
|
if abs(p.gamma - 1.0) > 1e-3:
|
|
x = np.power(np.clip(x, 0.0, 1.0), p.gamma)
|
|
|
|
if p.blur_enabled and p.blur_kernel >= 3:
|
|
x = cv2.GaussianBlur(x, (p.blur_kernel, p.blur_kernel), 0)
|
|
|
|
if p.noise_enabled:
|
|
# Converte sigma DN RAW10 para escala 0..1.
|
|
sigma01 = float(p.noise_sigma_dn) / 1023.0
|
|
noise = rng_np.normal(0.0, sigma01, size=x.shape).astype(np.float32)
|
|
x += noise
|
|
|
|
return np.clip(x, 0.0, 1.0).astype(np.float32, copy=False)
|
|
|
|
|
|
def apply_radiometry_raw(raw: np.ndarray, role: str, p: AugParams, cam: CameraBin, rng_np: np.random.Generator, light_map_cache: Dict[Tuple[int, int], np.ndarray]) -> np.ndarray:
|
|
h, w = raw.shape[:2]
|
|
x = raw.astype(np.float32)
|
|
|
|
# Base global de exposição: muda todo mundo junto.
|
|
x *= p.exposure_mult
|
|
|
|
# Variação por papel espectral/canal.
|
|
# OBS: RGB não deve passar por aqui, porque RGB Bayer precisa ser demosaicado
|
|
# antes de qualquer warp/radiometria por canal.
|
|
if role == "re":
|
|
x *= p.re_gain
|
|
elif role == "nir":
|
|
x *= p.nir_gain
|
|
|
|
# Mesmo mapa espacial de luz por tamanho, para manter coerência física.
|
|
key = (h, w)
|
|
if key not in light_map_cache:
|
|
light_map_cache[key] = make_spatial_light_map((h, w), p)
|
|
x *= light_map_cache[key]
|
|
|
|
# Gamma leve, opcional. Em RAW científico puro, gamma seria discutível.
|
|
# Mantemos bem pequeno para simular resposta/exposição não ideal.
|
|
x = apply_gamma_raw(x, p.gamma)
|
|
|
|
if p.blur_enabled and p.blur_kernel >= 3:
|
|
x = cv2.GaussianBlur(x, (p.blur_kernel, p.blur_kernel), 0)
|
|
|
|
if p.noise_enabled:
|
|
noise = rng_np.normal(0.0, p.noise_sigma_dn, size=x.shape).astype(np.float32)
|
|
x += noise
|
|
|
|
return np.clip(x, 0, 1023).astype(np.uint16)
|
|
|
|
|
|
# ============================================================
|
|
# Processamento de amostra/grupo
|
|
# ============================================================
|
|
|
|
def validate_mask_colors(mask_before: np.ndarray, mask_after: np.ndarray, name: str) -> None:
|
|
if mask_before.ndim == 2 or mask_after.ndim == 2:
|
|
before_colors = int(np.unique(mask_before.reshape(-1)).size)
|
|
after_colors = int(np.unique(mask_after.reshape(-1)).size)
|
|
else:
|
|
before_colors = len(set(map(tuple, mask_before.reshape(-1, mask_before.shape[2]))))
|
|
after_colors = len(set(map(tuple, mask_after.reshape(-1, mask_after.shape[2]))))
|
|
|
|
if after_colors > max(before_colors * 3, 64):
|
|
print(
|
|
f"[WARN] {name}: máscara ganhou muitas cores. "
|
|
f"antes={before_colors}, depois={after_colors}. Confira interpolação NEAREST."
|
|
)
|
|
|
|
|
|
def save_aug_meta(
|
|
meta: Dict[str, Any],
|
|
out_path: Path,
|
|
base: str,
|
|
new_base: str,
|
|
params: AugParams,
|
|
cameras: Dict[str, CameraBin],
|
|
preview_method: Optional[str] = None,
|
|
) -> None:
|
|
out = copy.deepcopy(meta)
|
|
out["synthetic"] = True
|
|
out["augmentation"] = {
|
|
"script": "_5_augmentation_raw_oak.py",
|
|
"parent_base": base,
|
|
"new_base": new_base,
|
|
"params": asdict(params),
|
|
}
|
|
|
|
# Deixa o JSON compatível com o check_saved_files / validador.
|
|
out["saved_payload_type"] = "raw_native_multi"
|
|
out["saved_preview_method"] = preview_method or "unknown"
|
|
out["saved_payload_paths"] = {}
|
|
out["saved_payload_shapes"] = {}
|
|
out["saved_payload_dtypes"] = {}
|
|
|
|
stream_meta = copy.deepcopy(out.get("stream_meta") if isinstance(out.get("stream_meta"), dict) else {})
|
|
camera_info = stream_meta.get("camera_info") if isinstance(stream_meta.get("camera_info"), dict) else {}
|
|
camera_frames = stream_meta.get("camera_frames") if isinstance(stream_meta.get("camera_frames"), dict) else {}
|
|
|
|
out["augmented_bins"] = {}
|
|
|
|
for role, cam in cameras.items():
|
|
packed_w = int(math.ceil(int(cam.width) * 10 / 8))
|
|
fname = f"{new_base}_{cam.cam_key}.bin"
|
|
|
|
out["saved_payload_paths"][cam.cam_key] = fname
|
|
out["saved_payload_shapes"][cam.cam_key] = [int(cam.height), int(packed_w)]
|
|
out["saved_payload_dtypes"][cam.cam_key] = "uint8"
|
|
|
|
cam_info = dict(camera_info.get(cam.cam_key, {}) or {})
|
|
cam_info.update({
|
|
"role": role,
|
|
"width": int(cam.width),
|
|
"height": int(cam.height),
|
|
"bit_depth": int(cam.bit_depth),
|
|
"raw_format": cam.raw_format or DEFAULT_RAW_FORMAT,
|
|
"packed": True,
|
|
"channels": 1,
|
|
"bayer_pattern": cam.bayer_pattern or DEFAULT_BAYER_PATTERN,
|
|
})
|
|
camera_info[cam.cam_key] = cam_info
|
|
camera_frames[cam.cam_key] = dict(cam_info)
|
|
|
|
out["augmented_bins"][role] = {
|
|
"cam_key": cam.cam_key,
|
|
"filename": fname,
|
|
"role": role,
|
|
"width": int(cam.width),
|
|
"height": int(cam.height),
|
|
"bit_depth": int(cam.bit_depth),
|
|
"raw_format": cam.raw_format,
|
|
"bayer_pattern": cam.bayer_pattern,
|
|
}
|
|
|
|
stream_meta["frame_type"] = "RAW_BRUTO"
|
|
stream_meta["camera_info"] = camera_info
|
|
stream_meta["camera_frames"] = camera_frames
|
|
stream_meta["payload_sources"] = [cameras[r].cam_key for r in ("rgb", "re", "nir") if r in cameras]
|
|
out["stream_meta"] = stream_meta
|
|
|
|
save_json(out_path, out)
|
|
|
|
|
|
def process_one_sample(
|
|
base: str,
|
|
group_name: str,
|
|
src_group: Path,
|
|
dst_group: Path,
|
|
mask_path: Path,
|
|
meta_path: Path,
|
|
copies: int,
|
|
rng_global: random.Random,
|
|
dry_run: bool = False,
|
|
) -> Tuple[int, int]:
|
|
bins_dir = src_group / "bins"
|
|
meta = load_json(meta_path)
|
|
cameras = discover_sample_cameras(meta, bins_dir, base)
|
|
|
|
required = ["rgb", "re", "nir"]
|
|
missing = [r for r in required if r not in cameras]
|
|
if missing:
|
|
print(f"[WARN] [{group_name}] {base}: câmeras ausentes {missing}. Pulando.")
|
|
return 0, 1
|
|
|
|
mask = load_mask(mask_path)
|
|
|
|
# Carrega todos os RAWs uma vez.
|
|
for role, cam in cameras.items():
|
|
cam.data = read_raw_bin(cam.path, cam.width, cam.height, cam.raw_format)
|
|
|
|
generated = 0
|
|
errors = 0
|
|
|
|
out_bins = dst_group / "bins"
|
|
out_masks = dst_group / "masks"
|
|
out_metas = dst_group / "metas"
|
|
out_previews = dst_group / "previews"
|
|
|
|
for i in range(copies):
|
|
seed_value = rng_global.randint(0, 2**31 - 1)
|
|
rng = random.Random(seed_value)
|
|
rng_np = np.random.default_rng(seed_value)
|
|
params = sample_aug_params(rng, seed_value)
|
|
new_base = f"{base}_aug_{i:02d}"
|
|
|
|
try:
|
|
if dry_run:
|
|
generated += 1
|
|
continue
|
|
|
|
# Mesma geometria para mask e raws.
|
|
mask_aug = apply_geometry_mask(mask, params, random.Random(seed_value + 1000))
|
|
validate_mask_colors(mask, mask_aug, f"{new_base}_mask")
|
|
|
|
light_map_cache: Dict[Tuple[int, int], np.ndarray] = {}
|
|
cams_out: Dict[str, CameraBin] = {}
|
|
|
|
for role in required:
|
|
cam = cameras[role]
|
|
assert cam.data is not None
|
|
|
|
if role == "rgb":
|
|
# RGB RAW é mosaico Bayer. Não podemos aplicar warp direto nele,
|
|
# porque isso mistura pixels R/G/B antes do demosaic e cria ruído colorido.
|
|
# Fluxo correto offline:
|
|
# RAW Bayer -> RGB linear -> geometria/radiometria -> mosaico Bayer -> RAW10
|
|
rgb01 = demosaic_raw10_to_rgb01(
|
|
cam.data,
|
|
bayer_pattern=cam.bayer_pattern,
|
|
algorithm="ea",
|
|
)
|
|
|
|
rgb_geom = rgb01
|
|
if params.flip_h:
|
|
rgb_geom = cv2.flip(rgb_geom, 1)
|
|
|
|
h_rgb, w_rgb = rgb_geom.shape[:2]
|
|
M_rgb = build_affine_matrix(w_rgb, h_rgb, params)
|
|
rgb_geom = warp_image(rgb_geom, M_rgb, (w_rgb, h_rgb), is_mask=False)
|
|
|
|
H_rgb = build_perspective_matrix(w_rgb, h_rgb, params, random.Random(seed_value + 1000))
|
|
if H_rgb is not None:
|
|
rgb_geom = warp_perspective(rgb_geom, H_rgb, (w_rgb, h_rgb), is_mask=False)
|
|
|
|
rgb_aug = apply_radiometry_rgb01(
|
|
rgb_geom,
|
|
params,
|
|
rng_np,
|
|
light_map_cache,
|
|
)
|
|
|
|
raw_aug = remosaic_rgb01_to_bayer_raw10(
|
|
rgb_aug,
|
|
bayer_pattern=cam.bayer_pattern,
|
|
)
|
|
|
|
else:
|
|
raw_geom = apply_geometry_raw(cam.data, params, random.Random(seed_value + 1000))
|
|
raw_aug = apply_radiometry_raw(raw_geom, role, params, cam, rng_np, light_map_cache)
|
|
|
|
out_cam = copy.deepcopy(cam)
|
|
out_cam.path = out_bins / f"{new_base}_{cam.cam_key}.bin"
|
|
out_cam.data = raw_aug
|
|
cams_out[role] = out_cam
|
|
|
|
write_raw_bin(out_cam.path, raw_aug, cam.raw_format)
|
|
|
|
# Máscara com mesmo nome-base.
|
|
mask_ext = mask_path.suffix.lower() if mask_path.suffix.lower() in IMG_EXTS else ".png"
|
|
save_mask(out_masks / f"{new_base}{mask_ext}", mask_aug)
|
|
|
|
# Preview visual usando o mesmo caminho do validador/check_saved_files:
|
|
# bins augmentados -> RawProcessorCore -> MULTISPEC final -> RGB final.
|
|
# Se o core falhar, cai no fallback CAM_A e avisa no console.
|
|
preview_method = save_preview_from_augmented_raws(
|
|
out_previews / f"{new_base}.png",
|
|
cams_out,
|
|
meta=meta,
|
|
)
|
|
|
|
# Meta novo, com mesmo stem do preview/mask para manter o layout dataset limpo.
|
|
save_aug_meta(
|
|
meta=meta,
|
|
out_path=out_metas / f"{new_base}.json",
|
|
base=base,
|
|
new_base=new_base,
|
|
params=params,
|
|
cameras=cams_out,
|
|
preview_method=preview_method,
|
|
)
|
|
|
|
generated += 1
|
|
|
|
except Exception as e:
|
|
errors += 1
|
|
print(f"[ERRO] [{group_name}] {new_base}: {e}")
|
|
|
|
return generated, errors
|
|
|
|
|
|
def parse_group_copies(text: Optional[str]) -> Dict[str, int]:
|
|
"""
|
|
Ex: chao:1,chao_cana:3,chao_erva:3,chao_cana_erva:5
|
|
"""
|
|
out: Dict[str, int] = {}
|
|
if not text:
|
|
return out
|
|
|
|
for part in text.split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
if ":" not in part:
|
|
raise ValueError(f"group-copies inválido: {part}. Use grupo:N")
|
|
g, n = part.split(":", 1)
|
|
out[g.strip()] = int(n.strip())
|
|
return out
|
|
|
|
|
|
def process_group(
|
|
src_root: Path,
|
|
dst_root: Path,
|
|
group_name: str,
|
|
copies: int,
|
|
limit: Optional[int],
|
|
seed: int,
|
|
dry_run: bool,
|
|
) -> Tuple[int, int, int]:
|
|
src_group = src_root / group_name
|
|
dst_group = dst_root / group_name
|
|
|
|
mask_map = map_files_by_base(src_group / "masks", IMG_EXTS)
|
|
meta_map = map_files_by_base(src_group / "metas", META_EXTS)
|
|
|
|
bases = sorted(set(mask_map.keys()) & set(meta_map.keys()))
|
|
|
|
if limit is not None and limit > 0 and limit < len(bases):
|
|
rng_select = random.Random(seed)
|
|
bases = sorted(rng_select.sample(bases, limit))
|
|
print(f"[INFO] [{group_name}] limit={limit}, amostras selecionadas={len(bases)}")
|
|
|
|
if not bases:
|
|
print(f"[WARN] [{group_name}] Nenhum par mask/meta encontrado.")
|
|
return 0, 0, 0
|
|
|
|
ensure_dir(dst_group / "bins")
|
|
ensure_dir(dst_group / "masks")
|
|
ensure_dir(dst_group / "metas")
|
|
ensure_dir(dst_group / "previews")
|
|
|
|
rng_global = random.Random(seed + abs(hash(group_name)) % 1000000)
|
|
|
|
total_generated = 0
|
|
total_errors = 0
|
|
|
|
for idx, base in enumerate(bases, start=1):
|
|
gen, err = process_one_sample(
|
|
base=base,
|
|
group_name=group_name,
|
|
src_group=src_group,
|
|
dst_group=dst_group,
|
|
mask_path=mask_map[base],
|
|
meta_path=meta_map[base],
|
|
copies=copies,
|
|
rng_global=rng_global,
|
|
dry_run=dry_run,
|
|
)
|
|
total_generated += gen
|
|
total_errors += err
|
|
|
|
if idx % 25 == 0 or idx == len(bases):
|
|
print(
|
|
f"[INFO] [{group_name}] {idx}/{len(bases)} amostras | "
|
|
f"gerados={total_generated} | erros={total_errors}"
|
|
)
|
|
|
|
print(
|
|
f"[OK] Grupo '{group_name}' concluído: "
|
|
f"originais={len(bases)} | copies={copies} | gerados={total_generated} | erros={total_errors}"
|
|
)
|
|
|
|
return len(bases), total_generated, total_errors
|
|
|
|
|
|
# ============================================================
|
|
# CLI
|
|
# ============================================================
|
|
|
|
def main() -> None:
|
|
dataset_base = Path("dataset")
|
|
|
|
default_src = dataset_base / "original" / "group"
|
|
if not default_src.exists():
|
|
# Compatibilidade com script antigo que usava "originals".
|
|
alt = dataset_base / "originals" / "group"
|
|
if alt.exists():
|
|
default_src = alt
|
|
|
|
default_dst = dataset_base / "augmented" / "group"
|
|
|
|
ap = argparse.ArgumentParser(
|
|
description="Augmentation multiespectral RAW10 packed para dataset OAK-FCC-3."
|
|
)
|
|
|
|
ap.add_argument("--copies", type=int, default=3, help="Cópias augmentadas por amostra, se --group-copies não sobrescrever.")
|
|
ap.add_argument("--group-copies", type=str, default=None, help="Cópias por grupo. Ex: chao:1,chao_cana:3,chao_erva:3,chao_cana_erva:5")
|
|
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula.")
|
|
ap.add_argument("--src-root", type=str, default=str(default_src), help="Raiz dos grupos originais.")
|
|
ap.add_argument("--dst-root", type=str, default=str(default_dst), help="Raiz dos grupos augmentados.")
|
|
ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de amostras originais por grupo.")
|
|
ap.add_argument("--seed", type=int, default=42, help="Seed geral para reproducibilidade.")
|
|
ap.add_argument("--clear-dst", action="store_true", help="Apaga dst-root antes de gerar.")
|
|
ap.add_argument("--dry-run", action="store_true", help="Simula sem salvar arquivos.")
|
|
|
|
args = ap.parse_args()
|
|
|
|
src_root = Path(args.src_root)
|
|
dst_root = Path(args.dst_root)
|
|
group_copies = parse_group_copies(args.group_copies)
|
|
|
|
print("==============================================")
|
|
print("Augmentation RAW OAK-FCC-3")
|
|
print(f"MODEL/DATASET : {dataset_base}")
|
|
print(f"SRC_ROOT : {src_root}")
|
|
print(f"DST_ROOT : {dst_root}")
|
|
print(f"copies : {args.copies}")
|
|
print(f"group_copies : {group_copies if group_copies else '{}'}")
|
|
print(f"limit : {args.limit}")
|
|
print(f"seed : {args.seed}")
|
|
print(f"dry_run : {args.dry_run}")
|
|
print("==============================================")
|
|
|
|
if not src_root.exists():
|
|
raise SystemExit(f"[ERRO] src-root não encontrado: {src_root}")
|
|
|
|
if args.clear_dst and not args.dry_run:
|
|
print(f"[INFO] Limpando destino: {dst_root}")
|
|
clear_dir(dst_root)
|
|
else:
|
|
ensure_dir(dst_root)
|
|
|
|
groups = list_groups(src_root)
|
|
|
|
if args.groups:
|
|
wanted = {g.strip() for g in args.groups.split(",") if g.strip()}
|
|
groups = [g for g in groups if g in wanted]
|
|
|
|
if not groups:
|
|
raise SystemExit("[WARN] Nenhum grupo válido encontrado.")
|
|
|
|
print(f"Grupos encontrados: {', '.join(groups)}")
|
|
|
|
total_originals = 0
|
|
total_generated = 0
|
|
total_errors = 0
|
|
|
|
for group_name in groups:
|
|
copies = group_copies.get(group_name, args.copies)
|
|
if copies <= 0:
|
|
print(f"[INFO] [{group_name}] copies={copies}. Pulando.")
|
|
continue
|
|
|
|
n_orig, n_gen, n_err = process_group(
|
|
src_root=src_root,
|
|
dst_root=dst_root,
|
|
group_name=group_name,
|
|
copies=copies,
|
|
limit=args.limit,
|
|
seed=args.seed,
|
|
dry_run=args.dry_run,
|
|
)
|
|
total_originals += n_orig
|
|
total_generated += n_gen
|
|
total_errors += n_err
|
|
|
|
print("\n==============================================")
|
|
print("Resumo final")
|
|
print(f"Originais processados : {total_originals}")
|
|
print(f"Amostras geradas : {total_generated}")
|
|
print(f"Erros/Pulos : {total_errors}")
|
|
print(f"Destino : {dst_root}")
|
|
print("==============================================")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|