zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Update datasets.py

This commit is contained in:
Glenn Jocher 2020-06-14 10:32:59 -07:00 committed by GitHub
parent 1374581929
commit 335c607a24
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
utils/datasets.py
View File

@ -261,10 +261,6 @@ class LoadImagesAndLabels(Dataset): # for training/testing
try: try:
path = str(Path(path)) # os-agnostic path = str(Path(path)) # os-agnostic
parent = str(Path(path).parent) + os.sep parent = str(Path(path).parent) + os.sep
joined_path = os.path.join(os.getcwd(), parent)
print('Parent folder override for loading image files: %s' % (os.path.abspath(joined_path)))
if os.path.isfile(path): # file if os.path.isfile(path): # file
with open(path, 'r') as f: with open(path, 'r') as f:
f = f.read().splitlines() f = f.read().splitlines()
@ -278,10 +274,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
raise Exception('Error loading data from %s. See %s' % (path, help_url)) raise Exception('Error loading data from %s. See %s' % (path, help_url))
n = len(self.img_files) n = len(self.img_files)
print('Loaded %g images' % (n))
assert n > 0, 'No images found in %s. See %s' % (path, help_url) assert n > 0, 'No images found in %s. See %s' % (path, help_url)
bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
nb = bi[-1] + 1 # number of batches nb = bi[-1] + 1 # number of batches
@ -298,8 +291,6 @@ class LoadImagesAndLabels(Dataset): # for training/testing
self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt') self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt')
for x in self.img_files] for x in self.img_files]
# Rectangular Training https://github.com/ultralytics/yolov3/issues/232
if self.rect:
# Read image shapes (wh) # Read image shapes (wh)
sp = path.replace('.txt', '') + '.shapes' # shapefile path sp = path.replace('.txt', '') + '.shapes' # shapefile path
try: try:
@ -310,8 +301,12 @@ class LoadImagesAndLabels(Dataset): # for training/testing
s = [exif_size(Image.open(f)) for f in tqdm(self.img_files, desc='Reading image shapes')] s = [exif_size(Image.open(f)) for f in tqdm(self.img_files, desc='Reading image shapes')]
np.savetxt(sp, s, fmt='%g') # overwrites existing (if any) np.savetxt(sp, s, fmt='%g') # overwrites existing (if any)
self.shapes = np.array(s, dtype=np.float64)
# Rectangular Training https://github.com/ultralytics/yolov3/issues/232
if self.rect:
# Sort by aspect ratio # Sort by aspect ratio
s = np.array(s, dtype=np.float64) s = self.shapes # wh
ar = s[:, 1] / s[:, 0] # aspect ratio ar = s[:, 1] / s[:, 0] # aspect ratio
irect = ar.argsort() irect = ar.argsort()
self.img_files = [self.img_files[i] for i in irect] self.img_files = [self.img_files[i] for i in irect]