Spaces:
Runtime error
Runtime error
| import os | |
| import cv2 | |
| import argparse | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| from torchvision import datasets, transforms | |
| from torch.autograd import Variable | |
| from network_v0.model import PointModel | |
| from datasets.hp_loader import PatchesDataset | |
| from torch.utils.data import DataLoader | |
| from evaluation.evaluate import evaluate_keypoint_net | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Testing") | |
| parser.add_argument("--device", default=0, type=int, help="which gpu to run on.") | |
| parser.add_argument("--test_dir", required=True, type=str, help="Test data path.") | |
| opt = parser.parse_args() | |
| torch.manual_seed(0) | |
| use_gpu = torch.cuda.is_available() | |
| if use_gpu: | |
| torch.cuda.set_device(opt.device) | |
| # Load data in 320x240 | |
| hp_dataset_320x240 = PatchesDataset( | |
| root_dir=opt.test_dir, use_color=True, output_shape=(320, 240), type="all" | |
| ) | |
| data_loader_320x240 = DataLoader( | |
| hp_dataset_320x240, | |
| batch_size=1, | |
| pin_memory=False, | |
| shuffle=False, | |
| num_workers=4, | |
| worker_init_fn=None, | |
| sampler=None, | |
| ) | |
| # Load data in 640x480 | |
| hp_dataset_640x480 = PatchesDataset( | |
| root_dir=opt.test_dir, use_color=True, output_shape=(640, 480), type="all" | |
| ) | |
| data_loader_640x480 = DataLoader( | |
| hp_dataset_640x480, | |
| batch_size=1, | |
| pin_memory=False, | |
| shuffle=False, | |
| num_workers=4, | |
| worker_init_fn=None, | |
| sampler=None, | |
| ) | |
| # Load model | |
| model = PointModel(is_test=True) | |
| ckpt = torch.load("./checkpoints/PointModel_v0.pth") | |
| model.load_state_dict(ckpt["model_state"]) | |
| model = model.eval() | |
| if use_gpu: | |
| model = model.cuda() | |
| print("Evaluating in 320x240, 300 points") | |
| rep, loc, c1, c3, c5, mscore = evaluate_keypoint_net( | |
| data_loader_320x240, model, output_shape=(320, 240), top_k=300 | |
| ) | |
| print("Repeatability: {0:.3f}".format(rep)) | |
| print("Localization Error: {0:.3f}".format(loc)) | |
| print("H-1 Accuracy: {:.3f}".format(c1)) | |
| print("H-3 Accuracy: {:.3f}".format(c3)) | |
| print("H-5 Accuracy: {:.3f}".format(c5)) | |
| print("Matching Score: {:.3f}".format(mscore)) | |
| print("\n") | |
| print("Evaluating in 640x480, 1000 points") | |
| rep, loc, c1, c3, c5, mscore = evaluate_keypoint_net( | |
| data_loader_640x480, model, output_shape=(640, 480), top_k=1000 | |
| ) | |
| print("Repeatability: {0:.3f}".format(rep)) | |
| print("Localization Error: {0:.3f}".format(loc)) | |
| print("H-1 Accuracy: {:.3f}".format(c1)) | |
| print("H-3 Accuracy: {:.3f}".format(c3)) | |
| print("H-5 Accuracy: {:.3f}".format(c5)) | |
| print("Matching Score: {:.3f}".format(mscore)) | |
| print("\n") | |
| if __name__ == "__main__": | |
| main() | |