189 lines
7.4 KiB
Python
189 lines
7.4 KiB
Python
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("--modelo", default=MODELO, help="Nome do modelo (base da árvore de pastas).")
|
|
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)
|