polys_to_centerlines.py

#!/usr/bin/env python 
#
## Copyright (c) 2013
##
## Joachim Ungar ([email protected])
##
## Permission is hereby granted, free of charge, to any person obtaining a copy 
## of this software and associated documentation files (the "Software"), to deal 
## in the Software without restriction, including without limitation the rights 
## to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
## copies of the Software, and to permit persons to whom the Software is furnished
##  to do so, subject to the following conditions:
##
## The above copyright notice and this permission notice shall be included in all 
## copies or substantial portions of the Software.
##
## THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
## IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
## FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
## AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
## LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
## OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
## SOFTWARE.


import ogr
import osr
import sys
import argparse
import numpy as np
from scipy.spatial import Voronoi, voronoi_plot_2d
import shapely
import networkx as nx
import math
import decimal
import os

def main(args):

    parser = argparse.ArgumentParser()
    parser.add_argument("input_shp", nargs=1, type=str)
    parser.add_argument("output_shp", nargs=1, type=str)
    parser.add_argument("segmentize", nargs=1, type=float)
    parser.add_argument("simplify", nargs=1, type=float)
    parsed = parser.parse_args(args)
    input_shp = parsed.input_shp[0]
    output_shp = parsed.output_shp[0]
    segmentize_threshold = parsed.segmentize[0]
    simplify_threshold = parsed.segmentize[0]


    # read source shapefile
    # ---------------------------------------------------
    # Get the driver
    driver = ogr.GetDriverByName('ESRI Shapefile')
    # Open a shapefile
    input_dataset = driver.Open(input_shp, 0)    
    # create the data source
    if os.path.exists(output_shp):
        driver.DeleteDataSource(output_shp)
    target_dataset = driver.CreateDataSource(output_shp)


    # create target shapefile    
    # ---------------------------------------------------
    
    # create the spatial reference, WGS84
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    # create the layer
    target_layer = target_dataset.CreateLayer("regions", srs, ogr.wkbLineString)


    # calculate geometry feature by feature
    # ---------------------------------------------------
    source_layer = input_dataset.GetLayer()

    # create target layer attributes
    source_layer_defn = source_layer.GetLayerDefn()
    target_layer_defn = source_layer_defn
    for i in range(target_layer_defn.GetFieldCount()):
        target_layer.CreateField(target_layer_defn.GetFieldDefn(i))
    
    
    for index in xrange(source_layer.GetFeatureCount()):
        try:
            source_feature = source_layer.GetFeature(index)
            geometry = source_feature.GetGeometryRef()
            linestring = geometry.GetGeometryRef(0)

            # create target feature
            target_feature = ogr.Feature(target_layer_defn)
            
            # generalize
            linestring.Simplify(simplify_threshold)
    
            # segmentize
            linestring.Segmentize(segmentize_threshold)
            
            # extract points        
            points = extract_points(linestring)
    
            # calculate voronoi
            vor = Voronoi(points)
            #print vor.vertices
    
            ridge_lines = ogr.Geometry(ogr.wkbMultiLineString)
    
            # create graph & add edge length as attribute
            G = nx.Graph()
    
            for i in vor.ridge_vertices:
                if i[0]>-1 and i[1]>-1:
                    templine = ogr.Geometry(ogr.wkbLineString)
                    p1 = vor.vertices[i][0]
                    p2 = vor.vertices[i][1]
                    p1lon = p1[0]
                    p1lat = p1[1]
                    p2lon = p2[0]
                    p2lat = p2[1]
                    if point_inside_polygon(p1lon,p1lat,points) and point_inside_polygon(p2lon,p2lat,points):
                        # create points
                        point1 = ogr.Geometry(ogr.wkbPoint)
                        point1.AddPoint(p1lon, p1lat)
                        point2 = ogr.Geometry(ogr.wkbPoint)
                        point2.AddPoint(p2lon, p2lat)
                        #create ogr layer
                        templine.AddPoint(p1lon, p1lat)
                        templine.AddPoint(p2lon, p2lat)
                        ridge_lines.AddGeometry(templine)
                        # create graph
                        dist = point1.Distance(point2)
                        G.add_nodes_from([i[0], i[1]])
                        G.add_edge(i[0], i[1], weight=dist )
    
            # get all end nodes
            end_nodes = get_end_nodes(G)
            #print end_nodes
    
            # iterate through nodes and get shortest path
            longest_path = get_longest_path(end_nodes,G,vor.vertices)
    
            templine2 = ogr.Geometry(ogr.wkbLineString)
            for i in longest_path:
                templine2.AddPoint(vor.vertices[i][0], vor.vertices[i][1])

            #print source_feature.GetField('name')
            #print source_feature.GetFieldCount()

            
            #feat_defn = source_layer.GetLayerDefn()
            #for i in range(feat_defn.GetFieldCount()):
            #    field_defn = feat_defn.GetFieldDefn(i)
            #    print field_defn

            #print templine2.ExportToWkt()
            
            
            # write output to target shapefile
            # ---------------------------------------------------
 
            # Copy attributes
            for i in range(source_layer_defn.GetFieldCount()):
                target_feature.SetField(source_layer_defn.GetFieldDefn(i).GetName(), source_feature.GetField(i))
            # get WKT
            target_line = ogr.CreateGeometryFromWkt(templine2.ExportToWkt())
            # Set the feature geometry using the line
            target_feature.SetGeometry(target_line)
            # Create the feature in the layer (shapefile)
            target_layer.CreateFeature(target_feature)
            # Destroy the feature to free resources
            target_feature.Destroy()
    
            # go through paths and get angle penalty
            #for i in path1:
                #print vor.vertices[i-1]
                #print vor.vertices[i]
                #print vor.vertices[i+1]
                #angle = math.sqrt((vor.vertices[i-1][0] - vor.vertices[i+1][0])**2 + (vor.vertices[i-1][0] - vor.vertices[i+1][0])**2)
                #print str(angle)
    
            # multiply each path with angle penalty
    
            #print pointlayer.ExportToWkt()
            #print ridge_lines.ExportToWkt()
    
        except Exception,e: print str(e)


    # clean up
    # ---------------------------------------------------
    target_dataset.Destroy()


