#!/usr/bin/env python
#
# heatmap.py - Generates heat map images and animations from geographic data
# Copyright 2010 Seth Golub
# http://www.sethoscope.net/heatmap/
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
from __future__ import print_function
import sys
import logging
import math
from PIL import Image
from PIL import ImageColor
from itertools import chain
import tempfile
import os.path
import shutil
import subprocess
import platform
import glob
from collections import defaultdict
import xml.etree.cElementTree as ET
from colorsys import hsv_to_rgb
try:
import cPickle as pickle
except ImportError:
import pickle
__version__ = '1.13'
class Coordinate(object):
def __init__(self, x, y):
self.x = x
self.y = y
first = property(lambda self: self.x)
second = property(lambda self: self.y)
def copy(self):
return self.__class__(self.first, self.second)
def __str__(self):
return '(%s, %s)' % (str(self.x), str(self.y))
def __hash__(self):
return hash((self.x, self.y))
def __eq__(self, o):
return True if self.x == o.x and self.y == o.y else False
def __sub__(self, o):
return self.__class__(self.first - o.first, self.second - o.second)
class LatLon(Coordinate):
def __init__(self, lat, lon):
self.lat = lat
self.lon = lon
def get_lat(self):
return self.y
def set_lat(self, lat):
self.y = lat
def get_lon(self):
return self.x
def set_lon(self, lon):
self.x = lon
lat = property(get_lat, set_lat)
lon = property(get_lon, set_lon)
first = property(get_lat)
second = property(get_lon)
class TrackLog:
class Trkseg(list): # for GPX tags
pass
class Trkpt: # for GPX tags
def __init__(self, lat, lon):
self.coords = LatLon(float(lat), float(lon))
def __str__(self):
return str(self.coords)
def _parse(self, filename):
self._segments = []
for event, elem in ET.iterparse(filename, ('start', 'end')):
elem.tag = elem.tag[elem.tag.rfind('}') + 1:] # remove namespace
if elem.tag == "trkseg":
if event == 'start':
self._segments.append(TrackLog.Trkseg())
else: # event == 'end'
yield self._segments.pop()
elem.clear() # delete contents from parse tree
elif elem.tag == 'trkpt' and event == 'end':
point = TrackLog.Trkpt(elem.attrib['lat'], elem.attrib['lon'])
self._segments[-1].append(point)
elem.clear() # clear the trkpt node to minimize memory usage
def __init__(self, filename):
self.filename = filename
def segments(self):
'''Parse file and yield segments containing points'''
logging.info('reading GPX track from %s' % self.filename)
return self._parse(self.filename)
class Projection(object):
# For guessing scale, we pretend the earth is a sphere with this
# radius in meters, as in Web Mercator (the projection all the
# online maps use).
EARTH_RADIUS = 6378137 # in meters
def get_pixels_per_degree(self):
try:
return self._pixels_per_degree
except AttributeError:
raise AttributeError('projection scale was never set')
def set_pixels_per_degree(self, val):
self._pixels_per_degree = val
logging.info('scale: %f meters/pixel (%f pixels/degree)'
% (self.meters_per_pixel, val))
def get_meters_per_pixel(self):
return 2 * math.pi * self.EARTH_RADIUS / 360 / self.pixels_per_degree
def set_meters_per_pixel(self, val):
self.pixels_per_degree = 2 * math.pi * self.EARTH_RADIUS / 360 / val
return val
pixels_per_degree = property(get_pixels_per_degree, set_pixels_per_degree)
meters_per_pixel = property(get_meters_per_pixel, set_meters_per_pixel)
def is_scaled(self):
return hasattr(self, '_pixels_per_degree')
def project(self, coords):
raise NotImplementedError
def inverse_project(self, coords): # Not all projections can do this.
raise NotImplementedError
def auto_set_scale(self, extent_in, padding, width=None, height=None):
# We need to choose a scale at which the data's bounding box,
# once projected onto the map, will fit in the specified height
# and/or width. The catch is that we can't project until we
# have a scale, so what we'll do is set a provisional scale,
# project the bounding box onto the map, then adjust the scale
# appropriately. This way we don't need to know anything about
# the projection.
#
# Projection subclasses are free to override this method with
# something simpler that just solves for scale given the lat/lon
# and x/y bounds.
# We'll work large to minimize roundoff error.
SCALE_FACTOR = 1000000.0
self.pixels_per_degree = SCALE_FACTOR
extent_out = extent_in.map(self.project)
padding *= 2 # padding-per-edge -> padding-in-each-dimension
try:
if height:
self.pixels_per_degree = pixels_per_lat = (
float(height - padding) /
extent_out.size().y * SCALE_FACTOR)
if width:
self.pixels_per_degree = (
float(width - padding) /
extent_out.size().x * SCALE_FACTOR)
if height:
self.pixels_per_degree = min(self.pixels_per_degree,
pixels_per_lat)
except ZeroDivisionError:
raise ZeroDivisionError(
'You need at least two data points for auto scaling. '
'Try specifying the scale explicitly (or extent + '
'height or width).')
assert(self.pixels_per_degree > 0)
# Treats Lat/Lon as a square grid.
class EquirectangularProjection(Projection):
# http://en.wikipedia.org/wiki/Equirectangular_projection
def project(self, coord):
x = coord.lon * self.pixels_per_degree
y = -coord.lat * self.pixels_per_degree
return Coordinate(x, y)
def inverse_project(self, coord):
lat = -coord.y / self.pixels_per_degree
lon = coord.x / self.pixels_per_degree
return LatLon(lat, lon)
class MercatorProjection(Projection):
def set_pixels_per_degree(self, val):
super(MercatorProjection, self).set_pixels_per_degree(val)
self._pixels_per_radian = val * (180 / math.pi)
pixels_per_degree = property(Projection.get_pixels_per_degree,
set_pixels_per_degree)
def project(self, coord):
x = coord.lon * self.pixels_per_degree
y = -self._pixels_per_radian * math.log(
math.tan((math.pi/4 + math.pi/360 * coord.lat)))
return Coordinate(x, y)
def inverse_project(self, coord):
lat = (360 / math.pi *
math.atan(math.exp(-coord.y / self._pixels_per_radian)) - 90)
lon = coord.x / self.pixels_per_degree
return LatLon(lat, lon)
class Extent():
def __init__(self, coords=None, shapes=None):
if coords:
coords = tuple(coords) # if it's a generator, slurp them all
self.min = coords[0].__class__(min(c.first for c in coords),
min(c.second for c in coords))
self.max = coords[0].__class__(max(c.first for c in coords),
max(c.second for c in coords))
elif shapes:
self.from_shapes(shapes)
else:
raise ValueError('Extent must be initialized')
def __str__(self):
return '%s,%s,%s,%s' % (self.min.y, self.min.x, self.max.y, self.max.x)
def update(self, other):
'''grow this bounding box so that it includes the other'''
self.min.x = min(self.min.x, other.min.x)
self.min.y = min(self.min.y, other.min.y)
self.max.x = max(self.max.x, other.max.x)
self.max.y = max(self.max.y, other.max.y)
def from_bounding_box(self, other):
self.min = other.min.copy()
self.max = other.max.copy()
def from_shapes(self, shapes):
shapes = iter(shapes)
self.from_bounding_box(next(shapes).extent)
for s in shapes:
self.update(s.extent)
def corners(self):
return (self.min, self.max)
def size(self):
return self.max.__class__(self.max.x - self.min.x,
self.max.y - self.min.y)
def grow(self, pad):
self.min.x -= pad
self.min.y -= pad
self.max.x += pad
self.max.y += pad
def resize(self, width=None, height=None):
if width:
self.max.x += float(width - self.size().x) / 2
self.min.x = self.max.x - width
if height:
self.max.y += float(height - self.size().y) / 2
self.min.y = self.max.y - height
def is_inside(self, coord):
return (coord.x >= self.min.x and coord.x <= self.max.x and
coord.y >= self.min.y and coord.y <= self.max.y)
def map(self, func):
'''Returns a new Extent whose corners are a function of the
corners of this one. The expected use is to project a Extent
onto a map. For example: bbox_xy = bbox_ll.map(projector.project)'''
return Extent(coords=(func(self.min), func(self.max)))
class Matrix(defaultdict):
'''An abstract sparse matrix, with data stored as {coord : value}.'''
@staticmethod
def matrix_factory(decay):
# If decay is 0 or 1, we can accumulate as we go and save lots of
# memory.
if decay == 1.0:
logging.info('creating a summing matrix')
return SummingMatrix()
elif decay == 0.0:
logging.info('creating a maxing matrix')
return MaxingMatrix()
logging.info('creating an appending matrix')
return AppendingMatrix(decay)
def __init__(self, default_factory=float):
self.default_factory = default_factory
def add(self, coord, val):
raise NotImplementedError
def extent(self):
return(Extent(coords=self.keys()))
def finalized(self):
return self
class SummingMatrix(Matrix):
def add(self, coord, val):
self[coord] += val
class MaxingMatrix(Matrix):
def add(self, coord, val):
self[coord] = max(val, self.get(coord, val))
class AppendingMatrix(Matrix):
def __init__(self, decay):
self.default_factory = list
self.decay = decay
def add(self, coord, val):
self[coord].append(val)
def finalized(self):
logging.info('combining coincident points')
m = Matrix()
for (coord, values) in self.items():
m[coord] = self.reduce(self.decay, values)
return m
@staticmethod
def reduce(decay, values):
'''
Returns a weighted sum of the values, where weight N is
pow(decay,N). This means the largest value counts fully, but
additional values have diminishing contributions. decay=0 makes
the reduction equivalent to max(), which makes each data point
visible, but says nothing about their relative magnitude.
decay=1 makes this like sum(), which makes the relative
magnitude of the points more visible, but could make smaller
values hard to see. Experiment with values between 0 and 1.
Values outside that range will give weird results.
'''
# It would be nice to do this on the fly, while accumulating data, but
# it needs to be insensitive to data order.
weight = 1.0
total = 0.0
values.sort(reverse=True)
for value in values:
total += value * weight
weight *= decay
return total
class Point:
def __init__(self, coord, weight=1.0):
self.coord = coord
self.weight = weight
def __str__(self):
return 'P(%s)' % str(self.coord)
@staticmethod
def general_distance(x, y):
# assumes square units, which causes distortion in some projections
return math.hypot(x, y)
@property
def extent(self):
if not hasattr(self, '_extent'):
self._extent = Extent(coords=(self.coord,))
return self._extent
# From a modularity standpoint, it would be reasonable to cache
# distances, not heat values, and let the kernel cache the
# distance to heat map, but this is substantially faster.
heat_cache = {}
@classmethod
def _initialize_heat_cache(cls, kernel):
cache = {}
for x in range(kernel.radius + 1):
for y in range(kernel.radius + 1):
cache[(x, y)] = kernel.heat(cls.general_distance(x, y))
cls.heat_cache[kernel] = cache
def add_heat_to_matrix(self, matrix, kernel):
if kernel not in Point.heat_cache:
Point._initialize_heat_cache(kernel)
cache = Point.heat_cache[kernel]
x = int(self.coord.x)
y = int(self.coord.y)
for dx in range(-kernel.radius, kernel.radius + 1):
for dy in range(-kernel.radius, kernel.radius + 1):
matrix.add(Coordinate(x + dx, y + dy),
self.weight * cache[(abs(dx), abs(dy))])
def map(self, func):
return Point(func(self.coord), self.weight)
class LineSegment:
def __init__(self, start, end, weight=1.0):
self.start = start
self.end = end
self.weight = weight
self.length_squared = float((self.end.x - self.start.x) ** 2 +
(self.end.y - self.start.y) ** 2)
self.extent = Extent(coords=(start, end))
def __str__(self):
return 'LineSegment(%s, %s)' % (self.start, self.end)
def distance(self, coord):
# http://stackoverflow.com/questions/849211/shortest-distance-between-a-point-and-a-line-segment
# http://www.topcoder.com/tc?d1=tutorials&d2=geometry1&module=Static#line_point_distance
# http://local.wasp.uwa.edu.au/~pbourke/geometry/pointline/
try:
dx = (self.end.x - self.start.x)
dy = (self.end.y - self.start.y)
u = ((coord.x - self.start.x) * dx +
(coord.y - self.start.y) * dy) / self.length_squared
if u < 0:
u = 0
elif u > 1:
u = 1
except ZeroDivisionError:
u = 0 # Our line is zero-length. That's ok.
dx = self.start.x + u * dx - coord.x
dy = self.start.y + u * dy - coord.y
return math.hypot(dx, dy)
def add_heat_to_matrix(self, matrix, kernel):
# Iterate over every point in a bounding box around this, with an
# extra margin given by the kernel's self-reported maximum range.
# TODO: There is probably a more clever iteration that skips more
# of the empty space.
for x in range(int(self.extent.min.x - kernel.radius),
int(self.extent.max.x + kernel.radius + 1)):
for y in range(int(self.extent.min.y - kernel.radius),
int(self.extent.max.y + kernel.radius + 1)):
coord = Coordinate(x, y)
heat = kernel.heat(self.distance(coord))
if heat:
matrix.add(coord, self.weight * heat)
def map(self, func):
return LineSegment(func(self.start), func(self.end))
class LinearKernel:
'''Uses a linear falloff, essentially turning a point into a cone.'''
def __init__(self, radius):
self.radius = radius # in pixels
self.radius_float = float(radius) # worthwhile time saver
def heat(self, distance):
if distance >= self.radius:
return 0.0
return 1.0 - (distance / self.radius_float)
class GaussianKernel:
def __init__(self, radius):
'''radius is the distance beyond which you should not bother.'''
self.radius = radius
# We set the scale such that the heat value drops to 1/256 of
# the peak at a distance of radius.
self.scale = math.log(256) / radius
def heat(self, distance):
'''Returns 1.0 at center, 1/e at radius pixels from center.'''
if distance >= self.radius:
return 0.0
return math.e ** (-distance * self.scale)
class ColorMap:
DEFAULT_HSVA_MIN_STR = '000ffff00'
DEFAULT_HSVA_MAX_STR = '02affffff'
@staticmethod
def _str_to_float(string, base=16, maxval=256):
return float(int(string, base)) / maxval
@staticmethod
def str_to_hsva(string):
'''
Returns a 4-tuple of ints from a hex string color specification,
such that AAABBCCDD becomes AAA, BB, CC, DD. For example,
str2hsva('06688bbff') returns (102, 136, 187, 255). Note that
the first number is 3 digits.
'''
if string.startswith('#'):
string = string[1:] # Leading "#" is now optional.
return tuple(ColorMap._str_to_float(s) for s in (string[0:3],
string[3:5],
string[5:7],
string[7:9]))
def __init__(self, hsva_min=None, hsva_max=None, image=None, steps=256):
'''
Create a color map based on a progression in the specified
range, or using pixels in a provided image.
If supplied, hsva_min and hsva_max must each be a 4-tuple of
(hue, saturation, value, alpha), where each is a float from
0.0 to 1.0. The gradient will be a linear progression from
hsva_min to hsva_max, including both ends of the range.
The optional steps argument specifies how many discrete steps
there should be in the color gradient when using hsva_min
and hsva_max.
'''
# TODO: do the interpolation in Lab space instead of HSV
self.values = []
if image:
assert image.mode == 'RGBA', (
'Gradient image must be RGBA. Yours is %s.' % image.mode)
num_rows = image.size[1]
self.values = [image.getpixel((0, row)) for row in range(num_rows)]
self.values.reverse()
if self.values[0][3] != 0:
logging.warn('In gradient image %s, the bottom-left pixel is '
'not fully transparent. If the output appears '
'blocky, make sure your gradient image '
'transitions to fully transparent at the bottom.'
% os.path.basename(image.filename))
if self.values[-1][3] != 255:
logging.warn('In gradient image %s, the top-left pixel is '
'not fully opaque. If the output appears '
'dim, try increasing the opacity of the '
'upper region of your gradient image.'
% os.path.basename(image.filename))
else:
if not hsva_min:
hsva_min = ColorMap.str_to_hsva(self.DEFAULT_HSVA_MIN_STR)
if not hsva_max:
hsva_max = ColorMap.str_to_hsva(self.DEFAULT_HSVA_MAX_STR)
# Turn (h1,s1,v1,a1), (h2,s2,v2,a2) into (h2-h1,s2-s1,v2-v1,a2-a1)
hsva_range = list(map(lambda min, max: max - min,
hsva_min, hsva_max))
for value in range(0, steps):
hsva = list(map(
lambda range, min: value / float(steps - 1) * range + min,
hsva_range, hsva_min))
hsva[0] = hsva[0] % 1 # in case hue is out of range
rgba = tuple(
[int(x * 255)
for x in hsv_to_rgb(*hsva[0:3]) + (hsva[3],)])
self.values.append(rgba)
def get(self, floatval):
return self.values[int(floatval * (len(self.values) - 1))]
class ImageMaker():
def __init__(self, config):
'''Each argument to the constructor should be a 4-tuple of (hue,
saturaton, value, alpha), one to use for minimum data values and
one for maximum. Each should be in [0,1], however because hue is
circular, you may specify hue in any range and it will be shifted
into [0,1] as needed. This is so you can wrap around the color
wheel in either direction.'''
self.config = config
if config.background and not config.background_image:
self.background = ImageColor.getrgb(config.background)
else:
self.background = None
@staticmethod
def _blend_pixels(a, b):
# a is RGBA, b is RGB; we could write this more generically,
# but why complicate things?
alpha = a[3] / 255.0
return tuple(
map(lambda aa, bb: int(aa * alpha + bb * (1 - alpha)), a[:3], b))
def make_image(self, matrix):
extent = self.config.extent_out
if not extent:
extent = matrix.extent()
extent.resize((self.config.width or 1) - 1,
(self.config.height or 1) - 1)
size = extent.size()
size.x = int(size.x) + 1
size.y = int(size.y) + 1
logging.info('saving image (%d x %d)' % (size.x, size.y))
if self.background:
img = Image.new('RGB', (size.x, size.y), self.background)
else:
img = Image.new('RGBA', (size.x, size.y))
maxval = max(matrix.values())
logging.info('maximum density: %f' % maxval)
pixels = img.load()
for (coord, val) in matrix.items():
if val == 0.0:
continue
x = int(coord.x - extent.min.x)
y = int(coord.y - extent.min.y)
if extent.is_inside(coord):
color = self.config.colormap.get(val / maxval)
if self.background:
pixels[x, y] = ImageMaker._blend_pixels(color,
self.background)
else:
pixels[x, y] = color
if self.config.background_image:
img = Image.composite(img, self.config.background_image,
img.split()[3])
return img
class ImageSeriesMaker():
'''Creates a movie showing the data appearing on the heatmap.'''
def __init__(self, config):
self.config = config
self.image_maker = ImageMaker(config)
self.tmpdir = tempfile.mkdtemp()
self.imgfile_template = os.path.join(self.tmpdir, 'frame-%05d.png')
def _save_image(self, matrix):
self.frame_count += 1
logging.info('Frame %d' % (self.frame_count))
matrix = matrix.finalized()
image = self.image_maker.make_image(matrix)
image.save(self.imgfile_template % self.frame_count)
def maybe_save_image(self, matrix):
self.inputs_since_output += 1
if self.inputs_since_output >= self.config.frequency:
self._save_image(matrix)
self.inputs_since_output = 0
@staticmethod
def create_movie(infiles, outfile, ffmpegopts):
command = ['ffmpeg', '-i', infiles]
if ffmpegopts:
# I hope they don't have spaces in their arguments
command.extend(ffmpegopts.split())
command.append(outfile)
logging.info('Encoding video: %s' % ' '.join(command))
subprocess.call(command)
def run(self):
logging.info('Putting animation frames in %s' % self.tmpdir)
self.inputs_since_output = 0
self.frame_count = 0
matrix = process_shapes(self.config, self.maybe_save_image)
if ((not self.frame_count or
self.inputs_since_output >= self.config.straggler_threshold)):
self._save_image(matrix)
self.create_movie(self.imgfile_template,
self.config.output,
self.config.ffmpegopts)
if self.config.keepframes:
logging.info('The animation frames are in %s' % self.tmpdir)
else:
shutil.rmtree(self.tmpdir)
return matrix
def _get_osm_image(bbox, zoom, osm_base):
# Just a wrapper for osm.createOSMImage to translate coordinate schemes
try:
from osmviz.manager import PILImageManager, OSMManager
osm = OSMManager(
image_manager=PILImageManager('RGB'),
server=osm_base)
(c1, c2) = bbox.corners()
image, bounds = osm.createOSMImage((c1.lat, c2.lat, c1.lon, c2.lon),
zoom)
(lat1, lat2, lon1, lon2) = bounds
return image, Extent(coords=(LatLon(lat1, lon1),
LatLon(lat2, lon2)))
except ImportError as e:
logging.error(
"ImportError: %s.\n"
"The --osm option depends on the osmviz module, available from\n"
"http://cbick.github.com/osmviz/\n\n" % str(e))
sys.exit(1)
def _scale_for_osm_zoom(zoom):
return 256 * pow(2, zoom) / 360.0
def choose_osm_zoom(config, padding):
# Since we know we're only going to do this with Mercator, we could do
# a bit more math and solve this directly, but as a first pass method,
# we instead project the bounding box into pixel-land at a high zoom
# level, then see the power of two we're off by.
if config.zoom:
return config.zoom
if not (config.width or config.height):
raise ValueError('For OSM, you must specify height, width, or zoom')
crazy_zoom_level = 30
proj = MercatorProjection()
scale = _scale_for_osm_zoom(crazy_zoom_level)
proj.pixels_per_degree = scale
bbox_crazy_xy = config.extent_in.map(proj.project)
if config.width:
size_ratio = width_ratio = (
float(bbox_crazy_xy.size().x) / (config.width - 2 * padding))
if config.height:
size_ratio = (
float(bbox_crazy_xy.size().y) / (config.height - 2 * padding))
if config.width:
size_ratio = max(size_ratio, width_ratio)
# TODO: We use --height and --width as upper bounds, choosing a zoom
# level that lets our image be no larger than the specified size.
# It might be desirable to use them as lower bounds or to get as close
# as possible, whether larger or smaller (where "close" probably means
# in pixels, not scale factors).
# TODO: This is off by a little bit at small scales.
zoom = int(crazy_zoom_level - math.log(size_ratio, 2))
logging.info('Choosing OSM zoom level %d' % zoom)
return zoom
def get_osm_background(config, padding):
zoom = choose_osm_zoom(config, padding)
proj = MercatorProjection()
proj.pixels_per_degree = _scale_for_osm_zoom(zoom)
bbox_xy = config.extent_in.map(proj.project)
# We're not checking that the padding fits within the specified size.
bbox_xy.grow(padding)
bbox_ll = bbox_xy.map(proj.inverse_project)
image, img_bbox_ll = _get_osm_image(bbox_ll, zoom, config.osm_base)
img_bbox_xy = img_bbox_ll.map(proj.project)
# TODO: this crops to our data extent, which means we're not making
# an image of the requested dimensions. Perhaps we should let the
# user specify whether to treat the requested size as min,max,exact.
offset = bbox_xy.min - img_bbox_xy.min
image = image.crop((
int(offset.x),
int(offset.y),
int(offset.x + bbox_xy.size().x + 1),
int(offset.y + bbox_xy.size().y + 1)))
config.background_image = image
config.extent_in = bbox_ll
config.projection = proj
(config.width, config.height) = image.size
return image, bbox_ll, proj
def process_shapes(config, hook=None):
matrix = Matrix.matrix_factory(config.decay)
logging.info('processing data')
for shape in config.shapes:
shape = shape.map(config.projection.project)
# TODO: skip shapes outside map extent
shape.add_heat_to_matrix(matrix, config.kernel)
if hook:
hook(matrix)
return matrix
class FileReader():
'''Abstract base class for file readers.'''
def __init__(self, filenames=[], extras={}):
self.filenames = filenames
self.config = extras
logging.debug('%s for %s' % (self.__class__.__name__, str(filenames)))
def __iter__(self):
'''Readers can be iterated over, yielding all their shapes.'''
filenames = self.filenames
if platform.system() == 'Windows': # On Windows, apps do the globbing
filenames = chain.from_iterable(glob.iglob(f) for f in filenames)
return chain.from_iterable(self.read_file(f) for f in filenames)
def read_file(self, filename):
'''a simple file opener, for simple subclasses'''
logging.info('reading from %s' % filename)
return self.parse(open(filename, 'rU'))
class GPXFileReader(FileReader):
'''GPX track file reader'''
def read_file(self, filename):
track = TrackLog(filename)
for trkseg in track.segments():
for i, p1 in enumerate(trkseg[:-1]):
p2 = trkseg[i + 1]
yield LineSegment(p1.coords, p2.coords)
class PlainFileReader(FileReader):
'''
Reads files containing one space-separated coordinate pair per
line, with optional point weight as third term.
'''
@staticmethod
def parse(lines):
count = 0
for line in lines:
line = line.strip()
if len(line) > 0: # ignore blank lines
values = [float(x) for x in line.split()]
assert len(values) == 2 or len(values) == 3, (
'input lines must have two or three values: %s' % line)
(lat, lon) = values[0:2]
weight = 1.0 if len(values) == 2 else values[2]
count += 1
yield Point(LatLon(lat, lon), weight)
logging.info('read %d points' % count)
class CSVFileReader(FileReader):
def parse(self, lines):
import csv
reader = csv.reader(lines)
if self.config['ignore_csv_header']:
next(reader) # Skip header line
count = 0
for row in reader:
(lat, lon) = (float(row[0]), float(row[1]))
count += 1
yield Point(LatLon(lat, lon))
logging.info('read %d points' % count)
class SQLiteReader(FileReader):
def read_file(self, filename):
import sqlite3
sqlite_file = filename # name of the sqlite database file
sqlite_table = self.config['sqlite_table'] # name of the table to be queried
# Connecting to the database file
conn = sqlite3.connect(sqlite_file)
c = conn.cursor()
c.execute('SELECT "lat", "lon" FROM %s' % (sqlite_table))
all_rows = c.fetchall()
conn.close()
count = 0
for row in all_rows:
(lat, lon) = (float(row[0]), float(row[1]))
count += 1
yield Point(LatLon(lat, lon))
logging.info('read %d points' % count)
class ShapeFileReader(FileReader):
'''ESRI Shapefile reader'''
@staticmethod
def read_file(filename):
try:
import ogr
except ImportError:
try:
from osgeo import ogr
except ImportError:
raise ImportError("You need to have python-gdal bindings "
"installed")
driver = ogr.GetDriverByName("ESRI Shapefile")
dataSource = driver.Open(filename, 0)
if dataSource is None:
raise Exception("Not a valid shape file")
layer = dataSource.GetLayer()
if layer.GetGeomType() != 1:
raise Exception("Only point layers are supported")
spatial_reference = layer.GetSpatialRef()
if spatial_reference is None:
raise Exception("The shapefile doesn't have spatial reference")
spatial_reference.AutoIdentifyEPSG()
auth_code = spatial_reference.GetAuthorityCode(None)
if auth_code == '':
raise Exception("The input shapefile projection could not be "
"recognized")
if auth_code != '4326':
# TODO: implement reproject layer
# (maybe geometry by geometry is easier)
raise Exception("Currently only Lng-Lat WGS84 is supported "
"(EPSG 4326)")
count = 0
for feature in layer:
geom = feature.GetGeometryRef()
lat = geom.GetY()
lon = geom.GetX()
count += 1
yield Point(LatLon(lat, lon))
logging.info('read %d points' % count)
class AutoFileReader(FileReader):
'''Delegates a reader based on each filename extension.'''
def read_file(self, filename):
types = {'.shp': ShapeFileReader,
'.gpx': GPXFileReader,
'.csv': CSVFileReader,
'.sqlite': SQLiteReader,
'.db': SQLiteReader,
}
try:
_, ext = os.path.splitext(filename)
reader_class = types[ext]
except KeyError:
reader_class = PlainFileReader
reader = reader_class([filename], self.config)
return reader.read_file(filename)
class Configuration(object):
'''
This object holds the settings for creating a heatmap as well as
an iterator for the input data.
Most of the command line processing is about settings and data, so
the command line arguments are also processed with this object.
This happens in two phases.
First the settings are parsed and turned into more useful objects
in set_from_options(). Command line flags go in, and the
Configuration object is populated with the specified values and
defaults.
In the second phase, various other parameters are computed. These
are things we set automatically based on the other settings or on
the data. You can skip this if you set everything manually, but
The idea is that someone could import this module, populate a
Configuration instance manually, and run the process themselves.
Where possible, this object contains instances, rather than option
strings (e.g. for projection, kernel, colormap, etc).
Every parameter is explained in the glossary dictionary, and only
documented parameters are allowed. Parameters default to None.
'''
glossary = {
# Many of these are exactly the same as the command line option.
# In those cases, the documentation is left blank.
# Many have default values based on the command line defaults.
'output': '',
'width': '',
'height': '',
'margin': '',
'files': 'input files',
'shapes': 'unprojected iterable of shapes (Points and LineSegments)',
'projection': 'Projection instance',
'colormap': 'ColorMap instance',
'decay': '',
'kernel': 'kernel instance',
'extent_in': 'extent in original space',
'extent_out': 'extent in projected space',
'background': '',
'background_image': '',
'background_brightness': '',
# OpenStreetMap background tiles
'osm': 'True/False; see command line options',
'osm_base': '',
'zoom': '',
# These are for making an animation, ignored otherwise.
'ffmpegopts': '',
'keepframes': '',
'frequency': '',
'straggler_threshold': '',
# We always instantiate an ArgumentParser in order to set up
# default values. You can use this ArgumentParser in your own
# script, perhaps adding your own arguments.
'argparser': 'ArgumentParser instance for command line processing',
'sqlite_table': 'Table name to read from',
}
_kernels = {'linear': LinearKernel,
'gaussian': GaussianKernel, }
_projections = {'equirectangular': EquirectangularProjection,
'mercator': MercatorProjection, }
_filetypes = {'plain': PlainFileReader,
'gpx': GPXFileReader,
'csv': CSVFileReader,
'shp': ShapeFileReader,
'db': SQLiteReader,
'sqlite': SQLiteReader,
'auto': AutoFileReader, }
def __init__(self, use_defaults=True):
for k in self.glossary.keys():
setattr(self, k, None)
self.argparser = self._make_argparser()
if use_defaults:
self.set_defaults()
def set_defaults(self):
args = self.argparser.parse_args([])
self.set_from_options(args)
def _make_argparser(self):
'''Return a an ArgumentParser set up for our command line options.'''
from argparse import (ArgumentParser, ArgumentDefaultsHelpFormatter,
SUPPRESS)
description = 'plot a heatmap from coordinate data'
parser = ArgumentParser(description=description,
formatter_class=ArgumentDefaultsHelpFormatter)
inputs = parser.add_mutually_exclusive_group()
inputs.add_argument('-p', '--points', help=SUPPRESS)
inputs.add_argument('--csv', metavar='FILE', help=SUPPRESS)
parser.add_argument('--shp_file', help=SUPPRESS)
inputs.add_argument('-g', '--gpx', help=SUPPRESS)
parser.add_argument(
'--filetype',
choices=list(self._filetypes.keys()), default='auto',
help=('Treat all input files as this type. ("auto" will guess '
'based on the filename extension.); default: %(default)s'))
parser.add_argument(
'--ignore_csv_header', action='store_true',
help='ignore first line of CSV input files')
parser.add_argument(
'--sqlite_table', metavar='NAME',
help='sqlite table name to read from. lat and lon columns assumed')
parser.add_argument(
'-s', '--scale', type=float,
help='meters per pixel, approximate'),
parser.add_argument(
'-W', '--width', type=int,
help='width of output image'),
parser.add_argument(
'-H', '--height', type=int,
help='height of output image'),
parser.add_argument(
'-P', '--projection', metavar='NAME',
choices=list(self._projections.keys()), default='mercator',
help='default: %(default)s')
parser.add_argument(
'-e', '--extent', metavar='RANGE',
help=(
'Clip results to RANGE, which is specified as '
'lat1,lon1,lat2,lon2;'
' (for square mercator: -85.0511,-180,85.0511,180)'))
parser.add_argument(
'-R', '--margin', type=int, default=0,
help=(
'Try to keep data at least this many pixels away from image '
'border.'))
parser.add_argument(
'-r', '--radius', type=int, default=5,
help='pixel radius of point blobs; default: %(default)s')
parser.add_argument(
'-d', '--decay', type=float, default=0.95,
help=(
'float in [0,1]; Larger values give more weight to data '
'magnitude. Smaller values are more democratic. default:'
'%(default)s'))
parser.add_argument(
'-S', '--save', metavar='FILE', help='save processed data to FILE')
parser.add_argument(
'-L', '--load', metavar='FILE',
help='load processed data from FILE')
parser.add_argument(
'-o', '--output', metavar='FILE',
help='name of output file (image or video)')
parser.add_argument(
'-a', '--animate', action='store_true',
help='Make an animation instead of a static image')
parser.add_argument(
'--frequency', type=int, default=1,
help='input points per animation frame; default: %(default)s')
parser.add_argument(
'--straggler_threshold', type=int, default=1,
help='add one more animation frame if >= this many inputs remain')
parser.add_argument(
'-F', '--ffmpegopts', metavar='STR',
help='extra options to pass to ffmpeg when making an animation')
parser.add_argument(
'-K', '--keepframes', action='store_true',
help='keep intermediate images after creating an animation')
parser.add_argument(
'-b', '--background', metavar='COLOR',
help='composite onto this background (color name or #rrggbb)')
parser.add_argument(
'-I', '--background_image', metavar='FILE',
help='composite onto this image')
parser.add_argument(
'-B', '--background_brightness', type=float,
help='Multiply each pixel in background image by this.')
parser.add_argument(
'-m', '--hsva_min', metavar='HEX',
default=ColorMap.DEFAULT_HSVA_MIN_STR,
help='hhhssvvaa hex for minimum data values; default: %(default)s')
parser.add_argument(
'-M', '--hsva_max', metavar='HEX',
default=ColorMap.DEFAULT_HSVA_MAX_STR,
help='hhhssvvaa hex for maximum data values; default: %(default)s')
parser.add_argument(
'-G', '--gradient', metavar='FILE',
help=(
'Take color gradient from this the first column of pixels in '
'this image. Overrides -m and -M.'))
parser.add_argument(
'-k', '--kernel',
default='linear',
choices=list(self._kernels.keys()),
help=('Kernel to use for the falling-off function; '
'default: %(default)s'))
parser.add_argument(
'--osm', action='store_true',
help='Composite onto OpenStreetMap tiles')
parser.add_argument(
'--osm_base', metavar='URL',
default='http://tile.openstreetmap.org',
help='Base URL for map tiles; default %(default)s')
parser.add_argument(
'-z', '--zoom', type=int,
help='Zoom level for OSM; 0 (the default) means autozoom')
parser.add_argument('-v', '--verbose', action='store_true')
parser.add_argument('--debug', action='store_true')
parser.add_argument('--version', action='version',
version='%(prog)s ' + str(__version__))
parser.add_argument(
'files', nargs='*', help='input files', metavar='FILE')
return parser
def legacy_cmdline_input(self, options):
flag_to_filetype = {'points': 'plain',
'gpx': 'gpx',
'csv': 'csv',
'shp_file': 'shp'}
for flag, filetype in flag_to_filetype.items():
filename = getattr(options, flag)
if filename:
logging.warn(
'--%s is deprecated; now you can just put the input file '
'at the end of the command line. This legacy support '
'will be removed someday.' % flag)
return filetype, [filename]
raise ValueError
def set_files_from_legacy_cmdline(self, options):
try:
self.filetype, self.files = self.legacy_cmdline_input(options)
except ValueError:
pass
def set_from_options(self, options):
for k in self.glossary.keys():
try:
setattr(self, k, getattr(options, k))
except AttributeError:
pass
self.kernel = self._kernels[options.kernel](options.radius)
self.projection = self._projections[options.projection]()
if options.scale:
self.projection.meters_per_pixel = options.scale
if options.gradient:
self.colormap = ColorMap(image=Image.open(options.gradient))
else:
self.colormap = ColorMap(
hsva_min=ColorMap.str_to_hsva(options.hsva_min),
hsva_max=ColorMap.str_to_hsva(options.hsva_max))
self.set_files_from_legacy_cmdline(options)
if self.files:
if self.shapes:
logging.warning('We have both input files and shapes')
logging.warning('Ignoring the input files!')
else:
logging.debug('reading input files')
self.shapes = self._filetypes[options.filetype](
self.files, options.__dict__)
if options.extent:
(lat1, lon1, lat2, lon2) = \
[float(f) for f in options.extent.split(',')]
self.extent_in = Extent(coords=(LatLon(lat1, lon1),
LatLon(lat2, lon2)))
if options.background_image:
self.background_image = Image.open(options.background_image)
(self.width, self.height) = self.background_image.size
def fill_missing(self):
if not self.shapes:
raise ValueError('no input specified')
padding = self.margin + self.kernel.radius
if not self.extent_in:
logging.debug('reading input data')
self.shapes = list(self.shapes)
logging.debug('read %d shapes' % len(self.shapes))
self.extent_in = Extent(shapes=self.shapes)
if self.osm:
get_osm_background(self, padding)
else:
if not self.projection.is_scaled():
self.projection.auto_set_scale(self.extent_in, padding,
self.width, self.height)
if not (self.width or self.height or self.background_image):
raise ValueError('You must specify width or height or '
'scale or background_image or both osm '
'and zoom.')
if self.background_brightness is not None:
if self.background_image:
self.background_image = self.background_image.point(
lambda x: x * self.background_brightness)
self.background_brightness = None # idempotence
else:
logging.warning(
'background brightness specified, but no background image')
if not self.extent_out:
self.extent_out = self.extent_in.map(self.projection.project)
self.extent_out.grow(padding)
logging.info('input extent: %s' % str(self.extent_out.map(
self.projection.inverse_project)))
logging.info('output extent: %s' % str(self.extent_out))
def main():
logging.basicConfig(format='%(relativeCreated)8d ms // %(message)s')
config = Configuration(use_defaults=False)
args = config.argparser.parse_args()
if args.verbose:
logging.getLogger().setLevel(logging.INFO)
if args.debug:
logging.getLogger().setLevel(logging.DEBUG)
elif args.load:
logging.info('loading data')
matrix = pickle.load(open(args.load, 'rb'))
config, matrix['config'].argparser = matrix['config'], config.argparser
del matrix['config']
config.set_from_options(args)
config.fill_missing()
else:
config.set_from_options(args)
config.fill_missing()
if args.animate:
animator = ImageSeriesMaker(config)
matrix = animator.run()
else:
matrix = process_shapes(config)
matrix = matrix.finalized()
if args.output and not args.animate:
image = ImageMaker(config).make_image(matrix)
image.save(args.output)
if args.save:
logging.info('saving data')
matrix['config'] = config
del config.argparser # can't pickle an ArgumentParser
pickle.dump(matrix, open(args.save, 'wb'), 2)
logging.info('end')
if __name__ == '__main__':
main()