uncached label removal

This commit is contained in:
Glenn Jocher 2020-04-28 11:07:26 -07:00
parent b1d385a8de
commit 15f1343dfc
1 changed files with 75 additions and 86 deletions

View File

@ -257,7 +257,7 @@ class LoadStreams: # multiple IP or RTSP cameras
class LoadImagesAndLabels(Dataset): # for training/testing class LoadImagesAndLabels(Dataset): # for training/testing
def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False, def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
cache_labels=True, cache_images=False, single_cls=False): cache_images=False, single_cls=False):
path = str(Path(path)) # os-agnostic path = str(Path(path)) # os-agnostic
assert os.path.isfile(path), 'File not found %s. See %s' % (path, help_url) assert os.path.isfile(path), 'File not found %s. See %s' % (path, help_url)
with open(path, 'r') as f: with open(path, 'r') as f:
@ -315,11 +315,9 @@ class LoadImagesAndLabels(Dataset): # for training/testing
self.batch_shapes = np.ceil(np.array(shapes) * img_size / 64.).astype(np.int) * 64 self.batch_shapes = np.ceil(np.array(shapes) * img_size / 64.).astype(np.int) * 64
# Preload labels (required for weighted CE training) # Cache labels
self.imgs = [None] * n self.imgs = [None] * n
self.labels = [None] * n self.labels = [np.zeros((0, 5), dtype=np.float32)] * n
if cache_labels or image_weights: # cache labels for faster training
self.labels = [np.zeros((0, 5))] * n
extract_bounding_boxes = False extract_bounding_boxes = False
create_datasubset = False create_datasubset = False
pbar = tqdm(self.label_files, desc='Caching labels') pbar = tqdm(self.label_files, desc='Caching labels')
@ -432,7 +430,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
# Load labels # Load labels
labels = [] labels = []
x = self.labels[index] x = self.labels[index]
if x is not None and 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()
labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0] # pad width labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0] # pad width
@ -502,9 +500,9 @@ def load_image(self, index):
# loads 1 image from dataset, returns img, original hw, resized hw # loads 1 image from dataset, returns img, original hw, resized hw
img = self.imgs[index] img = self.imgs[index]
if img is None: # not cached if img is None: # not cached
img_path = self.img_files[index] path = self.img_files[index]
img = cv2.imread(img_path) # BGR img = cv2.imread(path) # BGR
assert img is not None, 'Image Not Found ' + img_path assert img is not None, 'Image Not Found ' + path
h0, w0 = img.shape[:2] # orig hw h0, w0 = img.shape[:2] # orig hw
r = self.img_size / max(h0, w0) # resize image to img_size r = self.img_size / max(h0, w0) # resize image to img_size
if r < 1 or (self.augment and r != 1): # always resize down, only resize up if training with augmentation if r < 1 or (self.augment and r != 1): # always resize down, only resize up if training with augmentation
@ -557,23 +555,14 @@ def load_mosaic(self, index):
padw = x1a - x1b padw = x1a - x1b
padh = y1a - y1b padh = y1a - y1b
# Load labels # Labels
label_path = self.label_files[index]
if os.path.isfile(label_path):
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)
if x.size > 0:
# Normalized xywh to pixel xyxy format
labels = x.copy() labels = x.copy()
if x.size > 0: # Normalized xywh to pixel xyxy format
labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw
labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + padh labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + padh
labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padw labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padw
labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + padh labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + padh
else:
labels = np.zeros((0, 5), dtype=np.float32)
labels4.append(labels) labels4.append(labels)
# Concat/clip labels # Concat/clip labels
@ -585,10 +574,10 @@ def load_mosaic(self, index):
# Augment # Augment
# img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)] # center crop (WARNING, requires box pruning) # img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)] # center crop (WARNING, requires box pruning)
img4, labels4 = random_affine(img4, labels4, img4, labels4 = random_affine(img4, labels4,
degrees=self.hyp['degrees'] * 1, degrees=self.hyp['degrees'],
translate=self.hyp['translate'] * 1, translate=self.hyp['translate'],
scale=self.hyp['scale'] * 1, scale=self.hyp['scale'],
shear=self.hyp['shear'] * 1, shear=self.hyp['shear'],
border=-s // 2) # border to remove border=-s // 2) # border to remove
return img4, labels4 return img4, labels4
@ -688,7 +677,7 @@ def random_affine(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10,
area = w * h area = w * h
area0 = (targets[:, 3] - targets[:, 1]) * (targets[:, 4] - targets[:, 2]) area0 = (targets[:, 3] - targets[:, 1]) * (targets[:, 4] - targets[:, 2])
ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16)) # aspect ratio ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16)) # aspect ratio
i = (w > 4) & (h > 4) & (area / (area0 + 1e-16) > 0.2) & (ar < 10) i = (w > 4) & (h > 4) & (area / (area0 * s + 1e-16) > 0.2) & (ar < 10)
targets = targets[i] targets = targets[i]
targets[:, 1:5] = xy[i] targets[:, 1:5] = xy[i]