diff --git a/detect.py b/detect.py index 63feaa6ad..5935367d3 100644 --- a/detect.py +++ b/detect.py @@ -81,12 +81,12 @@ def detect(save_img=False): # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 - p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy() + p, s, im0, frame = Path(path[i]), '%g: ' % i, im0s[i].copy(), dataset.count else: - p, s, im0 = Path(path), '', im0s + p, s, im0, frame = Path(path), '', im0s, getattr(dataset, 'frame', 0) save_path = str(save_dir / p.name) - txt_path = str(save_dir / 'labels' / p.stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '') + txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh if len(det): @@ -96,7 +96,7 @@ def detect(save_img=False): # Print results for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class - s += '%g %ss, ' % (n, names[int(c)]) # add to string + s += f'{n} {names[int(c)]}s, ' # add to string # Write results for *xyxy, conf, cls in reversed(det): @@ -107,11 +107,11 @@ def detect(save_img=False): f.write(('%g ' * len(line)).rstrip() % line + '\n') if save_img or view_img: # Add bbox to image - label = '%s %.2f' % (names[int(cls)], conf) + label = f'{names[int(cls)]} {conf:.2f}' plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) # Print time (inference + NMS) - print('%sDone. (%.3fs)' % (s, t2 - t1)) + print(f'{s}Done. ({t2 - t1:.3f}s)') # Stream results if view_img: @@ -121,9 +121,9 @@ def detect(save_img=False): # Save results (image with detections) if save_img: - if dataset.mode == 'images': + if dataset.mode == 'image': cv2.imwrite(save_path, im0) - else: + else: # 'video' if vid_path != save_path: # new video vid_path = save_path if isinstance(vid_writer, cv2.VideoWriter): @@ -140,7 +140,7 @@ def detect(save_img=False): s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' print(f"Results saved to {save_dir}{s}") - print('Done. (%.3fs)' % (time.time() - t0)) + print(f'Done. ({time.time() - t0:.3f}s)') if __name__ == '__main__': diff --git a/utils/datasets.py b/utils/datasets.py index dc3c8fb74..7d10d1c0d 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -138,7 +138,7 @@ class LoadImages: # for inference self.files = images + videos self.nf = ni + nv # number of files self.video_flag = [False] * ni + [True] * nv - self.mode = 'images' + self.mode = 'image' if any(videos): self.new_video(videos[0]) # new video else: @@ -256,7 +256,7 @@ class LoadWebcam: # for inference class LoadStreams: # multiple IP or RTSP cameras def __init__(self, sources='streams.txt', img_size=640): - self.mode = 'images' + self.mode = 'stream' self.img_size = img_size if os.path.isfile(sources):