Browse Source

Hybrid auto-labelling support (#1646)

5.0
Glenn Jocher GitHub 3 years ago
parent
commit
86f4247515
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 4 deletions
  1. +6
    -4
      test.py

+ 6
- 4
test.py View File

dataloader=None, dataloader=None,
save_dir=Path(''), # for saving images save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling save_txt=False, # for auto-labelling
save_conf=False,
save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences
plots=True, plots=True,
log_imgs=0): # number of logged images log_imgs=0): # number of logged images


else: # called directly else: # called directly
set_logging() set_logging()
device = select_device(opt.device, batch_size=batch_size) device = select_device(opt.device, batch_size=batch_size)
save_txt = opt.save_txt # save *.txt labels


# Directories # Directories
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run


# Run NMS # Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_txt else [] # for autolabelling
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_synchronized() t = time_synchronized()
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb) output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb)
t1 += time_synchronized() - t t1 += time_synchronized() - t
parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--verbose', action='store_true', help='report mAP by class') parser.add_argument('--verbose', action='store_true', help='report mAP by class')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
parser.add_argument('--project', default='runs/test', help='save to project/name') parser.add_argument('--project', default='runs/test', help='save to project/name')
opt.single_cls, opt.single_cls,
opt.augment, opt.augment,
opt.verbose, opt.verbose,
save_txt=opt.save_txt,
save_txt=opt.save_txt | opt.save_hybrid,
save_hybrid=opt.save_hybrid,
save_conf=opt.save_conf, save_conf=opt.save_conf,
) )



Loading…
Cancel
Save