glenn-jocher commited on
Commit
eae3330
·
1 Parent(s): 0afbb8d

update datasets.py LoadImages() path improvements and Mixup

Browse files
Files changed (1) hide show
  1. utils/datasets.py +25 -16
utils/datasets.py CHANGED
@@ -68,35 +68,37 @@ def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=Fa
68
 
69
  class LoadImages: # for inference
70
  def __init__(self, path, img_size=640):
71
- path = str(Path(path)) # os-agnostic
72
- files = []
73
- if os.path.isdir(path):
74
- files = sorted(glob.glob(os.path.join(path, '*.*')))
75
- elif os.path.isfile(path):
76
- files = [path]
 
 
77
 
78
  images = [x for x in files if os.path.splitext(x)[-1].lower() in img_formats]
79
  videos = [x for x in files if os.path.splitext(x)[-1].lower() in vid_formats]
80
- nI, nV = len(images), len(videos)
81
 
82
  self.img_size = img_size
83
  self.files = images + videos
84
- self.nF = nI + nV # number of files
85
- self.video_flag = [False] * nI + [True] * nV
86
  self.mode = 'images'
87
  if any(videos):
88
  self.new_video(videos[0]) # new video
89
  else:
90
  self.cap = None
91
- assert self.nF > 0, 'No images or videos found in %s. Supported formats are:\nimages: %s\nvideos: %s' % \
92
- (path, img_formats, vid_formats)
93
 
94
  def __iter__(self):
95
  self.count = 0
96
  return self
97
 
98
  def __next__(self):
99
- if self.count == self.nF:
100
  raise StopIteration
101
  path = self.files[self.count]
102
 
@@ -107,7 +109,7 @@ class LoadImages: # for inference
107
  if not ret_val:
108
  self.count += 1
109
  self.cap.release()
110
- if self.count == self.nF: # last video
111
  raise StopIteration
112
  else:
113
  path = self.files[self.count]
@@ -115,14 +117,14 @@ class LoadImages: # for inference
115
  ret_val, img0 = self.cap.read()
116
 
117
  self.frame += 1
118
- print('video %g/%g (%g/%g) %s: ' % (self.count + 1, self.nF, self.frame, self.nframes, path), end='')
119
 
120
  else:
121
  # Read image
122
  self.count += 1
123
  img0 = cv2.imread(path) # BGR
124
  assert img0 is not None, 'Image Not Found ' + path
125
- print('image %g/%g %s: ' % (self.count, self.nF, path), end='')
126
 
127
  # Padded resize
128
  img = letterbox(img0, new_shape=self.img_size)[0]
@@ -140,7 +142,7 @@ class LoadImages: # for inference
140
  self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
141
 
142
  def __len__(self):
143
- return self.nF # number of files
144
 
145
 
146
  class LoadWebcam: # for inference
@@ -470,6 +472,13 @@ class LoadImagesAndLabels(Dataset): # for training/testing
470
  img, labels = load_mosaic(self, index)
471
  shapes = None
472
 
 
 
 
 
 
 
 
473
  else:
474
  # Load image
475
  img, (h0, w0), (h, w) = load_image(self, index)
 
68
 
69
  class LoadImages: # for inference
70
  def __init__(self, path, img_size=640):
71
+ p = str(Path(path)) # os-agnostic
72
+ p = os.path.abspath(p) # absolute path
73
+ if os.path.isdir(p):
74
+ files = sorted(glob.glob(os.path.join(p, '*.*')))
75
+ elif os.path.isfile(p):
76
+ files = [p]
77
+ else:
78
+ raise Exception('ERROR: %s does not exist' % p)
79
 
80
  images = [x for x in files if os.path.splitext(x)[-1].lower() in img_formats]
81
  videos = [x for x in files if os.path.splitext(x)[-1].lower() in vid_formats]
82
+ ni, nv = len(images), len(videos)
83
 
84
  self.img_size = img_size
85
  self.files = images + videos
86
+ self.nf = ni + nv # number of files
87
+ self.video_flag = [False] * ni + [True] * nv
88
  self.mode = 'images'
89
  if any(videos):
90
  self.new_video(videos[0]) # new video
91
  else:
92
  self.cap = None
93
+ assert self.nf > 0, 'No images or videos found in %s. Supported formats are:\nimages: %s\nvideos: %s' % \
94
+ (p, img_formats, vid_formats)
95
 
96
  def __iter__(self):
97
  self.count = 0
98
  return self
99
 
100
  def __next__(self):
101
+ if self.count == self.nf:
102
  raise StopIteration
103
  path = self.files[self.count]
104
 
 
109
  if not ret_val:
110
  self.count += 1
111
  self.cap.release()
112
+ if self.count == self.nf: # last video
113
  raise StopIteration
114
  else:
115
  path = self.files[self.count]
 
117
  ret_val, img0 = self.cap.read()
118
 
119
  self.frame += 1
120
+ print('video %g/%g (%g/%g) %s: ' % (self.count + 1, self.nf, self.frame, self.nframes, path), end='')
121
 
122
  else:
123
  # Read image
124
  self.count += 1
125
  img0 = cv2.imread(path) # BGR
126
  assert img0 is not None, 'Image Not Found ' + path
127
+ print('image %g/%g %s: ' % (self.count, self.nf, path), end='')
128
 
129
  # Padded resize
130
  img = letterbox(img0, new_shape=self.img_size)[0]
 
142
  self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
143
 
144
  def __len__(self):
145
+ return self.nf # number of files
146
 
147
 
148
  class LoadWebcam: # for inference
 
472
  img, labels = load_mosaic(self, index)
473
  shapes = None
474
 
475
+ # MixUp https://arxiv.org/pdf/1710.09412.pdf
476
+ # if random.random() < 0.5:
477
+ # img2, labels2 = load_mosaic(self, random.randint(0, len(self.labels) - 1))
478
+ # r = np.random.beta(0.3, 0.3) # mixup ratio, alpha=beta=0.3
479
+ # img = (img * r + img2 * (1 - r)).astype(np.uint8)
480
+ # labels = np.concatenate((labels, labels2), 0)
481
+
482
  else:
483
  # Load image
484
  img, (h0, w0), (h, w) = load_image(self, index)