videos_to_colmap.py 21.4 KB
Newer Older
Clement Pinard's avatar
Clement Pinard committed
1
2
3
4
5
6
from colmap_util import read_model as rm, database as db
import anafi_metadata as am
from wrappers import FFMpeg, PDraw
from sklearn.cluster import KMeans
from sklearn.metrics import pairwise_distances_argmin_min
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
nicolas's avatar
nicolas committed
7
from edit_exif import set_gps_location
Clement Pinard's avatar
Clement Pinard committed
8
9
10
11
12
from path import Path
import pandas as pd
import numpy as np
from pyproj import Proj
from tqdm import tqdm
13
import tempfile
Clement Pinard's avatar
Clement Pinard committed
14
15
16
17
18
19
20

parser = ArgumentParser(description='Take all the drone videos of a folder and put the frame '
                                    'location in a COLMAP file for vizualisation',
                        formatter_class=ArgumentDefaultsHelpFormatter)

parser.add_argument('--video_folder', metavar='DIR',
                    help='path to videos', type=Path)
Clément Pinard's avatar
Clément Pinard committed
21
22
23
24
25
26
27
28
29
30
31
parser.add_argument('--system', default='epsg:2154',
                    help='coordinates system used for GPS, should be the same as the LAS files used')
parser.add_argument('--centroid_path', default=None, help="path to centroid generated in las2ply.py")
parser.add_argument('--colmap_img_root', metavar='DIR', type=Path,
                    help="folder that will be used as \"image_path\" parameter when using COLMAP", required=True)
parser.add_argument('--output_format', metavar='EXT', default="bin", choices=["bin", "txt"],
                    help='format of the COLMAP file that will be outputed, used for vizualisation only')
parser.add_argument('--vid_ext', nargs='+', default=[".mp4", ".MP4"],
                    help="format of video files that will be scraped from input folder")
parser.add_argument('--pic_ext', nargs='+', default=[".jpg", ".JPG", ".png", ".PNG"],
                    help='format of images that will be scraped from already existing images in colmap image_path folder')
Clement Pinard's avatar
Clement Pinard committed
32
parser.add_argument('--nw', default='',
Clément Pinard's avatar
Clément Pinard committed
33
                    help="native-wrapper.sh file location (see Anafi SDK documentation)")
Clement Pinard's avatar
Clement Pinard committed
34
35
parser.add_argument('--fps', default=1, type=int,
                    help="framerate at which videos will be scanned WITH reconstruction")
Clément Pinard's avatar
Clément Pinard committed
36
37
38
39
40
41
42
43
44
parser.add_argument('--total_frames', default=200, type=int, help="number of frames used for thorough photogrammetry")
parser.add_argument('--max_sequence_length', default=1000, help='Number max of frames for a chunk. '
                    'This is for RAM purpose, as loading feature matches of thousands of frames can take up GBs of RAM')
parser.add_argument('--orientation_weight', default=1, type=float,
                    help="Weight applied to orientation during optimal sample. "
                    "Higher means two pictures with same location but different orientation will be considered farer apart")
parser.add_argument('--resolution_weight', default=1, type=float, help="same as orientation, but with image size")
parser.add_argument('--save_space', action="store_true", help="if selected, will only extract from ffmpeg frames used for thorough photogrammetry")
parser.add_argument('--thorough_db', type=Path, help="output db file which will be used by COLMAP for photogrammetry")
Clément Pinard's avatar
Clément Pinard committed
45
parser.add_argument('-v', '--verbose', action="count", default=0)
Clement Pinard's avatar
Clement Pinard committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64


def world_coord_from_frame(frame_qvec, frame_tvec):
    '''
    frame_qvec is written in the NED system (north east down)
    frame_tvec is already is the world system (east norht up)
    '''
    world2NED = np.float32([[0, 1, 0],
                            [1, 0, 0],
                            [0, 0, -1]])
    NED2cam = np.float32([[0, 1, 0],
                          [0, 0, 1],
                          [1, 0, 0]])
    world2cam = NED2cam @ rm.qvec2rotmat(frame_qvec).T @ world2NED
    cam_tvec = - world2cam  @ frame_tvec
    cam_qvec = rm.rotmat2qvec(world2cam)
    return cam_qvec, cam_tvec


