updates
This commit is contained in:
parent
6b222df35d
commit
2244c72a1b
|
@ -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()
|
||||||
|
|
Loading…
Reference in New Issue