From 26ef99809bf19e75c6e588754698ba9895b42576 Mon Sep 17 00:00:00 2001 From: Timothy Guo Date: Fri, 16 Jul 2021 09:19:47 +1000 Subject: [PATCH] Fixed issue #6 --- write_datasets.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/write_datasets.py b/write_datasets.py index 68037c9..d3bb3b9 100644 --- a/write_datasets.py +++ b/write_datasets.py @@ -1,6 +1,5 @@ import os import argparse -import random from collections import Counter ''' @@ -30,12 +29,10 @@ def write_datasets(args): number_of_images = len(images) index_valid = round(number_of_images * args.valid_pct / 100) - validfiles = random.sample(images, index_valid) - traintestfiles = list(set(images).difference(set(validfiles))) - index_train = round(len(traintestfiles) * - (args.train_pct / (100 - args.valid_pct))) - trainfiles = random.sample(traintestfiles, index_train) - testfiles = list(set(traintestfiles).difference(set(trainfiles))) + index_train = round(number_of_images * args.train_pct / 100) + validfiles = images[:index_valid] + trainfiles = images[index_valid:(index_valid+index_train)] + testfiles = images[(index_valid+index_train):] print('Number of images:', number_of_images)