TKK_E32231767/Setup_Dataset/split_dataset.py

43 lines
891 B
Python

import os
import shutil
import random
base_path = "datasetfix"
output_path = "dataset_split_v2"
classes = ["A", "B", "C", "TL"]
train_ratio = 0.8
for c in classes:
src = os.path.join(base_path, c)
files = os.listdir(src)
random.shuffle(files)
split_index = int(len(files) * train_ratio)
train_files = files[:split_index]
val_files = files[split_index:]
# TRAIN
train_dst = os.path.join(output_path, "train", c)
os.makedirs(train_dst, exist_ok=True)
for f in train_files:
shutil.copy(
os.path.join(src, f),
os.path.join(train_dst, f)
)
# VALIDATION
val_dst = os.path.join(output_path, "val", c)
os.makedirs(val_dst, exist_ok=True)
for f in val_files:
shutil.copy(
os.path.join(src, f),
os.path.join(val_dst, f)
)
print("Split dataset selesai!")