260 lines
8.8 KiB
Python
260 lines
8.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
"""
|
||
|
|
Augmenta imagens e máscaras *por grupo*.
|
||
|
|
|
||
|
|
Entrada (via config.json -> MODELO):
|
||
|
|
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
|
||
|
|
|
||
|
|
Se "original/group" não existir, faz fallback para:
|
||
|
|
MODELO/dataset/original/{images,masks}
|
||
|
|
MODELO/dataset/augmented/{images,masks}
|
||
|
|
|
||
|
|
Transf. geométricas (aplicam a img e máscara) e fotométricas (apenas imagem).
|
||
|
|
|
||
|
|
Uso:
|
||
|
|
python _3_augmentation_grouped.py --copies 5
|
||
|
|
python _3_augmentation_grouped.py --copies 5 --groups chao,chao_erva,cana
|
||
|
|
"""
|
||
|
|
import os
|
||
|
|
import json
|
||
|
|
import cv2
|
||
|
|
from PIL import Image
|
||
|
|
import albumentations as A
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
# ⚙️ Configurações
|
||
|
|
with open("config.json", "r", encoding="utf-8") as f:
|
||
|
|
config = json.load(f)
|
||
|
|
MODELO = config.get("camera", ".")
|
||
|
|
|
||
|
|
# 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")
|
||
|
|
AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images")
|
||
|
|
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
|
||
|
|
|
||
|
|
# Extensões aceitas
|
||
|
|
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
||
|
|
MSK_EXTS = (".png", ".jpg", ".jpeg") # manter prioridade PNG quando possível
|
||
|
|
|
||
|
|
def garantir_dir(p):
|
||
|
|
os.makedirs(p, exist_ok=True)
|
||
|
|
|
||
|
|
# Pipeline de augmentations
|
||
|
|
train_tf = A.Compose([
|
||
|
|
A.HorizontalFlip(p=0.5),
|
||
|
|
|
||
|
|
# Geométricas (aplicam em imagem e máscara)
|
||
|
|
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
|
||
|
|
),
|
||
|
|
|
||
|
|
# Fotométricas (somente imagem)
|
||
|
|
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),
|
||
|
|
], additional_targets={'mask':'mask'})
|
||
|
|
|
||
|
|
def load_rgb(path):
|
||
|
|
# cv2 lê BGR → converte para RGB
|
||
|
|
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 list_groups(root):
|
||
|
|
"""Lista grupos válidos (que contêm subpastas images e masks)."""
|
||
|
|
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
|
||
|
|
if os.path.isdir(os.path.join(gdir, "images")) and os.path.isdir(os.path.join(gdir, "masks")):
|
||
|
|
grupos.append(name)
|
||
|
|
return grupos
|
||
|
|
|
||
|
|
def map_masks_by_base(msk_dir):
|
||
|
|
"""Mapeia máscaras 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(MSK_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:
|
||
|
|
# mantém .png se disponível
|
||
|
|
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 ensure_aug_dirs(group_name=None):
|
||
|
|
"""Cria diretórios de saída para o grupo ou modo antigo."""
|
||
|
|
if group_name:
|
||
|
|
img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images")
|
||
|
|
msk_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks")
|
||
|
|
else:
|
||
|
|
img_out = AUG_OLD_IMG
|
||
|
|
msk_out = AUG_OLD_MSK
|
||
|
|
garantir_dir(img_out)
|
||
|
|
garantir_dir(msk_out)
|
||
|
|
return img_out, msk_out
|
||
|
|
|
||
|
|
def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies):
|
||
|
|
base_img, img_ext = os.path.splitext(os.path.basename(img_path))
|
||
|
|
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path))
|
||
|
|
|
||
|
|
# padroniza pelo base da imagem
|
||
|
|
base = base_img
|
||
|
|
|
||
|
|
img = load_rgb(img_path)
|
||
|
|
msk = load_rgb(msk_path)
|
||
|
|
|
||
|
|
gen = 0
|
||
|
|
for i in range(copies):
|
||
|
|
aug = train_tf(image=img, mask=msk)
|
||
|
|
img_aug = aug["image"]
|
||
|
|
msk_aug = aug["mask"]
|
||
|
|
|
||
|
|
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_aug)
|
||
|
|
save_rgb(out_msk, msk_aug)
|
||
|
|
gen += 1
|
||
|
|
return gen
|
||
|
|
|
||
|
|
def process_group(group_name, copies):
|
||
|
|
"""Processa um grupo único (images/masks dentro de ORIG_GROUP_ROOT/<group_name>/)."""
|
||
|
|
img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images")
|
||
|
|
msk_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks")
|
||
|
|
if not (os.path.isdir(img_dir) and os.path.isdir(msk_dir)):
|
||
|
|
print(f"[WARN] Grupo '{group_name}' inválido (sem images/masks). Pulando.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
|
||
|
|
msk_map = map_masks_by_base(msk_dir)
|
||
|
|
img_out_dir, msk_out_dir = ensure_aug_dirs(group_name)
|
||
|
|
|
||
|
|
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
|
||
|
|
try:
|
||
|
|
count += augment_pair(
|
||
|
|
os.path.join(img_dir, img_file),
|
||
|
|
msk_file,
|
||
|
|
img_out_dir,
|
||
|
|
msk_out_dir,
|
||
|
|
copies=copies
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[ERRO] [{group_name}] {img_file}: {e}")
|
||
|
|
print(f"[OK] Grupo '{group_name}' → {count} pares gerados.")
|
||
|
|
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_masks_by_base(ORIG_OLD_MSK)
|
||
|
|
img_out_dir, msk_out_dir = ensure_aug_dirs(group_name=None)
|
||
|
|
|
||
|
|
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
|
||
|
|
try:
|
||
|
|
count += augment_pair(
|
||
|
|
os.path.join(ORIG_OLD_IMG, img_file),
|
||
|
|
msk_file,
|
||
|
|
img_out_dir,
|
||
|
|
msk_out_dir,
|
||
|
|
copies=copies
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[ERRO] (legacy) {img_file}: {e}")
|
||
|
|
print(f"[OK] Legacy → {count} pares gerados.")
|
||
|
|
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:
|
||
|
|
# filtra pelos grupos desejados
|
||
|
|
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:
|
||
|
|
# sem estrutura de grupos
|
||
|
|
total += process_legacy(copies)
|
||
|
|
|
||
|
|
print(f"\nAugmentation completed! Total: {total} pares gerados.")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
ap = argparse.ArgumentParser(description="Augmentação por grupos (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)
|