wm/layer2img.py

186 lines
5.6 KiB
Python
Raw Normal View History

2021-05-19 22:57:46 +03:00
#!/usr/bin/python3
import argparse
import geopandas
import psycopg2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
INCH = 25.4 # mm
BOUNDS = ('xmin', 'ymin', 'xmax', 'ymax')
GREEN, ORANGE, PURPLE = '#1b9e77', '#d95f02', '#7570b3'
PSQL_CREDS="host=127.0.0.1 dbname=osm user=osm password=osm"
2021-05-19 22:57:46 +03:00
def arrowplot(axes, x, y, narrs=30, dspace=0.1, direc='pos', \
hl=0.1, hw=5, c='black'):
2021-05-19 22:57:46 +03:00
''' narrs : Number of arrows that will be drawn along the curve
dspace : Shift the position of the arrows along the curve.
Should be between 0. and 1.
direc : can be 'pos' or 'neg' to select direction of the arrows
2021-05-19 22:57:46 +03:00
hl : length of the arrow head
2021-05-19 22:57:46 +03:00
2021-05-19 22:57:46 +03:00
hw : width of the arrow head
2021-05-19 22:57:46 +03:00
2021-05-19 22:57:46 +03:00
c : color of the edge and face of the arrow head
2021-05-19 22:57:46 +03:00
https://stackoverflow.com/questions/8247973
'''
# r is the distance spanned between pairs of points
r = [0]
for i in range(1,len(x)):
2021-05-19 22:57:46 +03:00
dx = x[i]-x[i-1]
dy = y[i]-y[i-1]
2021-05-19 22:57:46 +03:00
r.append(np.sqrt(dx*dx+dy*dy))
r = np.array(r)
# rtot is a cumulative sum of r, it's used to save time
rtot = []
for i in range(len(r)):
rtot.append(r[0:i].sum())
rtot.append(r.sum())
# based on narrs set the arrow spacing
aspace = r.sum() / narrs
if direc is 'neg':
2021-05-19 22:57:46 +03:00
dspace = -1.*abs(dspace)
2021-05-19 22:57:46 +03:00
else:
dspace = abs(dspace)
arrowData = [] # will hold tuples of x,y,theta for each arrow
arrowPos = aspace*(dspace) # current point on walk along data
# could set arrowPos to 0 if you want
# an arrow at the beginning of the curve
ndrawn = 0
2021-05-19 22:57:46 +03:00
rcount = 1
2021-05-19 22:57:46 +03:00
while arrowPos < r.sum() and ndrawn < narrs:
x1,x2 = x[rcount-1],x[rcount]
y1,y2 = y[rcount-1],y[rcount]
da = arrowPos-rtot[rcount]
theta = np.arctan2((x2-x1),(y2-y1))
ax = np.sin(theta)*da+x1
ay = np.cos(theta)*da+y1
arrowData.append((ax,ay,theta))
ndrawn += 1
arrowPos+=aspace
2021-05-19 22:57:46 +03:00
while arrowPos > rtot[rcount+1]:
2021-05-19 22:57:46 +03:00
rcount+=1
if arrowPos > rtot[-1]:
break
# could be done in above block if you want
for ax,ay,theta in arrowData:
# use aspace as a guide for size and length of things
# scaling factors were chosen by experimenting a bit
dx0 = np.sin(theta)*hl/2. + ax
dy0 = np.cos(theta)*hl/2. + ay
dx1 = -1.*np.sin(theta)*hl/2. + ax
dy1 = -1.*np.cos(theta)*hl/2. + ay
if direc is 'neg' :
2021-05-19 22:57:46 +03:00
ax0 = dx0
2021-05-19 22:57:46 +03:00
ay0 = dy0
ax1 = dx1
2021-05-19 22:57:46 +03:00
ay1 = dy1
2021-05-19 22:57:46 +03:00
else:
2021-05-19 22:57:46 +03:00
ax0 = dx1
2021-05-19 22:57:46 +03:00
ay0 = dy1
ax1 = dx0
2021-05-19 22:57:46 +03:00
ay1 = dy0
2021-05-19 22:57:46 +03:00
axes.annotate('', xy=(ax0, ay0), xycoords='data',
xytext=(ax1, ay1), textcoords='data',
arrowprops=dict( headwidth=hw, frac=1., ec=c, fc=c))
def plt_size(string):
if not string:
return None
try:
w, h = string.split("x")
return float(w) / INCH, float(h) / INCH
except Exception as e:
raise argparse.ArgumentTypeError from e
def parse_args():
parser = argparse.ArgumentParser(
description='Convert geopackage to an image')
group1 = parser.add_mutually_exclusive_group()
group1.add_argument('--group1-infile')
group1.add_argument('--group1-table')
2021-05-19 22:57:46 +03:00
parser.add_argument('--group1-arrows', type=bool)
2021-05-19 22:57:46 +03:00
group2 = parser.add_mutually_exclusive_group()
group2.add_argument('--group2-infile', type=str)
group2.add_argument('--group2-table', type=str)
2021-05-19 22:57:46 +03:00
parser.add_argument('--group2-arrows', type=bool)
2021-05-19 22:57:46 +03:00
group3 = parser.add_mutually_exclusive_group()
group3.add_argument('--group3-infile', type=str)
group3.add_argument('--group3-table', type=str)
2021-05-19 22:57:46 +03:00
parser.add_argument('--group3-arrows', type=bool)
parser.add_argument('-o', '--outfile', metavar='<file>')
parser.add_argument(
'--size', type=plt_size, help='Figure size in mm (WWxHH)')
parser.add_argument( '--clip', type=float, nargs=4, metavar=BOUNDS)
2021-05-19 22:57:46 +03:00
return parser.parse_args()
def read_layer(maybe_table, maybe_file):
if maybe_table:
conn = psycopg2.connect(PSQL_CREDS)
sql = "SELECT way FROM %s" % maybe_table
return geopandas.read_postgis(sql, con=conn, geom_col='way')
elif maybe_file:
return geopandas.read_file(maybe_file)
def add_lines(ax, group):
for g in group.to_dict()['way'].values():
for geom in getattr(g, 'geoms', [g]):
x, y = zip(*geom.coords)
2021-05-19 22:57:46 +03:00
narrs = geom.length / 25
arrowplot(ax, np.array(x), np.array(y), narrs=narrs)
2021-05-19 22:57:46 +03:00
def main():
args = parse_args()
group1 = read_layer(args.group1_table, args.group1_infile)
group2 = read_layer(args.group2_table, args.group2_infile)
group3 = read_layer(args.group3_table, args.group3_infile)
rc('text', usetex=True)
fig, ax = plt.subplots()
if args.size:
fig.set_size_inches(args.size)
if c := args.clip:
ax.set_xlim(left=c[0], right=c[2])
ax.set_ylim(bottom=c[1], top=c[3])
if group1 is not None:
2021-05-19 22:57:46 +03:00
group1.plot(ax=ax, color=ORANGE)
2021-05-19 22:57:46 +03:00
args.group1_arrows and add_lines(ax, group1)
2021-05-19 22:57:46 +03:00
if group2 is not None:
2021-05-19 22:57:46 +03:00
group2.plot(ax=ax, color=PURPLE)
2021-05-19 22:57:46 +03:00
args.group2_arrows and add_lines(ax, group1)
2021-05-19 22:57:46 +03:00
if group3 is not None:
group3.plot(ax=ax, color=GREEN)
2021-05-19 22:57:46 +03:00
args.group3_arrows and add_lines(ax, group1)
2021-05-19 22:57:46 +03:00
ax.axis('off')
ax.margins(0, 0)
#fig.tight_layout(0)
if args.outfile:
fig.savefig(args.outfile, bbox_inches=0, dpi=600)
else:
plt.show()
if __name__ == '__main__':
main()