62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
import os
|
|
import shutil
|
|
import random
|
|
|
|
# ⚙️ Configurações
|
|
MODELO = "oak-1"
|
|
RESOLUCAO = (384, 384)
|
|
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}")
|
|
pasta_destino = os.path.join(MODELO, "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(os.path.join(pasta_masks, 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!")
|