61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
|
import os
|
||
|
|
import shutil
|
||
|
|
import random
|
||
|
|
|
||
|
|
# === CONFIGURAÇÕES ===
|
||
|
|
RESOLUCAO_BASE = "512x512"
|
||
|
|
PASTA_ORIGEM = os.path.join("dataset", RESOLUCAO_BASE)
|
||
|
|
PASTA_DESTINO = os.path.join("dataset", "split")
|
||
|
|
|
||
|
|
PERCENT_TRAIN = 0.7
|
||
|
|
PERCENT_VAL = 0.2
|
||
|
|
PERCENT_TEST = 0.1
|
||
|
|
|
||
|
|
SEED = 42
|
||
|
|
random.seed(SEED)
|
||
|
|
|
||
|
|
# === Coleta imagens ===
|
||
|
|
pasta_rgb = os.path.join(PASTA_ORIGEM, "images")
|
||
|
|
pasta_masks = os.path.join(PASTA_ORIGEM, "masks")
|
||
|
|
|
||
|
|
arquivos = sorted([f for f in os.listdir(pasta_rgb) if f.endswith(".jpg") or f.endswith(".jpeg")])
|
||
|
|
|
||
|
|
# Embaralha
|
||
|
|
random.shuffle(arquivos)
|
||
|
|
|
||
|
|
# Divide
|
||
|
|
total = len(arquivos)
|
||
|
|
n_train = int(total * PERCENT_TRAIN)
|
||
|
|
n_val = int(total * PERCENT_VAL)
|
||
|
|
|
||
|
|
arquivos_train = arquivos[:n_train]
|
||
|
|
arquivos_val = arquivos[n_train:n_train+n_val]
|
||
|
|
arquivos_test = arquivos[n_train+n_val:]
|
||
|
|
|
||
|
|
conjuntos = {
|
||
|
|
"train": arquivos_train,
|
||
|
|
"val": arquivos_val,
|
||
|
|
"test": arquivos_test
|
||
|
|
}
|
||
|
|
|
||
|
|
# === Função auxiliar ===
|
||
|
|
def copiar(imagens, conjunto):
|
||
|
|
path_img_dest = os.path.join(PASTA_DESTINO, conjunto, "images")
|
||
|
|
path_mask_dest = os.path.join(PASTA_DESTINO, conjunto, "masks")
|
||
|
|
os.makedirs(path_img_dest, exist_ok=True)
|
||
|
|
os.makedirs(path_mask_dest, exist_ok=True)
|
||
|
|
|
||
|
|
for nome in imagens:
|
||
|
|
nome_mask = nome.replace(".jpg", ".png").replace(".jpeg", ".png")
|
||
|
|
if not os.path.exists(nome_mask):
|
||
|
|
continue
|
||
|
|
shutil.copy2(os.path.join(pasta_rgb, nome), os.path.join(path_img_dest, nome))
|
||
|
|
shutil.copy2(os.path.join(pasta_masks, nome_mask), os.path.join(path_mask_dest, nome_mask))
|
||
|
|
|
||
|
|
# === Executa cópia ===
|
||
|
|
for conjunto, lista in conjuntos.items():
|
||
|
|
print(f"[{conjunto}] {len(lista)} arquivos")
|
||
|
|
copiar(lista, conjunto)
|
||
|
|
|
||
|
|
print("\n✅ Dataset dividido com sucesso!")
|