| 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 |
| 202 | self.count += 1 |
| 203 | img0 = cv2.imread(path) # BGR |
| 204 | assert img0 is not None, 'Image Not Found ' + path |
| 205 | print('image %g/%g %s: ' % (self.count, self.nf, path), end='') |
| 206 | |
| 207 | # Padded resize |
| 208 | img = letterbox(img0, new_shape=self.img_size, auto_size=self.auto_size)[0] |
| 209 | |
| 210 | # Convert |
| 211 | img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 |
| 212 | img = np.ascontiguousarray(img) |
| 213 | |
| 214 | return path, img, img0, self.cap |
| 215 | |
| 216 | def new_video(self, path): |
| 217 | self.frame = 0 |