nicolas's avatar
nicolas committed
65
66
67
68
69
70
def set_gps(frames_list, metadata, image_path):
    for frame in frames_list:
        relative = str(frame.relpath(image_path))
        row = metadata[metadata["image_path"] == relative]
        if len(row) > 0:
            row = row.iloc[0]
71
72
73
74
75
            if row["location_valid"]:
                set_gps_location(frame,
                                 lat=row["location_latitude"],
                                 lng=row["location_longitude"],
                                 altitude=row["location_altitude"])
nicolas's avatar
nicolas committed
76
77
78
79
80
81


def get_georef(metadata):
    relevant_data = metadata[["location_valid", "image_path", "x", "y", "z"]]
    path_list = []
    georef_list = []
Clément Pinard's avatar
Clément Pinard committed
82
    for _, (loc_valid, path, x, y, alt) in relevant_data.iterrows():
nicolas's avatar
nicolas committed
83
        path_list.append(path)
Clément Pinard's avatar
Clément Pinard committed
84
        if loc_valid:
nicolas's avatar
nicolas committed
85
86
87
88
            georef_list.append("{} {} {} {}\n".format(path, x, y, alt))
    return georef_list, path_list


Clement Pinard's avatar
Clement Pinard committed
89
def optimal_sample(metadata, num_frames, orientation_weight, resolution_weight):
90
    valid_metadata = metadata[~metadata["sampled"]].dropna()
Clément Pinard's avatar
Clément Pinard committed
91
92
    if len(valid_metadata) == 0:
        return metadata
93
94
95
96
97
98
    XYZ = valid_metadata[["x", "y", "z"]].values
    axis_angle = valid_metadata[["frame_quat_x", "frame_quat_y", "frame_quat_z"]].values
    if True in valid_metadata["indoor"].unique():
        # We have indoor videos, without absolute positions. We assume each video is very far
        # from the other ones. As such we will have an optimal subsampling of each video
        # It won't leverage video proximity from each other but it's better than nothing
Clement Pinard's avatar
Clement Pinard committed
99
        diameter = (XYZ.max(axis=0) - XYZ.min(axis=0))
100
101
102
103
        indoor_videos = valid_metadata.loc[valid_metadata["indoor"]]["video"].unique()
        new_centroids = 2 * diameter * np.linspace(0, 10, len(indoor_videos)).reshape(-1, 1)
        for centroid, v in zip(new_centroids, indoor_videos):
            video_index = (valid_metadata["video"] == v).values
Clement Pinard's avatar
Clement Pinard committed
104
105
            XYZ[video_index] += centroid

106
    frame_size = valid_metadata["video_quality"].values
Clement Pinard's avatar
Clement Pinard committed
107
108
109
110
111
112
113
114
    weighted_point_cloud = np.concatenate([XYZ, orientation_weight * axis_angle], axis=1)

    if resolution_weight == 0:
        weights = None
    else:
        weights = frame_size ** resolution_weight
    km = KMeans(n_clusters=num_frames).fit(weighted_point_cloud, sample_weight=weights)
    closest, _ = pairwise_distances_argmin_min(km.cluster_centers_, weighted_point_cloud)
115
    metadata.at[valid_metadata.index[closest], "sampled"] = True
Clement Pinard's avatar
Clement Pinard committed
116
117
118
    return metadata


119
def register_new_cameras(cameras_dataframe, database, camera_dict):
Clement Pinard's avatar
Clement Pinard committed
120
    camera_ids = []
Clément Pinard's avatar
Clément Pinard committed
121
122
123
124
125
    for _, row in cameras_dataframe.iterrows():
        w, h, hfov, vfov, camera_model = row.reindex(["width", "height", "picture_hfov", "picture_vfov", "camera_model"])
        prior_focal_length = False
        single_focal = ('SIMPLE' in camera_model) or ('RADIAL' in camera_model)
        if hfov != 0:
