agrobot_base/Python/OAK/datasets/_5_augmentation.py

#!/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_oak.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")  # manter prioridade PNG quando possível
MSK2_EXTS = (".png", ".jpg", ".jpeg")

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={
    'mask2': '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 map_masks2_by_base(msk2_dir):
    """Mapeia máscaras2 por base (prioriza .png)."""
    by_base = {}
    if not os.path.isdir(msk2_dir):
        return by_base
    for fname in os.listdir(msk2_dir):
        f_lower = fname.lower()
        if not f_lower.endswith(MSK2_EXTS):
            continue
        base, ext = os.path.splitext(fname)
        cand = os.path.join(msk2_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 ensure_aug_dirs(group_name=None, use_masks2=False):
    """Cria diretórios de saída para o grupo ou modo antigo. Se use_masks2, cria masks2."""
    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")
        msk2_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks2") if use_masks2 else None
    else:
        img_out = AUG_OLD_IMG
        msk_out = AUG_OLD_MSK
        msk2_out = AUG_OLD_MSK2 if use_masks2 else None
    garantir_dir(img_out)
    garantir_dir(msk_out)
    if use_masks2 and msk2_out:
        garantir_dir(msk2_out)
    return img_out, msk_out, msk2_out

def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies, msk2_path=None, msk2_out_dir=None):
    base_img, img_ext = os.path.splitext(os.path.basename(img_path))
    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

    # padroniza pelo base da imagem
    base = base_img

    img = load_rgb(img_path)
    msk = load_rgb(msk_path)
    msk2 = load_rgb(msk2_path) if msk2_path else None

    gen = 0
    for i in range(copies):
        if msk2 is not None and msk2_out_dir:
            aug = train_tf(image=img, mask=msk, mask2=msk2)
        else:
            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)

        if msk2 is not None and msk2_out_dir:
            msk2_aug = aug["mask2"]
            out_msk2 = os.path.join(msk2_out_dir, f"{base}_aug_{i:02d}{msk2_ext}")
            save_rgb(out_msk2, msk2_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")
    msk2_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks2")
    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)
    use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir)
    msk2_map = map_masks2_by_base(msk2_dir) if use_masks2 else {}
    img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name, use_masks2=use_masks2)

    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
        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.")
        try:
            count += augment_pair(
                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
            )
        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)
    use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2)
    msk2_map = map_masks2_by_base(ORIG_OLD_MSK2) if use_masks2 else {}
    img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name=None, use_masks2=use_masks2)

    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_pair(
                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
            )
        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)