(predictor, vis_folder, current_time, args)
| 176 | |
| 177 | |
| 178 | def image_demo(predictor, vis_folder, current_time, args): |
| 179 | if osp.isdir(args.path): |
| 180 | files = get_image_list(args.path) |
| 181 | else: |
| 182 | files = [args.path] |
| 183 | files.sort() |
| 184 | tracker = BYTETracker(args, frame_rate=args.fps) |
| 185 | timer = Timer() |
| 186 | results = [] |
| 187 | |
| 188 | for frame_id, img_path in enumerate(files, 1): |
| 189 | outputs, img_info = predictor.inference(img_path, timer) |
| 190 | if outputs[0] is not None: |
| 191 | online_targets = tracker.update(outputs[0], [img_info['height'], img_info['width']], exp.test_size) |
| 192 | online_tlwhs = [] |
| 193 | online_ids = [] |
| 194 | online_scores = [] |
| 195 | for t in online_targets: |
| 196 | tlwh = t.tlwh |
| 197 | tid = t.track_id |
| 198 | vertical = tlwh[2] / tlwh[3] > args.aspect_ratio_thresh |
| 199 | if tlwh[2] * tlwh[3] > args.min_box_area and not vertical: |
| 200 | online_tlwhs.append(tlwh) |
| 201 | online_ids.append(tid) |
| 202 | online_scores.append(t.score) |
| 203 | # save results |
| 204 | results.append( |
| 205 | f"{frame_id},{tid},{tlwh[0]:.2f},{tlwh[1]:.2f},{tlwh[2]:.2f},{tlwh[3]:.2f},{t.score:.2f},-1,-1,-1\n" |
| 206 | ) |
| 207 | timer.toc() |
| 208 | online_im = plot_tracking( |
| 209 | img_info['raw_img'], online_tlwhs, online_ids, frame_id=frame_id, fps=1. / timer.average_time |
| 210 | ) |
| 211 | else: |
| 212 | timer.toc() |
| 213 | online_im = img_info['raw_img'] |
| 214 | |
| 215 | # result_image = predictor.visual(outputs[0], img_info, predictor.confthre) |
| 216 | if args.save_result: |
| 217 | timestamp = time.strftime("%Y_%m_%d_%H_%M_%S", current_time) |
| 218 | save_folder = osp.join(vis_folder, timestamp) |
| 219 | os.makedirs(save_folder, exist_ok=True) |
| 220 | cv2.imwrite(osp.join(save_folder, osp.basename(img_path)), online_im) |
| 221 | |
| 222 | if frame_id % 20 == 0: |
| 223 | logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1. / max(1e-5, timer.average_time))) |
| 224 | |
| 225 | ch = cv2.waitKey(0) |
| 226 | if ch == 27 or ch == ord("q") or ch == ord("Q"): |
| 227 | break |
| 228 | |
| 229 | if args.save_result: |
| 230 | res_file = osp.join(vis_folder, f"{timestamp}.txt") |
| 231 | with open(res_file, 'w') as f: |
| 232 | f.writelines(results) |
| 233 | logger.info(f"save results to {res_file}") |
| 234 | |
| 235 |
no test coverage detected