This commit is contained in:
Glenn Jocher 2019-06-25 11:45:38 +02:00
parent 6b222df35d
commit 2244c72a1b
1 changed files with 39 additions and 33 deletions

View File

@ -187,38 +187,41 @@ 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 = [np.zeros((0, 5))] * n self.labels = [None] * n
iter = tqdm(self.label_files, desc='Reading labels') if n > 10 else self.label_files preload_labels = False
extract_bounding_boxes = False if preload_labels:
for i, file in enumerate(iter): self.labels = [np.zeros((0, 5))] * n
try: iter = tqdm(self.label_files, desc='Reading labels') if n > 10 else self.label_files
with open(file, 'r') as f: extract_bounding_boxes = False
l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32) for i, file in enumerate(iter):
if l.shape[0]: try:
assert l.shape[1] == 5, '> 5 label columns: %s' % file with open(file, 'r') as f:
assert (l >= 0).all(), 'negative labels: %s' % file l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32)
assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels: %s' % file if l.shape[0]:
self.labels[i] = l assert l.shape[1] == 5, '> 5 label columns: %s' % file
assert (l >= 0).all(), 'negative labels: %s' % file
assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels: %s' % file
self.labels[i] = l
# 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])
img = cv2.imread(str(p)) img = cv2.imread(str(p))
h, w, _ = img.shape h, w, _ = img.shape
for j, x in enumerate(l): for j, x in enumerate(l):
f = '%s%sclassification%s%g_%g_%s' % (p.parent.parent, os.sep, os.sep, x[0], j, p.name) f = '%s%sclassification%s%g_%g_%s' % (
if not os.path.exists(Path(f).parent): p.parent.parent, os.sep, os.sep, x[0], j, p.name)
os.makedirs(Path(f).parent) # make new output folder if not os.path.exists(Path(f).parent):
box = xywh2xyxy(x[1:].reshape(-1, 4)).ravel() os.makedirs(Path(f).parent) # make new output folder
box = np.clip(box, 0, 1) # clip boxes outside of image box = xywh2xyxy(x[1:].reshape(-1, 4)).ravel()
result = cv2.imwrite(f, img[int(box[1] * h):int(box[3] * h), box = np.clip(box, 0, 1) # clip boxes outside of image
int(box[0] * w):int(box[2] * w)]) result = cv2.imwrite(f, img[int(box[1] * h):int(box[3] * h),
if not result: int(box[0] * w):int(box[2] * w)])
print('stop') if not result:
print('stop')
except: except:
pass # print('Warning: missing labels for %s' % self.img_files[i]) # missing label file pass # print('Warning: missing labels for %s' % self.img_files[i]) # missing label file
assert len(np.concatenate(self.labels, 0)) > 0, 'No labels found. Incorrect label paths provided.' assert len(np.concatenate(self.labels, 0)) > 0, 'No labels found. Incorrect label paths provided.'
def __len__(self): def __len__(self):
return len(self.img_files) return len(self.img_files)
@ -274,9 +277,12 @@ class LoadImagesAndLabels(Dataset): # for training/testing
# Load labels # Load labels
labels = [] labels = []
if os.path.isfile(label_path): if os.path.isfile(label_path):
# with open(label_path, 'r') as f:
# x = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32)
x = self.labels[index] x = self.labels[index]
if x is None: # labels not preloaded
with open(label_path, 'r') as f:
x = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32)
self.labels[index] = x # save for next time
if x.size > 0: if x.size > 0:
# Normalized xywh to pixel xyxy format # Normalized xywh to pixel xyxy format
labels = x.copy() labels = x.copy()