source: subversion/applications/utils/osm-extract/polygons/polybuffer.py

Last change on this file was 19035, checked in by fsteggink, 8 years ago

more advanced db connection, parsing improvements

File size: 6.3 KB
Line 
1#!/usr/bin/python
2# polybuffer.py
3# version 0.2
4
5# History
6# - 12/01/2009 0.1   First release
7# - 12/10/2009 0.2   More advanced DB connection, improvements in polygon parsing
8
9
10# Usage: ./polybuffer.py -h
11# Example use with stdin and stdout: cat poly/quebec2pts.txt | ./polybuffer.py >out.txt
12
13# Requirements:
14# * Recent Python version (tested with 2.6)
15# * PostGIS (tested with 1.3)
16# * PyGreSQL (http://www.pygresql.org/)
17
18# Notes:
19# * Only tested on Ubuntu, but should run everywhere
20# * In the input file, holes in polygons are not (yet?) recognized, but they are generated in the output file
21# * Poly file is assumed to be in WGS84
22# * TODO Use a different SRS in PostGIS, eventually make it an option. For now, WGS84 is sufficient
23
24
25import getpass
26import random
27import re
28import sys
29import _pg
30from optparse import OptionParser
31
32
33# Do not change any of the following constants
34VERSION = "0.2"
35
36
37# Read a polygon from file
38# NB: holes aren't supported yet
39def read_polygon(f):
40
41        coords = []
42        first_coord = True
43        while True:
44                line = f.readline()
45                if not(line):
46                        break;
47                       
48                line = line.strip()
49                if line == "END":
50                        break
51               
52                # NOTE: needs to happen before not(line).
53                # There can be a blank line in the poly file if the centroid wasn't calculated!
54                if first_coord:
55                        first_coord = False
56                        continue       
57               
58                if not(line):
59                        continue
60               
61                ords = line.split()
62                coords.append("%f %f" % (float(ords[0]), float(ords[1])))
63       
64        if len(coords) < 3:
65                return None
66       
67        polygon = "((" + ",".join(coords) + "))"
68       
69        return polygon
70
71
72# Read a multipolygon from the file
73# First line: name (discarded)
74# Polygon: numeric identifier, list of lon, lat, END
75# Last line: END
76def read_multipolygon(f):
77       
78        polygons = []
79        while True:
80                dummy = f.readline()
81                if not(dummy):
82                        break
83               
84                polygon = read_polygon(f)
85                if polygon != None:
86                        polygons.append(polygon)
87               
88        wkt = "MULTIPOLYGON(" + ",".join(polygons) + ")"
89       
90        return wkt
91
92
93# Write a polygon to the file
94def write_polygon(f, wkt, p):
95
96        match = re.search("^\(\((?P<pdata>.*)\)\)$", wkt)
97        pdata = match.group("pdata")
98        rings = re.split("\),\(", pdata)
99       
100        first_ring = True
101        for ring in rings:
102                coords = re.split(",", ring)                   
103       
104                p = p + 1
105                if first_ring:
106                        f.write(str(p) + "\n")
107                        first_ring = False
108                else:
109                        f.write("!" + str(p) + "\n")
110               
111                for coord in coords:
112                        ords = coord.split()
113                        f.write("\t%E\t%E\n" % (float(ords[0]), float(ords[1])))
114               
115                f.write("END\n")
116       
117        return p
118
119
120# Write a multipolygon to the file
121def write_multipolygon(f, wkt):
122
123        match = re.search("^MULTIPOLYGON\((?P<mpdata>.*)\)$", wkt)
124       
125        if match:
126                mpdata = match.group("mpdata")
127                polygons = re.split("(?<=\)\)),(?=\(\()", mpdata)
128       
129                p = 0
130                for polygon in polygons:
131                        p = write_polygon(f, polygon, p)
132               
133                return
134               
135        match = re.search("^POLYGON(?P<pdata>.*)$", wkt)
136        if match:
137                pdata = match.group("pdata")
138                write_polygon(f, pdata, 0)
139       
140
141# Calculate a buffer around the polygon, using PostGIS
142def buffer_polygon(wkt, db, buffer_distance):
143
144        # Generate random table name
145        table_name = "polybuffer_"
146        for i in random.sample("abcdefghijklmnopqrstuvwxyz", 10):
147                table_name += i
148       
149        # Create a table in the DB
150        sql = "CREATE TABLE %s ( id integer );" % (table_name)
151        db.query(sql)
152        sql = "SELECT AddGeometryColumn('', '%s', 'the_geom', '-1', 'GEOMETRY' ,2);" % (table_name)
153        db.query(sql)
154
155        # Load data into it
156        sql = "INSERT INTO %s (id, the_geom) VALUES (1, GeomFromText('%s'));" % (table_name, wkt)
157        db.query(sql)
158       
159        # Convert buffer distance from meter to degree
160        dist = float(buffer_distance) * 9 / 1000000
161       
162        # Simplify geometry
163        sql = "UPDATE %s SET the_geom = ST_Simplify(the_geom, %f);" % (table_name, dist / 10)
164        db.query(sql)
165       
166        # Calculate buffer
167        sql = "UPDATE %s SET the_geom = ST_Buffer(the_geom, %f);" % (table_name, dist)
168        db.query(sql)
169       
170        # Extract data from it
171        sql = "SELECT AsText(the_geom) AS wkt FROM %s;" % (table_name)
172        result = db.query(sql).dictresult()
173       
174        for record in result:
175                wkt = record["wkt"]
176                break
177       
178        # Delete table
179        sql = "SELECT DropGeometryColumn('', '%s', 'the_geom')" % (table_name)
180        db.query(sql)
181        sql = "DROP TABLE %s;" % (table_name)
182        db.query(sql)
183       
184        return wkt
185
186
187# Perform all the conversion steps
188def convert_poly (input_file, output_file, buffer_distance, db):
189
190        # Steps:
191        # 1. Convert poly to WKT
192        # 2. PostGIS actions:
193        #    a. Create temp table
194        #    b. Load data
195        #    c. Buffer data
196        #    d. Get data as WKT
197        # 3. Convert WKT to poly
198       
199        # Note that holes aren't supported yet on import. They are on export.
200       
201        # Read input file, and convert polygon to WKT
202        if input_file:
203                f = open(input_file, "r")
204        else:
205                f = sys.stdin
206               
207        name = f.readline().strip()
208        wkt = read_multipolygon(f)     
209        f.close()
210
211        # Buffer the polygon
212        wkt = buffer_polygon(wkt, db, buffer_distance)
213
214        # Convert WKT to polygon, and write output file
215        if output_file:
216                f = open(output_file, "w")
217        else:
218                f = sys.stdout
219       
220        name = "%s_%f" % (name, float(buffer_distance))
221        f.write(name + "\n")           
222        write_multipolygon(f, wkt)
223        f.write("END\n")
224        f.close()       
225
226        return
227
228
229def main():
230
231        default_user = getpass.getuser()
232
233        # Parse arguments
234        arg_list = "-i <input> -o <output> [-b <distance>] [-d <db> ...]"
235        usage = "Usage: %prog " + arg_list
236        version = "%prog " + VERSION
237        parser = OptionParser(usage, version=version)
238        parser.add_option("-i", "--input", dest="input", help="input file (default stdin)")
239        parser.add_option("-o", "--output", dest="output", help="output file (default stdout)")
240        parser.add_option("-b", "--buffer", type="float", dest="buffer", help="buffer distance in m (default '1000')", default=1000)
241        parser.add_option("-d", "--dbname", dest="dbname", help="database")
242        parser.add_option("--host", dest="host", help="host (default 'localhost')", default="localhost")
243        parser.add_option("--port", type="int", dest="port", help="port (default '5432')", default="5432")
244        parser.add_option("-u", "--user", dest="user", help="user name (default '%s')" % (default_user), default=default_user)
245        parser.add_option("-w", "--password", action="store_true", dest="passwd", help="password", default=False)
246        (options, args) = parser.parse_args()
247       
248        passwd = ""
249        if options.passwd:
250                passwd = getpass.getpass("Please enter your password: ")
251       
252        db = _pg.connect(dbname=options.dbname,
253                                         host=options.host,
254                                         port=options.port, 
255                                         user=options.user, 
256                                         passwd=passwd)
257                                         
258        convert_poly(options.input, options.output, options.buffer, db)
259
260
261if __name__ == "__main__":
262        main()
263
Note: See TracBrowser for help on using the repository browser.