From 348108c6e57acf63a75fe3cb4ce789cfb28da9d4 Mon Sep 17 00:00:00 2001 From: sharkAndshar Date: Mon, 21 Oct 2024 15:11:11 +0800 Subject: [PATCH] Add basic cog support --- Cargo.toml | 4 + martin/Cargo.toml | 9 +- martin/src/cog/errors.rs | 40 ++++ martin/src/cog/mod.rs | 414 ++++++++++++++++++++++++++++++++++++++ martin/src/config.rs | 17 ++ martin/src/file_config.rs | 4 + martin/src/lib.rs | 2 + martin/src/utils/error.rs | 4 + 8 files changed, 492 insertions(+), 2 deletions(-) create mode 100644 martin/src/cog/errors.rs create mode 100644 martin/src/cog/mod.rs diff --git a/Cargo.toml b/Cargo.toml index 27b11a0c1..f469d72a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ anyhow = "1.0" approx = "0.5.1" async-trait = "0.1" bit-set = "0.8" +bytemuck = "1.19.0" brotli = ">=5, <8" cargo-husky = { version = "1", features = ["user-hooks"], default-features = false } clap = { version = "4", features = ["derive"] } @@ -45,6 +46,7 @@ env_logger = "0.11" flate2 = "1" flume = "0.11" futures = "0.3" +image = "0.25.5" indoc = "2" insta = "1" itertools = "0.13" @@ -59,6 +61,7 @@ moka = { version = "0.12", features = ["future"] } num_cpus = "1" pbf_font_tools = { version = "2.5.1", features = ["freetype"] } pmtiles = { version = "0.11", features = ["http-async", "mmap-async-tokio", "tilejson", "reqwest-rustls-tls-native-roots"] } +png = "0.17.14" postgis = "0.9" postgres = { version = "0.19", features = ["with-time-0_3", "with-uuid-1", "with-serde_json-1"] } postgres-protocol = "0.6" @@ -85,6 +88,7 @@ static-files = "0.2" subst = { version = "0.3", features = ["yaml"] } testcontainers-modules = { version = "0.11.3", features = ["postgres"] } thiserror = "1" +tiff = "0.9.1" tile-grid = "0.6" tilejson = "0.4" tokio = { version = "1", features = ["macros"] } diff --git a/martin/Cargo.toml b/martin/Cargo.toml index 28dc1c4d8..fd5d0b598 100644 --- a/martin/Cargo.toml +++ b/martin/Cargo.toml @@ -59,12 +59,13 @@ name = "bench" harness = false [features] -default = ["webui", "fonts", "lambda", "mbtiles", "pmtiles", "postgres", "sprites"] +default = ["webui", "fonts", "lambda", "mbtiles", "pmtiles", "cog", "postgres", "sprites"] webui = ["dep:actix-web-static-files", "dep:static-files"] fonts = ["dep:bit-set", "dep:pbf_font_tools"] lambda = ["dep:lambda-web"] mbtiles = ["dep:mbtiles"] pmtiles = ["dep:pmtiles"] +cog = ["dep:tiff", "dep:png", "dep:bytemuck"] postgres = ["dep:deadpool-postgres", "dep:json-patch", "dep:postgis", "dep:postgres", "dep:postgres-protocol", "dep:semver", "dep:tokio-postgres-rustls"] sprites = ["dep:spreet", "tokio/fs"] bless-tests = [] @@ -77,12 +78,14 @@ actix-web-static-files = { workspace = true, optional = true } actix-web.workspace = true async-trait.workspace = true bit-set = { workspace = true, optional = true } +bytemuck = { workspace = true, optional = true } brotli.workspace = true clap.workspace = true deadpool-postgres = { workspace = true, optional = true } enum-display.workspace = true env_logger.workspace = true futures.workspace = true +image = {workspace = true, optional=true } itertools.workspace = true json-patch = { workspace = true, optional = true } lambda-web = { workspace = true, optional = true } @@ -93,6 +96,7 @@ moka.workspace = true num_cpus.workspace = true pbf_font_tools = { workspace = true, optional = true } pmtiles = { workspace = true, optional = true } +png= { workspace = true, optional=true } postgis = { workspace = true, optional = true } postgres = { workspace = true, optional = true } postgres-protocol = { workspace = true, optional = true } @@ -101,14 +105,15 @@ rustls-native-certs.workspace = true rustls-pemfile.workspace = true rustls.workspace = true semver = { workspace = true, optional = true } -serde.workspace = true serde_json.workspace = true serde_with.workspace = true serde_yaml.workspace = true +serde.workspace = true spreet = { workspace = true, optional = true } static-files = { workspace = true, optional = true } subst.workspace = true thiserror.workspace = true +tiff= { workspace = true, optional=true } tilejson.workspace = true tokio = { workspace = true, features = ["io-std"] } tokio-postgres-rustls = { workspace = true, optional = true } diff --git a/martin/src/cog/errors.rs b/martin/src/cog/errors.rs new file mode 100644 index 000000000..6501cfeff --- /dev/null +++ b/martin/src/cog/errors.rs @@ -0,0 +1,40 @@ +use std::path::PathBuf; + +use png::EncodingError; +use tiff::TiffError; + +#[derive(thiserror::Error, Debug)] +pub enum CogError { + #[error("Couldn't decoded {1} as tiff file: {0}")] + InvalidTifFile(TiffError, PathBuf), + + #[error("Requested zoom level:{0} from file {1} is out of range, the zoom level is from {2} to {3}")] + ZoomOutOfRange(u8, PathBuf, u8, u8), + + #[error("Couldn't find any image(the tiff tag newsubfile is not mask) in the tiff file: {0}")] + NoImagesFound(PathBuf), + + #[error("Couldn't seek to ifd number {1} (0 based indexing) in tiff file {2} : {0}")] + IfdSeekFailed(TiffError, usize, PathBuf), + + #[error("Too many images in the tiff file: {0}")] + TooManyImages(PathBuf), + + #[error("Couldn't find tags {1:?} at ifd {2} of tiff file {3} : {0}")] + TagsNotFound(TiffError, Vec, usize, PathBuf), + + #[error("Planar configuration not equals to 1 is not supported, the tiff file is {2}")] + PlanaConfigurationNotSupported(PathBuf, usize, u16), + + #[error("Failed to transform {1}th tile data at ifd {2} from {3} to png: {0}")] + ToPngBytesFailed(TiffError, u32, usize, PathBuf), + + #[error("Failed to read {1}th chunk(0 based index) at ifd {2} from tiff file {3}: {0}")] + ReadChunkFailed(TiffError, u32, usize, PathBuf), + + #[error("Failed to write header of png file at {0}, the color type is {1:?}, the bit depth is {2:?}: {3}")] + WritePngHeaderFailed(PathBuf, png::ColorType, png::BitDepth, EncodingError), + + #[error("Failed to write pixel bytes to png file at {0}: {1}")] + WriteToPngFailed(PathBuf, EncodingError), +} diff --git a/martin/src/cog/mod.rs b/martin/src/cog/mod.rs new file mode 100644 index 000000000..ad4493fbc --- /dev/null +++ b/martin/src/cog/mod.rs @@ -0,0 +1,414 @@ +mod errors; + +pub use errors::CogError; + +use std::collections::HashMap; +use std::fs::File; +use std::path::Path; +use std::vec; +use std::{fmt::Debug, path::PathBuf}; + +use std::io::BufWriter; +use tiff::decoder::{Decoder, DecodingResult}; +use tiff::tags::Tag; +use tiff::ColorType; + +use async_trait::async_trait; +use martin_tile_utils::{Format, TileCoord, TileInfo}; +use serde::{Deserialize, Serialize}; +use tilejson::{tilejson, TileJSON}; +use url::Url; + +extern crate bytemuck; +extern crate tilejson; + +use crate::file_config::FileError; +use crate::{ + config::UnrecognizedValues, + file_config::{ConfigExtras, FileResult, SourceConfigExtras}, + MartinResult, Source, TileData, UrlQuery, +}; + +#[derive(Clone, Debug)] +pub struct CogSource { + id: String, + path: PathBuf, + meta: Meta, + tilejson: TileJSON, + tileinfo: TileInfo, +} + +#[derive(Clone, Debug)] +struct Meta { + min_zoom: u8, + max_zoom: u8, + zoom_and_ifd: HashMap, + zoom_and_tile_across_down: HashMap, +} + +#[async_trait] +impl Source for CogSource { + fn get_id(&self) -> &str { + &self.id + } + + fn get_tilejson(&self) -> &TileJSON { + &self.tilejson + } + + fn get_tile_info(&self) -> TileInfo { + self.tileinfo + } + + fn clone_source(&self) -> Box { + Box::new(self.clone()) + } + + async fn get_tile( + &self, + xyz: TileCoord, + _url_query: Option<&UrlQuery>, + ) -> MartinResult { + let tif_file = + File::open(&self.path).map_err(|e| FileError::IoError(e, self.path.clone()))?; + let mut decoder = + Decoder::new(tif_file).map_err(|e| CogError::InvalidTifFile(e, self.path.clone()))?; + decoder = decoder.with_limits(tiff::decoder::Limits::unlimited()); + + let ifd = self.meta.zoom_and_ifd.get(&(xyz.z)).ok_or_else(|| { + CogError::ZoomOutOfRange( + xyz.z, + self.path.clone(), + self.meta.min_zoom, + self.meta.max_zoom, + ) + })?; + + decoder + .seek_to_image(*ifd) + .map_err(|e| CogError::IfdSeekFailed(e, *ifd, self.path.clone()))?; + + let tiles_across = self + .meta + .zoom_and_tile_across_down + .get(&(xyz.z)) + .ok_or_else(|| { + CogError::ZoomOutOfRange( + xyz.z, + self.path.clone(), + self.meta.min_zoom, + self.meta.max_zoom, + ) + })? + .0; + let tile_idx = xyz.y * tiles_across + xyz.x; + let decode_result = decoder + .read_chunk(tile_idx) + .map_err(|e| CogError::ReadChunkFailed(e, tile_idx, *ifd, self.path.clone()))?; + let color_type = decoder + .colortype() + .map_err(|e| CogError::InvalidTifFile(e, self.path.clone()))?; + + let tile_width = decoder.chunk_dimensions().0; + let tile_height = decoder.chunk_dimensions().1; + let (data_width, data_height) = decoder.chunk_data_dimensions(tile_idx); + let png_bytes = match color_type { + ColorType::Gray(_) => todo!(), + ColorType::RGB(_) => rgb_to_png( + decode_result, + tile_width, + tile_height, + data_width, + data_height, + &self.path, + ), + ColorType::Palette(_) => todo!(), + ColorType::GrayA(_) => todo!(), + ColorType::RGBA(_) => todo!(), + ColorType::CMYK(_) => todo!(), + ColorType::YCbCr(_) => todo!(), + }?; + + Ok(png_bytes) + } +} + +fn rgb_to_png( + data: DecodingResult, + tile_width: u32, + tile_height: u32, + data_width: u32, + data_height: u32, + path: &PathBuf, +) -> Result, CogError> { + let is_padded = data_width != tile_width; + match data { + DecodingResult::U8(vec) => { + let mut buffer = Vec::new(); + { + let w = BufWriter::new(&mut buffer); + let mut encoder = png::Encoder::new(w, tile_width, tile_height); + encoder.set_color(png::ColorType::Rgb); + encoder.set_depth(png::BitDepth::Eight); + let mut writer = encoder.write_header().map_err(|e| { + CogError::WritePngHeaderFailed( + path.clone(), + png::ColorType::Rgb, + png::BitDepth::Eight, + e, + ) + })?; + if is_padded { + //todo the no_data value should read from the tiff file, or from configuration + let arr = pad_data(0, &vec, tile_width, tile_height, data_width, data_height); + writer + .write_image_data(&arr) + .map_err(|e| CogError::WriteToPngFailed(path.clone(), e))?; + } else { + writer + .write_image_data(&vec) + .map_err(|e| CogError::WriteToPngFailed(path.clone(), e))?; + } + } + Ok(buffer) + } + DecodingResult::U16(vec) => { + let mut buffer = Vec::new(); + let w = BufWriter::new(&mut buffer); + let mut encoder = png::Encoder::new(w, tile_width, tile_height); + encoder.set_color(png::ColorType::Rgb); + encoder.set_depth(png::BitDepth::Sixteen); + + if is_padded { + //todo the no_data value should read from the tiff file + let arr = pad_data(0, &vec, tile_width, tile_height, data_width, data_height); + let u8_vec: &[u8] = bytemuck::cast_slice(&arr); + let mut writer = encoder + .write_header() + .map_err(|e| CogError::WriteToPngFailed(path.clone(), e))?; + writer + .write_image_data(u8_vec) + .map_err(|e| CogError::WriteToPngFailed(path.clone(), e))?; + } else { + let u8_vec: &[u8] = bytemuck::cast_slice(&vec); + let mut writer = encoder + .write_header() + .map_err(|e| CogError::WriteToPngFailed(path.clone(), e))?; + writer + .write_image_data(u8_vec) + .map_err(|e| CogError::WriteToPngFailed(path.clone(), e))?; + } + Ok(buffer) + } + _ => todo!(), + } +} + +fn pad_data( + no_data: T, + vec: &[T], + tile_width: u32, + tile_height: u32, + data_width: u32, + data_height: u32, +) -> Vec { + let mut arr = vec![no_data; (tile_width * tile_height * 3) as usize]; + for row in 0..data_height { + for col in 0..data_width { + let idx = row * data_width * 3 + col * 3; + let arr_idx = row * tile_width * 3 + col * 3; + arr[arr_idx as usize] = vec[idx as usize]; // r + arr[(arr_idx + 1) as usize] = vec[(idx + 1) as usize]; // g + arr[(arr_idx + 2) as usize] = vec[(idx + 2) as usize]; // b + } + } + arr +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct CogConfig { + #[serde(flatten)] + pub unrecognized: UnrecognizedValues, +} + +impl ConfigExtras for CogConfig { + fn get_unrecognized(&self) -> &UnrecognizedValues { + &self.unrecognized + } +} + +impl SourceConfigExtras for CogConfig { + async fn new_sources(&self, id: String, path: PathBuf) -> FileResult> { + let tilejson = get_tilejson(); + let tileinfo = TileInfo::new(Format::Png, martin_tile_utils::Encoding::Uncompressed); + let meta = get_meta(&path)?; + Ok(Box::new(CogSource { + id, + path, + meta, + tilejson, + tileinfo, + })) + } + + #[allow(clippy::no_effect_underscore_binding)] + async fn new_sources_url(&self, _id: String, _url: Url) -> FileResult> { + unreachable!() + } + + fn parse_urls() -> bool { + false + } +} + +//todo add more to tileJson +fn get_tilejson() -> TileJSON { + tilejson! {tiles: vec![] } +} + +fn get_meta(path: &PathBuf) -> Result { + let tif_file = File::open(path).map_err(|e| FileError::IoError(e, path.clone()))?; + let mut decoder = Decoder::new(tif_file) + .map_err(|e| CogError::InvalidTifFile(e, path.clone()))? + .with_limits(tiff::decoder::Limits::unlimited()); + + let images_ifd = get_images_ifd(&mut decoder); + + let mut zoom_and_ifd: HashMap = HashMap::new(); + let mut zoom_and_tile_across_down: HashMap = HashMap::new(); + + for image_ifd in &images_ifd { + decoder + .seek_to_image(*image_ifd) + .map_err(|e| CogError::IfdSeekFailed(e, *image_ifd, path.clone()))?; + + let zoom = u8::try_from(images_ifd.len() - (image_ifd + 1)) + .map_err(|_| CogError::TooManyImages(path.clone()))?; + + let planar_configuration: u16 = decoder + .get_tag_unsigned(Tag::PlanarConfiguration) + .map_err(|e| { + CogError::TagsNotFound( + e, + vec![Tag::PlanarConfiguration.to_u16()], + *image_ifd, + path.clone(), + ) + })?; + + if planar_configuration != 1 { + Err(CogError::PlanaConfigurationNotSupported( + path.clone(), + *image_ifd, + planar_configuration, + ))?; + } + + let (tiles_across, tiles_down) = get_across_down(&mut decoder, path, *image_ifd)?; + + zoom_and_ifd.insert(zoom, *image_ifd); + zoom_and_tile_across_down.insert(zoom, (tiles_across, tiles_down)); + } + + let min_zoom = zoom_and_ifd + .keys() + .min() + .ok_or_else(|| CogError::NoImagesFound(path.clone()))?; + + let max_zoom = zoom_and_ifd + .keys() + .max() + .ok_or_else(|| CogError::NoImagesFound(path.clone()))?; + Ok(Meta { + min_zoom: *min_zoom, + max_zoom: *max_zoom, + zoom_and_ifd, + zoom_and_tile_across_down, + }) +} + +fn get_across_down( + decoder: &mut Decoder, + path: &Path, + image_ifd: usize, +) -> Result<(u32, u32), FileError> { + let (tile_width, tile_height) = (decoder.chunk_dimensions().0, decoder.chunk_dimensions().1); + let (image_width, image_length) = get_image_width_length(decoder, path, image_ifd)?; + let tiles_across = (image_width + tile_width - 1) / tile_width; + let tiles_down = (image_length + tile_height - 1) / tile_height; + + Ok((tiles_across, tiles_down)) +} + +fn get_image_width_length( + decoder: &mut Decoder, + path: &Path, + image_ifd: usize, +) -> Result<(u32, u32), FileError> { + let (image_width, image_length) = decoder.dimensions().map_err(|e| { + CogError::TagsNotFound( + e, + vec![Tag::ImageWidth.to_u16(), Tag::ImageLength.to_u16()], + image_ifd, + path.to_path_buf(), + ) + })?; + + Ok((image_width, image_length)) +} + +fn get_images_ifd(decoder: &mut Decoder) -> Vec { + let mut res = vec![]; + let mut ifd_idx = 0; + loop { + let is_image = decoder + .get_tag_u32(Tag::NewSubfileType) //based on the tiff6.0 spec, it's 32-bit(4-byte)unsigned integer + .map_or_else(|_| true, |v| v & 4 != 4); + if is_image { + //todo We should not ignore mask in the future + res.push(ifd_idx); + } + + ifd_idx += 1; + + let next_res = decoder.seek_to_image(ifd_idx); + if next_res.is_err() { + break; + } + } + res +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use martin_tile_utils::{TileCoord, TileInfo}; + use tilejson::tilejson; + + use crate::Source; + + use super::get_meta; + + #[actix_rt::test] + async fn can_get_tile() -> () { + let path = + PathBuf::from("/home/zettakit/repos/martin/tests/fixtures/cog/tiled_cog_chengdu.tif"); + let meta = get_meta(&path).unwrap(); + let source = super::CogSource { + id: "test".to_string(), + path, + meta, + tilejson: tilejson! {tiles: vec![] }, + tileinfo: TileInfo { + format: martin_tile_utils::Format::Png, + encoding: martin_tile_utils::Encoding::Uncompressed, + }, + }; + let query = None; + let _tile = source.get_tile(TileCoord { z: 2, x: 1, y: 1 }, query).await; + let _bytes = _tile.unwrap(); + + todo!() + } +} diff --git a/martin/src/config.rs b/martin/src/config.rs index eff279dc8..4e9ca02f8 100644 --- a/martin/src/config.rs +++ b/martin/src/config.rs @@ -54,6 +54,10 @@ pub struct Config { #[serde(default, skip_serializing_if = "FileConfigEnum::is_none")] pub mbtiles: FileConfigEnum, + #[cfg(feature = "cog")] + #[serde(default, skip_serializing_if = "FileConfigEnum::is_none")] + pub cog: FileConfigEnum, + #[cfg(feature = "sprites")] #[serde(default, skip_serializing_if = "FileConfigEnum::is_none")] pub sprites: FileConfigEnum, @@ -86,6 +90,9 @@ impl Config { #[cfg(feature = "mbtiles")] res.extend(self.mbtiles.finalize("mbtiles.")?); + #[cfg(feature = "cog")] + res.extend(self.cog.finalize("cog.")?); + #[cfg(feature = "sprites")] res.extend(self.sprites.finalize("sprites.")?); @@ -103,6 +110,9 @@ impl Config { #[cfg(feature = "mbtiles")] let is_empty = is_empty && self.mbtiles.is_empty(); + #[cfg(feature = "cog")] + let is_empty = is_empty && self.cog.is_empty(); + #[cfg(feature = "sprites")] let is_empty = is_empty && self.sprites.is_empty(); @@ -179,6 +189,13 @@ impl Config { sources.push(Box::pin(val)); } + #[cfg(feature = "cog")] + if !self.cog.is_empty() { + let cfg = &mut self.cog; + let val = crate::file_config::resolve_files(cfg, idr, cache.clone(), "tif"); + sources.push(Box::pin(val)); + } + Ok(TileSources::new(try_join_all(sources).await?)) } diff --git a/martin/src/file_config.rs b/martin/src/file_config.rs index 695ef9d01..9b0c46dfa 100644 --- a/martin/src/file_config.rs +++ b/martin/src/file_config.rs @@ -45,6 +45,10 @@ pub enum FileError { #[cfg(feature = "pmtiles")] #[error(r#"PMTiles error {0} processing {1}"#)] PmtError(pmtiles::PmtError, String), + + #[cfg(feature = "cog")] + #[error(transparent)] + CogError(#[from] crate::cog::CogError), } pub trait ConfigExtras: Clone + Debug + Default + PartialEq + Send { diff --git a/martin/src/lib.rs b/martin/src/lib.rs index 59548426a..bddd5fab1 100644 --- a/martin/src/lib.rs +++ b/martin/src/lib.rs @@ -14,6 +14,8 @@ pub use utils::{ }; pub mod args; +#[cfg(feature = "cog")] +pub mod cog; pub mod file_config; #[cfg(feature = "fonts")] pub mod fonts; diff --git a/martin/src/utils/error.rs b/martin/src/utils/error.rs index d59eef77c..b4298d627 100644 --- a/martin/src/utils/error.rs +++ b/martin/src/utils/error.rs @@ -64,6 +64,10 @@ pub enum MartinError { #[error(transparent)] MbtilesError(#[from] mbtiles::MbtError), + #[cfg(feature = "cog")] + #[error(transparent)] + CogError(#[from] crate::cog::CogError), + #[error(transparent)] FileError(#[from] crate::file_config::FileError),