agrobot_base/Python/OAK/datasets/_2_add_new_images.py

# copy_pairs.py
import json
import os
import cv2
import csv
import shutil
import argparse

# ⚙️ Configurações (MODELO via config.json, pode sobrescrever na CLI)
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config.get("camera", ".")

# Pastas
PASTA_NEW_IMAGES = os.path.join(MODELO, "dataset", "original", "new_images")
PASTA_NEW_MASKS  = os.path.join(MODELO, "dataset", "original", "new_masks")

PASTA_FINAL_IMAGES = os.path.join(MODELO, "dataset", "original", "images")
PASTA_FINAL_MASKS  = os.path.join(MODELO, "dataset", "original", "masks")

# Extensões aceitas
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
EXT_MASKS   = (".png", ".jpg", ".jpeg")  # prioridade será .png quando houver

MANIFESTO = "manifest.csv"

# ===================================================

def garantir_pasta(p):
    os.makedirs(p, exist_ok=True)

def nome_disponivel(dest_dir, base_name, ext):
    """
    Retorna um caminho disponível em dest_dir, garantindo unicidade com sufixos _001, _002, ...
    """
    cand = os.path.join(dest_dir, base_name + ext)
    if not os.path.exists(cand):
        return cand

    i = 1
    while True:
        cand = os.path.join(dest_dir, f"{base_name}_{i:03d}{ext}")
        if not os.path.exists(cand):
            return cand
        i += 1

def mapear_masks_por_base(pasta_masks):
    """
    Cria um dicionário {base: caminho_mask} escolhendo .png com prioridade
    quando houver múltiplas máscaras para o mesmo base.
    """
    mapa = {}
    for nome in os.listdir(pasta_masks):
        lower = nome.lower()
        if not lower.endswith(EXT_MASKS):
            continue
        base, ext = os.path.splitext(nome)
        caminho = os.path.join(pasta_masks, nome)
        # Prioriza PNG se houver mais de uma
        if base not in mapa:
            mapa[base] = caminho
        else:
            atual_ext = os.path.splitext(mapa[base])[1].lower()
            if atual_ext != ".png" and ext.lower() == ".png":
                mapa[base] = caminho
    return mapa

def copiar_com_pareamento(caminho_img_src, caminho_mask_src, dest_img_dir, dest_mask_dir):
    """
    Copia imagem e máscara mantendo mesmo nome-base. Resolve colisões com sufixo _NNN.
    Retorna (dst_img_path, dst_mask_path).
    """
    garantir_pasta(dest_img_dir)
    garantir_pasta(dest_mask_dir)

    base_img_src = os.path.splitext(os.path.basename(caminho_img_src))[0]
    img_ext = os.path.splitext(caminho_img_src)[1].lower()
    mask_ext = os.path.splitext(caminho_mask_src)[1].lower()

    # 1) escolhe nome disponível para a imagem
    dst_img_path = nome_disponivel(dest_img_dir, base_img_src, img_ext)
    new_base = os.path.splitext(os.path.basename(dst_img_path))[0]

    # 2) tenta a máscara com o mesmo base
    dst_mask_path = os.path.join(dest_mask_dir, new_base + mask_ext)

    # 3) se já existir uma máscara com esse nome, gera um base novo e sincroniza a imagem
    if os.path.exists(dst_mask_path):
        dst_mask_path = nome_disponivel(dest_mask_dir, new_base, mask_ext)
        new_base = os.path.splitext(os.path.basename(dst_mask_path))[0]
        # sincroniza imagem com o mesmo base
        dst_img_path = os.path.join(dest_img_dir, new_base + img_ext)
        if os.path.exists(dst_img_path):
            dst_img_path = nome_disponivel(dest_img_dir, new_base, img_ext)

    # 4) copia
    shutil.copy2(caminho_img_src, dst_img_path)
    shutil.copy2(caminho_mask_src, dst_mask_path)

    print(f"[COPIADO] {os.path.basename(dst_img_path)}  |  {os.path.basename(dst_mask_path)}")
    return dst_img_path, dst_mask_path