126
            fx = w / (2 * np.tan(hfov * np.pi/360))
Clément Pinard's avatar
Clément Pinard committed
127
128
129
130
131
            # If the model is not single focal, only knowing hfov is not enough, you also need to know vfov
            prior_focal_length = single_focal
        else:
            fx = w / 2  # As if hfov was 90 degrees
        if vfov != 0:
132
            fy = h / (2 * np.tan(vfov * np.pi/360))
Clément Pinard's avatar
Clément Pinard committed
133
            prior_focal_length = True
134
        else:
Clément Pinard's avatar
Clément Pinard committed
135
            fy = w / 2  # As if vfov was 90 degrees
136
137
        model_id = rm.CAMERA_MODEL_NAMES[camera_model].model_id
        num_params = rm.CAMERA_MODEL_NAMES[camera_model].num_params
Clément Pinard's avatar
Clément Pinard committed
138
139
140
141
142
        if ('SIMPLE' in camera_model) or ('RADIAL' in camera_model):
            params = np.array([fx, w/2, h/2] + [0] * (num_params - 3))
        else:
            params = np.array([fx, fy, w/2, h/2] + [0] * (num_params - 4))
        db_id = database.add_camera(model_id, int(w), int(h), params, prior_focal_length=prior_focal_length)
Clement Pinard's avatar
Clement Pinard committed
143
144
        camera_ids.append(db_id)
        camera_dict[db_id] = rm.Camera(id=db_id,
145
                                       model=camera_model,
Clement Pinard's avatar
Clement Pinard committed
146
147
148
149
150
151
152
153
                                       width=int(w),
                                       height=int(h),
                                       params=params)
    ids_series = pd.Series(camera_ids)
    return cameras_dataframe.set_index(ids_series)


def process_video_folder(videos_list, existing_pictures, output_video_folder, image_path, system, centroid,
Clément Pinard's avatar
Clément Pinard committed
154
                         thorough_db, fps=1, total_frames=500, orientation_weight=1, resolution_weight=1,
155
                         output_colmap_format="bin", save_space=False, include_lowfps_thorough=False,
Clément Pinard's avatar
Clément Pinard committed
156
                         max_sequence_length=1000, num_neighbours=10, existing_georef=False, **env):
Clement Pinard's avatar
Clement Pinard committed
157
158
159
160
161
    proj = Proj(system)
    final_metadata = []
    video_output_folders = {}
    images = {}
    colmap_cameras = {}
162
    tempfile_database = Path(tempfile.NamedTemporaryFile().name)
Clément Pinard's avatar
Clément Pinard committed
163
164
    if thorough_db.isfile():
        thorough_db.copy(thorough_db.stripext() + "_backup.db")
Clement Pinard's avatar
Clement Pinard committed
165
    path_lists_output = {}
166
    database = db.COLMAPDatabase.connect(thorough_db)
Clement Pinard's avatar
Clement Pinard committed
167
168
169
    database.create_tables()

    print("extracting metadata for {} videos...".format(len(videos_list)))
170
    videos_summary = {"anafi": {"indoor": 0, "outdoor": 0}, "generic": 0}
Clément Pinard's avatar
Clément Pinard committed
171
172

    indoor_video_diameters = {}
Clement Pinard's avatar
Clement Pinard committed
173
    for v in tqdm(videos_list):
174
        width, height, framerate, num_frames = env["ffmpeg"].get_size_and_framerate(v)
Clément Pinard's avatar
Clément Pinard committed
175
        video_output_folder = output_video_folder / "{}x{}".format(width, height) / v.stem
Clement Pinard's avatar
Clement Pinard committed
176
177
178
        video_output_folder.makedirs_p()
        video_output_folders[v] = video_output_folder

179
180
181
182
183
        try:
            metadata = am.extract_metadata(v.parent, v, env["pdraw"], proj,
                                           width, height, framerate)
            metadata["model"] = "anafi"
            metadata["camera_model"] = "PINHOLE"
