Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,911 Bytes
9e15541 |
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 |
import argparse
import sys
sys.path.append(".")
from pathlib import Path
import subprocess
import cv2
import numpy as np
from tqdm import tqdm
from datasets.kitti_360.kitti_360_dataset import Kitti360Dataset
DRY_RUN = False
def main():
parser = argparse.ArgumentParser("KITTI 360 Preprocessing")
parser.add_argument("--data_path", "-d", type=str)
parser.add_argument("--resolution", "-r", default=(192, 640))
parser.add_argument("--fisheye_rotation", "-f", default=(0, -15))
parser.add_argument("--only_fisheye", "-o", action="store_true")
args = parser.parse_args()
data_path = Path(args.data_path)
resolution = args.resolution
rotation = args.fisheye_rotation
only_fisheye = args.only_fisheye
print("Setting up dataset")
dataset = Kitti360Dataset(
data_path=data_path,
pose_path=data_path / "data_poses",
split_path=None,
return_stereo=True,
frame_count=1,
fisheye_rotation=rotation,
color_aug=False,
return_segmentation=False,
)
print("Setting up folders...")
for i in tqdm(range(len(dataset))):
sequence, id, is_right = dataset._datapoints[i]
if is_right:
continue
image_00 = data_path / "data_2d_raw" / sequence / "image_00" / f"data_{resolution[0]}x{resolution[1]}"
image_01 = data_path / "data_2d_raw" / sequence / "image_01" / f"data_{resolution[0]}x{resolution[1]}"
image_02 = data_path / "data_2d_raw" / sequence / "image_02" / f"data_{resolution[0]}x{resolution[1]}_{rotation[0]}x{rotation[1]}"
image_03 = data_path / "data_2d_raw" / sequence / "image_03" / f"data_{resolution[0]}x{resolution[1]}_{rotation[0]}x{rotation[1]}"
img_id = dataset._img_ids[sequence][id]
if (image_00 / f"{img_id:010d}.png").exists():
continue
data = dataset[i]
image_00.mkdir(exist_ok=True, parents=True)
image_01.mkdir(exist_ok=True, parents=True)
image_02.mkdir(exist_ok=True, parents=True)
image_03.mkdir(exist_ok=True, parents=True)
img_00 = (np.transpose(data["imgs"][0].numpy(), (1, 2, 0)) * .5 + .5) * 255.
img_01 = (np.transpose(data["imgs"][1].numpy(), (1, 2, 0)) * .5 + .5) * 255.
img_02 = (np.transpose(data["imgs"][2].numpy(), (1, 2, 0)) * .5 + .5) * 255.
img_03 = (np.transpose(data["imgs"][3].numpy(), (1, 2, 0)) * .5 + .5) * 255.
if not only_fisheye:
cv2.imwrite(str(image_00 / f"{img_id:010d}.png"), cv2.cvtColor(img_00, cv2.COLOR_RGB2BGR))
cv2.imwrite(str(image_01 / f"{img_id:010d}.png"), cv2.cvtColor(img_01, cv2.COLOR_RGB2BGR))
cv2.imwrite(str(image_02 / f"{img_id:010d}.png"), cv2.cvtColor(img_02, cv2.COLOR_RGB2BGR))
cv2.imwrite(str(image_03 / f"{img_id:010d}.png"), cv2.cvtColor(img_03, cv2.COLOR_RGB2BGR))
if __name__ == "__main__":
main()
|