Spaces:
Running
Running
| #!/usr/bin/env python | |
| # coding: utf-8 | |
| import os, glob, cv2 | |
| import argparse | |
| from argparse import Namespace | |
| import yaml | |
| from tqdm import tqdm | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader, SequentialSampler | |
| from src.datasets.custom_dataloader import TestDataLoader | |
| from src.utils.dataset import read_img_gray | |
| from configs.data.base import cfg as data_cfg | |
| import viz | |
| def get_model_config(method_name, dataset_name, root_dir='viz'): | |
| config_file = f'{root_dir}/configs/{method_name}.yml' | |
| with open(config_file, 'r') as f: | |
| model_conf = yaml.load(f, Loader=yaml.FullLoader)[dataset_name] | |
| return model_conf | |
| class DemoDataset(Dataset): | |
| def __init__(self, dataset_dir, img_file=None, resize=0, down_factor=16): | |
| self.dataset_dir = dataset_dir | |
| if img_file is None: | |
| self.list_img_files = glob.glob(os.path.join(dataset_dir, "*.*")) | |
| self.list_img_files.sort() | |
| else: | |
| with open(img_file) as f: | |
| self.list_img_files = [os.path.join(dataset_dir, img_file.strip()) for img_file in f.readlines()] | |
| self.resize = resize | |
| self.down_factor = down_factor | |
| def __len__(self): | |
| return len(self.list_img_files) | |
| def __getitem__(self, idx): | |
| img_path = self.list_img_files[idx] #os.path.join(self.dataset_dir, self.list_img_files[idx]) | |
| img, scale = read_img_gray(img_path, resize=self.resize, down_factor=self.down_factor) | |
| return {"img": img, "id": idx, "img_path": img_path} | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='Visualize matches') | |
| parser.add_argument('--gpu', '-gpu', type=str, default='0') | |
| parser.add_argument('--method', type=str, default=None) | |
| parser.add_argument('--dataset_dir', type=str, default='data/aachen-day-night') | |
| parser.add_argument('--pair_dir', type=str, default=None) | |
| parser.add_argument( | |
| '--dataset_name', type=str, choices=['megadepth', 'scannet', 'aachen_v1.1', 'inloc'], default='megadepth' | |
| ) | |
| parser.add_argument('--measure_time', action="store_true") | |
| parser.add_argument('--no_viz', action="store_true") | |
| parser.add_argument('--compute_eval_metrics', action="store_true") | |
| parser.add_argument('--run_demo', action="store_true") | |
| args = parser.parse_args() | |
| model_cfg = get_model_config(args.method, args.dataset_name) | |
| class_name = model_cfg["class"] | |
| model = viz.__dict__[class_name](model_cfg) | |
| # all_args = Namespace(**vars(args), **model_cfg) | |
| if not args.run_demo: | |
| if args.dataset_name == 'megadepth': | |
| from configs.data.megadepth_test_1500 import cfg | |
| data_cfg.merge_from_other_cfg(cfg) | |
| elif args.dataset_name == 'scannet': | |
| from configs.data.scannet_test_1500 import cfg | |
| data_cfg.merge_from_other_cfg(cfg) | |
| elif args.dataset_name == 'aachen_v1.1': | |
| data_cfg.merge_from_list(["DATASET.TEST_DATA_SOURCE", "aachen_v1.1", | |
| "DATASET.TEST_DATA_ROOT", os.path.join(args.dataset_dir, "images/images_upright"), | |
| "DATASET.TEST_LIST_PATH", args.pair_dir, | |
| "DATASET.TEST_IMGSIZE", model_cfg["imsize"]]) | |
| elif args.dataset_name == 'inloc': | |
| data_cfg.merge_from_list(["DATASET.TEST_DATA_SOURCE", "inloc", | |
| "DATASET.TEST_DATA_ROOT", args.dataset_dir, | |
| "DATASET.TEST_LIST_PATH", args.pair_dir, | |
| "DATASET.TEST_IMGSIZE", model_cfg["imsize"]]) | |
| has_ground_truth = str(data_cfg.DATASET.TEST_DATA_SOURCE).lower() in ["megadepth", "scannet"] | |
| dataloader = TestDataLoader(data_cfg) | |
| with torch.no_grad(): | |
| for data_dict in tqdm(dataloader): | |
| for k, v in data_dict.items(): | |
| if isinstance(v, torch.Tensor): | |
| data_dict[k] = v.cuda() if torch.cuda.is_available() else v | |
| img_root_dir = data_cfg.DATASET.TEST_DATA_ROOT | |
| model.match_and_draw(data_dict, root_dir=img_root_dir, ground_truth=has_ground_truth, | |
| measure_time=args.measure_time, viz_matches=(not args.no_viz)) | |
| if args.measure_time: | |
| print("Running time for each image is {} miliseconds".format(model.measure_time())) | |
| if args.compute_eval_metrics and has_ground_truth: | |
| model.compute_eval_metrics() | |
| else: | |
| demo_dataset = DemoDataset(args.dataset_dir, img_file=args.pair_dir, resize=640) | |
| sampler = SequentialSampler(demo_dataset) | |
| dataloader = DataLoader(demo_dataset, batch_size=1, sampler=sampler) | |
| writer = cv2.VideoWriter('topicfm_demo.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 15, (640 * 2 + 5, 480 * 2 + 10)) | |
| model.run_demo(iter(dataloader), writer) #, output_dir="demo", no_display=True) | |