glenn-jocher commited on
Commit
54043a9
·
unverified ·
1 Parent(s): bc52ea2

Streaming --save-txt bug fix (#1672)

Browse files

* Streaming --save-txt bug fix

* cleanup

Files changed (2) hide show
  1. detect.py +9 -9
  2. utils/datasets.py +2 -2
detect.py CHANGED
@@ -81,12 +81,12 @@ def detect(save_img=False):
81
  # Process detections
82
  for i, det in enumerate(pred): # detections per image
83
  if webcam: # batch_size >= 1
84
- p, s, im0 = Path(path[i]), '%g: ' % i, im0s[i].copy()
85
  else:
86
- p, s, im0 = Path(path), '', im0s
87
 
88
  save_path = str(save_dir / p.name)
89
- txt_path = str(save_dir / 'labels' / p.stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
90
  s += '%gx%g ' % img.shape[2:] # print string
91
  gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
92
  if len(det):
@@ -96,7 +96,7 @@ def detect(save_img=False):
96
  # Print results
97
  for c in det[:, -1].unique():
98
  n = (det[:, -1] == c).sum() # detections per class
99
- s += '%g %ss, ' % (n, names[int(c)]) # add to string
100
 
101
  # Write results
102
  for *xyxy, conf, cls in reversed(det):
@@ -107,11 +107,11 @@ def detect(save_img=False):
107
  f.write(('%g ' * len(line)).rstrip() % line + '\n')
108
 
109
  if save_img or view_img: # Add bbox to image
110
- label = '%s %.2f' % (names[int(cls)], conf)
111
  plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
112
 
113
  # Print time (inference + NMS)
114
- print('%sDone. (%.3fs)' % (s, t2 - t1))
115
 
116
  # Stream results
117
  if view_img:
@@ -121,9 +121,9 @@ def detect(save_img=False):
121
 
122
  # Save results (image with detections)
123
  if save_img:
124
- if dataset.mode == 'images':
125
  cv2.imwrite(save_path, im0)
126
- else:
127
  if vid_path != save_path: # new video
128
  vid_path = save_path
129
  if isinstance(vid_writer, cv2.VideoWriter):
@@ -140,7 +140,7 @@ def detect(save_img=False):
140
  s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
141
  print(f"Results saved to {save_dir}{s}")
142
 
143
- print('Done. (%.3fs)' % (time.time() - t0))
144
 
145
 
146
  if __name__ == '__main__':
 
81
  # Process detections
82
  for i, det in enumerate(pred): # detections per image
83
  if webcam: # batch_size >= 1
84
+ p, s, im0, frame = Path(path[i]), '%g: ' % i, im0s[i].copy(), dataset.count
85
  else:
86
+ p, s, im0, frame = Path(path), '', im0s, getattr(dataset, 'frame', 0)
87
 
88
  save_path = str(save_dir / p.name)
89
+ txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')
90
  s += '%gx%g ' % img.shape[2:] # print string
91
  gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
92
  if len(det):
 
96
  # Print results
97
  for c in det[:, -1].unique():
98
  n = (det[:, -1] == c).sum() # detections per class
99
+ s += f'{n} {names[int(c)]}s, ' # add to string
100
 
101
  # Write results
102
  for *xyxy, conf, cls in reversed(det):
 
107
  f.write(('%g ' * len(line)).rstrip() % line + '\n')
108
 
109
  if save_img or view_img: # Add bbox to image
110
+ label = f'{names[int(cls)]} {conf:.2f}'
111
  plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
112
 
113
  # Print time (inference + NMS)
114
+ print(f'{s}Done. ({t2 - t1:.3f}s)')
115
 
116
  # Stream results
117
  if view_img:
 
121
 
122
  # Save results (image with detections)
123
  if save_img:
124
+ if dataset.mode == 'image':
125
  cv2.imwrite(save_path, im0)
126
+ else: # 'video'
127
  if vid_path != save_path: # new video
128
  vid_path = save_path
129
  if isinstance(vid_writer, cv2.VideoWriter):
 
140
  s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
141
  print(f"Results saved to {save_dir}{s}")
142
 
143
+ print(f'Done. ({time.time() - t0:.3f}s)')
144
 
145
 
146
  if __name__ == '__main__':
utils/datasets.py CHANGED
@@ -138,7 +138,7 @@ class LoadImages: # for inference
138
  self.files = images + videos
139
  self.nf = ni + nv # number of files
140
  self.video_flag = [False] * ni + [True] * nv
141
- self.mode = 'images'
142
  if any(videos):
143
  self.new_video(videos[0]) # new video
144
  else:
@@ -256,7 +256,7 @@ class LoadWebcam: # for inference
256
 
257
  class LoadStreams: # multiple IP or RTSP cameras
258
  def __init__(self, sources='streams.txt', img_size=640):
259
- self.mode = 'images'
260
  self.img_size = img_size
261
 
262
  if os.path.isfile(sources):
 
138
  self.files = images + videos
139
  self.nf = ni + nv # number of files
140
  self.video_flag = [False] * ni + [True] * nv
141
+ self.mode = 'image'
142
  if any(videos):
143
  self.new_video(videos[0]) # new video
144
  else:
 
256
 
257
  class LoadStreams: # multiple IP or RTSP cameras
258
  def __init__(self, sources='streams.txt', img_size=640):
259
+ self.mode = 'stream'
260
  self.img_size = img_size
261
 
262
  if os.path.isfile(sources):