Clément Pinard's avatar
Clément Pinard committed
184
185
            raw_positions = metadata[["x", "y", "z"]]
            video_displacement_diameter = np.linalg.norm(raw_positions.values.max(axis=0) - raw_positions.values.min(axis=0))
186
187
            if metadata["indoor"].iloc[0]:
                videos_summary["anafi"]["indoor"] += 1
Clément Pinard's avatar
Clément Pinard committed
188
                indoor_video_diameters[video_displacement_diameter] = v
189
190
191
192
193
194
            else:
                videos_summary["anafi"]["outdoor"] += 1
                raw_positions = metadata[["x", "y", "z"]]
                if centroid is None:
                    '''No centroid (possibly because there was no georeferenced lidar model in the first place)
                    set it as the first valid GPS position of the first outdoor video'''
Clément Pinard's avatar
Clément Pinard committed
195
                    centroid = raw_positions[metadata["location_valid"]].iloc[0].values
196
                zero_centered_positions = raw_positions.values - centroid
Clément Pinard's avatar
bug fix    
Clément Pinard committed
197
198
199
200
                radius = np.max(np.abs(zero_centered_positions))
                if radius > 1000:
                    print("Warning, your positions coordinates are most likely too high, have you configured the right GPS system ?")
                    print("It should be the same as the one used for the Lidar point cloud")
201
202
203
204
205
206
207
208
209
210
211
212
                metadata["x"], metadata["y"], metadata["z"] = zero_centered_positions.transpose()
        except Exception:
            # No metadata found, construct a simpler dataframe without location
            metadata = pd.DataFrame({"video": [v] * num_frames})
            metadata["height"] = height
            metadata["width"] = width
            metadata["framerate"] = framerate
            metadata["video_quality"] = height * width / framerate
            metadata['frame'] = metadata.index + 1
            # timestemp is in microseconds
            metadata['time'] = 1e6 * metadata.index / framerate
            metadata['indoor'] = True
Clément Pinard's avatar
Clément Pinard committed
213
            metadata['location_valid'] = False
214
            metadata["model"] = "generic"
Clément Pinard's avatar
Clément Pinard committed
215
216
217
            metadata["camera_model"] = "PINHOLE"
            metadata["picture_hfov"] = 0
            metadata["picture_vfov"] = 0
Clément Pinard's avatar
bug fix    
Clément Pinard committed
218
219
220
221
            metadata["frame_quat_w"] = np.NaN
            metadata["frame_quat_x"] = np.NaN
            metadata["frame_quat_y"] = np.NaN
            metadata["frame_quat_z"] = np.NaN
Clément Pinard's avatar
Clément Pinard committed
222
223
224
            metadata["x"] = np.NaN
            metadata["y"] = np.NaN
            metadata["z"] = np.NaN
225
            videos_summary["generic"] += 1
Clément Pinard's avatar
Clément Pinard committed
226
        metadata["num_frames"] = num_frames
227
228
229
230
        if include_lowfps_thorough:
            by_time = metadata.set_index(pd.to_datetime(metadata["time"], unit="us"))
            by_time_lowfps = by_time.resample("{:.3f}S".format(1/fps)).first()
            metadata["sampled"] = by_time["time"].isin(by_time_lowfps["time"]).values
Clément Pinard's avatar
Clément Pinard committed
231
        else:
232
233
            metadata["sampled"] = False
        final_metadata.append(metadata)
Clement Pinard's avatar
Clement Pinard committed
234
    final_metadata = pd.concat(final_metadata, ignore_index=True)
235
236
237
    print("{} outdoor anafi videos".format(videos_summary["anafi"]["outdoor"]))
    print("{} indoor anafi videos".format(videos_summary["anafi"]["indoor"]))
    print("{} generic videos".format(videos_summary["generic"]))
Clement Pinard's avatar
Clement Pinard committed
238

Clément Pinard's avatar
Clément Pinard committed
239
    if(not existing_georef and videos_summary["anafi"]["outdoor"] == 0 and videos_summary["anafi"]["indoor"] > 0):
