agrobot_base/Python/OAK/datasets/_2_create_full_mask.py

import json
import os
import cv2
import csv
import shutil
import argparse
import numpy as np

# ⚙️ Configurações
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config["camera"]

# Pastas (ajuste conforme sua estrutura)
PASTA_NEW_IMAGES = os.path.join(MODELO, "dataset", "original", "new_images")  # novas fotos
PASTA_NEW_MASKS  = os.path.join(MODELO, "dataset", "original", "new_masks")   # onde salvar as novas máscaras sólidas

PASTA_FINAL_IMAGES = os.path.join(MODELO, "dataset", "original", "images")    # dataset final (imagens)
PASTA_FINAL_MASKS  = os.path.join(MODELO, "dataset", "original", "masks")     # dataset final (máscaras)

# Classe-alvo (RGB) para preencher a máscara sólida
COR_CLASSE_RGB = (128, 0, 0)  # (R, G, B)
COPIAR_IMAGENS = True

# Controle
EXT_IMAGENS = (".jpg", ".jpeg", ".png")
MASK_EXT_OUT = ".png"        # saída das máscaras sempre PNG
FORCAR_SOBRESCRITA_NEW_MASK = False   # sobrescrever máscara em new_masks se já existir
VALIDAR_DIM_MASK_EXISTENTE = True     # se existir, validar dimensões

MANIFESTO = "manifest.csv"
# ===================================================

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

def criar_mask_solida(dim_h, dim_w, cor_rgb):
    r, g, b = cor_rgb
    mask_bgr = np.zeros((dim_h, dim_w, 3), dtype=np.uint8)
    mask_bgr[:] = (b, g, r)  # OpenCV usa BGR
    return mask_bgr

def caminho_mask_new_para_img(caminho_img):
    base = os.path.splitext(os.path.basename(caminho_img))[0]
    return os.path.join(PASTA_NEW_MASKS, base + MASK_EXT_OUT)

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

def salvar_mask_solidaria(caminho_img, cor_rgb, forcar=False, validar_dim=True):
    (h, w), _ = ler_dim(caminho_img)
    caminho_mask = caminho_mask_new_para_img(caminho_img)

    if os.path.exists(caminho_mask) and not forcar:
        if validar_dim:
            m = cv2.imread(caminho_mask, cv2.IMREAD_COLOR)
            if m is None or m.shape[0] != h or m.shape[1] != w:
                print(f"[ALERTA] Máscara existente com dimensão diferente, recriando: {caminho_mask}")
            else:
                print(f"[OK] Já existe (mantida): {os.path.basename(caminho_mask)}")
                return caminho_mask
        else:
            print(f"[OK] Já existe (mantida): {os.path.basename(caminho_mask)}")
            return caminho_mask

    mask_bgr = criar_mask_solida(h, w, cor_rgb)
    garantir_pasta(os.path.dirname(caminho_mask))
    cv2.imwrite(caminho_mask, mask_bgr)
    print(f"[CRIADA] {os.path.basename(caminho_mask)}  ({w}x{h})")
    return caminho_mask

def nome_disponivel(dest_dir, base_name, ext):
    """
    Retorna um nome disponível. Se 'base_name.ext' não existir, usa ele.
    Caso exista, tenta base_name_001.ext, base_name_002.ext, ...
    """
    candidate = os.path.join(dest_dir, base_name + ext)
    if not os.path.exists(candidate):
        return candidate

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

def copiar_com_pareamento(caminho_img_src, caminho_mask_src, dest_img_dir, dest_mask_dir):
    """
    Copia imagem e máscara mantendo o mesmo nome-base (com renome em caso de conflito).
    Retorna (dst_img_path, dst_mask_path).
    """
    garantir_pasta(dest_img_dir)
    garantir_pasta(dest_mask_dir)

    # base sem extensão (da imagem)
    base = os.path.splitext(os.path.basename(caminho_img_src))[0]

    # ext de saída: imagem mantém sua extensão original; máscara = PNG
    img_ext = os.path.splitext(caminho_img_src)[1].lower()
    mask_ext = ".png"

    # escolhe nomes finais não colidentes
    dst_img_path = nome_disponivel(dest_img_dir, base, img_ext)
    new_base = os.path.splitext(os.path.basename(dst_img_path))[0]  # pode ter ganhado _001
    dst_mask_path = os.path.join(dest_mask_dir, new_base + mask_ext)

    # se máscara final já existe por acaso, procura próximo disponível e sincroniza com 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]
        # ajustar a imagem para manter o mesmo base
        dst_img_path = os.path.join(dest_img_dir, new_base + img_ext)

        # se por acaso colidiu de novo com imagem, volta a encontrar disponível
        if os.path.exists(dst_img_path):
            dst_img_path = nome_disponivel(dest_img_dir, new_base, img_ext)

    # copia arquivos
    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 processar_novas_imagens(fazer_copia_final=True, manifesto_csv=None):
    garantir_pasta(PASTA_NEW_IMAGES)
    garantir_pasta(PASTA_NEW_MASKS)
    garantir_pasta(PASTA_FINAL_IMAGES)
    garantir_pasta(PASTA_FINAL_MASKS)

    registros = []
    total, criadas_mask, copiados, puladas, erros = 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)
        try:
            # 1) criar máscara sólida em new_masks
            antes = os.path.exists(caminho_mask_new_para_img(caminho_img))
            caminho_mask_new = salvar_mask_solidaria(
                caminho_img,
                COR_CLASSE_RGB,
                forcar=FORCAR_SOBRESCRITA_NEW_MASK,
                validar_dim=VALIDAR_DIM_MASK_EXISTENTE
            )
            if caminho_mask_new and not antes:
                criadas_mask += 1

            # 2) copiar imagem + máscara para as pastas finais (com renome se necessário)
            if fazer_copia_final:
                dst_img, dst_mask = copiar_com_pareamento(
                    caminho_img, caminho_mask_new, PASTA_FINAL_IMAGES, PASTA_FINAL_MASKS
                )
                copiados += 1
                registros.append([caminho_img, caminho_mask_new, dst_img, dst_mask])
            else:
                puladas += 1

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

    # salvar manifesto (opcional)
    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} | masks_criadas={criadas_mask} | copiados={copiados} | puladas={puladas} | erros={erros}")

def build_cli():
    ap = argparse.ArgumentParser(description="Gera máscaras sólidas para novas imagens e copia para dataset final com dedup.")
    ap.add_argument("--no-copy", action="store_true", help="Não copia para as pastas finais (só cria masks em new_masks).")
    ap.add_argument("--manifest", default=MANIFESTO, help="Caminho do CSV de manifesto a gerar (ou vazio para não gerar).")
    return ap

if __name__ == "__main__":
    processar_novas_imagens(fazer_copia_final=COPIAR_IMAGENS, manifesto_csv=None)