Spaces:
Build error
Build error
File size: 2,783 Bytes
c939ae6 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import shutil
import os
import random
import math
import yaml
def split_dataset(dataset_path,train_percentage, validation_percentage,classes=[],mode="move",output_path=""):
def shift(source,destination,mode):
if mode == "copy":
shutil.copyfile(source,destination)
else:
shutil.move(source,destination)
def diff_lists(l1,l2):
from collections import Counter
return list((Counter(l1) - Counter(l2)).elements())
if mode == "copy" and output_path == "":
raise Exception("Cannot copy files on the same directory")
if validation_percentage > train_percentage:
raise Exception("validation_percentage must be lower than train_percentage")
other_files = None
if classes == []:
classes = os.listdir(dataset_path)
else:
other_files = diff_lists(os.listdir(dataset_path),classes)
test_percentage = 1 - train_percentage
if not dataset_path.endswith("/"):
dataset_path = dataset_path + "/"
if output_path == "":
output_path = dataset_path
train_dir = output_path + "/train"
test_dir = output_path + "/test"
validation_dir = output_path + "/valid"
seed = 42
for _class in classes:
os.makedirs(train_dir + "/" + _class,exist_ok=True)
os.makedirs(test_dir + "/" + _class,exist_ok=True)
os.makedirs(validation_dir + "/" + _class,exist_ok=True)
data = sorted(os.listdir(dataset_path + _class + "/"))
random.Random(seed).shuffle(data)
data_length = len(data)
test_size = math.floor(data_length * test_percentage)
validation_size = math.floor(data_length * validation_percentage)
for i,single_data in enumerate(data):
single_data_path = dataset_path + _class + "/" + single_data
if i < test_size:
shift(single_data_path, test_dir + "/" + _class + "/" + single_data, mode)
elif test_size < i <= test_size + validation_size:
shift(single_data_path, validation_dir + "/" + _class + "/" + single_data, mode)
else:
shift(single_data_path, train_dir + "/" + _class + "/" + single_data, mode)
if mode == "move":
shutil.rmtree(dataset_path + _class)
if other_files is not None:
for file in other_files:
shift(dataset_path + file, output_path + "/" + file, mode)
os.system("unzip -n dataset.zip")
with open("params.yaml", 'r') as fd:
params = yaml.safe_load(fd)
train_percentage = params['preparation']['train_percentage']
validation_percentage = params['preparation']['validation_percentage']
split_dataset(dataset_path="dataset/yolo",train_percentage=train_percentage,validation_percentage=validation_percentage,classes=["images","labels"])
|