Spaces:
Running
Running
File size: 4,218 Bytes
9665c2c b0cf684 9665c2c b0cf684 9665c2c b0cf684 9665c2c b0cf684 9665c2c 38e60a1 9665c2c |
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 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
import argparse
import shutil
import zipfile
from pathlib import Path
import numpy as np
from tqdm.auto import tqdm
from ... import logger
from ...osm.tiling import TileManager
from ...osm.viz import GeoPlotter
from ...utils.geo import BoundaryBox, Projection
from ...utils.io import DATA_URL, download_file
from .dataset import KittiDataModule
from .utils import parse_gps_file
split_files = ["test1_files.txt", "test2_files.txt", "train_files.txt"]
def prepare_osm(
data_dir,
osm_path,
output_path,
tile_margin=512,
ppm=2,
):
all_latlon = []
for gps_path in data_dir.glob("2011_*/*/oxts/data/*.txt"):
all_latlon.append(parse_gps_file(gps_path)[0])
if not all_latlon:
raise ValueError(f"Cannot find any GPS file in {data_dir}.")
all_latlon = np.stack(all_latlon)
projection = Projection.from_points(all_latlon)
all_xy = projection.project(all_latlon)
bbox_map = BoundaryBox(all_xy.min(0), all_xy.max(0)) + tile_margin
plotter = GeoPlotter()
plotter.points(all_latlon, "red", name="GPS")
plotter.bbox(projection.unproject(bbox_map), "blue", "tiling bounding box")
plotter.fig.write_html(data_dir / "split_kitti.html")
tile_manager = TileManager.from_bbox(
projection,
bbox_map,
ppm,
path=osm_path,
)
tile_manager.save(output_path)
return tile_manager
def download(data_dir: Path):
data_dir.mkdir(exist_ok=True, parents=True)
this_dir = Path(__file__).parent
seqs = set()
for f in split_files:
shutil.copy(this_dir / f, data_dir)
with open(this_dir / f, "r") as fid:
info = fid.read()
for line in info.split("\n"):
if line:
seq = line.split()[0].split("/")[1][: -len("_sync")]
seqs.add(seq)
dates = {"_".join(s.split("_")[:3]) for s in seqs}
logger.info("Downloading data for %d sequences in %d dates", len(seqs), len(dates))
for seq in tqdm(seqs):
logger.info("Working on %s.", seq)
date = "_".join(seq.split("_")[:3])
url = f"https://s3.eu-central-1.amazonaws.com/avg-kitti/raw_data/{seq}/{seq}_sync.zip" # noqa E501
seq_dir = data_dir / date / f"{seq}_sync"
if seq_dir.exists():
continue
zip_path = download_file(url, data_dir)
with zipfile.ZipFile(zip_path, "r") as z:
z.extractall(data_dir)
# Delete unused files to save space.
for image_index in [0, 1, 3]:
shutil.rmtree(seq_dir / f"image_0{image_index}")
shutil.rmtree(seq_dir / "velodyne_points")
Path(zip_path).unlink()
for date in tqdm(dates):
url = (
f"https://s3.eu-central-1.amazonaws.com/avg-kitti/raw_data/{date}_calib.zip"
)
zip_path = download_file(url, data_dir)
with zipfile.ZipFile(zip_path, "r") as z:
z.extractall(data_dir)
Path(zip_path).unlink()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--data_dir", type=Path, default=Path(KittiDataModule.default_cfg["data_dir"])
)
parser.add_argument("--pixel_per_meter", type=int, default=2)
parser.add_argument("--generate_tiles", action="store_true")
args = parser.parse_args()
args.data_dir.mkdir(exist_ok=True, parents=True)
download(args.data_dir)
tiles_path = args.data_dir / KittiDataModule.default_cfg["tiles_filename"]
if args.generate_tiles:
logger.info("Generating the map tiles.")
osm_filename = "karlsruhe.osm"
osm_path = args.data_dir / osm_filename
if not osm_path.exists():
logger.info("Downloading OSM raw data.")
download_file(DATA_URL + f"/osm/{osm_filename}", osm_path)
if not osm_path.exists():
raise FileNotFoundError(f"No OSM data file at {osm_path}.")
prepare_osm(args.data_dir, osm_path, tiles_path, ppm=args.pixel_per_meter)
(args.data_dir / ".downloaded").touch()
else:
logger.info("Downloading pre-generated map tiles.")
download_file(DATA_URL + "/tiles/kitti.pkl", tiles_path)
|