A method for performing building polygonization by moving vertices to optimize angles between lines.

polygonize_buildings.py

import argparse
import os

import fiona
import numpy as np
import rasterio
import rasterio.crs
import rasterio.features
import rasterio.mask
import rasterio.warp
import shapely
import shapely.geometry
import utm
from scipy.optimize import minimize
from tqdm import tqdm


def get_utm_zone_from_lat_lon(lat, lon):
    """Calculates the UTM CRS for a given lat lon (EPSG:4326) point

    Returns:
        a rasterio CRS object
    """
    projection_dict = {"proj": "utm", "zone": utm.latlon_to_zone_number(lat, lon)}
    if lat < 0:
        projection_dict["south"] = True
    return rasterio.crs.CRS.from_dict(projection_dict)


def get_azimuth(pt1, pt2):
    a = np.arctan2(pt1[0] - pt2[0], pt1[1] - pt2[1])
    return np.degrees(a) if a >= 0 else np.degrees(a) + 360


def get_angle(pt1, pt2, pt3):
    a1 = get_azimuth(pt1, pt2)
    a2 = get_azimuth(pt3, pt2)
    angle = a2 - a1
    return angle if angle > 0 else angle + 360


def angle_loss(angles_in_degrees, dcoords):
    val = 0
    for r in np.radians(angles_in_degrees):
        val += -np.cos(r * 4) + 1
    val += np.linalg.norm(dcoords, ord=2)
    return val


def get_loss_function(shape):
    coords = np.array(list(shape.exterior.coords)[:-1])

    def loss(dcoords):
        dcoords = dcoords.reshape(-1, 2)
        new_coords = coords + dcoords
        angles = []
        for i in range(-2, coords.shape[0] - 2):
            angles.append(
                get_angle(new_coords[i], new_coords[i + 1], new_coords[i + 2])
            )
        val = angle_loss(angles, dcoords)
        return val

    return loss, coords


def optimize(shape, tolerance=0.5, maxiter=50):
    shape = shape.simplify(tolerance)

    loss_fn, coords = get_loss_function(shape)
    result = minimize(
        loss_fn, np.zeros_like(coords).ravel(), options={"maxiter": maxiter}
    )
    dcoords = result.x.reshape(-1, 2)
    shape = shapely.geometry.Polygon(shell=coords + dcoords)
    shape = shape.simplify(tolerance)

    return shape


def set_up_parser():
    """Set up the argument parser.

    Returns:
        the argument parser
    """
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "-e",
        "--epsilon",
        default=0.5,
        type=float,
        help="epsilon value used in Douglas-Peucker simplification",
    )
    parser.add_argument(
        "--input_fn",
        type=str,
        required=True,
        help="input mask filename (must be a GeoTIFF where the value of '1' indicates a predicted building)",
    )
    parser.add_argument(
        "--output_fn",
        type=str,
        required=True,
        help="output filename (must be a gpkg)",
    )
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="print results to stdout"
    )

    return parser


def main(args):

    assert args.input_fn.lower().endswith(".tif")
    assert os.path.exists(args.input_fn)
    assert args.output_fn.lower().endswith(".gpkg")
    assert not os.path.exists(args.output_fn)

    with rasterio.open(args.input_fn) as f:
        src_crs = f.crs
        dst_crs = get_utm_zone_from_lat_lon(f.bounds.top, f.bounds.left)
        features = list(
            rasterio.features.dataset_features(
                f,
                bidx=1,
                sampling=1,
                band=True,
                as_mask=False,
                with_nodata=False,
                geographic=False,
                precision=-1,
            )
        )

    shapes = []
    for feature in features:
        if feature["properties"]["val"] == 1:
            geom = rasterio.warp.transform_geom(src_crs, dst_crs, feature["geometry"])
            shapes.append(shapely.geometry.shape(geom))

    schema = {"geometry": "Polygon", "properties": {"id": "int"}}

    with fiona.open(
        args.output_fn, "w", driver="GPKG", crs="EPSG:4326", schema=schema
    ) as f:
        for i, shape in enumerate(tqdm(shapes)):
            shape = shape.simplify(args.epsilon)

            new_shape = optimize(shape, args.epsilon)

            geom = shapely.geometry.mapping(new_shape)
            geom = rasterio.warp.transform_geom(dst_crs, src_crs, geom)
            row = {"type": "Feature", "geometry": geom, "properties": {"id": i}}
            f.write(row)


if __name__ == "__main__":
    parser = set_up_parser()
    args = parser.parse_args()

    main(args)