Clément Pinard's avatar
Clément Pinard committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
        # We have no GPS data but we have navdata, which will help rescale the colmap model
        # Take the longest video and do as if the GPS was valid
        longest_video = indoor_video_diameters[max(indoor_video_diameters)]
        print("Only indoor videos used, will use {} for COLMAP rescaling".format(longest_video))
        video_index = final_metadata["video"] == longest_video
        if include_lowfps_thorough:
            # We already added frames to be sampled so we just copy the boolean to the "location_valid" column
            final_metadata.loc[video_index, "location_valid"] = final_metadata.loc[video_index, "sampled"]
        else:
            # Take frames at lowfps, add it to the thorough photogrammetry and mark their location as valid
            video_md = final_metadata[video_index]
            by_time = video_md.set_index(pd.to_datetime(video_md["time"], unit="us"))
            by_time_lowfps = by_time.resample("{:.3f}S".format(1/fps)).first()
            to_georef = by_time["time"].isin(by_time_lowfps["time"]).values
            final_metadata.loc[video_index, "sampled"] = to_georef
            final_metadata.loc[video_index, "location_valid"] = to_georef

Clement Pinard's avatar
Clement Pinard committed
257
258
    print("{} frames in total".format(len(final_metadata)))

259
260
261
    cam_fields = ["width", "height", "framerate", "picture_hfov", "picture_vfov", "camera_model"]
    cameras_dataframe = final_metadata[final_metadata["model"] == "anafi"][cam_fields].drop_duplicates()
    cameras_dataframe = register_new_cameras(cameras_dataframe, database, colmap_cameras)
Clement Pinard's avatar
Clement Pinard committed
262
263
264
    final_metadata["camera_id"] = 0
    for cam_id, row in cameras_dataframe.iterrows():
        final_metadata.loc[(final_metadata[cam_fields] == row).all(axis=1), "camera_id"] = cam_id
265
266
267
    if any(final_metadata["model"] == "generic"):
        print("Undefined remaining cameras, assigning generic models to them")
        generic_frames = final_metadata[final_metadata["model"] == "generic"]
Clément Pinard's avatar
Clément Pinard committed
268
269
        generic_cam_fields = cam_fields + ["video"]
        generic_cameras_dataframe = generic_frames[generic_cam_fields]
270
271
272
273
274
275
        fixed_camera = True
        if fixed_camera:
            generic_cameras_dataframe = generic_cameras_dataframe.drop_duplicates()
        generic_cameras_dataframe = register_new_cameras(generic_cameras_dataframe, database, colmap_cameras)
        if fixed_camera:
            for cam_id, row in generic_cameras_dataframe.iterrows():
Clément Pinard's avatar
Clément Pinard committed
276
                final_metadata.loc[(final_metadata[generic_cam_fields] == row).all(axis=1), "camera_id"] = cam_id
277
278
279
280
281
282
283
        else:
            final_metadata.loc[generic_frames.index, "camera_id"] = generic_cameras_dataframe.index
        cameras_dataframe = cameras_dataframe.append(generic_cameras_dataframe)
    print("Cameras : ")
    print(cameras_dataframe)

    to_extract = total_frames - len(existing_pictures) - sum(final_metadata["sampled"])
Clement Pinard's avatar
Clement Pinard committed
284

Clément Pinard's avatar
Clément Pinard committed
285
    if to_extract <= 0:
286
        pass
Clément Pinard's avatar
Clément Pinard committed
287
    elif to_extract < len(final_metadata):
Clement Pinard's avatar
Clement Pinard committed
288
289
290
291
292
293
        print("subsampling based on K-Means, to get {}"
              " frames from videos, for a total of {} frames".format(to_extract, total_frames))
        final_metadata = optimal_sample(final_metadata, total_frames - len(existing_pictures),
                                        orientation_weight,
                                        resolution_weight)
        print("Done.")
Clément Pinard's avatar
Clément Pinard committed
294
295
    else:
        final_metadata["sampled"] = True
Clement Pinard's avatar
Clement Pinard committed
296

297
    print("Constructing COLMAP model with {:,} frames".format(sum(final_metadata["sampled"])))
Clement Pinard's avatar
Clement Pinard committed
298

