Spaces:
Sleeping
Sleeping
File size: 1,942 Bytes
097f176 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import os
import shutil
from sklearn.model_selection import train_test_split
def split_dataset(images_dir, labels_dir, val_images_dir, val_labels_dir, val_ratio=0.2):
"""
Divide o conjunto de dados em treino e validação, excluindo imagens "_rotated".
:param images_dir: Diretório das imagens originais.
:param labels_dir: Diretório das labels originais.
:param val_images_dir: Diretório para salvar imagens de validação.
:param val_labels_dir: Diretório para salvar labels de validação.
:param val_ratio: Proporção de imagens a serem usadas na validação.
"""
# Cria pastas de validação, se não existirem
os.makedirs(val_images_dir, exist_ok=True)
os.makedirs(val_labels_dir, exist_ok=True)
# Lista de imagens normais (sem _rotated)
normal_images = [f for f in os.listdir(images_dir) if f.endswith(".jpg") and "_rotated" not in f]
# Associa labels existentes
normal_images = [f for f in normal_images if os.path.exists(os.path.join(labels_dir, f.replace(".jpg", ".txt")))]
# Divide em treino e validação
train_images, val_images = train_test_split(normal_images, test_size=val_ratio, random_state=42)
# Move as imagens e labels para o diretório de validação
for image in val_images:
label = image.replace(".jpg", ".txt")
# Move a imagem
shutil.move(os.path.join(images_dir, image), os.path.join(val_images_dir, image))
# Move a label
shutil.move(os.path.join(labels_dir, label), os.path.join(val_labels_dir, label))
print(f"Divisão concluída! {len(val_images)} imagens movidas para validação.")
# Exemplos de uso
split_dataset(
images_dir="dataset/images/train",
labels_dir="dataset/labels/train",
val_images_dir="dataset/images/val",
val_labels_dir="dataset/labels/val",
val_ratio=0.2
)
|