stud

study spacejunk
Log | Files | Refs | LICENSE

greedy-points-to-multiline.py (2363B) - Raw

      1 #!/usr/bin/env python3
      2 """
      3 This script takes a point layer and connects them all in a "greedy" way: it
      4 will take a random point, connect its closest one, and repeat the loop until
      5 all of the points are connected.
      6 
      7 This is useful if you have point area boudnaries and want to make a real
      8 connected polygon.
      9 """
     10 
     11 import sys
     12 import rtree
     13 import os
     14 from osgeo import ogr
     15 
     16 
     17 def main():
     18     driver = ogr.GetDriverByName('gpkg')
     19 
     20     inDataSource = driver.Open(sys.argv[1], 0)
     21     inLayer = inDataSource.GetLayer()
     22     srs = inLayer.GetSpatialRef()
     23     outDataSource = driver.CreateDataSource(sys.argv[2])
     24     outLayerName, _ = os.path.splitext(os.path.basename(sys.argv[2]))
     25     outLayer = outDataSource.CreateLayer(
     26             outLayerName,
     27             srs,
     28             geom_type=ogr.wkbLineString
     29     )
     30 
     31     points = extract_layer_points(inLayer)
     32     path = greedy_path(points)
     33     write_path(srs, outLayer, path)
     34 
     35     inDataSource.Destroy()
     36     outDataSource.Destroy()
     37 
     38 
     39 def extract_layer_points(layer):
     40     """Extracts points from layer in a format:
     41         dict(
     42           i : int
     43           points: (x, y : int)
     44         )
     45     """
     46     points = {}
     47     for (i, pt) in enumerate(layer):
     48         p = pt.GetGeometryRef().GetPoint()
     49         points[i] = (p[0], p[1])
     50     return points
     51 
     52 
     53 def greedy_path(points):
     54     """Given a dict of points, return a "greedy" path between them.
     55 
     56     It will start with a random point, and keep navigating to the closest one,
     57     until all the points have been visited.
     58     """
     59     idx = rtree.index.Index(interleaved=False)
     60     for (i, pt) in points.items():
     61         idx.insert(i, pt)
     62 
     63     def _bounds(_pts):
     64         return (_pts[0], _pts[0], _pts[1], _pts[1])
     65 
     66     ret = [0]
     67     idx.delete(0, _bounds(points[0]))
     68     while len(ret) < len(points):
     69         pt = _bounds(points[ret[-1]])
     70         nearest = next(idx.nearest(pt))
     71         ret.append(nearest)
     72         idx.delete(nearest, _bounds(points[nearest]))
     73     ret.append(0)
     74     return [points[point] for point in ret]
     75 
     76 
     77 def write_path(srs, outLayer, path):
     78     "writes the linestring (path) to the outLayer."
     79     line = ogr.Geometry(ogr.wkbLineString)
     80     for point in path:
     81         line.AddPoint(point[0], point[1])
     82 
     83     outFeature = ogr.Feature(outLayer.GetLayerDefn())
     84     outFeature.SetGeometry(line)
     85     outLayer.CreateFeature(outFeature)
     86 
     87 
     88 if __name__ == '__main__':
     89     main()