add_video_to_db.py 2.67 KB
Newer Older
Clement Pinard's avatar
Clement Pinard committed
1
2
3
4
5
6
7
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
Clément Pinard's avatar
Clément Pinard committed
8
9
parser = ArgumentParser(description='Add video frames to a colmap db file. '
                        'Needs original colmap db file to link the frames to the right cameras',
Clement Pinard's avatar
Clement Pinard committed
10
11
12
13
                        formatter_class=ArgumentDefaultsHelpFormatter)

parser.add_argument('--frame_list', metavar='PATH',
                    help='path to list with relative path to images', type=Path, default=None)
Clément Pinard's avatar
Clément Pinard committed
14
parser.add_argument('--metadata', metavar='PATH', required=True,
Clement Pinard's avatar
Clement Pinard committed
15
16
17
18
19
                    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
20
def add_to_db(db_path, metadata, frame_list_path, **env):
Clement Pinard's avatar
Clement Pinard committed
21
22
23
    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
29
        metadata = metadata[metadata["image_path"].isin(frame_list)]

Clément Pinard's avatar
Clément Pinard committed
30
    for _, row in tqdm(metadata.iterrows(), total=len(metadata)):
Clement Pinard's avatar
Clement Pinard committed
31
32
        image_path = row["image_path"]
        camera_id = row["camera_id"]
33
        if "location_valid" in row.keys() and row["location_valid"]:
Clement Pinard's avatar
Clement Pinard committed
34
35
36
37
            frame_gps = row[["location_longitude", "location_latitude", "location_altitude"]]
        else:
            frame_gps = np.full(3, np.NaN)
        try:
38
            frame_ids.append(database.add_image(image_path, int(camera_id), prior_t=frame_gps, image_id=row["db_id"]))
Clement Pinard's avatar
Clement Pinard committed
39
        except IntegrityError:
Clément Pinard's avatar
Clément Pinard committed
40
41
42
            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
43
            assert(existing_camera_id == camera_id)
Clément Pinard's avatar
Clément Pinard committed
44
            frame_ids.append(sql_output[1])
Clément Pinard's avatar
Clément Pinard committed
45
46
    database.commit()
    database.close()
nicolas's avatar
nicolas committed
47
48
49
50
51
52
53
54
55
56
57
58
59
    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
60
61
62
63


def main():
    args = parser.parse_args()
Clément Pinard's avatar
Clément Pinard committed
64
    add_to_db(args.database, pd.read_csv(args.metadata), args.frame_list)
Clement Pinard's avatar
Clement Pinard committed
65
66
67
68
69
70

    return


if __name__ == '__main__':
    main()