korrel/korrel.py

#!/usr/bin/python3
"""
Python tool to tag photos along a GPX track.
"""

import argparse
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from pathlib import Path
import sys

import gpx
import exif


def engrave_gpx_data(
    exif_data: exif.Image, photopath: Path, wpt: gpx.waypoint.Waypoint
):
    """Engrave the GPS data from a Waypoint into the image"""
    # GPSVersionID
    exif_data["gps_version_id"] = 2

    lat = wpt.lat
    # GPSLatitudeRef
    exif_data["gps_latitude_ref"] = "N" if lat >= 0 else "S"
    lat = abs(lat)
    # GPSLatitude
    intlat = int(lat)
    declat = int((lat - intlat) * 60)
    decdeclat = (lat - intlat - declat / Decimal(60)) * 3600
    exif_data["gps_latitude"] = (intlat, declat, decdeclat)

    long = wpt.lon
    # GPSLongitudeRef
    exif_data["gps_longitude_ref"] = "E" if long >= 0 else "W"
    # GPSLongitude
    long = abs(long)
    intlong = int(long)
    declong = int((long - intlong) * 60)
    decdeclong = (long - intlong - declong / Decimal(60)) * 3600
    exif_data["gps_longitude"] = (intlong, declong, decdeclong)

    ele = wpt.ele
    if ele is not None:
        # GPSAltitudeRef
        exif_data["gps_altitude_ref"] = 0 if ele >= 0 else 1
        # GPSAltitude
        exif_data["gps_altitude"] = abs(ele)

    if wpt.hdop:
        # GPSDOP
        exif_data["gps_dop"] = wpt.hdop

    utctime = wpt.time.replace(tzinfo = timezone.utc)
    # GPSTimeStamp
    exif_data["gps_datestamp"] = utctime.strftime("%Y:%m:%d")
    # GPSDateStamp
    exif_data["gps_timestamp"] = (utctime.hour, utctime.minute, utctime.second)

    with photopath.open(mode="wb") as photo:
        then = utctime.strftime("%Y:%m:%d %H:%M:%S")
        print(
            f"[{then}] {photopath.name}: {intlat}°{declat}'{decdeclat:02.6}\"N ({lat}) {intlong}°{declong}'{decdeclong:02.6}\"W ({long})"
        )
        photo.write(exif_data.get_file())


def run(gpxf: str, photodir: str):
    """Run the program"""
    # Get the file, potentially
    gpxobj = gpx.GPX.from_file(gpxf)
    # Open the directory
    photos = Path(photodir)
    if not photos.exists():
        print(f'Unable to open photographs folder "{photodir}" !', file=sys.stderr)
        return

    # Get the time bounds of the GPX segments
    gpx_time_bounds = compute_gpx_time_bounds(gpxobj)

    time_photo_dict = find_applicable_images(gpx_time_bounds, photos)
    times_to_find = list(time_photo_dict.keys())
    times_to_find.sort()

    # Find the earliest point that is okay for all photos
    last_point = None
    for track in gpxobj.tracks:
        for seg in track.segments:
            for point in seg.points:
                if last_point is None:
                    last_point = point
                    continue

                if len(times_to_find) == 0:
                    # Job's done
                    break

                while (
                    len(times_to_find) > 0
                    and last_point.time <= times_to_find[0] < point.time
                ):
                    # Oh, `last_point` was what we wanted
                    engrave_gpx_data(*time_photo_dict[times_to_find[0]], last_point)
                    times_to_find.pop(0)
                last_point = point


def find_applicable_images(
    gpx_time_bounds: list[tuple[datetime, timedelta]], photos: Path
) -> dict[datetime, (exif.Image, Path)]:
    """Find images in the photo path within the given GPX segment time bounds"""
    result = {}
    local_tz = datetime.now(timezone.utc).astimezone().tzinfo
    for root, _, candidates in photos.walk(
        top_down=True,
        on_error=lambda x: print(f"Could not open {x}"),
        follow_symlinks=True,
    ):
        for candidate in candidates:
            with open(root / candidate, mode="rb") as candidate_file:
                exif_data = exif.Image(candidate_file)
                if len(exif_data.list_all()) == 0:
                    continue
                candidate_time = datetime.strptime(
                    exif_data.datetime_original, "%Y:%m:%d %H:%M:%S"
                ).replace(tzinfo=local_tz)

                # Find an applicable time bound
                for bound_min, bound_dur in gpx_time_bounds:
                    if bound_min <= candidate_time < bound_min + bound_dur:
                        result[candidate_time] = (exif_data, Path(root / candidate))

    print(f"Found {len(result)} candidate pictures")

    return result


def compute_gpx_time_bounds(gpxobj: gpx.gpx.GPX) -> list[tuple[datetime, timedelta]]:
    """Compute the time bounds of GPX segments"""
    return [
        (seg.points[0].time, seg.duration)
        for track in gpxobj.tracks
        for seg in track.segments
    ]


def main():
    """The main function"""
    args = parse_arguments()
    run(args.gpx_file, args.photo_older)


def parse_arguments() -> argparse.Namespace:
    """Parse arguments to the programs"""
    parser = argparse.ArgumentParser(
        prog="korrel",
        description="Correlate and tag images using a GPX track",
        usage="%(prog)s [-h] <gpx_file> <photo_folder>",
    )
    parser.add_argument("gpx_file")
    parser.add_argument("photo_older")

    return parser.parse_args()


if __name__ == "__main__":
    main()