@torch.no_grad()
def main(conf, image_dir, export_dir, as_half=False):
    logging.info(f'''Extracting local features with configuration:
{pprint.pformat(conf)}''')
    device = ('cuda' if torch.cuda.is_available() else 'cpu')
    Model = dynamic_load(extractors, conf['model']['name'])
    model = Model(conf['model']).eval().to(device)
    loader = ImageDataset(image_dir, conf['preprocessing'])
    loader = torch.utils.data.DataLoader(loader, num_workers=1)
    feature_path = Path(export_dir, (conf['output'] + '.h5'))
    feature_path.parent.mkdir(exist_ok=True, parents=True)
    feature_file = h5py.File(str(feature_path), 'a')
    for data in tqdm(loader):
        pred = model(map_tensor(data, (lambda x: x.to(device))))
        pred = {k: v[0].cpu().numpy() for (k, v) in pred.items()}
        pred['image_size'] = original_size = data['original_size'][0].numpy()
        if ('keypoints' in pred):
            size = np.array(data['image'].shape[(- 2):][::(- 1)])
            scales = (original_size / size).astype(np.float32)
            pred['keypoints'] = (((pred['keypoints'] + 0.5) * scales[None]) - 0.5)
        if as_half:
            for k in pred:
                dt = pred[k].dtype
                if ((dt == np.float32) and (dt != np.float16)):
                    pred[k] = pred[k].astype(np.float16)
        grp = feature_file.create_group(data['name'][0])
        for (k, v) in pred.items():
            grp.create_dataset(k, data=v)
        del pred
    feature_file.close()
    logging.info('Finished exporting features.')
@torch.no_grad()
def main(conf, image_dir, export_dir, as_half=False):
    logging.info(f'Extracting local features with configuration:
{pprint.pformat(conf)}')
    device = ('cuda' if torch.cuda.is_available() else 'cpu')
    Model = dynamic_load(extractors, conf['model']['name'])
    model = Model(conf['model']).eval().to(device)
    loader = ImageDataset(image_dir, conf['preprocessing'])
    loader = torch.utils.data.DataLoader(loader, num_workers=1)
    feature_path = Path(export_dir, (conf['output'] + '.h5'))
    feature_path.parent.mkdir(exist_ok=True, parents=True)
    feature_file = h5py.File(str(feature_path), 'a')
    for data in tqdm(loader):
        pred = model(map_tensor(data, (lambda x: x.to(device))))
        pred = {k: v[0].cpu().numpy() for (k, v) in pred.items()}
        pred['image_size'] = original_size = data['original_size'][0].numpy()
        if ('keypoints' in pred):
            size = np.array(data['image'].shape[(- 2):][::(- 1)])
            scales = (original_size / size).astype(np.float32)
            pred['keypoints'] = (((pred['keypoints'] + 0.5) * scales[None]) - 0.5)
        if as_half:
            for k in pred:
                dt = pred[k].dtype
                if ((dt == np.float32) and (dt != np.float16)):
                    pred[k] = pred[k].astype(np.float16)
        grp = feature_file.create_group(data['name'][0])
        for (k, v) in pred.items():
            grp.create_dataset(k, data=v)
        del pred
    feature_file.close()
    logging.info('Finished exporting features.')