def eval_best(cfg, models, device, test_loader, interp, fixed_test_size, verbose): assert len(models) == 1, 'Not yet supported multi models in this mode' assert osp.exists(cfg.TEST.SNAPSHOT_DIR[0]), 'SNAPSHOT_DIR is not found' start_iter = cfg.TEST.SNAPSHOT_STEP step = cfg.TEST.SNAPSHOT_STEP max_iter = cfg.TEST.SNAPSHOT_MAXITER cache_path = osp.join(cfg.TEST.SNAPSHOT_DIR[0], 'all_res.pkl') if osp.exists(cache_path): all_res = pickle_load(cache_path) else: all_res = {} cur_best_miou = -1 cur_best_model = '' for i_iter in range(start_iter, max_iter + 1, step): restore_from = osp.join(cfg.TEST.SNAPSHOT_DIR[0], f'model_{i_iter}.pth') if not osp.exists(restore_from): # continue if cfg.TEST.WAIT_MODEL: print('Waiting for model..!') while not osp.exists(restore_from): time.sleep(5) print("Evaluating model", restore_from) if i_iter not in all_res.keys(): load_checkpoint_for_evaluation(models[0], restore_from, device) # eval hist = np.zeros((cfg.NUM_CLASSES, cfg.NUM_CLASSES)) # for index, batch in enumerate(test_loader): # image, _, _, name = batch test_iter = iter(test_loader) for index in tqdm(range(len(test_loader))): #TODO: need to check image, _, label, _, name, _ = next(test_iter) image, label, _, name = next(test_iter) if not fixed_test_size: interp = nn.Upsample(size=(label.shape[1], label.shape[2]), mode='bilinear', align_corners=True) with torch.no_grad(): pred_main = models[0](image.cuda(device))[1] output = interp(pred_main).cpu().data[0].numpy() output = output.transpose(1, 2, 0) output = np.argmax(output, axis=2) label = label.numpy()[0] hist += fast_hist(label.flatten(), output.flatten(), cfg.NUM_CLASSES) if verbose and index > 0 and index % 100 == 0: print('{:d} / {:d}: {:0.2f}'.format( index, len(test_loader), 100 * np.nanmean(per_class_iu(hist)))) inters_over_union_classes = per_class_iu(hist) all_res[i_iter] = inters_over_union_classes pickle_dump(all_res, cache_path) else: inters_over_union_classes = all_res[i_iter] computed_miou = round(np.nanmean(inters_over_union_classes) * 100, 2) if cur_best_miou < computed_miou: cur_best_miou = computed_miou cur_best_model = restore_from print('\tCurrent mIoU:', computed_miou) print('\tCurrent best model:', cur_best_model) print('\tCurrent best mIoU:', cur_best_miou) if verbose: display_stats(cfg, test_loader.dataset.class_names, inters_over_union_classes)
def eval_best(cfg, models, device, test_loader, interp, fixed_test_size, verbose): # -------------------------------------------------------- # # codes to initialize wandb for storing logs on its cloud wandb.init(project='FDA_integration_to_INTRA_DA') for key, val in cfg.items(): wandb.config.update({key: val}) # -------------------------------------------------------- # assert len(models) == 1, 'Not yet supported multi models in this mode' assert osp.exists(cfg.TEST.SNAPSHOT_DIR[0]), 'SNAPSHOT_DIR is not found' start_iter = cfg.TEST.SNAPSHOT_STEP step = cfg.TEST.SNAPSHOT_STEP max_iter = cfg.TEST.SNAPSHOT_MAXITER cache_path = osp.join(cfg.TEST.SNAPSHOT_DIR[0], 'all_res.pkl') if osp.exists(cache_path): cache_path = pickle_load(cache_path) else: all_res = {} cur_best_miou = -1 cur_best_model = '' for i_iter in range(start_iter, max_iter + 1, step): restore_from = osp.join(cfg.TEST.SNAPSHOT_DIR[0], f'model_{i_iter}.pth') if not osp.exists(restore_from): # continue if cfg.TEST.WAIT_MODEL: print('Waiting for model..!') while not osp.exists(restore_from): time.sleep(5) print("Evaluating model", restore_from) if i_iter not in all_res.keys(): load_checkpoint_for_evaluation(models[0], restore_from, device) # eval hist = np.zeros((cfg.NUM_CLASSES, cfg.NUM_CLASSES)) # for index, batch in enumerate(test_loader): # image, _, _, name = batch test_iter = iter(test_loader) for index in tqdm(range(len(test_loader))): image, label, _, name = next(test_iter) if not fixed_test_size: interp = nn.Upsample(size=(label.shape[1], label.shape[2]), mode='bilinear', align_corners=True) with torch.no_grad(): pred_main = models[0](image.cuda(device))[1] output = interp(pred_main).cpu().data[0].numpy() output = output.transpose(1, 2, 0) output = np.argmax(output, axis=2) label = label.numpy()[0] hist += fast_hist(label.flatten(), output.flatten(), cfg.NUM_CLASSES) if verbose and index > 0 and index % 100 == 0: print('{:d} / {:d}: {:0.2f}'.format( index, len(test_loader), 100 * np.nanmean(per_class_iu(hist)))) inters_over_union_classes = per_class_iu(hist) all_res[i_iter] = inters_over_union_classes pickle_dump(all_res, cache_path) # -------------------------------------------------------- # # save logs at weight and biases IoU_classes = {} for idx in range(cfg.NUM_CLASSES): IoU_classes[test_loader.dataset.class_names[idx]] = round( inters_over_union_classes[idx] * 100, 2) wandb.log(IoU_classes, step=(i_iter)) wandb.log( { 'mIoU19': round( np.nanmean(inters_over_union_classes) * 100, 2) }, step=(i_iter)) wandb.log( { 'val_prediction': wandb.Image( colorize_mask(np.asarray( output, dtype=np.uint8)).convert('RGB')) }, step=(i_iter)) # -------------------------------------------------------- # else: inters_over_union_classes = all_res[i_iter] computed_miou = round(np.nanmean(inters_over_union_classes) * 100, 2) if cur_best_miou < computed_miou: cur_best_miou = computed_miou cur_best_model = restore_from print('\tCurrent mIoU:', computed_miou) print('\tCurrent best model:', cur_best_model) print('\tCurrent best mIoU:', cur_best_miou) wandb.log({'best mIoU': cur_best_miou}, step=(i_iter)) if verbose: display_stats(cfg, test_loader.dataset.class_names, inters_over_union_classes)