299
300
301
302
    database.commit()
    thorough_db.copy(tempfile_database)
    temp_database = db.COLMAPDatabase.connect(tempfile_database)

Clement Pinard's avatar
Clement Pinard committed
303
    final_metadata["image_path"] = ""
304
305
    final_metadata["db_id"] = -1
    for current_id, row in tqdm(final_metadata.iterrows(), total=len(final_metadata)):
Clement Pinard's avatar
Clement Pinard committed
306
307
308
        video = row["video"]
        frame = row["frame"]
        camera_id = row["camera_id"]
Clément Pinard's avatar
Clément Pinard committed
309
        current_image_path = video_output_folders[video].relpath(image_path) / video.stem + "_{:05d}.jpg".format(frame)
Clement Pinard's avatar
Clement Pinard committed
310

311
312
313
        final_metadata.at[current_id, "image_path"] = current_image_path
        db_image_id = temp_database.add_image(current_image_path, int(camera_id))
        final_metadata.at[current_id, "db_id"] = db_image_id
Clement Pinard's avatar
Clement Pinard committed
314
315
316
317
318
319

        if row["sampled"]:
            frame_qvec = row[["frame_quat_w",
                              "frame_quat_x",
                              "frame_quat_y",
                              "frame_quat_z"]].values
320
321
            if True in pd.isnull(frame_qvec):
                frame_qvec = np.array([1, 0, 0, 0])
nicolas's avatar
nicolas committed
322
323
            x, y, z = row[["x", "y", "z"]]
            frame_tvec = np.array([x, y, z])
Clement Pinard's avatar
Clement Pinard committed
324
325
326
327
328
329
            if row["location_valid"]:
                frame_gps = row[["location_longitude", "location_latitude", "location_altitude"]]
            else:
                frame_gps = np.full(3, np.NaN)

            world_qvec, world_tvec = world_coord_from_frame(frame_qvec, frame_tvec)
330
331
332
333
            database.add_image(current_image_path, int(camera_id), prior_t=frame_gps, image_id=db_image_id)
            images[db_image_id] = rm.Image(id=db_image_id, qvec=world_qvec, tvec=world_tvec,
                                           camera_id=camera_id, name=current_image_path,
                                           xys=[], point3D_ids=[])
Clement Pinard's avatar
Clement Pinard committed
334
335
336

    database.commit()
    database.close()
337
338
    temp_database.commit()
    temp_database.close()
Clement Pinard's avatar
Clement Pinard committed
339
340
341
    rm.write_model(colmap_cameras, images, {}, output_video_folder, "." + output_colmap_format)
    print("COLMAP model created")

nicolas's avatar
nicolas committed
342
343
344
345
    thorough_georef, thorough_paths = get_georef(final_metadata[final_metadata["sampled"]])
    path_lists_output["thorough"] = {}
    path_lists_output["thorough"]["frames"] = thorough_paths
    path_lists_output["thorough"]["georef"] = thorough_georef
Clement Pinard's avatar
Clement Pinard committed
346
347
348
349
350

    print("Extracting frames from videos")

    for v in tqdm(videos_list):
        video_metadata = final_metadata[final_metadata["video"] == v]
nicolas's avatar
nicolas committed
351
        by_time = video_metadata.set_index(pd.to_datetime(video_metadata["time"], unit="us"))
Clement Pinard's avatar
Clement Pinard committed
352
353
        video_folder = video_output_folders[v]
        video_metadata.to_csv(video_folder/"metadata.csv")
nicolas's avatar
nicolas committed
354
355
356
357
358
        path_lists_output[v] = {}
        video_metadata_1fps = by_time.resample("{:.3f}S".format(1/fps)).first()
        georef, frame_paths = get_georef(video_metadata_1fps)
        path_lists_output[v]["frames_lowfps"] = frame_paths
        path_lists_output[v]["georef_lowfps"] = georef
Clément Pinard's avatar
Clément Pinard committed
359
        num_chunks = len(video_metadata) // max_sequence_length + 1
