| 294 | |
| 295 | @torch.no_grad() |
| 296 | def inference(model, |
| 297 | criterion, |
| 298 | postprocessors, |
| 299 | data_loader, |
| 300 | device, |
| 301 | output_dir, |
| 302 | wo_class_error=False, |
| 303 | tmpdir=None, |
| 304 | gpu_collect=False, |
| 305 | args=None, |
| 306 | logger=None): |
| 307 | try: |
| 308 | need_tgt_for_training = args.use_dn |
| 309 | except: |
| 310 | need_tgt_for_training = False |
| 311 | model.eval() |
| 312 | criterion.eval() |
| 313 | |
| 314 | metric_logger = utils.MetricLogger(delimiter=' ') |
| 315 | if not wo_class_error: |
| 316 | metric_logger.add_meter( |
| 317 | 'class_error', utils.SmoothedValue(window_size=1, |
| 318 | fmt='{value:.2f}')) |
| 319 | header = 'Test:' |
| 320 | iou_types = tuple(k for k in ('bbox', 'keypoints')) |
| 321 | try: |
| 322 | useCats = args.useCats |
| 323 | except: |
| 324 | useCats = True |
| 325 | if not useCats: |
| 326 | print('useCats: {} !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'.format( |
| 327 | useCats)) |
| 328 | |
| 329 | dataset = data_loader.dataset |
| 330 | rank, world_size = get_dist_info() |
| 331 | time.sleep(2) |
| 332 | for data_batch in metric_logger.log_every( |
| 333 | data_loader, 10, header, logger=logger): |
| 334 | # i = i+1 |
| 335 | with torch.cuda.amp.autocast(enabled=args.amp): |
| 336 | if need_tgt_for_training: |
| 337 | # outputs = model(samples, targets) |
| 338 | outputs, targets, data_batch_nc = model(data_batch) |
| 339 | else: |
| 340 | outputs,targets, data_batch_nc = model(data_batch) |
| 341 | |
| 342 | orig_target_sizes = torch.stack([t["size"] for t in targets], dim=0) |
| 343 | result = postprocessors['bbox'](outputs, orig_target_sizes, targets, data_batch_nc) |
| 344 | |
| 345 | dataset.inference(result) |
| 346 | |
| 347 | if rank == 0 and args.to_vid: |
| 348 | if hasattr(dataset,'result_img_dir'): |
| 349 | images_to_video(dataset.result_img_dir, os.path.join(dataset.output_path, dataset.img_name+'_demo.mp4'),remove_raw_file=False, fps=30) |
| 350 | # shutil.rmtree(dataset.result_img_dir) |
| 351 | # shutil.rmtree(dataset.tmp_dir) |
| 352 | |
| 353 | |