(data,
weights=None,
batch_size=16,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_conf=False,
plots=True,
log_imgs=0)
| 28 | |
| 29 | |
| 30 | def test(data, |
| 31 | weights=None, |
| 32 | batch_size=16, |
| 33 | imgsz=640, |
| 34 | conf_thres=0.001, |
| 35 | iou_thres=0.6, # for NMS |
| 36 | save_json=False, |
| 37 | single_cls=False, |
| 38 | augment=False, |
| 39 | verbose=False, |
| 40 | model=None, |
| 41 | dataloader=None, |
| 42 | save_dir=Path(''), # for saving images |
| 43 | save_txt=False, # for auto-labelling |
| 44 | save_conf=False, |
| 45 | plots=True, |
| 46 | log_imgs=0): # number of logged images |
| 47 | |
| 48 | # Initialize/load model and set device |
| 49 | training = model is not None |
| 50 | if training: # called by train.py |
| 51 | device = next(model.parameters()).device # get model device |
| 52 | |
| 53 | else: # called directly |
| 54 | set_logging() |
| 55 | device = select_device(opt.device, batch_size=batch_size) |
| 56 | save_txt = opt.save_txt # save *.txt labels |
| 57 | |
| 58 | # Directories |
| 59 | save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run |
| 60 | (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir |
| 61 | |
| 62 | # Load model |
| 63 | model = Darknet(opt.cfg).to(device) |
| 64 | |
| 65 | # load model |
| 66 | try: |
| 67 | ckpt = torch.load(weights[0], map_location=device) # load checkpoint |
| 68 | ckpt['model'] = {k: v for k, v in ckpt['model'].items() if model.state_dict()[k].numel() == v.numel()} |
| 69 | model.load_state_dict(ckpt['model'], strict=False) |
| 70 | except: |
| 71 | load_darknet_weights(model, weights[0]) |
| 72 | imgsz = check_img_size(imgsz, s=64) # check img_size |
| 73 | |
| 74 | # Half |
| 75 | half = device.type != 'cpu' # half precision only supported on CUDA |
| 76 | if half: |
| 77 | model.half() |
| 78 | |
| 79 | # Configure |
| 80 | model.eval() |
| 81 | is_coco = data.endswith('coco.yaml') # is COCO dataset |
| 82 | with open(data) as f: |
| 83 | data = yaml.load(f, Loader=yaml.FullLoader) # model dict |
| 84 | check_dataset(data) # check |
| 85 | nc = 1 if single_cls else int(data['nc']) # number of classes |
| 86 | iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 |
| 87 | niou = iouv.numel() |
no test coverage detected