-
Notifications
You must be signed in to change notification settings - Fork 5
/
extract_pictures_from_model.py
66 lines (52 loc) · 2.69 KB
/
extract_pictures_from_model.py
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
from colmap_util import read_model as rm
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from path import Path
import pandas as pd
parser = ArgumentParser(description='create a new colmap model with only the frames of selected video',
formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('--input_model', metavar='DIR', type=Path, required=True,
help='folder where the cameras.bin and images.bin are located')
parser.add_argument('--output_model', metavar='DIR', type=Path, required=True,
help='Output folder where the modified COLMAP model will be saved')
parser.add_argument('--output_format', choices=['.txt', '.bin'], default='.txt')
group = parser.add_mutually_exclusive_group()
group.add_argument('--metadata_path', metavar="CSV", type=Path, default=None,
help='Path to metadata CSV file of the desired video. '
'Usually in /pictures/Videos/<size>/<video_name>/metadata.csv')
group.add_argument('--picture_list_path', type=Path, default=None,
help='Path to list of picture to extract from model. '
'Picture paths must be relatvie to colmap root')
def extract_pictures(input, output, picture_list, output_format='.bin'):
cameras = rm.read_cameras_binary(input / "cameras.bin")
images = rm.read_images_binary(input / "images.bin")
images_per_name = {}
camera_ids = []
def add_image(image):
images_per_name[image.name] = image._replace(xys=[], point3D_ids=[])
cam_id = image.camera_id
if cam_id not in camera_ids:
camera_ids.append(cam_id)
for id, image in images.items():
if image.name in picture_list:
add_image(image)
if len(images_per_name) == 1:
# Add also first picture so that we have multiple pictures.
# Otherwise, GourndTruth Creator will error
for id, image in images.items():
if image.name not in picture_list:
add_image(image)
break
output_cameras = {cid: cameras[cid] for cid in camera_ids if cid in cameras.keys()}
rm.write_model(output_cameras, images_per_name, {}, output, output_format)
return len(images_per_name) > 1
def main():
args = parser.parse_args()
if args.metadata_path is not None:
picture_list = pd.read_csv(args.metadata_path)["image_path"].values
elif args.picture_list_path is not None:
with open(args.picture_list_path, 'r') as f:
picture_list = [line[:-1] for line in f.readlines()]
extract_pictures(args.input_model, args.output_model, picture_list, args.output_format)
return
if __name__ == '__main__':
main()