Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| # @Author : xuelun | |
| import os | |
| import torch | |
| import warnings | |
| import numpy as np | |
| from tqdm import tqdm | |
| from os.path import join | |
| from pathlib import Path | |
| from argparse import ArgumentParser | |
| from hloc import pairs_from_exhaustive | |
| from hloc import extract_features, match_features, match_dense, reconstruction | |
| from hloc.utils import segment | |
| from hloc.utils.io import read_image | |
| from hloc.match_dense import ImagePairDataset | |
| from networks.lightglue.superpoint import SuperPoint | |
| from networks.lightglue.models.matchers.lightglue import LightGlue | |
| from networks.mit_semseg.models import ModelBuilder, SegmentationModule | |
| def segmentation(images, segment_root, matcher_conf): | |
| # initial device | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| # initial segmentation mode | |
| net_encoder = ModelBuilder.build_encoder( | |
| arch='resnet50dilated', | |
| fc_dim=2048, | |
| weights='weights/encoder_epoch_20.pth') | |
| net_decoder = ModelBuilder.build_decoder( | |
| arch='ppm_deepsup', | |
| fc_dim=2048, | |
| num_class=150, | |
| weights='weights/decoder_epoch_20.pth', | |
| use_softmax=True) | |
| crit = torch.nn.NLLLoss(ignore_index=-1) | |
| segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) | |
| segmentation_module = segmentation_module.to(device).eval() | |
| # initial data reader | |
| dataset = ImagePairDataset(None, matcher_conf["preprocessing"], None) | |
| # Segment images | |
| image_list = sorted(os.listdir(images)) | |
| with torch.no_grad(): | |
| for img in tqdm(image_list): | |
| segment_path = join(segment_root, '{}.npy'.format(img[:-4])) | |
| if not os.path.exists(segment_path): | |
| rgb = read_image(images / img, dataset.conf.grayscale) | |
| mask = segment(rgb, 1920, device, segmentation_module) | |
| np.save(segment_path, mask) | |
| def main(scene_name, version): | |
| # Setup | |
| images = Path('inputs') / scene_name / 'images' | |
| outputs = Path('outputs') / scene_name / version | |
| outputs.mkdir(parents=True, exist_ok=True) | |
| os.environ['GIMRECONSTRUCTION'] = str(outputs) | |
| segment_root = Path('outputs') / scene_name / 'segment' | |
| segment_root.mkdir(parents=True, exist_ok=True) | |
| sfm_dir = outputs / 'sparse' | |
| mvs_path = outputs / 'dense' | |
| database_path = sfm_dir / 'database.db' | |
| image_pairs = outputs / 'pairs-near.txt' | |
| feature_conf = matcher_conf = None | |
| if version == 'gim_dkm': | |
| feature_conf = None | |
| matcher_conf = match_dense.confs[version] | |
| elif version == 'gim_lightglue': | |
| feature_conf = extract_features.confs['gim_superpoint'] | |
| matcher_conf = match_features.confs[version] | |
| # Find image pairs via pair-wise image | |
| exhaustive_pairs = pairs_from_exhaustive.main(image_pairs, image_list=images) | |
| segmentation(images, segment_root, matcher_conf) | |
| # Extract and match local features | |
| with warnings.catch_warnings(): | |
| warnings.filterwarnings("ignore", category=UserWarning) | |
| if version == 'gim_dkm': | |
| feature_path, match_path = match_dense.main(matcher_conf, image_pairs, | |
| images, outputs) | |
| elif version == 'gim_lightglue': | |
| checkpoints_path = join('weights', 'gim_lightglue_100h.ckpt') | |
| detector = SuperPoint({ | |
| 'max_num_keypoints': 2048, | |
| 'force_num_keypoints': True, | |
| 'detection_threshold': 0.0, | |
| 'nms_radius': 3, | |
| 'trainable': False, | |
| }) | |
| state_dict = torch.load(checkpoints_path, map_location='cpu') | |
| if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] | |
| for k in list(state_dict.keys()): | |
| if k.startswith('model.'): | |
| state_dict.pop(k) | |
| if k.startswith('superpoint.'): | |
| state_dict[k.replace('superpoint.', '', 1)] = state_dict.pop(k) | |
| detector.load_state_dict(state_dict) | |
| model = LightGlue({ | |
| 'filter_threshold': 0.1, | |
| 'flash': False, | |
| 'checkpointed': True, | |
| }) | |
| state_dict = torch.load(checkpoints_path, map_location='cpu') | |
| if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] | |
| for k in list(state_dict.keys()): | |
| if k.startswith('superpoint.'): | |
| state_dict.pop(k) | |
| if k.startswith('model.'): | |
| state_dict[k.replace('model.', '', 1)] = state_dict.pop(k) | |
| model.load_state_dict(state_dict) | |
| feature_path = extract_features.main(feature_conf, images, outputs, | |
| model=detector) | |
| match_path = match_features.main(matcher_conf, image_pairs, | |
| feature_conf['output'], outputs, | |
| model=model) | |
| # sparse reconstruction | |
| reconstruction.main(sfm_dir, images, image_pairs, feature_path, match_path) | |
| if __name__ == '__main__': | |
| parser = ArgumentParser() | |
| parser.add_argument('--scene_name', type=str) | |
| parser.add_argument('--version', type=str, choices={'gim_dkm', 'gim_lightglue'}, | |
| default='gim_dkm') | |
| args = parser.parse_args() | |
| main(args.scene_name, args.version) | |