| 142 | |
| 143 | |
| 144 | class LoadImages: # for inference |
| 145 | def __init__(self, path, img_size=640, auto_size=32): |
| 146 | p = str(Path(path)) # os-agnostic |
| 147 | p = os.path.abspath(p) # absolute path |
| 148 | if '*' in p: |
| 149 | files = sorted(glob.glob(p, recursive=True)) # glob |
| 150 | elif os.path.isdir(p): |
| 151 | files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir |
| 152 | elif os.path.isfile(p): |
| 153 | files = [p] # files |
| 154 | else: |
| 155 | raise Exception('ERROR: %s does not exist' % p) |
| 156 | |
| 157 | images = [x for x in files if x.split('.')[-1].lower() in img_formats] |
| 158 | videos = [x for x in files if x.split('.')[-1].lower() in vid_formats] |
| 159 | ni, nv = len(images), len(videos) |
| 160 | |
| 161 | self.img_size = img_size |
| 162 | self.auto_size = auto_size |
| 163 | self.files = images + videos |
| 164 | self.nf = ni + nv # number of files |
| 165 | self.video_flag = [False] * ni + [True] * nv |
| 166 | self.mode = 'images' |
| 167 | if any(videos): |
| 168 | self.new_video(videos[0]) # new video |
| 169 | else: |
| 170 | self.cap = None |
| 171 | assert self.nf > 0, 'No images or videos found in %s. Supported formats are:\nimages: %s\nvideos: %s' % \ |
| 172 | (p, img_formats, vid_formats) |
| 173 | |
| 174 | def __iter__(self): |
| 175 | self.count = 0 |
| 176 | return self |
| 177 | |
| 178 | def __next__(self): |
| 179 | if self.count == self.nf: |
| 180 | raise StopIteration |
| 181 | path = self.files[self.count] |
| 182 | |
| 183 | if self.video_flag[self.count]: |
| 184 | # Read video |
| 185 | self.mode = 'video' |
| 186 | ret_val, img0 = self.cap.read() |
| 187 | if not ret_val: |
| 188 | self.count += 1 |
| 189 | self.cap.release() |
| 190 | if self.count == self.nf: # last video |
| 191 | raise StopIteration |
| 192 | else: |
| 193 | path = self.files[self.count] |
| 194 | self.new_video(path) |
| 195 | ret_val, img0 = self.cap.read() |
| 196 | |
| 197 | self.frame += 1 |
| 198 | print('video %g/%g (%g/%g) %s: ' % (self.count + 1, self.nf, self.frame, self.nframes, path), end='') |
| 199 | |
| 200 | else: |
| 201 | # Read image |