Generate palette info from block statistics, and benchmark how well it compresses.

Cargo.lock

@@ -6143,6 +6143,7 @@ dependencies = [
  "arr_macro",
  "bincode",
  "bitvec",
+ "clap",
  "criterion",
  "csv",
  "deflate 0.9.1",

common/net/src/msg/compression.rs

@@ -127,8 +127,9 @@ pub trait VoxelImageEncoding: Copy {
     type Workspace;
     type Output;
     fn create(width: u32, height: u32) -> Self::Workspace;
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>);
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>);
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -176,13 +177,14 @@ impl<const N: u32> VoxelImageEncoding for QuadPngEncoding<N> {
     }
 
     #[inline(always)]
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.0.put_pixel(x, y, image::Luma([kind as u8]));
         ws.3.put_pixel(x / N, y / N, image::Rgb([rgb.r, rgb.g, rgb.b]));
     }
 
     #[inline(always)]
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -451,7 +453,7 @@ impl<const AVERAGE_PALETTE: bool> VoxelImageEncoding for TriPngEncoding<AVERAGE_
         )
     }
 
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.0.put_pixel(x, y, image::Luma([kind as u8]));
         ws.1.put_pixel(x, y, image::Luma([0]));
         ws.2.put_pixel(x, y, image::Luma([0]));
@@ -461,6 +463,7 @@ impl<const AVERAGE_PALETTE: bool> VoxelImageEncoding for TriPngEncoding<AVERAGE_
     }
 
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -666,7 +669,7 @@ pub fn image_terrain<
     P: PackingFormula,
     VIE: VoxelImageEncoding,
 >(
-    _: VIE,
+    vie: VIE,
     packing: P,
     vol: &V,
     lo: Vec3<u32>,
@@ -697,10 +700,10 @@ pub fn image_terrain<
                     .unwrap_or(&Block::empty());
                 match (block.get_color(), block.get_sprite()) {
                     (Some(rgb), None) => {
-                        VIE::put_solid(&mut image, i, j, *block, rgb);
+                        VIE::put_solid(&vie, &mut image, i, j, *block, rgb);
                     },
                     (None, Some(sprite)) => {
-                        VIE::put_sprite(&mut image, i, j, *block, sprite, block.get_ori());
+                        VIE::put_sprite(&vie, &mut image, i, j, *block, sprite, block.get_ori());
                     },
                     _ => panic!(
                         "attr being used for color vs sprite is mutually exclusive (and that's \

world/Cargo.toml

@@ -6,7 +6,7 @@ edition = "2018"
 
 [features]
 simd = ["vek/platform_intrinsics"]
-bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "kiddo"]
+bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "kiddo", "clap"]
 
 default = ["simd"]
 
@@ -43,6 +43,7 @@ flate2 = { version = "1.0.20", optional = true }
 num-traits = { version = "0.2", optional = true }
 fallible-iterator = { version = "0.2.0", optional = true }
 kiddo = { version = "0.1.4", optional = true }
+clap = { version = "2.33.3", optional = true }
 
 
 [dev-dependencies]

world/examples/chunk_compression_benchmarks.rs

@@ -173,11 +173,12 @@ impl VoxelImageEncoding for PngEncoding {
         ImageBuffer::<Rgba<u8>, Vec<u8>>::new(width, height)
     }
 
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.put_pixel(x, y, image::Rgba([rgb.r, rgb.g, rgb.b, 255 - kind as u8]));
     }
 
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -223,11 +224,12 @@ impl VoxelImageEncoding for JpegEncoding {
         ImageBuffer::<Rgba<u8>, Vec<u8>>::new(width, height)
     }
 
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.put_pixel(x, y, image::Rgba([rgb.r, rgb.g, rgb.b, 255 - kind as u8]));
     }
 
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -268,7 +270,7 @@ impl VoxelImageEncoding for MixedEncoding {
         )
     }
 
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.0.put_pixel(x, y, image::Luma([kind as u8]));
         ws.1.put_pixel(x, y, image::Luma([0]));
         ws.2.put_pixel(x, y, image::Luma([0]));
