construct_evaluation_metadata.py 4.02 KB
Newer Older
Clément Pinard's avatar
Clément Pinard committed
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
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from path import Path
import random
import pandas as pd
from tqdm import tqdm
import numpy as np


parser = ArgumentParser(description='Convert dataset to KITTI format, optionnally create a visualization video',
                        formatter_class=ArgumentDefaultsHelpFormatter)

parser.add_argument('--dataset_dir', metavar='DIR', type=Path, required=True,
                    help='folder containing the converted dataset')
parser.add_argument('--split', type=float, default=0,
                    help="proportion between train and test. By default, the whole dataset serves for evaluation")
parser.add_argument('--seed', type=int, default=0,
                    help='seed for random classification between train and val')
parser.add_argument('--verbose', '-v', action='count', default=0)
parser.add_argument('--max_num_samples', default=500, type=int)
parser.add_argument('--min_shift', default=0, type=int,
                    help='Minimum of former frames with valid odometry')
parser.add_argument('--allow_interpolated_frames', action='store_true',
                    help='If set, will consider frames with interpolated odometry to be valid')


def main():
    args = parser.parse_args()
    random.seed(args.seed)
    folder_tree = args.dataset_dir.walkdirs()
    video_sequences = []
    for f in folder_tree:
        if f.files('*.jpg'):
            video_sequences.append(f)

    random.shuffle(video_sequences)
    n = len(video_sequences)
    train_videos = video_sequences[:int(n*args.split)]
    test_videos = video_sequences[int(n*args.split):]

    total_valid_frames = []
    for v in tqdm(test_videos):
        metadata = pd.read_csv(v/"metadata.csv")
        valid_odometry_frames = metadata["registered"]
        if not args.allow_interpolated_frames:
            valid_odometry_frames = valid_odometry_frames & ~metadata["interpolated"]
        # Construct valid sequences
        valid_diff = valid_odometry_frames.astype(float).diff()
        invalidity_start = valid_diff.index[valid_diff == -1].tolist()
        validity_start = valid_diff.index[valid_diff == 1].tolist()
        if valid_odometry_frames.iloc[0]:
            validity_start = [0] + validity_start
        if valid_odometry_frames.iloc[-1]:
            invalidity_start.append(len(valid_odometry_frames))
        valid_sequences = [metadata.iloc[s:e].copy() for s, e in zip(validity_start, invalidity_start)]
        for s in valid_sequences:
            if len(s) <= args.min_shift:
                continue
            tvec = s[["pose03", "pose13", "pose23"]]
            displacement = np.zeros_like(tvec)
            max_shift = 3
            for j in range(1, max_shift):
                displacement += tvec.diff(j) / j
            s["fpv_x"] = s["fx"] * displacement["pose03"] / displacement["pose23"] + s["cx"]
            s["fpv_y"] = s["fy"] * displacement["pose13"] / displacement["pose23"] + s["cy"]

            valid_frames = s.iloc[args.min_shift:]
            valid_frames = valid_frames[~valid_frames["interpolated"]]
68
            valid_frames["image_path"] = [args.dataset_dir.relpathto(v) / Path(f).basename() for f in valid_frames["image_path"].values]
Clément Pinard's avatar
Clément Pinard committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

            total_valid_frames.append(valid_frames)
    total_valid_frames_df = pd.concat(total_valid_frames)
    if len(total_valid_frames_df) <= args.max_num_samples:
        final_frames = total_valid_frames_df
    else:
        final_frames = total_valid_frames_df.sample(args.max_num_samples, random_state=args.seed)
    train_dirs_list_path = args.dataset_dir / "train_folders.txt"
    image_list_path = args.dataset_dir / "test_files.txt"
    fpv_list_path = args.dataset_dir / "fpv.txt"
    with open(image_list_path, 'w') as f:
        f.writelines(line + "\n" for line in final_frames["image_path"].values)
    np.savetxt(fpv_list_path, final_frames[["fpv_x", "fpv_y"]].values)
    if len(train_videos) > 0:
        with open(train_dirs_list_path, 'w') as f:
            f.writelines([folder + "\n" for folder in train_videos])


if __name__ == '__main__':
    main()