Commit
·
eae3330
1
Parent(s):
0afbb8d
update datasets.py LoadImages() path improvements and Mixup
Browse files- 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 |
-
|
72 |
-
|
73 |
-
if os.path.isdir(
|
74 |
-
files = sorted(glob.glob(os.path.join(
|
75 |
-
elif os.path.isfile(
|
76 |
-
files = [
|
|
|
|
|
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 |
-
|
81 |
|
82 |
self.img_size = img_size
|
83 |
self.files = images + videos
|
84 |
-
self.
|
85 |
-
self.video_flag = [False] *
|
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.
|
92 |
-
(
|
93 |
|
94 |
def __iter__(self):
|
95 |
self.count = 0
|
96 |
return self
|
97 |
|
98 |
def __next__(self):
|
99 |
-
if self.count == self.
|
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.
|
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.
|
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.
|
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.
|
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)
|