From 7a353a9c70b3fedf9e7936d21aab9c0e15d9c3a9 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 4 Jul 2019 14:03:13 +0200 Subject: [PATCH] updates --- utils/datasets.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/utils/datasets.py b/utils/datasets.py index 7a11dde6..a91aae7c 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -13,12 +13,13 @@ from tqdm import tqdm from utils.utils import xyxy2xywh, xywh2xyxy +img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif'] +vid_formats = ['.mov', '.avi', '.mp4'] + class LoadImages: # for inference def __init__(self, path, img_size=416): self.height = img_size - img_formats = ['.jpg', '.jpeg', '.png', '.tif'] - vid_formats = ['.mov', '.avi', '.mp4'] files = [] if os.path.isdir(path): @@ -146,11 +147,11 @@ class LoadImagesAndLabels(Dataset): # for training/testing self.augment = augment self.image_weights = image_weights self.rect = False if image_weights else rect - self.label_files = [x.replace('images', 'labels'). - replace('.jpeg', '.txt'). - replace('.jpg', '.txt'). - replace('.bmp', '.txt'). - replace('.png', '.txt') for x in self.img_files] + + # Define labels + self.label_files = [x.replace('images', 'labels') for x in self.img_files] + for f in img_formats: + self.label_files = [x.replace(f, '.txt') for x in self.label_files] # Rectangular Training https://github.com/ultralytics/yolov3/issues/232 if self.rect: @@ -169,9 +170,9 @@ class LoadImagesAndLabels(Dataset): # for training/testing # Sort by aspect ratio ar = s[:, 1] / s[:, 0] # aspect ratio i = ar.argsort() - ar = ar[i] self.img_files = [self.img_files[i] for i in i] self.label_files = [self.label_files[i] for i in i] + ar = ar[i] # Set training image shapes shapes = [[1, 1]] * nb