360
361
362
363
        chunks = [list(frames) for frames in np.array_split(video_metadata["image_path"],
                                                            num_chunks)]
        # Add some overlap between chunks, in order to ease the model merging afterwards
        for chunk, next_chunk in zip(chunks, chunks[1:]):
Clément Pinard's avatar
Clément Pinard committed
364
            chunk.extend(next_chunk[:num_neighbours])
365
366
        path_lists_output[v]["frames_full"] = chunks

Clement Pinard's avatar
Clement Pinard committed
367
        if save_space:
Clément Pinard's avatar
Clément Pinard committed
368
369
            frame_ids = set(video_metadata[video_metadata["sampled"]]["frame"].values) | \
                set(video_metadata_1fps["frame"].values)
370
            frame_ids = sorted(list(frame_ids))
Clément Pinard's avatar
Clément Pinard committed
371
            if len(frame_ids) > 0:
nicolas's avatar
nicolas committed
372
                extracted_frames = env["ffmpeg"].extract_specific_frames(v, video_folder, frame_ids)
Clement Pinard's avatar
Clement Pinard committed
373
        else:
nicolas's avatar
nicolas committed
374
375
            extracted_frames = env["ffmpeg"].extract_images(v, video_folder)
        set_gps(extracted_frames, video_metadata, image_path)
Clement Pinard's avatar
Clement Pinard committed
376
377
378
379
380
381
382
383

    return path_lists_output, video_output_folders


if __name__ == '__main__':
    args = parser.parse_args()
    env = vars(args)
    env["videos_list"] = sum((list(args.video_folder.walkfiles('*{}'.format(ext))) for ext in args.vid_ext), [])
Clément Pinard's avatar
Clément Pinard committed
384
    output_video_folder = args.colmap_img_root / "Videos"
Clement Pinard's avatar
Clement Pinard committed
385
    output_video_folder.makedirs_p()
Clément Pinard's avatar
Clément Pinard committed
386
    env["image_path"] = args.colmap_img_root
Clement Pinard's avatar
Clement Pinard committed
387
    env["output_video_folder"] = output_video_folder
Clément Pinard's avatar
Clément Pinard committed
388
    env["existing_pictures"] = sum((list(args.colmap_img_root.walkfiles('*{}'.format(ext))) for ext in args.pic_ext), [])
Clément Pinard's avatar
Clément Pinard committed
389
390
391
    env["pdraw"] = PDraw(args.nw, verbose=args.verbose)
    env["ffmpeg"] = FFMpeg(verbose=args.verbose)
    env["output_colmap_format"] = args.output_format
Clement Pinard's avatar
Clement Pinard committed
392
393
394
395
396
397
398
399
400

    if args.centroid_path is not None:
        centroid = np.loadtxt(args.centroid_path)
    else:
        centroid = np.zeros(3)
    env["centroid"] = centroid
    lists, extracted_video_folders = process_video_folder(**env)

    if lists is not None:
Clément Pinard's avatar
Clément Pinard committed
401
402
403
        with open(args.colmap_img_root/"video_frames_for_thorough_scan.txt", "w") as f:
            f.write("\n".join(lists["thorough"]["frames"]) + "\n")
        with open(args.colmap_img_root/"georef.txt", "w") as f:
Clément Pinard's avatar
Clément Pinard committed
404
            f.write("\n".join(lists["thorough"]["georef"]))
Clement Pinard's avatar
Clement Pinard committed
405
        for v in env["videos_list"]:
Clément Pinard's avatar
Clément Pinard committed
406
407
408
409
410
411
412
413
414
            video_folder = extracted_video_folders[v]
            with open(video_folder / "lowfps.txt", "w") as f:
                f.write("\n".join(lists[v]["frames_lowfps"]) + "\n")
            with open(video_folder / "georef.txt", "w") as f:
                f.write("\n".join(lists["thorough"]["georef"]) + "\n")
                f.write("\n".join(lists[v]["georef_lowfps"]) + "\n")
            for j, l in enumerate(lists[v]["frames_full"]):
                with open(video_folder / "full_chunk_{}.txt".format(j), "w") as f:
                    f.write("\n".join(l) + "\n")