def ler_dim(caminho_img):
    img = cv2.imread(caminho_img, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise RuntimeError(f"Erro ao abrir: {caminho_img}")
    h, w = img.shape[:2]
    return (h, w)

def processar_copias(so_mov=False, manifesto_csv=None, validar_tamanho=True, estrito=False):
    """
    - so_mov=False: copia (mantém em new_*). True: move (remove de new_* após copiar).
    - validar_tamanho=True: avisa se (w,h) imagem != (w,h) máscara; estrito=True -> pula nesses casos.
    """
    garantir_pasta(PASTA_NEW_IMAGES)
    garantir_pasta(PASTA_NEW_MASKS)
    garantir_pasta(PASTA_FINAL_IMAGES)
    garantir_pasta(PASTA_FINAL_MASKS)

    mapa_masks = mapear_masks_por_base(PASTA_NEW_MASKS)

    registros = []
    total, copiados, pulados, sem_mask, erros, dim_mismatch = 0, 0, 0, 0, 0, 0

    for nome in os.listdir(PASTA_NEW_IMAGES):
        if not nome.lower().endswith(EXT_IMAGENS):
            continue
        total += 1
        caminho_img = os.path.join(PASTA_NEW_IMAGES, nome)
        base, _ = os.path.splitext(nome)

        caminho_mask = mapa_masks.get(base)
        if not caminho_mask:
            sem_mask += 1
            print(f"[SKIP] Sem máscara correspondente para: {nome}")
            continue

        try:
            if validar_tamanho:
                try:
                    hi, wi = ler_dim(caminho_img)
                    hm, wm = ler_dim(caminho_mask)
                    if (hi, wi) != (hm, wm):
                        dim_mismatch += 1
                        msg = f"[AVISO] Dimensões diferentes (img {wi}x{hi} vs mask {wm}x{hm}) em base '{base}'"
                        if estrito:
                            print(msg + "  → pulando.")
                            pulados += 1
                            continue
                        else:
                            print(msg + "  → copiando mesmo assim.")
                except Exception as e_dim:
                    print(f"[AVISO] Falha ao validar dimensões: {e_dim}  → copiando mesmo assim.")

            dst_img, dst_mask = copiar_com_pareamento(
                caminho_img, caminho_mask, PASTA_FINAL_IMAGES, PASTA_FINAL_MASKS
            )
            copiados += 1
            registros.append([caminho_img, caminho_mask, dst_img, dst_mask])

            if so_mov:
                try:
                    os.remove(caminho_img)
                except Exception:
                    pass
                try:
                    os.remove(caminho_mask)
                except Exception:
                    pass

        except Exception as e:
            erros += 1
            print(f"[ERRO] {nome}: {e}")

    # Manifesto
    if manifesto_csv and registros:
        with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
            w = csv.writer(f)
            w.writerow(["src_image", "src_mask", "dst_image", "dst_mask"])
            w.writerows(registros)
        print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")

    print(f"\nResumo: total_imgs={total} | copiados={copiados} | pulados={pulados} | sem_mask={sem_mask} | "
          f"dim_mismatch={dim_mismatch} | erros={erros}")

def build_cli():
    ap = argparse.ArgumentParser(
        description="Copia (ou move) pares IMG+MASK de new_* para images/masks com nomes únicos e pareados."
    )
    ap.add_argument("--move", action="store_true", help="Move em vez de copiar (remove de new_* após copiar).")
    ap.add_argument("--manifest", default=MANIFESTO, help="CSV de manifesto a gerar ('' para não gerar).")
    ap.add_argument("--no-validate", action="store_true", help="Não validar dimensões de IMG e MASK.")
    ap.add_argument("--strict", action="store_true", help="Se validar dimensões e forem diferentes, pular o par.")
    return ap

if __name__ == "__main__":
    ap = build_cli()
    args = ap.parse_args()

    manifesto_csv = None if (args.manifest.strip() == "") else args.manifest
    processar_copias(
        so_mov=args.move,
        manifesto_csv=manifesto_csv,
        validar_tamanho=not args.no_validate,
        estrito=args.strict
    )