56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
|
import os
|
||
|
|
import random
|
||
|
|
import shutil
|
||
|
|
|
||
|
|
# === CONFIG ===
|
||
|
|
PASTA_ORIGEM_IMG = "dataset/1280x720/images"
|
||
|
|
PASTA_ORIGEM_LBL = "dataset/1280x720/labels"
|
||
|
|
PASTA_SAIDA = "yolov7/train"
|
||
|
|
|
||
|
|
SPLIT = {
|
||
|
|
"train": 0.7,
|
||
|
|
"val": 0.2,
|
||
|
|
"test": 0.1
|
||
|
|
}
|
||
|
|
|
||
|
|
# === PREPARA PASTAS ===
|
||
|
|
for tipo in SPLIT.keys():
|
||
|
|
os.makedirs(os.path.join(PASTA_SAIDA, tipo, "images"), exist_ok=True)
|
||
|
|
os.makedirs(os.path.join(PASTA_SAIDA, tipo, "labels"), exist_ok=True)
|
||
|
|
|
||
|
|
# === LISTA E EMBARALHA ===
|
||
|
|
arquivos = [f for f in os.listdir(PASTA_ORIGEM_IMG) if f.endswith(".jpg")]
|
||
|
|
random.shuffle(arquivos)
|
||
|
|
|
||
|
|
total = len(arquivos)
|
||
|
|
qt_train = int(SPLIT["train"] * total)
|
||
|
|
qt_val = int(SPLIT["val"] * total)
|
||
|
|
|
||
|
|
splits = {
|
||
|
|
"train": arquivos[:qt_train],
|
||
|
|
"val": arquivos[qt_train:qt_train+qt_val],
|
||
|
|
"test": arquivos[qt_train+qt_val:]
|
||
|
|
}
|
||
|
|
|
||
|
|
# === COPIA ===
|
||
|
|
for tipo, lista in splits.items():
|
||
|
|
for nome_img in lista:
|
||
|
|
nome_lbl = nome_img.replace(".jpg", ".txt")
|
||
|
|
|
||
|
|
origem_img = os.path.join(PASTA_ORIGEM_IMG, nome_img)
|
||
|
|
origem_lbl = os.path.join(PASTA_ORIGEM_LBL, nome_lbl)
|
||
|
|
|
||
|
|
destino_img = os.path.join(PASTA_SAIDA, tipo, "images", nome_img)
|
||
|
|
destino_lbl = os.path.join(PASTA_SAIDA, tipo, "labels", nome_lbl)
|
||
|
|
|
||
|
|
shutil.copy2(origem_img, destino_img)
|
||
|
|
|
||
|
|
if os.path.exists(origem_lbl):
|
||
|
|
shutil.copy2(origem_lbl, destino_lbl)
|
||
|
|
else:
|
||
|
|
print(f"[!] Label ausente: {nome_lbl}")
|
||
|
|
|
||
|
|
print(f"[✔] {tipo.upper()} - {len(lista)} imagens")
|
||
|
|
|
||
|
|
print("\n✅ Split do dataset para YOLOv7 concluído!")
|