This commit is contained in:
Glenn Jocher 2019-09-21 02:46:16 +02:00
parent db49211d70
commit 7de6584a34
1 changed files with 21 additions and 1 deletions

View File

@ -313,8 +313,9 @@ class LoadImagesAndLabels(Dataset): # for training/testing
if cache_labels or image_weights: # cache labels for faster training if cache_labels or image_weights: # cache labels for faster training
self.labels = [np.zeros((0, 5))] * n self.labels = [np.zeros((0, 5))] * n
extract_bounding_boxes = False extract_bounding_boxes = False
create_datasubset = False
pbar = tqdm(self.label_files, desc='Reading labels') pbar = tqdm(self.label_files, desc='Reading labels')
nm, nf, ne = 0, 0, 0 # number missing, number found, number empty nm, nf, ne, ns = 0, 0, 0, 0 # number missing, number found, number empty, number datasubset
for i, file in enumerate(pbar): for i, file in enumerate(pbar):
try: try:
with open(file, 'r') as f: with open(file, 'r') as f:
@ -330,6 +331,18 @@ class LoadImagesAndLabels(Dataset): # for training/testing
self.labels[i] = l self.labels[i] = l
nf += 1 # file found nf += 1 # file found
# Create subdataset (a smaller dataset)
if create_datasubset and ns < 1E4:
if ns == 0:
create_folder(path='./datasubset')
os.makedirs('./datasubset/images')
exclude_classes = 43
if exclude_classes not in l[:, 0]:
ns += 1
# shutil.copy(src=self.img_files[i], dst='./datasubset/images/') # copy image
with open('./datasubset/images.txt', 'a') as f:
f.write(self.img_files[i] + '\n')
# Extract object detection boxes for a second stage classifier # Extract object detection boxes for a second stage classifier
if extract_bounding_boxes: if extract_bounding_boxes:
p = Path(self.img_files[i]) p = Path(self.img_files[i])
@ -669,3 +682,10 @@ def convert_images2bmp():
'/Users/glennjocher/PycharmProjects/', '../') '/Users/glennjocher/PycharmProjects/', '../')
with open(label_path.replace('5k', '5k_bmp'), 'w') as file: with open(label_path.replace('5k', '5k_bmp'), 'w') as file:
file.write(lines) file.write(lines)
def create_folder(path='./new_folder'):
# Create folder
if os.path.exists(path):
shutil.rmtree(path) # delete output folder
os.makedirs(path) # make new output folder