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)