Initial commit

4 files changed, 443 insertions, 0 deletions

diff --git a/src/gpx.rs b/src/gpx.rs
new file mode 100644
index 0000000..7760ca5
--- /dev/null
+++ b/src/gpx.rs
@@ -0,0 +1,69 @@
+//! GPX data extraction functions.
+//!
+//! We *could* use the [gpx](https://github.com/georust/gpx) crate, but we don't care about much
+//! other than the coordinates of the tracks. By implementing the little functionality ourselves,
+//! we can use a fast XML parser ([roxmltree](https://github.com/RazrFalcon/roxmltree)).
+use std::{f64::consts::PI, fs, path::Path};
+
+use color_eyre::eyre::{eyre, Result};
+use roxmltree::{Document, Node, NodeType};
+
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct Coordinates {
+    longitude: f64,
+    latitude: f64,
+}
+
+impl Coordinates {
+    /// Calculates the [Web Mercator
+    /// projection](https://en.wikipedia.org/wiki/Web_Mercator_projection) of the coordinates.
+    /// Returns the `(x, y)` coordinates.
+    pub fn web_mercator(self, zoom: u32) -> (u64, u64) {
+        let lambda = self.longitude.to_radians();
+        let phi = self.latitude.to_radians();
+        let x = 2u64.pow(zoom) as f64 / (2.0 * PI) * 256.0 * (lambda + PI);
+        let y = 2u64.pow(zoom) as f64 / (2.0 * PI) * 256.0 * (PI - (PI / 4.0 + phi / 2.0).tan().ln());
+        (x.floor() as u64, y.floor() as u64)
+    }
+}
+
+fn is_track_node(node: &Node) -> bool {
+    node.node_type() == NodeType::Element && node.tag_name().name() == "trk"
+}
+
+fn is_track_segment(node: &Node) -> bool {
+    node.node_type() == NodeType::Element && node.tag_name().name() == "trkseg"
+}
+
+fn is_track_point(node: &Node) -> bool {
+    node.node_type() == NodeType::Element && node.tag_name().name() == "trkpt"
+}
+
+pub fn extract_from_str(input: &str) -> Result<Vec<Coordinates>> {
+    let mut result = Vec::new();
+    let document = Document::parse(input)?;
+    for node in document.root_element().children().filter(is_track_node) {
+        for segment in node.children().filter(is_track_segment) {
+            for point in segment.children().filter(is_track_point) {
+                let latitude = point.attribute("lat")
+                    .and_then(|l| l.parse::<f64>().ok())
+                    .ok_or_else(|| eyre!("Invalid latitude"))?;
+                let longitude = point.attribute("lon")
+                    .and_then(|l| l.parse::<f64>().ok())
+                    .ok_or_else(|| eyre!("Invalid longitude"))?;
+                result.push(Coordinates { latitude, longitude });
+            }
+        }
+    }
+    Ok(result)
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum Compression {
+    None,
+}
+
+pub fn extract_from_file<P: AsRef<Path>>(path: P, _compression: Compression) -> Result<Vec<Coordinates>> {
+    let content = fs::read_to_string(path)?;
+    extract_from_str(&content)
+}
diff --git a/src/layer.rs b/src/layer.rs
new file mode 100644
index 0000000..465953d
--- /dev/null
+++ b/src/layer.rs
@@ -0,0 +1,154 @@
+//! Lazy tiled image.
+//!
+//! This supports OSM-style "tiled" images, but not all of the tiles have to be present. If a tile
+//! is not present, a default pixel is returned. The tile is allocated with the first call to a
+//! mutating operation.
+use std::{fs::{self, File}, io::BufWriter, path::Path};
+
+use color_eyre::eyre::{bail, Result};
+use fnv::FnvHashMap;
+use image::{
+    codecs::png::{CompressionType, FilterType, PngEncoder},
+    ColorType, ImageBuffer, ImageEncoder, Pixel, Rgba, RgbaImage,
+};
+
+pub const TILE_HEIGHT: u64 = 256;
+pub const TILE_WIDTH: u64 = 256;
+
+/// Main "lazy image buffer" struct.
+#[derive(Debug, Clone)]
+pub struct TileLayer<P: Pixel> {
+    tiles: FnvHashMap<(u64, u64), ImageBuffer<P, Vec<P::Subpixel>>>,
+    default_pixel: P,
+    width: u64,
+    height: u64,
+}
+
+impl<P: Pixel> TileLayer<P> {
+    pub fn from_pixel(width: u64, height: u64, pixel: P) -> Self {
+        TileLayer {
+            tiles: Default::default(),
+            default_pixel: pixel,
+            width,
+            height,
+        }
+    }
+
+    pub fn width(&self) -> u64 {
+        self.width
+    }
+
+    pub fn height(&self) -> u64 {
+        self.height
+    }
+
+    fn index(&self, x: u64, y: u64) -> ((u64, u64), (u32, u32)) {
+        (
+            (x / TILE_WIDTH, y / TILE_HEIGHT),
+            ((x % TILE_WIDTH).try_into().unwrap(), (y % TILE_HEIGHT).try_into().unwrap()),
+        )
+    }
+
+    pub fn enumerate_tiles(&self) -> impl Iterator<Item = (u64, u64, &ImageBuffer<P, Vec<P::Subpixel>>)> {
+        self.tiles.iter().map(|((x, y), t)| (*x, *y, t))
+    }
+
+    pub fn tile_mut(&mut self, tile_x: u64, tile_y: u64) -> &mut ImageBuffer<P, Vec<P::Subpixel>> {
+        self.tiles
+            .entry((tile_x, tile_y))
+            .or_insert_with(|| ImageBuffer::from_pixel(TILE_WIDTH as u32, TILE_HEIGHT as u32, self.default_pixel))
+    }
+
+    pub fn tile_for_mut(&mut self, x: u64, y: u64) -> &mut ImageBuffer<P, Vec<P::Subpixel>> {
+        let ((tx, ty), _) = self.index(x, y);
+        self.tile_mut(tx, ty)
+    }
+
+    pub fn get_pixel_checked(&self, x: u64, y: u64) -> Option<&P> {
+        if x >= self.width || y >= self.height {
+            return None;
+        }
+
+        let (outer_idx, (inner_x, inner_y)) = self.index(x, y);
+        self.tiles
+            .get(&outer_idx)
+            .map(|tile| tile.get_pixel(inner_x, inner_y))
+            .or_else(|| Some(&self.default_pixel))
+    }
+
+    pub fn get_pixel(&self, x: u64, y: u64) -> &P {
+        // This is kinda cheating, but we care about the API for now, not the speed. We can
+        // optimize this later.
+        self.get_pixel_checked(x, y).unwrap()
+    }
+
+    pub fn get_pixel_mut_checked(&mut self, x: u64, y: u64) -> Option<&mut P> {
+        if x >= self.width || y >= self.height {
+            return None;
+        }
+
+        let ((outer_x, outer_y), (inner_x, inner_y)) = self.index(x, y);
+        Some(
+            self.tile_mut(outer_x, outer_y)
+                .get_pixel_mut(inner_x, inner_y),
+        )
+    }
+
+    pub fn get_pixel_mut(&mut self, x: u64, y: u64) -> &mut P {
+        self.get_pixel_mut_checked(x, y).unwrap()
+    }
+
+    /// Enumerate all pixels that are explicitely set in this layer.
+    pub fn enumerate_pixels(&self) -> impl Iterator<Item = (u64, u64, &P)> {
+        self.tiles.iter().flat_map(|((tx, ty), tile)| {
+            tile.enumerate_pixels()
+                .map(move |(x, y, p)| (u64::from(x) + tx * TILE_WIDTH, u64::from(y) + ty * TILE_HEIGHT, p))
+        })
+    }
+
+    /// Mutably enumerate all pixels that are explicitely set in this layer.
+    pub fn enumerate_pixels_mut(&mut self) -> impl Iterator<Item = (u64, u64, &mut P)> {
+        self.tiles.iter_mut().flat_map(|((tx, ty), tile)| {
+            tile.enumerate_pixels_mut()
+                .map(move |(x, y, p)| (u64::from(x) + tx * TILE_WIDTH, u64::from(y) + ty * TILE_HEIGHT, p))
+        })
+    }
+
+    pub fn pixels(&self) -> impl Iterator<Item = &P> {
+        self.enumerate_pixels().map(|x| x.2)
+    }
+
+    pub fn pixels_mut(&mut self) -> impl Iterator<Item = &mut P> {
+        self.enumerate_pixels_mut().map(|x| x.2)
+    }
+}
+
+impl TileLayer<Rgba<u8>> {
+    pub fn save_to_directory<S: AsRef<Path>>(&self, path: S) -> Result<()> {
+        let path = path.as_ref();
+
+        for ((x, y), tile) in self.tiles.iter() {
+            let folder = path.join(&x.to_string());
+            let metadata = folder.metadata();
+            match metadata {
+                Err(_) => fs::create_dir(&folder)?,
+                Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
+                _ => {},
+            }
+            let file = folder.join(&format!("{y}.png"));
+            compress_png(tile, file)?;
+        }
+
+        Ok(())
+    }
+}
+
+pub fn compress_png<P: AsRef<Path>>(image: &RgbaImage, path: P) -> Result<()> {
+    let outstream = BufWriter::new(File::create(path)?);
+    let encoder =
+        PngEncoder::new_with_quality(outstream, CompressionType::Best, FilterType::Adaptive);
+
+    encoder.write_image(&image, image.width(), image.height(), ColorType::Rgba8)?;
+
+    Ok(())
+}
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..b5bbdb3
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,31 @@
+use std::{env, fs, path::PathBuf};
+
+use color_eyre::eyre::Result;
+
+mod gpx;
+mod layer;
+mod renderer;
+
+fn main() -> Result<()> {
+    color_eyre::install()?;
+
+    let gpx_folder = env::args().nth(1).unwrap();
+    println!("Reading from {gpx_folder}");
+
+    let mut tracks = Vec::new();
+    for file in fs::read_dir(gpx_folder).unwrap() {
+        let file = file.unwrap();
+        let data = gpx::extract_from_file(file.path(), gpx::Compression::None).unwrap();
+        tracks.push(data);
+    }
+
+    for zoom in 0..=19 {
+        println!("Doing level {zoom}");
+        let counter = renderer::render_heatcounter(zoom, &tracks);
+        let target = ["tiles", &zoom.to_string()].iter().collect::<PathBuf>();
+        fs::create_dir(&target)?;
+        renderer::lazy_colorization(&counter, &target)?;
+    }
+
+    Ok(())
+}
diff --git a/src/renderer.rs b/src/renderer.rs
new file mode 100644
index 0000000..14dfb6f
--- /dev/null
+++ b/src/renderer.rs
@@ -0,0 +1,189 @@
+use std::{fs, mem, path::Path};
+
+use color_eyre::eyre::{bail, Result};
+use image::{ImageBuffer, Luma, Pixel, Rgba, RgbaImage};
+use nalgebra::{vector, Vector2};
+use num_traits::identities::Zero;
+
+use super::{
+    gpx::Coordinates,
+    layer::{self, TileLayer},
+};
+
+pub type HeatCounter = TileLayer<Luma<u32>>;
+
+pub type HeatMap = TileLayer<Rgba<u8>>;
+
+/// Returns (a - b)**2, but ensures that no underflow happens (if b > a).
+fn diff_squared(a: u64, b: u64) -> u64 {
+    if a > b {
+        (a - b).pow(2)
+    } else {
+        (b - a).pow(2)
+    }
+}
+
+fn render_circle<P: Pixel>(layer: &mut TileLayer<P>, center: (u64, u64), radius: u64, pixel: P) {
+    let x_lower = center.0.saturating_sub(radius);
+    let x_upper = (layer.width() - 1).min(center.0 + radius);
+    let y_lower = center.1.saturating_sub(radius);
+    let y_upper = (layer.height() - 1).min(center.1 + radius);
+
+    for x in x_lower..=x_upper {
+        for y in y_lower..=y_upper {
+            if diff_squared(center.0, x) + diff_squared(center.1, y) <= radius * radius {
+                *layer.get_pixel_mut(x, y) = pixel;
+            }
+        }
+    }
+}
+
+fn render_line<P: Pixel>(
+    layer: &mut TileLayer<P>,
+    start: (u64, u64),
+    end: (u64, u64),
+    thickness: u64,
+    pixel: P,
+) {
+    use imageproc::point::Point;
+
+    if start == end {
+        return;
+    }
+
+    fn unsigned_add(a: Vector2<u64>, b: Vector2<i32>) -> Vector2<u64> {
+        let x = if b[0] < 0 {
+            a[0] - b[0].abs() as u64
+        } else {
+            a[0] + b[0] as u64
+        };
+        let y = if b[1] < 0 {
+            a[1] - b[0].abs() as u64
+        } else {
+            a[1] + b[1] as u64
+        };
+        vector![x, y]
+    }
+
+    let r = vector![end.0 as f64, end.1 as f64] - vector![start.0 as f64, start.1 as f64];
+    let r = r.normalize();
+    let normal = vector![r[1], -r[0]];
+
+    let start = vector![start.0, start.1];
+    let end = vector![end.0, end.1];
+
+    let displacement = normal * thickness as f64;
+    let displacement = displacement.map(|x| x as i32);
+    let polygon = [
+        unsigned_add(start, displacement),
+        unsigned_add(end, displacement),
+        unsigned_add(end, -displacement),
+        unsigned_add(start, -displacement),
+    ];
+    let min_x = polygon.iter().map(|p| p[0]).min().unwrap();
+    let min_y = polygon.iter().map(|p| p[1]).min().unwrap();
+    let max_x = polygon.iter().map(|p| p[0]).max().unwrap();
+    let max_y = polygon.iter().map(|p| p[1]).max().unwrap();
+
+    let mut overlay = ImageBuffer::<P, Vec<P::Subpixel>>::new(
+        (max_x - min_x).try_into().unwrap(),
+        (max_y - min_y).try_into().unwrap(),
+    );
+    let adjusted_poly = polygon
+        .into_iter()
+        .map(|p| Point::new((p[0] - min_x) as i32, (p[1] - min_y) as i32))
+        .collect::<Vec<_>>();
+    imageproc::drawing::draw_polygon_mut(&mut overlay, &adjusted_poly, pixel);
+
+    for (x, y, pixel) in overlay.enumerate_pixels() {
+        if pixel.channels()[0] > Zero::zero() {
+            *layer.get_pixel_mut(u64::from(x) + min_x, u64::from(y) + min_y) = *pixel;
+        }
+    }
+}
+
+fn merge_heat_counter(base: &mut HeatCounter, overlay: &HeatCounter) {
+    for (x, y, source) in overlay.enumerate_pixels() {
+        let target = base.get_pixel_mut(x, y);
+        target[0] += source[0];
+    }
+}
+
+fn colorize_tile(tile: &ImageBuffer<Luma<u32>, Vec<u32>>, max: u32) -> RgbaImage {
+    let gradient = colorgrad::turbo();
+    let mut result = ImageBuffer::from_pixel(tile.width(), tile.height(), [0, 0, 0, 0].into());
+    for (x, y, pixel) in tile.enumerate_pixels() {
+        if pixel[0] > 0 {
+            let alpha = pixel[0] as f64 / max as f64;
+            let color = gradient.at(alpha);
+            let target = result.get_pixel_mut(x, y);
+            *target = color.to_rgba8().into();
+        }
+    }
+    result
+}
+
+pub fn colorize_heatcounter(layer: &HeatCounter) -> HeatMap {
+    let max = layer.pixels().map(|l| l.0[0]).max().unwrap_or_default();
+    let mut result = TileLayer::from_pixel(layer.width(), layer.height(), [0, 0, 0, 0].into());
+    if max == 0 {
+        return result;
+    }
+    for (tile_x, tile_y, tile) in layer.enumerate_tiles() {
+        let colorized = colorize_tile(&tile, max);
+        *result.tile_mut(tile_x, tile_y) = colorized;
+    }
+    result
+}
+
+/// Lazily colorizes a [`HeatCounter`] by colorizing it tile-by-tile and saving a tile before
+/// rendering the next one.
+///
+/// This has a way lower memory usage than [`colorize_heatcounter`].
+pub fn lazy_colorization<P: AsRef<Path>>(layer: &HeatCounter, base_dir: P) -> Result<()> {
+    let base_dir = base_dir.as_ref();
+    let max = layer.pixels().map(|l| l.0[0]).max().unwrap_or_default();
+    if max == 0 {
+        return Ok(());
+    }
+    for (tile_x, tile_y, tile) in layer.enumerate_tiles() {
+        let colorized = colorize_tile(&tile, max);
+        let folder = base_dir.join(&tile_x.to_string());
+        let metadata = folder.metadata();
+        match metadata {
+            Err(_) => fs::create_dir(&folder)?,
+            Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
+            _ => {}
+        }
+        let file = folder.join(&format!("{tile_y}.png"));
+        layer::compress_png(&colorized, file)?;
+    }
+    Ok(())
+}
+
+/// Renders the heat counter for the given zoom level and track points.
+pub fn render_heatcounter(zoom: u32, tracks: &[Vec<Coordinates>]) -> HeatCounter {
+    let size = 256 * 2u64.pow(zoom);
+
+    let mut heatcounter = TileLayer::from_pixel(size, size, [0].into());
+
+    for track in tracks {
+        let mut layer = TileLayer::from_pixel(size, size, [0].into());
+
+        let points = track
+            .iter()
+            .map(|coords| coords.web_mercator(zoom))
+            .collect::<Vec<_>>();
+
+        for point in points.iter() {
+            render_circle(&mut layer, *point, 10, [1].into());
+        }
+
+        for (a, b) in points.iter().zip(points.iter().skip(1)) {
+            render_line(&mut layer, *a, *b, 10, [1].into());
+        }
+
+        merge_heat_counter(&mut heatcounter, &layer);
+    }
+    heatcounter
+}