709 lines
24 KiB
Python
709 lines
24 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Augmenta imagens/máscaras *por grupo* usando RAW4 (float32, 0..1, R G IR B).
|
||
|
||
Novo fluxo (GAL5000):
|
||
|
||
Entrada:
|
||
MODELO/dataset/original/group/<grupo>/raws (.raw, 4 canais float32)
|
||
MODELO/dataset/original/group/<grupo>/masks (.png/.jpg)
|
||
(opcional) MODELO/dataset/original/group/<grupo>/masks2
|
||
|
||
Saída:
|
||
MODELO/dataset/augmented/group/<grupo>/previews (RGB gerado do RAW blurado)
|
||
MODELO/dataset/augmented/group/<grupo>/raws (RAW4 augmentado: geom + blur)
|
||
MODELO/dataset/augmented/group/<grupo>/masks
|
||
MODELO/dataset/augmented/group/<grupo>/masks2
|
||
|
||
Regras:
|
||
- GEOMETRIA (flip, shift, rotate) aplicada em RAW + masks (+ masks2).
|
||
- BLUR / MOTION BLUR aplicado no RAW, canal por canal, com o mesmo kernel.
|
||
- Preview é SEMPRE derivado do RAW já transformado (geom + blur).
|
||
- Máscaras não sofrem blur, apenas geom.
|
||
|
||
Uso:
|
||
python _5_augmentation_raw.py --copies 5
|
||
python _5_augmentation_raw.py --copies 5 --groups chao,chao_erva,cana
|
||
"""
|
||
import os
|
||
import json
|
||
import cv2
|
||
from PIL import Image
|
||
import albumentations as A
|
||
from albumentations import ReplayCompose
|
||
import argparse
|
||
import numpy as np
|
||
|
||
# ⚙️ Configurações
|
||
with open("config.json", "r", encoding="utf-8") as f:
|
||
config = json.load(f)
|
||
MODELO = config.get("camera", ".")
|
||
RESOLUCAO = config.get("raw_size")
|
||
|
||
USE_MASKS2 = config.get("dual_head", False)
|
||
|
||
# Pastas base
|
||
DATASET_BASE = os.path.join(MODELO, "dataset")
|
||
ORIG_GROUP_ROOT = os.path.join(DATASET_BASE, "original", "group")
|
||
AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group")
|
||
|
||
# Fallback (modo antigo, sem grupos) – mantido por compat
|
||
ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images")
|
||
ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks")
|
||
ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2")
|
||
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
|
||
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
|
||
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
|
||
|
||
# Extensões aceitas
|
||
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||
MSK_EXTS = (".png", ".jpg", ".jpeg")
|
||
MSK2_EXTS = (".png", ".jpg", ".jpeg")
|
||
RAW_EXTS = (".raw",)
|
||
|
||
|
||
def garantir_dir(p):
|
||
os.makedirs(p, exist_ok=True)
|
||
|
||
|
||
# ===============================
|
||
# Pipelines de augmentations
|
||
# ===============================
|
||
|
||
# 1) Geométricas: aplicam em preview “temporário”, mask, mask2 e RAW.
|
||
# RAW é passado como 'raw' em additional_targets para receber só geom.
|
||
train_geo = ReplayCompose([
|
||
A.HorizontalFlip(p=0.5),
|
||
A.ShiftScaleRotate(
|
||
shift_limit=0.01,
|
||
scale_limit=0.10,
|
||
rotate_limit=5,
|
||
border_mode=cv2.BORDER_REFLECT_101,
|
||
interpolation=cv2.INTER_LINEAR,
|
||
p=0.30
|
||
),
|
||
], additional_targets={
|
||
'mask2': 'mask',
|
||
'raw': 'mask', # trata raw como “mask” para receber apenas geom
|
||
})
|
||
|
||
|
||
# ===============================
|
||
# Funções de blur físico no RAW
|
||
# ===============================
|
||
|
||
def motion_blur_kernel(ksize=5, angle=0.0):
|
||
"""
|
||
Gera um kernel de motion blur 2D com tamanho ksize e ângulo em graus.
|
||
Usado tanto para RAW quanto para preview.
|
||
"""
|
||
ksize = int(ksize)
|
||
if ksize < 3:
|
||
ksize = 3
|
||
if ksize % 2 == 0:
|
||
ksize += 1
|
||
|
||
kernel = np.zeros((ksize, ksize), dtype=np.float32)
|
||
# linha horizontal no centro
|
||
kernel[ksize // 2, :] = 1.0
|
||
|
||
# rotaciona a linha para o ângulo desejado
|
||
center = (ksize / 2.0 - 0.5, ksize / 2.0 - 0.5)
|
||
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
||
kernel = cv2.warpAffine(kernel, M, (ksize, ksize))
|
||
|
||
s = kernel.sum()
|
||
if s > 0:
|
||
kernel /= s
|
||
return kernel
|
||
|
||
|
||
def apply_coherent_blur_raw(raw, p_motion=0.20, p_gauss=0.20):
|
||
"""
|
||
Aplica blur/motion blur de forma COERENTE no RAW, canal por canal.
|
||
|
||
- Mesmo kernel em TODOS os canais.
|
||
- Máscaras não passam aqui.
|
||
- Se nada for sorteado, retorna o RAW original.
|
||
"""
|
||
if raw is None:
|
||
return None
|
||
|
||
r = np.random.rand()
|
||
if r >= (p_motion + p_gauss):
|
||
# sem blur nesta amostra
|
||
return raw
|
||
|
||
# escolhe tamanho de kernel
|
||
k_choices = [3, 5, 7]
|
||
ksize = int(np.random.choice(k_choices))
|
||
|
||
if r < p_motion:
|
||
# Motion blur com ângulo pequeno (vibração / movimento leve)
|
||
angle = float(np.random.uniform(-20.0, 20.0))
|
||
kernel = motion_blur_kernel(ksize=ksize, angle=angle)
|
||
else:
|
||
# Gaussian blur (leve desfoque de foco)
|
||
if ksize % 2 == 0:
|
||
ksize += 1
|
||
g1d = cv2.getGaussianKernel(ksize, 0)
|
||
kernel = g1d @ g1d.T # vira kernel 2D
|
||
|
||
# Aplica kernel canal a canal
|
||
if raw.ndim == 2:
|
||
raw_blur = cv2.filter2D(raw, ddepth=-1, kernel=kernel)
|
||
else:
|
||
raw_blur = np.empty_like(raw)
|
||
h, w, c = raw.shape
|
||
for ch in range(c):
|
||
raw_blur[..., ch] = cv2.filter2D(
|
||
raw[..., ch],
|
||
ddepth=-1,
|
||
kernel=kernel,
|
||
borderType=cv2.BORDER_REFLECT_101,
|
||
)
|
||
|
||
# Garante que continua em [0,1] para float
|
||
if np.issubdtype(raw_blur.dtype, np.floating):
|
||
raw_blur = np.clip(raw_blur, 0.0, 1.0)
|
||
|
||
return raw_blur
|
||
|
||
|
||
# ===============================
|
||
# IO de imagens / RAW
|
||
# ===============================
|
||
|
||
def raw_to_preview_rgb(raw_np: np.ndarray) -> np.ndarray:
|
||
"""
|
||
raw_np: (H,W,C) float32 em 0..1, C=4 ou 5.
|
||
|
||
Layout esperado (após _load_raw):
|
||
- C=4 -> [R,G,B,IR]
|
||
- C=5 -> [R,G,B,IR,NDVI] (se um dia tiver)
|
||
|
||
Usamos sempre R,G,B para o preview e devolvemos (H,W,3) uint8.
|
||
"""
|
||
assert raw_np.ndim == 3, "raw_np deve ser (H,W,C)"
|
||
|
||
H, W, C = raw_np.shape
|
||
if C not in (4, 5):
|
||
raise RuntimeError(f"Esperado C=4 ou 5, veio {C} em raw_to_preview_rgb")
|
||
|
||
# RAW já vem como [R,G,B,IR] do _load_raw
|
||
r = raw_np[..., 0].astype(np.float32)
|
||
g = raw_np[..., 1].astype(np.float32)
|
||
b = raw_np[..., 2].astype(np.float32)
|
||
|
||
# Monta RGB em (H,W,3) para trabalhar
|
||
rgb = np.stack([r, g, b], axis=-1) # (H,W,3)
|
||
|
||
# Auto white balance simples (referenciado no G)
|
||
means = rgb.reshape(-1, 3).mean(axis=0) # (3,)
|
||
ref = means[1] # canal G
|
||
gains = ref / (means + 1e-6)
|
||
gains = np.clip(gains, 0.5, 2.0)
|
||
rgb_wb = rgb * gains # broadcasting em (3,)
|
||
|
||
# Força do WB
|
||
alpha = 1.0 # teste 0.4..0.8 se quiser misturar com original
|
||
if alpha >= 0.999:
|
||
rgb = rgb_wb
|
||
else:
|
||
rgb_orig = rgb.copy()
|
||
rgb = (1 - alpha) * rgb_orig + alpha * rgb_wb
|
||
|
||
rgb = np.clip(rgb, 0.0, 1.0)
|
||
|
||
# Gamma mais suave
|
||
rgb = np.power(rgb, 1/1.9, out=rgb)
|
||
|
||
# Converte pra uint8 OpenCV-friendly
|
||
rgb8 = (rgb * 255.0).clip(0, 255).astype(np.uint8) # (H,W,3)
|
||
|
||
return rgb8
|
||
|
||
|
||
def load_rgb(path):
|
||
im = cv2.imread(path, cv2.IMREAD_COLOR)
|
||
if im is None:
|
||
raise FileNotFoundError(path)
|
||
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
||
|
||
|
||
def save_rgb(path, arr_rgb):
|
||
Image.fromarray(arr_rgb).save(path)
|
||
|
||
|
||
def _load_raw(path: str, raw_hw) -> np.ndarray:
|
||
ext = os.path.splitext(path)[1].lower()
|
||
if ext == ".npz":
|
||
z = np.load(path)
|
||
key = list(z.keys())[0]
|
||
arr = z[key]
|
||
elif ext == ".npy":
|
||
arr = np.load(path)
|
||
elif ext == ".raw":
|
||
if raw_hw is None:
|
||
raise RuntimeError(
|
||
f"Para ler .raw preciso de raw_hw=(H,W). "
|
||
f"Passe raw_hw ao criar o dataset."
|
||
)
|
||
H, W = raw_hw
|
||
size_bytes = os.path.getsize(path)
|
||
mosa_bytes = H * W # mosaico uint8
|
||
raw4_bytes = 4 * H * W * 4 # 4 canais * H * W * 4 bytes (float32)
|
||
if size_bytes == mosa_bytes:
|
||
# MODO ANTIGO: mosaico uint8
|
||
arr = np.fromfile(path, dtype=np.uint8)
|
||
raw2d = arr.reshape(H, W)
|
||
if (H % 2) != 0 or (W % 2) != 0:
|
||
raise RuntimeError(f"raw_hw deve ser par em H e W p/ mosaico 2x2 (R,G,IR,B), veio H={H}, W={W}")
|
||
H2, W2 = H // 2, W // 2
|
||
r_sub = raw2d[0::2, 0::2]
|
||
g_sub = raw2d[0::2, 1::2]
|
||
ir_sub = raw2d[1::2, 0::2]
|
||
b_sub = raw2d[1::2, 1::2]
|
||
from PIL import Image
|
||
def upsample(ch_2d: np.ndarray) -> np.ndarray:
|
||
im = Image.fromarray(ch_2d) # uint8
|
||
im = im.resize((W, H), resample=Image.BILINEAR)
|
||
return np.array(im)
|
||
r_full = upsample(r_sub)
|
||
g_full = upsample(g_sub)
|
||
ir_full = upsample(ir_sub)
|
||
b_full = upsample(b_sub)
|
||
# (H, W, 4) uint8 [R,G,IR,B]
|
||
arr = np.stack([r_full, g_full, ir_full, b_full], axis=-1).astype(np.uint8)
|
||
# REORGANIZA PARA CONTRATO INTERNO [R,G,B,IR]
|
||
arr = arr[..., [0, 1, 3, 2]]
|
||
elif size_bytes == raw4_bytes:
|
||
# NOVO MODO: RAW4 float32 salvo pelo normalize
|
||
arr_f32 = np.fromfile(path, dtype=np.float32)
|
||
raw4 = arr_f32.reshape(4, H, W) # (C,H,W) [R,G,IR,B]
|
||
arr = np.transpose(raw4, (1, 2, 0)) # (H,W,4) [R,G,IR,B]
|
||
# REORGANIZA PARA CONTRATO INTERNO: [R,G,B,IR]
|
||
arr = arr[..., [0, 1, 3, 2]]
|
||
else:
|
||
raise RuntimeError(
|
||
f"Tamanho inesperado em {path}: {size_bytes} bytes "
|
||
f"(esperado {mosa_bytes} ou {raw4_bytes})"
|
||
)
|
||
else:
|
||
raise RuntimeError(f"Extensão não suportada para RAW: {ext} ({path})")
|
||
# Até aqui queremos (H,W,4) base
|
||
if arr.ndim != 3:
|
||
raise RuntimeError(f"RAW precisa ser 3D, veio {arr.shape} em {path}")
|
||
if arr.shape[0] == 4 and arr.shape[-1] != 4:
|
||
arr = np.transpose(arr, (1, 2, 0)) # (H,W,4)
|
||
if arr.shape[-1] != 4:
|
||
raise RuntimeError(f"Esperava 4 canais base, veio shape={arr.shape} em {path}")
|
||
return arr # (H,W,4)
|
||
|
||
def _scale_to_float01(raw: np.ndarray, max_value = None) -> np.ndarray:
|
||
if raw.dtype == np.uint16:
|
||
mv = float(max_value) if max_value is not None else 65535.0
|
||
elif raw.dtype == np.uint8:
|
||
mv = float(max_value) if max_value is not None else 255.0
|
||
else:
|
||
# float já vem normalizado ou não; se passar max_value, aplica
|
||
mv = float(max_value) if max_value is not None else None
|
||
raw_f = raw.astype(np.float32)
|
||
if mv is not None and mv > 0:
|
||
raw_f = raw_f / mv
|
||
# clip pra não explodir
|
||
raw_f = np.clip(raw_f, 0.0, 1.0)
|
||
return raw_f
|
||
|
||
|
||
def save_raw_any(path, arr):
|
||
arr = np.asarray(arr)
|
||
|
||
if arr.ndim == 3 and arr.shape[-1] == 4:
|
||
# arr está no contrato interno [R,G,B,IR]
|
||
arr_cam = internal_rgbir_to_camera_raw4(arr) # [R,G,IR,B]
|
||
arr_cam = np.transpose(arr_cam, (2, 0, 1)) # (4,H,W)
|
||
arr_cam.astype(np.float32).tofile(path)
|
||
else:
|
||
arr.tofile(path)
|
||
|
||
|
||
# ===============================
|
||
# Helpers de diretório / mapeamento
|
||
# ===============================
|
||
|
||
def list_groups(root):
|
||
"""
|
||
Lista grupos válidos para o NOVO fluxo:
|
||
- group/<g>/raws
|
||
- group/<g>/masks
|
||
(previews/images não são mais obrigatórios)
|
||
"""
|
||
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
|
||
has_masks = os.path.isdir(os.path.join(gdir, "masks"))
|
||
has_raws = os.path.isdir(os.path.join(gdir, "raws"))
|
||
if has_masks and has_raws:
|
||
grupos.append(name)
|
||
return grupos
|
||
|
||
|
||
def map_by_base_priorizando_png(msk_dir, exts):
|
||
"""Mapeia arquivos por base (prioriza .png)."""
|
||
by_base = {}
|
||
if not os.path.isdir(msk_dir):
|
||
return by_base
|
||
for fname in os.listdir(msk_dir):
|
||
f_lower = fname.lower()
|
||
if not f_lower.endswith(exts):
|
||
continue
|
||
base, ext = os.path.splitext(fname)
|
||
cand = os.path.join(msk_dir, 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):
|
||
by_base = {}
|
||
if not os.path.isdir(raw_dir):
|
||
return by_base
|
||
for fname in os.listdir(raw_dir):
|
||
f_lower = fname.lower()
|
||
if not f_lower.endswith(RAW_EXTS):
|
||
continue
|
||
base, _ = os.path.splitext(fname)
|
||
by_base[base] = os.path.join(raw_dir, fname)
|
||
return by_base
|
||
|
||
|
||
def ensure_aug_dirs(group_name=None, use_masks2=False, use_raw=False, is_preview_mode=True):
|
||
"""
|
||
Cria diretórios de saída para o grupo ou modo antigo.
|
||
|
||
is_preview_mode:
|
||
True -> saídas em 'previews' + 'raws'
|
||
False -> saídas em 'images' + (sem raws)
|
||
"""
|
||
if group_name:
|
||
base = os.path.join(AUG_GROUP_ROOT, group_name)
|
||
if is_preview_mode:
|
||
img_out = os.path.join(base, "previews")
|
||
else:
|
||
img_out = os.path.join(base, "images")
|
||
msk_out = os.path.join(base, "masks")
|
||
msk2_out = os.path.join(base, "masks2") if use_masks2 else None
|
||
raw_out = os.path.join(base, "raws") if use_raw else None
|
||
else:
|
||
# modo legacy
|
||
img_out = AUG_OLD_IMG
|
||
msk_out = AUG_OLD_MSK
|
||
msk2_out = AUG_OLD_MSK2 if use_masks2 else None
|
||
raw_out = None
|
||
|
||
garantir_dir(img_out)
|
||
garantir_dir(msk_out)
|
||
if use_masks2 and msk2_out:
|
||
garantir_dir(msk2_out)
|
||
if use_raw and raw_out:
|
||
garantir_dir(raw_out)
|
||
return img_out, msk_out, msk2_out, raw_out
|
||
|
||
|
||
# ===============================
|
||
# Núcleo: augmentação de um sample
|
||
# ===============================
|
||
|
||
def augment_sample(
|
||
raw_path,
|
||
msk_path,
|
||
img_out_dir,
|
||
msk_out_dir,
|
||
copies,
|
||
msk2_path=None,
|
||
msk2_out_dir=None,
|
||
raw_out_dir=None,
|
||
aug_suffix="aug",
|
||
):
|
||
"""
|
||
Faz a augmentação a partir de UM RAW4 + máscara (e opcionalmente máscara2).
|
||
|
||
Passos:
|
||
1) Carrega máscara (para saber HxW).
|
||
2) Carrega RAW4 float32 usando HxW da máscara.
|
||
3) Gera preview RGB a partir do RAW4 (para servir de "image" no Albumentations).
|
||
4) Aplica GEOMETRIA (mesma em preview, mask, mask2 e raw).
|
||
5) Aplica BLUR/MOTION BLUR no RAW geométrico, canal a canal, com o mesmo kernel.
|
||
6) Gera preview final a partir do RAW blurado.
|
||
7) Salva:
|
||
- preview_n em img_out_dir
|
||
- mask_n em msk_out_dir
|
||
- raw_n em raw_out_dir
|
||
- mask2_n em msk2_out_dir (se houver)
|
||
"""
|
||
base_raw, raw_ext = os.path.splitext(os.path.basename(raw_path))
|
||
base = base_raw
|
||
|
||
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
|
||
if base_msk != base:
|
||
print(f"[WARN] RAW '{base}' e máscara '{base_msk}' com bases diferentes, mas seguindo mesmo assim.")
|
||
|
||
msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
|
||
|
||
# 1) Máscaras
|
||
msk = load_rgb(msk_path)
|
||
msk2 = load_rgb(msk2_path) if msk2_path else None
|
||
|
||
# 2) RAW4
|
||
raw = _load_raw(raw_path, raw_hw=(RESOLUCAO[1], RESOLUCAO[0])) # (H,W,4) uint8
|
||
raw = _scale_to_float01(raw) # float32 0..1
|
||
|
||
# 3) Preview inicial derivado do RAW (antes de geom/blur, só referência)
|
||
img = raw_to_preview_rgb(raw)
|
||
|
||
gen = 0
|
||
for i in range(copies):
|
||
# 4) GEOMETRIA: mesma transformação em img, mask, mask2 e raw
|
||
if msk2 is not None and msk2_out_dir:
|
||
aug_geo = train_geo(image=img, mask=msk, mask2=msk2, raw=raw)
|
||
else:
|
||
aug_geo = train_geo(image=img, mask=msk, raw=raw)
|
||
|
||
msk_g = aug_geo["mask"]
|
||
raw_g = aug_geo.get("raw", None)
|
||
msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None
|
||
|
||
# 5) BLUR FÍSICO: aplica (ou não) o MESMO blur no RAW geométrico
|
||
raw_b = apply_coherent_blur_raw(raw_g)
|
||
|
||
# 6) Preview final derivado do RAW blurado
|
||
img_b = raw_to_preview_rgb(raw_b)
|
||
|
||
# 7) Salva cópias
|
||
out_img = os.path.join(img_out_dir, f"{base}_{aug_suffix}_{i:02d}.png")
|
||
out_msk = os.path.join(msk_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk_ext}")
|
||
save_rgb(out_img, img_b)
|
||
save_rgb(out_msk, msk_g)
|
||
|
||
if msk2_g is not None and msk2_out_dir:
|
||
out_msk2 = os.path.join(msk2_out_dir, f"{base}_{aug_suffix}_{i:02d}{msk2_ext}")
|
||
save_rgb(out_msk2, msk2_g)
|
||
|
||
if raw_b is not None and raw_out_dir:
|
||
out_raw = os.path.join(raw_out_dir, f"{base}_{aug_suffix}_{i:02d}{raw_ext}")
|
||
save_raw_any(out_raw, raw_b)
|
||
|
||
gen += 1
|
||
|
||
return gen
|
||
|
||
|
||
# ===============================
|
||
# Processamento por grupo
|
||
# ===============================
|
||
|
||
def process_group(group_name, copies, limit=None, seed=42, aug_suffix="aug"):
|
||
"""
|
||
Processa um grupo único usando APENAS:
|
||
- group/<g>/raws
|
||
- group/<g>/masks (+ masks2 opcional)
|
||
|
||
Previews de entrada são ignoradas.
|
||
As de saída são geradas do RAW4 (geom + blur).
|
||
"""
|
||
gdir = os.path.join(ORIG_GROUP_ROOT, group_name)
|
||
|
||
raw_dir = os.path.join(gdir, "raws")
|
||
msk_dir = os.path.join(gdir, "masks")
|
||
msk2_dir = os.path.join(gdir, "masks2")
|
||
|
||
if not (os.path.isdir(raw_dir) and os.path.isdir(msk_dir)):
|
||
print(f"[WARN] Grupo '{group_name}' inválido (precisa de 'raws' e 'masks'). Pulando.")
|
||
return 0
|
||
|
||
use_raw = True
|
||
use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
|
||
|
||
raw_map = map_raws_by_base(raw_dir)
|
||
items = sorted(raw_map.items())
|
||
|
||
if limit is not None and limit > 0:
|
||
rng = np.random.default_rng(seed)
|
||
if limit < len(items):
|
||
idx = rng.choice(len(items), size=limit, replace=False)
|
||
idx = sorted(idx.tolist())
|
||
items = [items[i] for i in idx]
|
||
|
||
msk_map = map_by_base_priorizando_png(msk_dir, MSK_EXTS)
|
||
msk2_map = map_by_base_priorizando_png(msk2_dir, MSK2_EXTS) if use_masks2 else {}
|
||
|
||
img_out_dir, msk_out_dir, msk2_out_dir, raw_out_dir = ensure_aug_dirs(
|
||
group_name,
|
||
use_masks2=use_masks2,
|
||
use_raw=use_raw,
|
||
is_preview_mode=True
|
||
)
|
||
|
||
count = 0
|
||
for base, raw_file in items:
|
||
msk_file = msk_map.get(base)
|
||
if not msk_file:
|
||
print(f"[WARN] [{group_name}] Máscara não encontrada para RAW {base}, pulando.")
|
||
continue
|
||
|
||
msk2_file = msk2_map.get(base) if use_masks2 else None
|
||
if use_masks2 and not msk2_file:
|
||
print(f"[WARN] [{group_name}] mask2 não encontrada para {base}, gerando só RAW+preview+mask.")
|
||
|
||
try:
|
||
count += augment_sample(
|
||
raw_path=raw_file,
|
||
msk_path=msk_file,
|
||
img_out_dir=img_out_dir,
|
||
msk_out_dir=msk_out_dir,
|
||
copies=copies,
|
||
msk2_path=msk2_file,
|
||
msk2_out_dir=msk2_out_dir,
|
||
raw_out_dir=raw_out_dir,
|
||
aug_suffix=aug_suffix,
|
||
)
|
||
except Exception as e:
|
||
print(f"[ERRO] [{group_name}] {base}: {e}")
|
||
print(f"[OK] Grupo '{group_name}' → {count} amostras geradas.")
|
||
return count
|
||
|
||
|
||
# ===============================
|
||
# Fallback legacy (sem RAW)
|
||
# ===============================
|
||
|
||
def process_legacy(copies):
|
||
"""
|
||
Fallback: modo sem grupos (original/images e original/masks).
|
||
Mantido por compat, mas não usa RAW.
|
||
"""
|
||
if not (os.path.isdir(ORIG_OLD_IMG) and os.path.isdir(ORIG_OLD_MSK)):
|
||
print("[WARN] Modo legacy não encontrado. Nada a fazer.")
|
||
return 0
|
||
|
||
imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||
msk_map = map_by_base_priorizando_png(ORIG_OLD_MSK, MSK_EXTS)
|
||
use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
|
||
msk2_map = map_by_base_priorizando_png(ORIG_OLD_MSK2, MSK2_EXTS) if use_masks2 else {}
|
||
|
||
img_out_dir, msk_out_dir, msk2_out_dir, _ = ensure_aug_dirs(
|
||
group_name=None,
|
||
use_masks2=use_masks2,
|
||
use_raw=False,
|
||
is_preview_mode=False
|
||
)
|
||
|
||
count = 0
|
||
for img_file in sorted(imgs):
|
||
base, img_ext = os.path.splitext(img_file)
|
||
msk_file = msk_map.get(base)
|
||
if not msk_file:
|
||
print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
|
||
continue
|
||
msk2_file = msk2_map.get(base) if use_masks2 else None
|
||
if use_masks2 and not msk2_file:
|
||
print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.")
|
||
|
||
try:
|
||
img_path = os.path.join(ORIG_OLD_IMG, img_file)
|
||
img = load_rgb(img_path)
|
||
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_file))
|
||
msk = load_rgb(msk_file)
|
||
msk2 = load_rgb(msk2_file) if msk2_file else None
|
||
|
||
for i in range(copies):
|
||
if msk2 is not None and msk2_out_dir:
|
||
aug_geo = train_geo(image=img, mask=msk, mask2=msk2)
|
||
else:
|
||
aug_geo = train_geo(image=img, mask=msk)
|
||
|
||
img_g = aug_geo["image"]
|
||
msk_g = aug_geo["mask"]
|
||
msk2_g = aug_geo.get("mask2", None) if (msk2 is not None and msk2_out_dir) else None
|
||
|
||
out_img = os.path.join(img_out_dir, f"{base}_aug_{i:02d}{img_ext}")
|
||
out_msk = os.path.join(msk_out_dir, f"{base_msk}_aug_{i:02d}{msk_ext}")
|
||
save_rgb(out_img, img_g)
|
||
save_rgb(out_msk, msk_g)
|
||
|
||
if msk2_g is not None and msk2_out_dir:
|
||
msk2_ext = os.path.splitext(msk2_file)[1]
|
||
out_msk2 = os.path.join(msk2_out_dir, f"{base_msk}_aug_{i:02d}{msk2_ext}")
|
||
save_rgb(out_msk2, msk2_g)
|
||
|
||
count += 1
|
||
|
||
except Exception as e:
|
||
print(f"[ERRO] (legacy) {img_file}: {e}")
|
||
print(f"[OK] Legacy → {count} amostras geradas.")
|
||
return count
|
||
|
||
def internal_rgbir_to_camera_raw4(arr: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Converte do contrato interno (H,W,4) [R,G,B,IR]
|
||
para o contrato bruto da câmera (H,W,4) [R,G,IR,B].
|
||
"""
|
||
if arr is None:
|
||
return None
|
||
if arr.ndim != 3 or arr.shape[-1] != 4:
|
||
raise RuntimeError(f"Esperado array (H,W,4), veio {arr.shape}")
|
||
return arr[..., [0, 1, 3, 2]]
|
||
|
||
|
||
# ===============================
|
||
# main
|
||
# ===============================
|
||
|
||
def main(copies=5, groups_csv=None, limit=None, seed=42, suffix="aug"):
|
||
total = 0
|
||
if os.path.isdir(ORIG_GROUP_ROOT):
|
||
grupos = list_groups(ORIG_GROUP_ROOT)
|
||
if groups_csv:
|
||
want = {g.strip() for g in groups_csv.split(",") if g.strip()}
|
||
grupos = [g for g in grupos if g in want]
|
||
if not grupos:
|
||
print("[WARN] Nenhum grupo válido encontrado após filtro.")
|
||
if not grupos:
|
||
print("[WARN] Nenhum grupo encontrado em original/group. Tentando modo legacy...")
|
||
total += process_legacy(copies)
|
||
else:
|
||
print(f"Grupos encontrados: {', '.join(grupos)}")
|
||
for g in grupos:
|
||
total += process_group(g, copies, limit=limit, seed=seed, aug_suffix=suffix)
|
||
else:
|
||
total += process_legacy(copies)
|
||
|
||
print(f"\nAugmentation completed! Total: {total} amostras geradas.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
ap = argparse.ArgumentParser(description="Augmentação por grupos usando RAW4 (preview gerado do RAW).")
|
||
ap.add_argument("--copies", type=int, default=5, help="Número de cópias augmentadas por imagem (default=5).")
|
||
ap.add_argument("--groups", type=str, default=None, help="Lista de grupos separados por vírgula (ex: chao,erva_cana).")
|
||
ap.add_argument("--limit", type=int, default=None, help="Quantidade máxima de imagens originais do grupo a augmentar.")
|
||
ap.add_argument("--seed", type=int, default=42, help="Seed para seleção reproduzível quando usar --limit.")
|
||
ap.add_argument("--suffix", type=str, default="aug", help="Sufixo usado no nome dos arquivos gerados.")
|
||
args = ap.parse_args()
|
||
main(
|
||
copies=args.copies,
|
||
groups_csv=args.groups,
|
||
limit=args.limit,
|
||
seed=args.seed,
|
||
suffix=args.suffix,
|
||
)
|