470 lines
16 KiB
Python
470 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Augmenta imagens/máscaras *por grupo*, agora com suporte a RAW.
|
|
|
|
Modos suportados:
|
|
|
|
1) Modo novo (RAW):
|
|
MODELO/dataset/original/group/<grupo>/previews
|
|
MODELO/dataset/original/group/<grupo>/raws
|
|
MODELO/dataset/original/group/<grupo>/masks
|
|
(opcional) MODELO/dataset/original/group/<grupo>/masks2
|
|
|
|
Saída:
|
|
MODELO/dataset/augmented/group/<grupo>/previews
|
|
MODELO/dataset/augmented/group/<grupo>/raws
|
|
MODELO/dataset/augmented/group/<grupo>/masks
|
|
MODELO/dataset/augmented/group/<grupo>/masks2 (se existir)
|
|
|
|
2) Modo antigo (sem RAW, só images/masks):
|
|
MODELO/dataset/original/group/<grupo>/images
|
|
MODELO/dataset/original/group/<grupo>/masks
|
|
|
|
Saída:
|
|
MODELO/dataset/augmented/group/<grupo>/images
|
|
MODELO/dataset/augmented/group/<grupo>/masks
|
|
|
|
Fallback legacy:
|
|
MODELO/dataset/original/{images,masks}
|
|
MODELO/dataset/augmented/{images,masks}
|
|
|
|
Transformações:
|
|
- Geométricas (HorizontalFlip, ShiftScaleRotate) → aplicam em preview, mask, mask2 e RAW.
|
|
- Fotométricas (brightness/contrast, gamma, blur, noise, flare, etc.) → só em preview RGB.
|
|
|
|
Uso:
|
|
python _5_augmentation.py --copies 5
|
|
python _5_augmentation.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", ".")
|
|
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)
|
|
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") # prioriza PNG
|
|
MSK2_EXTS = (".png", ".jpg", ".jpeg")
|
|
RAW_EXTS = (".raw",)
|
|
|
|
# Config de RAW (mosaico 1 canal; R G / IR B é tratado depois no loader)
|
|
RAW_DTYPE = np.uint16
|
|
RAW_CHANNELS = 1 # tratamos como grayscale (H, W)
|
|
# Se RAW tiver tamanho diferente, ajuste no loader depois.
|
|
# Aqui assumimos que altura/largura batem com preview/mask.
|
|
|
|
|
|
def garantir_dir(p):
|
|
os.makedirs(p, exist_ok=True)
|
|
|
|
|
|
# ===============================
|
|
# Pipelines de augmentations
|
|
# ===============================
|
|
|
|
# Geométricas: aplicam em preview, mask, mask2 e RAW
|
|
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', # tratamos RAW como "mask" pra NÃO sofrer fotométricas
|
|
})
|
|
|
|
# Fotométricas: só preview RGB
|
|
train_photo = A.Compose([
|
|
A.OneOf([
|
|
A.RandomBrightnessContrast(0.2, 0.2, p=1.0),
|
|
A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1.0),
|
|
A.RandomGamma(gamma_limit=(90, 110), p=1.0),
|
|
], p=0.70),
|
|
|
|
A.OneOf([
|
|
A.MotionBlur(blur_limit=3, p=1.0),
|
|
A.GaussianBlur(blur_limit=3, p=1.0),
|
|
], p=0.20),
|
|
|
|
A.OneOf([
|
|
A.GaussNoise(var_limit=(5.0, 15.0), p=1.0),
|
|
A.ImageCompression(quality_lower=50, quality_upper=85, p=1.0),
|
|
], p=0.20),
|
|
|
|
A.RandomShadow(p=0.10),
|
|
A.RandomSunFlare(p=0.10),
|
|
A.ChannelShuffle(p=0.05),
|
|
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10),
|
|
])
|
|
|
|
|
|
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_grayscale(path, like_shape=None):
|
|
"""
|
|
Carrega RAW como (H, W) em uint8 ou uint16 detectando pelo tamanho do arquivo.
|
|
"""
|
|
if like_shape is None:
|
|
raise RuntimeError("like_shape é obrigatório para RAW (preciso do HxW).")
|
|
|
|
h, w = like_shape[:2]
|
|
npx = h * w
|
|
|
|
fsize = os.path.getsize(path)
|
|
if fsize == npx: # 1 byte por pixel
|
|
dtype = np.uint8
|
|
elif fsize == npx * 2: # 2 bytes por pixel
|
|
dtype = np.uint16
|
|
else:
|
|
raise RuntimeError(
|
|
f"Tamanho inesperado para RAW {path}: {fsize} bytes "
|
|
f"(esperado {npx} (u8) ou {npx*2} (u16) para {w}x{h})."
|
|
)
|
|
|
|
data = np.fromfile(path, dtype=dtype)
|
|
if data.size != npx:
|
|
raise RuntimeError(f"RAW {path}: size={data.size} != {npx} (HxW)")
|
|
|
|
return data.reshape((h, w))
|
|
|
|
|
|
def save_raw_grayscale(path, arr):
|
|
np.asarray(arr).tofile(path)
|
|
|
|
|
|
def list_groups(root):
|
|
"""Lista grupos válidos (tem subpasta 'masks' e pelo menos 'images' ou 'previews')."""
|
|
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_imgs = os.path.isdir(os.path.join(gdir, "images")) or os.path.isdir(os.path.join(gdir, "previews"))
|
|
if has_masks and has_imgs:
|
|
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
|
|
|
|
|
|
def augment_sample(
|
|
img_path,
|
|
msk_path,
|
|
img_out_dir,
|
|
msk_out_dir,
|
|
copies,
|
|
msk2_path=None,
|
|
msk2_out_dir=None,
|
|
raw_path=None,
|
|
raw_out_dir=None,
|
|
):
|
|
base_img, img_ext = os.path.splitext(os.path.basename(img_path))
|
|
base = base_img
|
|
|
|
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
|
|
msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None
|
|
raw_ext = os.path.splitext(os.path.basename(raw_path))[1] if raw_path else None
|
|
|
|
img = load_rgb(img_path)
|
|
msk = load_rgb(msk_path)
|
|
msk2 = load_rgb(msk2_path) if msk2_path else None
|
|
|
|
if raw_path is not None and raw_out_dir is not None:
|
|
# usa shape do preview pra validar RAW
|
|
raw = load_raw_grayscale(raw_path, like_shape=img.shape)
|
|
else:
|
|
raw = None
|
|
|
|
gen = 0
|
|
for i in range(copies):
|
|
# 1) GEOM: mesma transformação em img, mask, mask2 e raw
|
|
if raw is not None:
|
|
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)
|
|
else:
|
|
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"]
|
|
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
|
|
|
|
# 2) FOTO: só preview RGB
|
|
img_p = train_photo(image=img_g)["image"]
|
|
|
|
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}_aug_{i:02d}{msk_ext}")
|
|
save_rgb(out_img, img_p)
|
|
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_{i:02d}{msk2_ext}")
|
|
save_rgb(out_msk2, msk2_g)
|
|
|
|
if raw_g is not None and raw_out_dir:
|
|
out_raw = os.path.join(raw_out_dir, f"{base}_aug_{i:02d}{raw_ext}")
|
|
save_raw_grayscale(out_raw, raw_g)
|
|
|
|
gen += 1
|
|
|
|
return gen
|
|
|
|
|
|
def process_group(group_name, copies):
|
|
"""
|
|
Processa um grupo único.
|
|
Suporta:
|
|
- group/<g>/previews + raws + masks (+ masks2)
|
|
- group/<g>/images + masks (+ masks2) [modo antigo]
|
|
"""
|
|
gdir = os.path.join(ORIG_GROUP_ROOT, group_name)
|
|
|
|
# prioridade para novo modo (previews/raws)
|
|
img_dir_previews = os.path.join(gdir, "previews")
|
|
img_dir_images = os.path.join(gdir, "images")
|
|
raw_dir = os.path.join(gdir, "raws")
|
|
msk_dir = os.path.join(gdir, "masks")
|
|
msk2_dir = os.path.join(gdir, "masks2")
|
|
|
|
use_preview_mode = os.path.isdir(img_dir_previews)
|
|
img_dir = img_dir_previews if use_preview_mode else img_dir_images
|
|
|
|
if not (os.path.isdir(img_dir) and os.path.isdir(msk_dir)):
|
|
print(f"[WARN] Grupo '{group_name}' inválido (sem images/previews ou masks). Pulando.")
|
|
return 0
|
|
|
|
use_raw = os.path.isdir(raw_dir)
|
|
use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
|
|
|
|
imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
|
msk_map = map_by_base_priorizando_png(msk_dir, MSK_EXTS)
|
|
raw_map = map_raws_by_base(raw_dir) if use_raw else {}
|
|
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=use_preview_mode
|
|
)
|
|
|
|
count = 0
|
|
for img_file in sorted(imgs):
|
|
base, _ = os.path.splitext(img_file)
|
|
msk_file = msk_map.get(base)
|
|
if not msk_file:
|
|
print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.")
|
|
continue
|
|
|
|
raw_file = raw_map.get(base) if use_raw else None
|
|
if use_raw and not raw_file:
|
|
print(f"[WARN] [{group_name}] RAW não encontrado para {img_file}, gerando só preview/mask (+mask2).")
|
|
|
|
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 {img_file}, gerando só img+mask(+raw).")
|
|
|
|
try:
|
|
count += augment_sample(
|
|
os.path.join(img_dir, img_file),
|
|
msk_file,
|
|
img_out_dir,
|
|
msk_out_dir,
|
|
copies=copies,
|
|
msk2_path=msk2_file,
|
|
msk2_out_dir=msk2_out_dir,
|
|
raw_path=raw_file,
|
|
raw_out_dir=raw_out_dir,
|
|
)
|
|
except Exception as e:
|
|
print(f"[ERRO] [{group_name}] {img_file}: {e}")
|
|
print(f"[OK] Grupo '{group_name}' → {count} amostras geradas.")
|
|
return count
|
|
|
|
|
|
def process_legacy(copies):
|
|
"""Fallback: modo sem grupos (original/images e original/masks)."""
|
|
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, _ = 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:
|
|
count += augment_sample(
|
|
os.path.join(ORIG_OLD_IMG, img_file),
|
|
msk_file,
|
|
img_out_dir,
|
|
msk_out_dir,
|
|
copies=copies,
|
|
msk2_path=msk2_file,
|
|
msk2_out_dir=msk2_out_dir,
|
|
raw_path=None,
|
|
raw_out_dir=None,
|
|
)
|
|
except Exception as e:
|
|
print(f"[ERRO] (legacy) {img_file}: {e}")
|
|
print(f"[OK] Legacy → {count} amostras geradas.")
|
|
return count
|
|
|
|
|
|
def main(copies=5, groups_csv=None):
|
|
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)
|
|
else:
|
|
total += process_legacy(copies)
|
|
|
|
print(f"\nAugmentation completed! Total: {total} amostras geradas.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ap = argparse.ArgumentParser(description="Augmentação por grupos (preview/raw/masks ou images/masks)")
|
|
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).")
|
|
args = ap.parse_args()
|
|
main(copies=args.copies, groups_csv=args.groups)
|