prepare_images.py 3.38 KB
Newer Older
Clément Pinard's avatar
Clément Pinard committed
1
2
3
4
5
6
7
from pyproj import Proj
from edit_exif import get_gps_location
import numpy as np
import rawpy
import imageio
import generate_sky_masks as gsm
import videos_to_colmap as v2c
Clément Pinard's avatar
Clément Pinard committed
8
import colmap_util as ci
Clément Pinard's avatar
Clément Pinard committed
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


def extract_gps_and_path(existing_pictures, image_path, system, centroid=None, **env):
    proj = Proj(system)
    georef_list = []
    for img in existing_pictures:
        gps = get_gps_location(img)
        if gps is not None:
            lat, lng, alt = gps
            x, y = proj(lng, lat)
            if centroid is None:
                centroid = np.array([x, y, alt])
            x -= centroid[0]
            y -= centroid[1]
            alt -= centroid[2]
            georef_list.append("{} {} {} {}\n".format(img.relpath(image_path), x, y, alt))
    return georef_list, centroid


def extract_pictures_to_workspace(input_folder, image_path, workspace, colmap,
                                  raw_ext, pic_ext, more_sift_features, **env):
    picture_folder = input_folder / "Pictures"
    picture_folder.merge_tree(image_path)
    raw_files = sum((list(image_path.walkfiles('*{}'.format(ext))) for ext in raw_ext), [])
    for raw in raw_files:
        if not any((raw.stripext() + ext).isfile() for ext in pic_ext):
            raw_array = rawpy.imread(raw)
            rgb = raw_array.postprocess()
            imageio.imsave(raw.stripext() + ".jpg", rgb)
        raw.remove()
    gsm.process_folder(folder_to_process=image_path, image_path=image_path, pic_ext=pic_ext, **env)
    colmap.extract_features(per_sub_folder=True, more=more_sift_features)
    return sum((list(image_path.walkfiles('*{}'.format(ext))) for ext in pic_ext), [])


def extract_videos_to_workspace(video_path, video_frame_list_thorough, georef_frames_list, **env):
    existing_georef, env["centroid"] = extract_gps_and_path(**env)
Clément Pinard's avatar
Clément Pinard committed
46
47
48
    path_lists, extracted_video_folders = v2c.process_video_folder(output_video_folder=video_path,
                                                                   existing_georef=existing_georef,
                                                                   **env)
Clément Pinard's avatar
Clément Pinard committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    if path_lists is not None:
        with open(video_frame_list_thorough, "w") as f:
            f.write("\n".join(path_lists["thorough"]["frames"]))
        with open(georef_frames_list, "w") as f:
            f.write("\n".join(existing_georef) + "\n")
            f.write("\n".join(path_lists["thorough"]["georef"]) + "\n")
        for v in env["videos_list"]:
            video_folder = extracted_video_folders[v]
            with open(video_folder / "lowfps.txt", "w") as f:
                f.write("\n".join(path_lists[v]["frames_lowfps"]) + "\n")
            with open(video_folder / "georef.txt", "w") as f:
                f.write("\n".join(existing_georef) + "\n")
                f.write("\n".join(path_lists["thorough"]["georef"]) + "\n")
                f.write("\n".join(path_lists[v]["georef_lowfps"]) + "\n")
            for j, l in enumerate(path_lists[v]["frames_full"]):
                with open(video_folder / "full_chunk_{}.txt".format(j), "w") as f:
                    f.write("\n".join(l) + "\n")
    gsm.process_folder(folder_to_process=video_path, **env)
    return extracted_video_folders
Clément Pinard's avatar
Clément Pinard committed
68
69
70
71
72
73


def choose_biggest_model(dir):
    colmap_model_dirs = dir.dirs("[0-9]*")
    model_sizes = [len(ci.read_model.read_images_binary(d/"images.bin")) for d in colmap_model_dirs]
    return colmap_model_dirs[model_sizes.index(max((model_sizes)))]