@@ -276,6 +278,7 @@ impl VoxelImageEncoding for MixedEncoding {
     }
 
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -378,12 +381,13 @@ impl VoxelImageEncoding for MixedEncodingSparseSprites {
         )
     }
 
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.0.put_pixel(x, y, image::Luma([kind as u8]));
         ws.1.put_pixel(x, y, image::Rgb([rgb.r, rgb.g, rgb.b]));
     }
 
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -439,12 +443,13 @@ impl VoxelImageEncoding for MixedEncodingDenseSprites {
         )
     }
 
-    fn put_solid(ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
         ws.0.put_pixel(x, y, image::Luma([kind as u8]));
         ws.3.put_pixel(x, y, image::Rgb([rgb.r, rgb.g, rgb.b]));
     }
 
     fn put_sprite(
+        &self,
         ws: &mut Self::Workspace,
         x: u32,
         y: u32,
@@ -488,6 +493,100 @@ impl VoxelImageEncoding for MixedEncodingDenseSprites {
     }
 }
 
+use kiddo::KdTree;
+
+lazy_static::lazy_static! {
+    pub static ref PALETTE: HashMap<BlockKind, KdTree<f32, u8, 3>> = {
+        let ron_bytes = include_bytes!("palettes.ron");
+        let palettes: HashMap<BlockKind, Vec<Rgb<u8>>> =
+            ron::de::from_bytes(&*ron_bytes).expect("palette should parse");
+        palettes
+            .into_iter()
+            .map(|(k, v)| {
+                let mut tree = KdTree::new();
+                for (i, rgb) in v.into_iter().enumerate() {
+                    tree.add(&[rgb.r as f32, rgb.g as f32, rgb.b as f32], i as u8)
+                        .expect("kdtree insert should succeed");
+                }
+                (k, tree)
+            })
+            .collect()
+    };
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct PaletteEncoding<'a, const N: u32>(&'a HashMap<BlockKind, KdTree<f32, u8, 3>>);
+
+impl<'a, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, N> {
+    #[allow(clippy::type_complexity)]
+    type Output = CompressedData<(Vec<u8>, [usize; 4])>;
+    #[allow(clippy::type_complexity)]
+    type Workspace = (
+        ImageBuffer<image::Luma<u8>, Vec<u8>>,
+        ImageBuffer<image::Luma<u8>, Vec<u8>>,
+        ImageBuffer<image::Luma<u8>, Vec<u8>>,
+        ImageBuffer<image::Luma<u8>, Vec<u8>>,
+    );
+
+    fn create(width: u32, height: u32) -> Self::Workspace {
+        (
+            ImageBuffer::new(width, height),
+            ImageBuffer::new(width, height),
+            ImageBuffer::new(width, height),
+            ImageBuffer::new(width / N, height / N),
+        )
+    }
+
+    fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
+        ws.0.put_pixel(x, y, image::Luma([kind as u8]));
+        let i = self.0[&kind]
+            .nearest_one(
+                &[rgb.r as f32, rgb.g as f32, rgb.b as f32],
+                &kiddo::distance::squared_euclidean,
+            )
+            .map(|(_, i)| *i)
+            .unwrap_or(0);
+        ws.3.put_pixel(x / N, y / N, image::Luma([i]));
+    }
+
+    fn put_sprite(
+        &self,
+        ws: &mut Self::Workspace,
+        x: u32,
+        y: u32,
+        kind: BlockKind,
+        sprite: SpriteKind,
+        ori: Option<u8>,
+    ) {
+        ws.0.put_pixel(x, y, image::Luma([kind as u8]));
+        ws.1.put_pixel(x, y, image::Luma([sprite as u8]));
+        ws.2.put_pixel(x, y, image::Luma([ori.unwrap_or(0)]));
+    }
+
+    fn finish(ws: &Self::Workspace) -> Option<Self::Output> {
+        let mut buf = Vec::new();
+        use image::codecs::png::{CompressionType, FilterType};
+        let mut indices = [0; 4];
+        let mut f = |x: &ImageBuffer<_, Vec<u8>>, i| {
+            let png = image::codecs::png::PngEncoder::new_with_quality(
+                &mut buf,
+                CompressionType::Rle,
+                FilterType::Up,
+            );
+            png.encode(&*x.as_raw(), x.width(), x.height(), image::ColorType::L8)
+                .ok()?;
+            indices[i] = buf.len();
+            Some(())
+        };
+        f(&ws.0, 0)?;
+        f(&ws.1, 1)?;
+        f(&ws.2, 2)?;
+        f(&ws.3, 3)?;
+
+        Some(CompressedData::compress(&(buf, indices), 1))
+    }
+}
+
 #[allow(clippy::many_single_char_names)]
 fn histogram_to_dictionary(histogram: &HashMap<Vec<u8>, usize>, dictionary: &mut Vec<u8>) {
     let mut tmp: Vec<(Vec<u8>, usize)> = histogram.iter().map(|(k, v)| (k.clone(), *v)).collect();
@@ -898,6 +997,15 @@ fn main() {
                         .unwrap();
                 let tripngconst_post = Instant::now();
 
+                let palettepng_pre = Instant::now();
+                let palettepng = image_terrain_chonk(
+                    PaletteEncoding::<4>(&PALETTE),
+                    WidePacking::<true>(),
+                    &chunk,
+                )
+                .unwrap();
+                let palettepng_post = Instant::now();
+
                 #[rustfmt::skip]
                 sizes.extend_from_slice(&[
                     ("quadpngfull", quadpngfull.data.len() as f32 / n as f32),
@@ -906,6 +1014,7 @@ fn main() {
                     ("quadpngquartwide", quadpngquartwide.data.len() as f32 / n as f32),
                     ("tripngaverage", tripngaverage.data.len() as f32 / n as f32),
                     ("tripngconst", tripngconst.data.len() as f32 / n as f32),
+                    ("palettepng", palettepng.data.len() as f32 / n as f32),
                 ]);
                 let best_idx = sizes
                     .iter()
@@ -926,6 +1035,7 @@ fn main() {
                     ("quadpngquartwide", (quadpngquartwide_post - quadpngquartwide_pre).subsec_nanos()),
                     ("tripngaverage", (tripngaverage_post - tripngaverage_pre).subsec_nanos()),
                     ("tripngconst", (tripngconst_post - tripngconst_pre).subsec_nanos()),
+                    ("palettepng", (palettepng_post - palettepng_pre).subsec_nanos()),
                 ]);
                 if false {
                     let bucket = z_buckets
@@ -965,6 +1075,15 @@ fn main() {
                     bucket.0 += 1;
                     bucket.1 += (tripngconst_post - tripngconst_pre).subsec_nanos() as f32;
                 }
+                if true {
+                    let bucket = z_buckets
+                        .entry("palettepng")
+                        .or_default()
+                        .entry(chunk.get_max_z() - chunk.get_min_z())
+                        .or_insert((0, 0.0));
+                    bucket.0 += 1;
+                    bucket.1 += (palettepng_post - palettepng_pre).subsec_nanos() as f32;
+                }
                 trace!(
                     "{} {}: uncompressed: {}, {:?} {} {:?}",
                     spiralpos.x,

world/examples/world_block_statistics.rs

@@ -1,5 +1,6 @@
+use clap::{App, Arg, SubCommand};
 use common::{
-    terrain::TerrainChunkSize,
+    terrain::{BlockKind, TerrainChunkSize},
     vol::{IntoVolIterator, RectVolSize},
 };
 use fallible_iterator::FallibleIterator;
@@ -12,6 +13,9 @@ use rusqlite::{Connection, ToSql, Transaction, TransactionBehavior, NO_PARAMS};
 use std::{
     collections::{HashMap, HashSet},
     error::Error,
+    fs::File,
+    io::Write,
+    str::FromStr,
     sync::mpsc,
     time::{SystemTime, UNIX_EPOCH},
 };
@@ -21,8 +25,8 @@ use veloren_world::{
     World,
 };
 
-fn block_statistics_db() -> Result<Connection, Box<dyn Error>> {
-    let conn = Connection::open("block_statistics.sqlite")?;
+fn block_statistics_db(db_path: &str) -> Result<Connection, Box<dyn Error>> {
+    let conn = Connection::open(db_path)?;
     #[rustfmt::skip]
     conn.execute_batch("
     CREATE TABLE IF NOT EXISTS chunk (
@@ -54,7 +58,7 @@ fn block_statistics_db() -> Result<Connection, Box<dyn Error>> {
     Ok(conn)
 }
 
-fn main() -> Result<(), Box<dyn Error>> {
+fn generate(db_path: &str, ymin: Option<i32>, ymax: Option<i32>) -> Result<(), Box<dyn Error>> {
     common_frontend::init_stdout(None);
     println!("Loading world");
     let pool = ThreadPoolBuilder::new().build().unwrap();
@@ -68,7 +72,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     );
     println!("Loaded world");
 
-    let conn = block_statistics_db()?;
+    let conn = block_statistics_db(db_path)?;
 
     let existing_chunks: HashSet<(i32, i32)> = conn
         .prepare("SELECT xcoord, ycoord FROM chunk")?
@@ -79,13 +83,13 @@ fn main() -> Result<(), Box<dyn Error>> {
     let sz = world.sim().get_size();
     let (tx, rx) = mpsc::channel();
     rayon::spawn(move || {
-        let coords: Vec<_> = (1..sz.y)
+        let coords: Vec<_> = (ymin.unwrap_or(1)..ymax.unwrap_or(sz.y as i32))
             .into_iter()
             .flat_map(move |y| {
                 let tx = tx.clone();
-                (1..sz.x)
+                (1..sz.x as i32)
                     .into_iter()
-                    .map(move |x| (tx.clone(), x as i32, y as i32))
+                    .map(move |x| (tx.clone(), x, y))
             })
             .collect();
         coords.into_par_iter().for_each(|(tx, x, y)| {
@@ -107,7 +111,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                 let hi = TerrainChunkSize::RECT_SIZE.as_().with_z(chunk.get_max_z());
                 let height = chunk.get_max_z() - chunk.get_min_z();
                 for (_, block) in chunk.vol_iter(lo, hi) {
-                    let mut rgb = block.get_color().unwrap_or(Rgb::new(0, 0, 0));
+                    let mut rgb = block.get_color().unwrap_or_else(|| Rgb::new(0, 0, 0));
                     let color: [f32; 3] = [rgb.r as _, rgb.g as _, rgb.b as _];
                     if let Ok((dist, nearest)) =
                         block_colors.nearest_one(&color, &squared_euclidean)
@@ -152,7 +156,7 @@ fn main() -> Result<(), Box<dyn Error>> {
             REPLACE INTO chunk (xcoord, ycoord, height, start_time, end_time)
             VALUES (?1, ?2, ?3, ?4, ?5)
         ")?;
-        println!("Inserting results for chunk at ({}, {})", x, y);
+        println!("Inserting results for chunk at ({}, {}): {}", x, y, i);
         for ((kind, color), count) in block_counts.iter() {
             insert_block.execute(&[
                 &x as &dyn ToSql,
@@ -182,3 +186,99 @@ fn main() -> Result<(), Box<dyn Error>> {
     }
     Ok(())
 }
+
+fn palette(conn: Connection) -> Result<(), Box<dyn Error>> {
+    let mut stmt =
+        conn.prepare("SELECT kind, r, g, b, SUM(quantity) FROM block GROUP BY kind, r, g, b")?;
+    let mut block_colors: HashMap<BlockKind, Vec<(Rgb<u8>, i64)>> = HashMap::new();
+
+    let mut rows = stmt.query(NO_PARAMS)?;
+    while let Some(row) = rows.next()? {
+        let kind = BlockKind::from_str(&row.get::<_, String>(0)?)?;
+        let rgb: Rgb<u8> = Rgb::new(row.get(1)?, row.get(2)?, row.get(3)?);
+        let count: i64 = row.get(4)?;
+        block_colors
+            .entry(kind)
+            .or_insert_with(Vec::new)
+            .push((rgb, count));
+    }
+    for (_, v) in block_colors.iter_mut() {
+        v.sort_by(|a, b| b.1.cmp(&a.1));
+    }
+
+    let mut palettes: HashMap<BlockKind, Vec<Rgb<u8>>> = HashMap::new();
+    for (kind, colors) in block_colors.iter() {
+        let palette = palettes.entry(*kind).or_insert_with(Vec::new);
+        if colors.len() <= 256 {
+            for (color, _) in colors {
+                palette.push(*color);
+            }
+            println!("{:?}: {:?}", kind, palette);
+            continue;
+        }
+        let mut radius = 1024.0;
+        let mut tree = KdTree::<f32, Rgb<u8>, 3>::new();
+        while palette.len() < 256 {
+            if let Some((color, _)) = colors.iter().find(|(color, _)| {
+                tree.nearest_one(
+                    &[color.r as f32, color.g as f32, color.b as f32],
+                    &squared_euclidean,
+                )
+                .map(|(dist, _)| dist > radius)
+                .unwrap_or(true)
+            }) {
+                palette.push(*color);
+                tree.add(&[color.r as f32, color.g as f32, color.b as f32], *color)?;
+                println!("{:?}, {:?}: {:?}", kind, radius, *color);
+            } else {
+                radius -= 1.0;
+            }
+        }
+    }
+    let mut f = File::create("palettes.ron")?;
+    let pretty = ron::ser::PrettyConfig::default().with_depth_limit(2);
+    write!(f, "{}", ron::ser::to_string_pretty(&palettes, pretty)?)?;
+    Ok(())
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let mut app = App::new("world_block_statistics")
+        .version(common::util::DISPLAY_VERSION_LONG.as_str())
+        .author("The veloren devs <https://gitlab.com/veloren/veloren>")
+        .about("Compute and process block statistics on generated chunks")
+        .subcommand(
+            SubCommand::with_name("generate")
+                .about("Generate block statistics")
+                .args(&[
+                    Arg::with_name("database")
+                        .required(true)
+                        .help("File to generate/resume generation"),
+                    Arg::with_name("ymin").long("ymin").takes_value(true),
+                    Arg::with_name("ymax").long("ymax").takes_value(true),
+                ]),
+        )
+        .subcommand(
+            SubCommand::with_name("palette")
+                .about("Compute a palette from previously gathered statistics")
+                .args(&[Arg::with_name("database").required(true)]),
+        );
+
+    let matches = app.clone().get_matches();
+    match matches.subcommand() {
+        ("generate", Some(matches)) => {
+            let db_path = matches.value_of("database").expect("database is required");
+            let ymin = matches.value_of("ymin").and_then(|x| i32::from_str(x).ok());
+            let ymax = matches.value_of("ymax").and_then(|x| i32::from_str(x).ok());
+            generate(db_path, ymin, ymax)?;
+        },
+        ("palette", Some(matches)) => {
+            let conn =
+                Connection::open(&matches.value_of("database").expect("database is required"))?;
+            palette(conn)?;
+        },
+        _ => {
+            app.print_help()?;
+        },
+    }
+    Ok(())
+}