add_video_to_db.py 2.79 KB
Newer Older
Clement Pinard's avatar
Clement Pinard committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from colmap_util import database as db
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from path import Path
import pandas as pd
import numpy as np
from sqlite3 import IntegrityError
from tqdm import tqdm
parser = ArgumentParser(description='Create vizualisation for specified video',
                        formatter_class=ArgumentDefaultsHelpFormatter)

parser.add_argument('--frame_list', metavar='PATH',
                    help='path to list with relative path to images', type=Path, default=None)
parser.add_argument('--metadata', metavar='PATH',
                    help='path to metadata csv file', type=Path)
parser.add_argument('--database', metavar='DB', required=True,
                    help='path to colmap database file, to get the image ids right')


Clément Pinard's avatar
Clément Pinard committed
19
def add_to_db(db_path, metadata_path, frame_list_path, input_frame_ids=None, **env):
Clement Pinard's avatar
Clement Pinard committed
20
21
22
23
    metadata = pd.read_csv(metadata_path)
    database = db.COLMAPDatabase.connect(db_path)

    frame_list = []
nicolas's avatar
nicolas committed
24
    frame_ids = []
Clement Pinard's avatar
Clement Pinard committed
25
26
    if frame_list_path is not None:
        with open(frame_list_path, "r") as f:
Clément Pinard's avatar
Clément Pinard committed
27
            frame_list = [line[:-1] for line in f.readlines()]
Clement Pinard's avatar
Clement Pinard committed
28
        metadata = metadata[metadata["image_path"].isin(frame_list)]
Clément Pinard's avatar
Clément Pinard committed
29
30
31
    if input_frame_ids:
        assert(len(metadata) == len(input_frame_ids))
        metadata["input_frame_id"] = input_frame_ids
Clement Pinard's avatar
Clement Pinard committed
32

Clément Pinard's avatar
Clément Pinard committed
33
    for _, row in tqdm(metadata.iterrows(), total=len(metadata)):
Clement Pinard's avatar
Clement Pinard committed
34
35
36
37
38
39
40
        image_path = row["image_path"]
        camera_id = row["camera_id"]
        if row["location_valid"]:
            frame_gps = row[["location_longitude", "location_latitude", "location_altitude"]]
        else:
            frame_gps = np.full(3, np.NaN)
        try:
Clément Pinard's avatar
Clément Pinard committed
41
42
            input_id = row["input_frame_id"] if input_frame_ids else None
            frame_ids.append(database.add_image(image_path, int(camera_id), prior_t=frame_gps, image_id=input_id))
Clement Pinard's avatar
Clement Pinard committed
43
        except IntegrityError:
Clément Pinard's avatar
Clément Pinard committed
44
45
46
            sql_string = "SELECT camera_id, image_id FROM images WHERE name='{}'".format(image_path)
            sql_output = next(database.execute(sql_string))
            existing_camera_id = sql_output[0]
Clement Pinard's avatar
Clement Pinard committed
47
            assert(existing_camera_id == camera_id)
Clément Pinard's avatar
Clément Pinard committed
48
            frame_ids.append(sql_output[1])
Clément Pinard's avatar
Clément Pinard committed
49
50
    database.commit()
    database.close()
nicolas's avatar
nicolas committed
51
52
53
54
55
56
57
58
59
60
61
62
63
    return frame_ids


def get_frame_without_features(db_path):
    database = db.COLMAPDatabase.connect(db_path)
    first_string = "SELECT image_id FROM descriptors WHERE cols=0"
    descriptors = list(database.execute(first_string))
    for d in descriptors:
        second_string = "SELECT name FROM images WHERE image_id={}".format(d)
        row = list(database.execute(second_string))[0]
        print(row)

    database.close()
Clement Pinard's avatar
Clement Pinard committed
64
65
66
67
68
69
70
71
72
73
74


def main():
    args = parser.parse_args()
    add_to_db(args.database, args.metadata, args.frame_list)

    return


if __name__ == '__main__':
    main()