File size: 5,610 Bytes
6979012 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | import random
import os
from pathlib import Path
from superpoint_pruning.distillation.utils import rescale_image, load_grayscale_image
from superpoint_pruning.models.superpoint import SuperPoint
from superpoint_pruning.paths import DATA_ROOT, DEFAULT_DATASET_NAME
import numpy as np
import torch
from tqdm import tqdm
import urllib.request
import zipfile
import argparse
DEFAULT_DATASET_URL = "http://rpg.ifi.uzh.ch/datasets/uzh-fpv-newer-versions/v3/indoor_forward_3_snapdragon_with_gt.zip"
DEFAULT_GT_INDEX_FILE = "train_base.txt"
def download_dataset_from_url(url: str, output_dir: Path) -> None:
zip_path = DATA_ROOT / "download.zip"
if output_dir.exists():
print(
f"Dataset directory {output_dir.resolve()} already exists. Skipping download."
)
return
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Downloading dataset from {url} to {zip_path.resolve()}")
urllib.request.urlretrieve(url, zip_path)
print(f"Extracting files from {zip_path.resolve()} to {output_dir.resolve()}")
with zipfile.ZipFile(zip_path, "r") as archive:
archive.extractall(output_dir)
zip_path.unlink()
print(f"Files extracted to {output_dir.resolve()}")
def generate_data_split(
image_dir: Path, split_ratio: float = 0.8
) -> tuple[list[str], list[str]]:
random.seed(42)
images = [path.name for path in image_dir.glob("*.png")]
random.shuffle(images)
split_index = int(len(images) * split_ratio)
train_images = images[:split_index]
val_images = images[split_index:]
with open(DATA_ROOT / "train.txt", "w") as f:
for image in train_images:
f.write(image + "\n")
with open(DATA_ROOT / "val.txt", "w") as f:
for image in val_images:
f.write(image + "\n")
print(
f"Train images: {len(train_images)}, Val images: {len(val_images)} - Saved to {DATA_ROOT}"
)
@torch.inference_mode()
def generate_ground_truth(
image_dir: Path,
id_file: Path,
num_keypoints: int = 512,
image_size: tuple[int, int] = (640, 480),
num_images: int = 250,
) -> None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
images = [path.name for path in image_dir.glob("*.png")]
model = SuperPoint(num_keypoints=num_keypoints, return_dense=True)
model.eval().to(device)
ids = id_file.read_text().splitlines()
ids = ids[:num_images]
id_set = set(ids)
selected_images = [image_id for image_id in images if image_id in id_set]
Path(DATA_ROOT / f"gt_id_file_{Path(id_file).stem}_{num_images}.txt").write_text(
"\n".join(selected_images) + "\n"
)
desc = None
kpts = None
for index, image_id in enumerate(tqdm(selected_images)):
image_path = os.path.join(image_dir, image_id)
image = load_grayscale_image(image_path)
img, _ = rescale_image(image, image_size)
img = img[None, None].astype(np.float32)
img = torch.from_numpy(img).to(device)
keypoint_logits, descriptor_logits = model(img)
desc_np = descriptor_logits.detach().cpu().numpy()[0]
kpts_np = keypoint_logits.detach().cpu().numpy()[0]
if desc is None or kpts is None:
desc = np.lib.format.open_memmap(
DATA_ROOT / f"desc_{id_file.stem}_{num_images}.npy",
mode="w+",
dtype=desc_np.dtype,
shape=(len(selected_images), *desc_np.shape),
)
kpts = np.lib.format.open_memmap(
DATA_ROOT / f"kpts_{id_file.stem}_{num_images}.npy",
mode="w+",
dtype=kpts_np.dtype,
shape=(len(selected_images), *kpts_np.shape),
)
desc[index] = desc_np
kpts[index] = kpts_np
if desc is not None:
desc.flush()
if kpts is not None:
kpts.flush()
print(f"Ground truth files generated to {DATA_ROOT}")
def add_parser_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument("--download-url", type=str, default=DEFAULT_DATASET_URL)
parser.add_argument("--output-dir-name", type=str, default=DEFAULT_DATASET_NAME)
parser.add_argument("--gt-index-file", type=str, default=DEFAULT_GT_INDEX_FILE)
parser.add_argument("--generate-data-split", action="store_true", default=False)
parser.add_argument("--generate-gt", action="store_true", default=False)
parser.add_argument("--num-keypoints", type=int, default=512)
parser.add_argument("--width", type=int, default=640)
parser.add_argument("--height", type=int, default=480)
parser.add_argument("--num-images", type=int, default=250)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Setup the dataset for distillation and evaluation."
)
add_parser_args(parser)
return parser
def main(args: argparse.Namespace) -> None:
if args.download_url and args.output_dir_name:
download_dataset_from_url(
args.download_url, DATA_ROOT / "datasets" / args.output_dir_name
)
if args.generate_data_split:
generate_data_split(DATA_ROOT / "datasets" / args.output_dir_name / "img", 0.8)
if args.generate_gt:
generate_ground_truth(
DATA_ROOT / "datasets" / args.output_dir_name / "img",
DATA_ROOT / args.gt_index_file,
num_images=args.num_images,
image_size=(args.width, args.height),
num_keypoints=args.num_keypoints,
)
if __name__ == "__main__":
main(build_parser().parse_args())
|