def get_end_nodes(graph):
    nodelist = []
    for i in graph.nodes_iter():
        if len(graph.neighbors(i))==1:
            nodelist.append(i)
    return nodelist


def get_longest_path(nodes,graph,vertices):
    paths = []
    distances = []
    for node1 in nodes:
        for node2 in nodes:
            path = nx.shortest_path(graph,node1,node2,"weight")
            if len(path)>1:
                distance = get_path_distance(path,graph)
                #print get_angle_penalty(path,vertices)
                paths.append(path)
                distances.append(distance)
    paths_sorted = [x for (y,x) in sorted(zip(distances,paths),reverse=True)]
    longest_path = paths_sorted[0]
    return longest_path


def get_path_distance(path,graph):
    distance = 0
    for i,w in enumerate(path):
        j=i+1
        if j<len(path):
            distance = distance + round(graph.edge[path[i]][path[j]]["weight"],6)
    return distance
    

def get_angle_penalty(path,vertices):
    penalty = 0
    #print path
    for i,w in enumerate(path):
        if i>0:
            j=i+1
            h=i-1
            #print "%s, %s, %s" %(path[h], path[i], path[j])
            if j<len(path):
                p1=vertices[path[i]]+180
                p2=vertices[path[j]]+180
                p3=vertices[path[h]]+180
                #print "%s, %s, %s" %(p1, p2, p3)
                #B^2-A^2-C^2) / 2AC )
                v12 = round(math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2), 6)
                v13 = round(math.sqrt((p1[0] - p3[0])**2 + (p1[1] - p3[1])**2), 6)
                v23 = round(math.sqrt((p2[0] - p3[0])**2 + (p2[1] - p3[1])**2), 6)
                angle = round(math.acos((v12**2 + v13**2 - v23**2)/(2 * v12 * v13)), 6)
                #print angle
                #penalty=(1-(180-angle)/180)
                #penalty = penalty + angle
                #print penalty
    return penalty


def extract_points(inputline):
    pointsX = []; pointsY = [];
    points = np.array([[0,0]]);
    numpoints = inputline.GetPointCount()
    for p in range(numpoints):
        lon, lat, z = inputline.GetPoint(p)
        pointsX.append(lon)
        pointsY.append(lat)
        points_temp = [lon, lat]
        points_old = points
        points = np.vstack((points_old,points_temp))
    return points


# from http://www.ariel.com.au/a/python-point-int-poly.html
def point_inside_polygon(x,y,poly):

    n = len(poly)
    inside =False

    p1x,p1y = poly[0]
    for i in range(n+1):
        p2x,p2y = poly[i % n]
        if y > min(p1y,p2y):
            if y <= max(p1y,p2y):
                if x <= max(p1x,p2x):
                    if p1y != p2y:
                        xinters = (y-p1y)*(p2x-p1x)/(p2y-p1y)+p1x
                    if p1x == p2x or x <= xinters:
                        inside = not inside
        p1x,p1y = p2x,p2y

    return inside


if __name__ == "__main__":
        main(sys.argv[1:])