projectify.py

#!/usr/bin/env python3
"""
Create a GeoTIFF from a PNG, using the georeferencing of a DEM GeoTIFF.

Example:
    python png_to_geotiff.py dem.tif image.png output.tif
"""

import argparse
import numpy as np
from PIL import Image
import rasterio

def main():
    parser = argparse.ArgumentParser(
        description="Create a GeoTIFF from a PNG using a DEM GeoTIFF's projection and transform."
    )
    parser.add_argument("dem_tif", help="Path to input DEM GeoTIFF with desired projection/transform")
    parser.add_argument("png", help="Path to input PNG file (visuals)")
    parser.add_argument("output_tif", help="Path to output GeoTIFF")

    args = parser.parse_args()

    # Open DEM to get metadata (CRS, transform, etc.)
    with rasterio.open(args.dem_tif) as dem_src:
        profile = dem_src.profile.copy()
        dem_height = dem_src.height
        dem_width = dem_src.width

    # Open PNG and convert to array
    img = Image.open(args.png)

    # If needed, resize PNG to match DEM dimensions
    if img.size != (dem_width, dem_height):
        img = img.resize((dem_width, dem_height), Image.BILINEAR)

    # Convert PNG to numpy array
    arr = np.array(img)

    # Handle band layout:
    # - Grayscale: shape (H, W) -> (1, H, W)
    # - RGB: shape (H, W, 3) -> (3, H, W)
    # - RGBA: shape (H, W, 4) -> (4, H, W)
    if arr.ndim == 2:
        # Single band
        arr = arr[np.newaxis, ...]
        count = 1
        dtype = arr.dtype
    elif arr.ndim == 3:
        # Multiple bands
        arr = np.transpose(arr, (2, 0, 1))  # to (bands, H, W)
        count = arr.shape[0]
        dtype = arr.dtype
    else:
        raise ValueError("Unsupported PNG array shape: {}".format(arr.shape))

    # Update profile for the new raster
    profile.update(
        driver="GTiff",
        height=dem_height,
        width=dem_width,
        count=count,
        dtype=dtype,
        nodata=None  # or a specific nodata value if desired
    )

    # Write new GeoTIFF with DEM georeferencing and PNG visuals
    with rasterio.open(args.output_tif, "w", **profile) as dst:
        dst.write(arr)

if __name__ == "__main__":
    main()