This commit is contained in:
Glenn Jocher 2019-08-06 11:26:52 +02:00
parent 7462a95873
commit 68a5f8e207
1 changed files with 2 additions and 3 deletions

View File

@ -214,11 +214,10 @@ class LoadImagesAndLabels(Dataset): # for training/testing
# Preload labels (required for weighted CE training) # Preload labels (required for weighted CE training)
self.imgs = [None] * n self.imgs = [None] * n
self.labels = [None] * n self.labels = [None] * n
preload_labels = False if augment or image_weights: # cache labels for faster training
if preload_labels or image_weights:
self.labels = [np.zeros((0, 5))] * n self.labels = [np.zeros((0, 5))] * n
extract_bounding_boxes = False extract_bounding_boxes = False
for i, file in enumerate(tqdm(self.label_files, desc='Reading labels') if n > 10 else self.label_files): for i, file in enumerate(tqdm(self.label_files, desc='Caching labels')):
try: try:
with open(file, 'r') as f: with open(file, 'r') as f:
l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32) l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32)