108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
|
|
# === CONFIGURAÇÕES ===
|
|
PASTA_BASE = "dataset"
|
|
FONTE_DADOS = ["original", "augmented"] # <<-- novas fontes
|
|
LABELMAP_PATH = os.path.join(PASTA_BASE, "labelmap.txt")
|
|
|
|
RESOLUCOES = {
|
|
"512x512": (512, 512),
|
|
"768x768": (768, 768),
|
|
"384x384": (384, 384)
|
|
}
|
|
|
|
# === Função para ler o labelmap ===
|
|
def carregar_labelmap_completo(caminho):
|
|
cor_para_id = {}
|
|
id_para_nome = {}
|
|
cores_bgr = []
|
|
|
|
with open(caminho, 'r') as arquivo:
|
|
for linha in arquivo:
|
|
if linha.startswith("#") or not linha.strip():
|
|
continue
|
|
partes = linha.strip().split(':')
|
|
if len(partes) >= 2:
|
|
nome_classe, cor_rgb_str = partes[0], partes[1]
|
|
r, g, b = map(int, cor_rgb_str.split(','))
|
|
cor_bgr = (b, g, r) # Corrige para BGR
|
|
|
|
classe_id = len(cor_para_id)
|
|
cor_para_id[cor_bgr] = classe_id
|
|
cores_bgr.append(cor_bgr)
|
|
id_para_nome[classe_id] = nome_classe
|
|
|
|
return cor_para_id, cores_bgr, id_para_nome
|
|
|
|
|
|
# === Converte uma máscara RGB para índice de classes ===
|
|
def converter_mask_rgb_para_ids(img_rgb, mapa):
|
|
h, w, _ = img_rgb.shape
|
|
mask = np.zeros((h, w), dtype=np.uint8)
|
|
for cor, classe_id in mapa.items():
|
|
r, g, b = cor
|
|
cond = (img_rgb[:,:,0] == b) & (img_rgb[:,:,1] == g) & (img_rgb[:,:,2] == r)
|
|
mask[cond] = classe_id
|
|
return mask
|
|
|
|
# === Início do processamento ===
|
|
cor_para_id, _, _ = carregar_labelmap_completo(LABELMAP_PATH)
|
|
|
|
# Cria pastas de saída
|
|
for nome_res, dim in RESOLUCOES.items():
|
|
os.makedirs(os.path.join(PASTA_BASE, nome_res, "images"), exist_ok=True)
|
|
os.makedirs(os.path.join(PASTA_BASE, nome_res, "masks"), exist_ok=True)
|
|
|
|
# Processa cada fonte de dados (original + augmented)
|
|
for fonte in FONTE_DADOS:
|
|
if (not os.path.exists(os.path.join(PASTA_BASE, fonte))):
|
|
continue
|
|
pasta_rgb = os.path.join(PASTA_BASE, fonte, "images")
|
|
pasta_masks = os.path.join(PASTA_BASE, fonte, "masks")
|
|
|
|
nomes_arquivos = sorted([f for f in os.listdir(pasta_rgb) if f.endswith(".jpg") or f.endswith(".jpeg")])
|
|
total = len(nomes_arquivos)
|
|
|
|
for i, nome in enumerate(nomes_arquivos, 1):
|
|
caminho_rgb = os.path.join(pasta_rgb, nome)
|
|
caminho_mask = os.path.join(pasta_masks, nome.replace(".jpg", ".png").replace(".jpeg", ".png"))
|
|
|
|
img_rgb = cv2.imread(caminho_rgb)
|
|
if img_rgb is None:
|
|
print(f"[!] Erro ao ler imagem {nome}")
|
|
continue
|
|
|
|
# Tenta carregar a máscara RGB (se existir)
|
|
if os.path.exists(caminho_mask):
|
|
img_mask_rgb = cv2.imread(caminho_mask)
|
|
img_mask_rgb = cv2.cvtColor(img_mask_rgb, cv2.COLOR_BGR2RGB) # ← CORRIGE isso!
|
|
if img_mask_rgb is not None:
|
|
mask_ids = converter_mask_rgb_para_ids(img_mask_rgb, cor_para_id)
|
|
else:
|
|
print(f"[!] Erro ao ler máscara {caminho_mask}, ignorando.")
|
|
mask_ids = None
|
|
else:
|
|
mask_ids = None
|
|
|
|
for nome_res, dim in RESOLUCOES.items():
|
|
# Cria nomes únicos baseados na fonte
|
|
nome_saida_img = f"{fonte}_{nome}"
|
|
nome_saida_mask = nome_saida_img.replace(".jpg", ".png").replace(".jpeg", ".png")
|
|
|
|
# Redimensiona e salva imagem
|
|
img_resized = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)
|
|
path_img_saida = os.path.join(PASTA_BASE, nome_res, "images", nome_saida_img)
|
|
cv2.imwrite(path_img_saida, img_resized)
|
|
|
|
# Redimensiona e salva máscara (se existir)
|
|
if mask_ids is not None:
|
|
mask_resized = cv2.resize(mask_ids, dim, interpolation=cv2.INTER_NEAREST)
|
|
path_mask_saida = os.path.join(PASTA_BASE, nome_res, "masks", nome_saida_mask)
|
|
cv2.imwrite(path_mask_saida, mask_resized)
|
|
|
|
print(f"[{fonte}] [{i}/{total}] Redimensionado: {nome}")
|
|
|
|
print("\n✅ Concluído com sucesso! Todas as fontes foram processadas.")
|