質問編集履歴

1

丸投げな質問だったため詳細を追記しました!ご迷惑おかけします。。

2021/03/13 05:45

投稿

Pino02
Pino02

スコア1

test CHANGED
File without changes
test CHANGED
@@ -2,4 +2,442 @@
2
2
 
3
3
 
4
4
 
5
+ ### 前提・実現したいこと
6
+
5
- YOLOv5を使ってWebCamでリアルタイムで認識はしてくれいるんですが、コンソール上にしか結果が表示されないのでそれをtxtファイルに出力したいです。
7
+ コンソール上にしか結果が表示されないのでそれをtxtファイルに出力したいです。
8
+
9
+
10
+
11
+ ### 発生している問題・エラーメッセージ
12
+
13
+ yolov5-masterというディレクトリにYOLOv5のファイルを解凍してあります。
14
+
15
+ その中のdetect.pyを`python detect.py --source 0`で実行しました。
16
+
17
+ そうするとコマンドプロントには下記のものはその一部ですが結果が返ってきます。
18
+
19
+ ですがこれをtxtファイルに出力できなくて困っています。
20
+
21
+
22
+
23
+ ### 試したこと
24
+
25
+ とりあえずprintしているところが118行目の
26
+
27
+ ```python
28
+
29
+ print(f'{s}Done. ({t2 - t1:.3f}s)')
30
+
31
+ ```
32
+
33
+ だと分かったのでこれを変数に格納して
34
+
35
+ 強引に結果を切り取ったりしてみました。
36
+
37
+ ```python
38
+
39
+ result = f'{s}Done. ({t2 - t1:.3f}s)'
40
+
41
+ result = result[14:]
42
+
43
+ ```
44
+
45
+ ### Code
46
+
47
+ ```Python
48
+
49
+ detect.py
50
+
51
+
52
+
53
+ import argparse
54
+
55
+ import time
56
+
57
+ from pathlib import Path
58
+
59
+
60
+
61
+ import cv2
62
+
63
+ import torch
64
+
65
+ import torch.backends.cudnn as cudnn
66
+
67
+ from numpy import random
68
+
69
+
70
+
71
+ from models.experimental import attempt_load
72
+
73
+ from utils.datasets import LoadStreams, LoadImages
74
+
75
+ from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
76
+
77
+ scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
78
+
79
+ from utils.plots import plot_one_box
80
+
81
+ from utils.torch_utils import select_device, load_classifier, time_synchronized
82
+
83
+
84
+
85
+
86
+
87
+ def detect(save_img=False):
88
+
89
+ source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
90
+
91
+ webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
92
+
93
+ ('rtsp://', 'rtmp://', 'http://'))
94
+
95
+
96
+
97
+ # Directories
98
+
99
+ save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
100
+
101
+ (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
102
+
103
+
104
+
105
+ # Initialize
106
+
107
+ set_logging()
108
+
109
+ device = select_device(opt.device)
110
+
111
+ half = device.type != 'cpu' # half precision only supported on CUDA
112
+
113
+
114
+
115
+ # Load model
116
+
117
+ model = attempt_load(weights, map_location=device) # load FP32 model
118
+
119
+ stride = int(model.stride.max()) # model stride
120
+
121
+ imgsz = check_img_size(imgsz, s=stride) # check img_size
122
+
123
+ if half:
124
+
125
+ model.half() # to FP16
126
+
127
+
128
+
129
+ # Second-stage classifier
130
+
131
+ classify = False
132
+
133
+ if classify:
134
+
135
+ modelc = load_classifier(name='resnet101', n=2) # initialize
136
+
137
+ modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
138
+
139
+
140
+
141
+ # Set Dataloader
142
+
143
+ vid_path, vid_writer = None, None
144
+
145
+ if webcam:
146
+
147
+ view_img = check_imshow()
148
+
149
+ cudnn.benchmark = True # set True to speed up constant image size inference
150
+
151
+ dataset = LoadStreams(source, img_size=imgsz, stride=stride)
152
+
153
+ else:
154
+
155
+ save_img = True
156
+
157
+ dataset = LoadImages(source, img_size=imgsz, stride=stride)
158
+
159
+
160
+
161
+ # Get names and colors
162
+
163
+ names = model.module.names if hasattr(model, 'module') else model.names
164
+
165
+ colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
166
+
167
+
168
+
169
+ # Run inference
170
+
171
+ if device.type != 'cpu':
172
+
173
+ model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
174
+
175
+ t0 = time.time()
176
+
177
+ for path, img, im0s, vid_cap in dataset:
178
+
179
+ img = torch.from_numpy(img).to(device)
180
+
181
+ img = img.half() if half else img.float() # uint8 to fp16/32
182
+
183
+ img /= 255.0 # 0 - 255 to 0.0 - 1.0
184
+
185
+ if img.ndimension() == 3:
186
+
187
+ img = img.unsqueeze(0)
188
+
189
+
190
+
191
+ # Inference
192
+
193
+ t1 = time_synchronized()
194
+
195
+ pred = model(img, augment=opt.augment)[0]
196
+
197
+
198
+
199
+ # Apply NMS
200
+
201
+ pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
202
+
203
+ t2 = time_synchronized()
204
+
205
+
206
+
207
+ # Apply Classifier
208
+
209
+ if classify:
210
+
211
+ pred = apply_classifier(pred, modelc, img, im0s)
212
+
213
+
214
+
215
+ # Process detections
216
+
217
+ for i, det in enumerate(pred): # detections per image
218
+
219
+ if webcam: # batch_size >= 1
220
+
221
+ p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
222
+
223
+ else:
224
+
225
+ p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
226
+
227
+
228
+
229
+ p = Path(p) # to Path
230
+
231
+ save_path = str(save_dir / p.name) # img.jpg
232
+
233
+ txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
234
+
235
+ s += '%gx%g ' % img.shape[2:] # print string
236
+
237
+ gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
238
+
239
+ if len(det):
240
+
241
+ # Rescale boxes from img_size to im0 size
242
+
243
+ det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
244
+
245
+
246
+
247
+ # Print results
248
+
249
+ for c in det[:, -1].unique():
250
+
251
+ n = (det[:, -1] == c).sum() # detections per class
252
+
253
+ s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
254
+
255
+
256
+
257
+ # Write results
258
+
259
+ for *xyxy, conf, cls in reversed(det):
260
+
261
+ if save_txt: # Write to file
262
+
263
+ xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
264
+
265
+ line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
266
+
267
+ with open(txt_path + '.txt', 'a') as f:
268
+
269
+ f.write(('%g ' * len(line)).rstrip() % line + '\n')
270
+
271
+
272
+
273
+ if save_img or view_img: # Add bbox to image
274
+
275
+ label = f'{names[int(cls)]} {conf:.2f}'
276
+
277
+ plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
278
+
279
+
280
+
281
+ # Print time (inference + NMS)
282
+
283
+ print(f'{s}Done. ({t2 - t1:.3f}s)')
284
+
285
+
286
+
287
+ # Stream results
288
+
289
+ if view_img:
290
+
291
+ cv2.imshow(str(p), im0)
292
+
293
+ cv2.waitKey(1) # 1 millisecond
294
+
295
+
296
+
297
+ # Save results (image with detections)
298
+
299
+ if save_img:
300
+
301
+ if dataset.mode == 'image':
302
+
303
+ cv2.imwrite(save_path, im0)
304
+
305
+ else: # 'video'
306
+
307
+ if vid_path != save_path: # new video
308
+
309
+ vid_path = save_path
310
+
311
+ if isinstance(vid_writer, cv2.VideoWriter):
312
+
313
+ vid_writer.release() # release previous video writer
314
+
315
+
316
+
317
+ fourcc = 'mp4v' # output video codec
318
+
319
+ fps = vid_cap.get(cv2.CAP_PROP_FPS)
320
+
321
+ w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
322
+
323
+ h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
324
+
325
+ vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
326
+
327
+ vid_writer.write(im0)
328
+
329
+
330
+
331
+ if save_txt or save_img:
332
+
333
+ s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
334
+
335
+ print(f"Results saved to {save_dir}{s}")
336
+
337
+
338
+
339
+ print(f'Done. ({time.time() - t0:.3f}s)')
340
+
341
+
342
+
343
+
344
+
345
+ if __name__ == '__main__':
346
+
347
+ parser = argparse.ArgumentParser()
348
+
349
+ parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
350
+
351
+ parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
352
+
353
+ parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
354
+
355
+ parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
356
+
357
+ parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
358
+
359
+ parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
360
+
361
+ parser.add_argument('--view-img', action='store_true', help='display results')
362
+
363
+ parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
364
+
365
+ parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
366
+
367
+ parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
368
+
369
+ parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
370
+
371
+ parser.add_argument('--augment', action='store_true', help='augmented inference')
372
+
373
+ parser.add_argument('--update', action='store_true', help='update all models')
374
+
375
+ parser.add_argument('--project', default='runs/detect', help='save results to project/name')
376
+
377
+ parser.add_argument('--name', default='exp', help='save results to project/name')
378
+
379
+ parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
380
+
381
+ opt = parser.parse_args()
382
+
383
+ print(opt)
384
+
385
+ check_requirements()
386
+
387
+
388
+
389
+ with torch.no_grad():
390
+
391
+ if opt.update: # update all models (to fix SourceChangeWarning)
392
+
393
+ for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
394
+
395
+ detect()
396
+
397
+ strip_optimizer(opt.weights)
398
+
399
+ else:
400
+
401
+ detect()
402
+
403
+ ```
404
+
405
+
406
+
407
+ ```cmd
408
+
409
+ 0: 480x640 1 tv, Done. (0.219s)
410
+
411
+ 0: 480x640 1 tv, 1 book, Done. (0.224s)
412
+
413
+ 0: 480x640 1 person, Done. (0.215s)
414
+
415
+ 0: 480x640 Done. (0.217s)
416
+
417
+ 0: 480x640 Done. (0.212s)
418
+
419
+ 0: 480x640 1 tv, Done. (0.205s)
420
+
421
+ 0: 480x640 1 tie, Done. (0.208s)
422
+
423
+ 0: 480x640 Done. (0.211s)
424
+
425
+ 0: 480x640 1 person, Done. (0.219s)
426
+
427
+ 0: 480x640 Done. (0.204s)
428
+
429
+ 0: 480x640 1 chair, Done. (0.218s)
430
+
431
+ 0: 480x640 1 carrot, Done. (0.210s)
432
+
433
+ 0: 480x640 Done. (0.212s)
434
+
435
+ 0: 480x640 Done. (0.215s)
436
+
437
+ 0: 480x640 1 person, Done. (0.211s)
438
+
439
+ 0: 480x640 1 person, Done. (0.214s)
440
+
441
+ 0: 480x640 1 person, 1 tv, 1 laptop, Done. (0.208s)
442
+
443
+ ```