diff options
| -rw-r--r-- | src/lib.rs | 1 | ||||
| -rw-r--r-- | src/main.rs | 40 | ||||
| -rw-r--r-- | src/renderer.rs | 42 | ||||
| -rw-r--r-- | src/storage.rs | 67 |
4 files changed, 99 insertions, 51 deletions
@@ -2,3 +2,4 @@ pub mod gpx; pub mod layer; pub mod renderer; +pub mod storage; diff --git a/src/main.rs b/src/main.rs index 53a1301..7230dce 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,17 +1,14 @@ -use std::{ - fs, - io::{self, ErrorKind}, - path::{Path, PathBuf}, -}; +use std::{io, path::PathBuf}; use clap::Parser; use color_eyre::{ - eyre::{bail, eyre, Context, Result}, + eyre::{bail, eyre, Result}, Report, }; use hittekaart::{ gpx::{self, Compression}, renderer, + storage::{Folder, Storage}, }; use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; use is_terminal::IsTerminal; @@ -88,7 +85,8 @@ fn main() -> Result<()> { let tracks = tracks.into_iter().collect::<Result<Vec<_>>>()?; bar.finish(); - ensure_output_directory(&args.output_directory)?; + let mut storage = Folder::new(args.output_directory.clone()); + storage.prepare()?; let multibar = MultiProgress::new(); if !use_progress_bars { @@ -107,37 +105,23 @@ fn main() -> Result<()> { bar.finish(); multibar.remove(&bar); - let target = [&args.output_directory, &zoom.to_string().into()] - .iter() - .collect::<PathBuf>(); - fs::create_dir(&target)?; + storage.prepare_zoom(zoom)?; let bar = make_bar(counter.tile_count().try_into().unwrap()).with_style(progress_style.clone()); multibar.insert_from_back(1, bar.clone()); bar.set_prefix("Saving heat tiles"); - renderer::lazy_colorization(counter, &target, |x| bar.inc(x.try_into().unwrap()))?; + renderer::lazy_colorization(counter, |rendered_tile| { + storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?; + bar.inc(1); + Ok(()) + })?; bar.finish(); multibar.remove(&bar); zoom_bar.inc(1); } + storage.finish()?; zoom_bar.finish(); Ok(()) } - -fn ensure_output_directory<P: AsRef<Path>>(path: P) -> Result<()> { - let path = path.as_ref(); - let metadata = fs::metadata(path); - match metadata { - Err(e) if e.kind() == ErrorKind::NotFound => { - let parent = path.parent().unwrap_or_else(|| Path::new("/")); - fs::create_dir(path) - .context(format!("Could not create output directory at {parent:?}"))? - } - Err(e) => Err(e).context("Error while checking output directory")?, - Ok(m) if m.is_dir() => (), - Ok(_) => bail!("Output directory is not a directory"), - } - Ok(()) -} diff --git a/src/renderer.rs b/src/renderer.rs index bb33ddc..4840e73 100644 --- a/src/renderer.rs +++ b/src/renderer.rs @@ -1,9 +1,6 @@ -use std::{fs, path::Path, thread}; +use std::thread; -use color_eyre::{ - eyre::{bail, Result}, - Report, -}; +use color_eyre::{eyre::Result, Report}; use image::{ImageBuffer, Luma, Pixel, RgbaImage}; use nalgebra::{vector, Vector2}; use rayon::iter::ParallelIterator; @@ -13,6 +10,13 @@ use super::{ layer::{self, TileLayer}, }; +#[derive(Debug, Clone)] +pub struct RenderedTile { + pub x: u64, + pub y: u64, + pub data: Vec<u8>, +} + pub type HeatCounter = TileLayer<Luma<u8>>; fn render_circle<P: Pixel>(layer: &mut TileLayer<P>, center: (u64, u64), radius: u64, pixel: P) { @@ -135,32 +139,21 @@ fn colorize_tile(tile: &ImageBuffer<Luma<u8>, Vec<u8>>, max: u32) -> RgbaImage { /// rendering the next one. /// /// This has a way lower memory usage than [`colorize_heatcounter`]. -pub fn lazy_colorization<P: AsRef<Path>, F: Fn(usize) + Send + Sync>( +pub fn lazy_colorization<F: FnMut(RenderedTile) -> Result<()> + Send + Sync>( layer: HeatCounter, - base_dir: P, - progress_callback: F, + mut save_callback: F, ) -> 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(()); } - type Job = (u64, u64, Vec<u8>); - let (tx, rx) = crossbeam_channel::bounded::<Job>(30); + let (tx, rx) = crossbeam_channel::bounded::<RenderedTile>(30); thread::scope(|s| { let saver = s.spawn(move || loop { - let Ok((tile_x, tile_y, data)) = rx.recv() else { return Ok(()) }; - 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")); - fs::write(file, data)?; + let Ok(tile) = rx.recv() else { return Ok::<_, Report>(()) }; + save_callback(tile)?; }); layer @@ -168,8 +161,11 @@ pub fn lazy_colorization<P: AsRef<Path>, F: Fn(usize) + Send + Sync>( .try_for_each_with(tx, |tx, (tile_x, tile_y, tile)| { let colorized = colorize_tile(&tile, max.into()); let data = layer::compress_png_as_bytes(&colorized)?; - tx.send((tile_x, tile_y, data))?; - progress_callback(1); + tx.send(RenderedTile { + x: tile_x, + y: tile_y, + data, + })?; Ok::<(), Report>(()) })?; diff --git a/src/storage.rs b/src/storage.rs new file mode 100644 index 0000000..fef8194 --- /dev/null +++ b/src/storage.rs @@ -0,0 +1,67 @@ +use color_eyre::{eyre::{bail, WrapErr}, Result}; +use std::{ + fs, + io::ErrorKind, + path::{Path, PathBuf}, +}; + +pub trait Storage { + fn prepare(&mut self) -> Result<()>; + fn prepare_zoom(&mut self, zoom: u32) -> Result<()>; + fn store(&mut self, zoom: u32, x: u64, y: u64, data: &[u8]) -> Result<()>; + fn finish(&mut self) -> Result<()>; +} + +#[derive(Debug)] +pub struct Folder { + base_dir: PathBuf, +} + +impl Folder { + pub fn new(base_dir: PathBuf) -> Self { + Folder { base_dir } + } +} + +impl Storage for Folder { + fn prepare(&mut self) -> Result<()> { + let path = &self.base_dir; + let metadata = fs::metadata(path); + match metadata { + Err(e) if e.kind() == ErrorKind::NotFound => { + let parent = path.parent().unwrap_or_else(|| Path::new("/")); + fs::create_dir(path) + .context(format!("Could not create output directory at {parent:?}"))? + } + Err(e) => Err(e).context("Error while checking output directory")?, + Ok(m) if m.is_dir() => (), + Ok(_) => bail!("Output directory is not a directory"), + } + Ok(()) + } + + fn prepare_zoom(&mut self, zoom: u32) -> Result<()> { + let target = [&self.base_dir, &zoom.to_string().into()] + .iter() + .collect::<PathBuf>(); + fs::create_dir(target)?; + Ok(()) + } + + fn store(&mut self, zoom: u32, x: u64, y: u64, data: &[u8]) -> Result<()> { + let folder = self.base_dir.join(zoom.to_string()).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")); + fs::write(file, data)?; + Ok(()) + } + + fn finish(&mut self) -> Result<()> { + Ok(()) + } +} |