Streaming --save-txt bug fix (#1672)
Browse files* Streaming --save-txt bug fix
* cleanup
- detect.py +9 -9
- 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) + ('
|
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 += '
|
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 = '
|
111 |
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
|
112 |
|
113 |
# Print time (inference + NMS)
|
114 |
-
print('
|
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 == '
|
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. (
|
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 = '
|
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 = '
|
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):
|