tymelapse/main.py at main · paul-leroux/tymelapse · GitHub

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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
# main.py
# Tymelapse orchestrator script

import os
import cv2
from glob import glob
from config.config import *
from config.global_variables import *

from tymepkg import (
    detect_keypoints_and_descriptors,
    match_keypoints_flann,
    match_keypoints_bf,
    compute_homography,
    warp_image,
    get_date_taken,
    preprocess_gray,
    convert_to_bgra,
    delete_previous_outputs,
    log_message,
    compute_alignment_metrics,
    is_alignment_acceptable
)

def perform_alignment_pass(input_dir, output_dir, log_file):
    image_paths = sorted(glob(os.path.join(input_dir, '*.jpg')))
    if not image_paths:
        print("No input images found.")
        return

    selected = input(f"Enter index of reference image (0 to {len(image_paths) - 1}): ")
    try:
        ref_img_path = image_paths[int(selected)]
    except (IndexError, ValueError):
        print("Invalid selection. Using middle image by default.")
        ref_img_path = image_paths[len(image_paths) // 2]

    print(f"Using '{os.path.basename(ref_img_path)}' as reference image.")
    log_message(log_file, f"Reference image selected: {os.path.basename(ref_img_path)}")

    ref_img = cv2.imread(ref_img_path)
    ref_gray = preprocess_gray(ref_img)
    ref_bgra = convert_to_bgra(ref_img)

    ref_kp_sift, ref_desc_sift = detect_keypoints_and_descriptors(ref_gray, method='sift')

    success_log = []
    fail_log = []

    for img_path in image_paths:
        img_name = os.path.basename(img_path)
        aligned = None
        H = None

        date_taken = get_date_taken(img_path)
        out_name = f"{date_taken}.png" if date_taken else f"aligned_{img_name[:-4]}.png"
        out_path = os.path.join(output_dir, out_name)

        print(f"Processing '{img_name}'...")

        img = cv2.imread(img_path)
        img_gray = preprocess_gray(img)
        img_bgra = convert_to_bgra(img)

        kp, desc = detect_keypoints_and_descriptors(img_gray, method='sift')
        if desc is not None and len(kp) >= 10:
            matches = match_keypoints_flann(ref_desc_sift, desc)
            if len(matches) >= 10:
                H = compute_homography(ref_kp_sift, kp, matches)
                if H is not None:
                    metrics = compute_alignment_metrics(H)
                    if is_alignment_acceptable(metrics, {
                        'translation': TRANSLATION_MAX_THRESHOLD,
                        'scale_x': SCALE_MAX_THRESHOLD_X,
                        'scale_y': SCALE_MAX_THRESHOLD_Y,
                        'shear': SHEAR_MAX_THRESHOLD
                    }):
                        aligned = warp_image(img_bgra, H, (ref_bgra.shape[1], ref_bgra.shape[0]))

        if aligned is None:
            kp_ref, desc_ref = detect_keypoints_and_descriptors(ref_gray, method='akaze')
            kp, desc = detect_keypoints_and_descriptors(img_gray, method='akaze')
            if desc_ref is not None and desc is not None and len(kp) >= 10:
                matches = match_keypoints_bf(desc_ref, desc)[:100]
                H = compute_homography(kp_ref, kp, matches)
                if H is not None:
                    metrics = compute_alignment_metrics(H)
                    if is_alignment_acceptable(metrics, {
                        'translation': TRANSLATION_MAX_THRESHOLD,
                        'scale_x': SCALE_MAX_THRESHOLD_X,
                        'scale_y': SCALE_MAX_THRESHOLD_Y,
                        'shear': SHEAR_MAX_THRESHOLD
                    }):
                        aligned = warp_image(img_bgra, H, (ref_bgra.shape[1], ref_bgra.shape[0]))

        if aligned is not None:
            cv2.imwrite(out_path, aligned)
            log_message(log_file, f"Processed {img_name}: aligned and saved as {out_name}")
            success_log.append((img_name, out_name))
        else:
            log_message(log_file, f"Failed to align {img_name}: not enough good matches or homography failed.")
            fail_log.append(img_name)

    log_message(log_file, f"\nAlignment complete. Success: {len(success_log)}, Failures: {len(fail_log)}")
    if fail_log:
        log_message(log_file, "Failed images:")
        for fail in fail_log:
            log_message(log_file, f" - {fail}")

def main():
    delete_previous_outputs(OUTPUT_DIR_PASS_01)
    delete_previous_outputs(OUTPUT_DIR_PASS_02)

    os.makedirs(LOG_DIR, exist_ok=True)
    log_file_path = os.path.join(LOG_DIR, 'alignment_log.txt')

    with open(log_file_path, "a", encoding="utf-8") as log_file:
        pass_number = input("Which alignment pass to run? (1 or 2): ")
        if pass_number == "1":
            perform_alignment_pass(INPUT_DIR, OUTPUT_DIR_PASS_01, log_file)
        elif pass_number == "2":
            perform_alignment_pass(INPUT_DIR, OUTPUT_DIR_PASS_02, log_file)
        else:
            print("Invalid pass number.")

if __name__ == "__main__":
    main()