Add benchmark for color quantization using the rtree crate.

This commit is contained in:
Avi Weinstock 2021-07-01 13:41:37 -04:00
parent 47b3b5e5d9
commit f9c263faf6
4 changed files with 195 additions and 43 deletions

33
Cargo.lock generated
View File

@ -106,7 +106,7 @@ version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "01aa95c33b750236bd69526dd773d618c5c8ad5e156daad1fc760dfd4be8def3"
dependencies = [
"heapless",
"heapless 0.5.6",
"nom 4.2.3",
]
@ -2269,6 +2269,18 @@ dependencies = [
"stable_deref_trait",
]
[[package]]
name = "heapless"
version = "0.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "634bd4d29cbf24424d0a4bfcbf80c6960129dc24424752a7d1d1390607023422"
dependencies = [
"as-slice",
"generic-array 0.14.4",
"hash32",
"stable_deref_trait",
]
[[package]]
name = "heck"
version = "0.3.3"
@ -3690,6 +3702,12 @@ dependencies = [
"winapi 0.3.9",
]
[[package]]
name = "pdqselect"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ec91767ecc0a0bbe558ce8c9da33c068066c57ecc8bb8477ef8c1ad3ef77c27"
[[package]]
name = "peeking_take_while"
version = "0.1.2"
@ -4374,6 +4392,18 @@ dependencies = [
"petgraph 0.5.1",
]
[[package]]
name = "rstar"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ce61d743ebe516592df4dd542dfe823577b811299f7bee1106feb1bbb993dbac"
dependencies = [
"heapless 0.6.1",
"num-traits",
"pdqselect",
"smallvec",
]
[[package]]
name = "rusqlite"
version = "0.24.2"
@ -6167,6 +6197,7 @@ dependencies = [
"rand_chacha 0.3.0",
"rayon",
"ron",
"rstar",
"rusqlite",
"serde",
"structopt",

View File

@ -123,7 +123,7 @@ impl PackingFormula for GridLtrPacking {
}
}
pub trait VoxelImageEncoding: Copy {
pub trait VoxelImageEncoding {
type Workspace;
type Output;
fn create(width: u32, height: u32) -> Self::Workspace;
@ -154,6 +154,37 @@ pub fn image_from_bytes<'a, I: ImageDecoder<'a>, P: 'static + Pixel<Subpixel = u
ImageBuffer::from_raw(w, h, buf)
}
impl<'a, VIE: VoxelImageEncoding> VoxelImageEncoding for &'a VIE {
type Output = VIE::Output;
type Workspace = VIE::Workspace;
fn create(width: u32, height: u32) -> Self::Workspace { VIE::create(width, height) }
fn put_solid(&self, ws: &mut Self::Workspace, x: u32, y: u32, kind: BlockKind, rgb: Rgb<u8>) {
(*self).put_solid(ws, x, y, kind, rgb)
}
fn put_sprite(
&self,
ws: &mut Self::Workspace,
x: u32,
y: u32,
kind: BlockKind,
sprite: SpriteKind,
ori: Option<u8>,
) {
(*self).put_sprite(ws, x, y, kind, sprite, ori)
}
fn finish(ws: &Self::Workspace) -> Option<Self::Output> { VIE::finish(ws) }
}
impl<'a, VIE: VoxelImageDecoding> VoxelImageDecoding for &'a VIE {
fn start(ws: &Self::Output) -> Option<Self::Workspace> { VIE::start(ws) }
fn get_block(ws: &Self::Workspace, x: u32, y: u32, is_border: bool) -> Block {
VIE::get_block(ws, x, y, is_border)
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct QuadPngEncoding<const RESOLUTION_DIVIDER: u32>();
@ -626,7 +657,7 @@ impl<const AVERAGE_PALETTE: bool> VoxelImageDecoding for TriPngEncoding<AVERAGE_
}
pub fn image_terrain_chonk<S: RectVolSize, M: Clone, P: PackingFormula, VIE: VoxelImageEncoding>(
vie: VIE,
vie: &VIE,
packing: P,
chonk: &Chonk<Block, S, M>,
) -> Option<VIE::Output> {
@ -645,7 +676,7 @@ pub fn image_terrain_volgrid<
P: PackingFormula,
VIE: VoxelImageEncoding,
>(
vie: VIE,
vie: &VIE,
packing: P,
volgrid: &VolGrid2d<Chonk<Block, S, M>>,
) -> Option<VIE::Output> {
@ -669,7 +700,7 @@ pub fn image_terrain<
P: PackingFormula,
VIE: VoxelImageEncoding,
>(
vie: VIE,
vie: &VIE,
packing: P,
vol: &V,
lo: Vec3<u32>,
@ -700,10 +731,10 @@ pub fn image_terrain<
.unwrap_or(&Block::empty());
match (block.get_color(), block.get_sprite()) {
(Some(rgb), None) => {
VIE::put_solid(&vie, &mut image, i, j, *block, rgb);
VIE::put_solid(vie, &mut image, i, j, *block, rgb);
},
(None, Some(sprite)) => {
VIE::put_sprite(&vie, &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 \
@ -772,7 +803,7 @@ impl<VIE: VoxelImageEncoding + VoxelImageDecoding, P: PackingFormula, M: Clone,
WireChonk<VIE, P, M, S>
{
pub fn from_chonk(vie: VIE, packing: P, chonk: &Chonk<Block, S, M>) -> Option<Self> {
let data = image_terrain_chonk(vie, packing, chonk)?;
let data = image_terrain_chonk(&vie, packing, chonk)?;
Some(Self {
zmin: chonk.get_min_z(),
zmax: chonk.get_max_z(),
@ -791,7 +822,7 @@ impl<VIE: VoxelImageEncoding + VoxelImageDecoding, P: PackingFormula, M: Clone,
pub fn to_chonk(&self) -> Option<Chonk<Block, S, M>> {
let mut chonk = Chonk::new(self.zmin, self.below, self.above, self.meta.clone());
write_image_terrain(
self.vie,
&self.vie,
self.packing,
&mut chonk,
&self.data,

View File

@ -6,7 +6,7 @@ edition = "2018"
[features]
simd = ["vek/platform_intrinsics"]
bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "kiddo", "clap"]
bin_compression = ["lz-fear", "deflate", "flate2", "image/jpeg", "num-traits", "fallible-iterator", "kiddo", "clap", "rstar"]
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 }
rstar = { version = "0.8.3", optional = true }
clap = { version = "2.33.3", optional = true }

View File

@ -494,9 +494,61 @@ impl VoxelImageEncoding for MixedEncodingDenseSprites {
}
use kiddo::KdTree;
use rstar::{PointDistance, RTree, RTreeObject, RTreeParams};
#[derive(Debug)]
struct ColorPoint {
rgb: Rgb<u8>,
index: u8,
}
impl RTreeObject for ColorPoint {
type Envelope = <[i32; 3] as RTreeObject>::Envelope;
fn envelope(&self) -> Self::Envelope {
[self.rgb.r as i32, self.rgb.g as i32, self.rgb.b as i32].envelope()
}
}
impl PointDistance for ColorPoint {
fn distance_2(&self, other: &[i32; 3]) -> i32 {
(self.rgb.r as i32 - other[0]).pow(2)
+ (self.rgb.g as i32 - other[1]).pow(2)
+ (self.rgb.b as i32 - other[2]).pow(2)
}
fn contains_point(&self, other: &[i32; 3]) -> bool {
&[self.rgb.r as i32, self.rgb.g as i32, self.rgb.b as i32] == other
}
}
struct TestParams;
impl RTreeParams for TestParams {
type DefaultInsertionStrategy = rstar::RStarInsertionStrategy;
const MAX_SIZE: usize = 4;
const MIN_SIZE: usize = 2;
const REINSERTION_COUNT: usize = 1;
}
lazy_static::lazy_static! {
pub static ref PALETTE: HashMap<BlockKind, KdTree<f32, u8, 3>> = {
static ref PALETTE_RTREE: HashMap<BlockKind, RTree<ColorPoint, TestParams>> = {
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 tree = RTree::bulk_load_with_params(v.into_iter()
.enumerate()
.map(|(index, rgb)| ColorPoint { rgb, index: index as u8 })
.collect()
);
(k, tree)
})
.collect()
};
pub static ref PALETTE_KDTREE: 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");
@ -514,10 +566,32 @@ lazy_static::lazy_static! {
};
}
#[derive(Debug, Clone, Copy)]
pub struct PaletteEncoding<'a, const N: u32>(&'a HashMap<BlockKind, KdTree<f32, u8, 3>>);
pub trait NearestNeighbor {
fn nearest_neighbor(&self, x: &Rgb<u8>) -> Option<u8>;
}
impl<'a, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, N> {
impl NearestNeighbor for KdTree<f32, u8, 3> {
fn nearest_neighbor(&self, x: &Rgb<u8>) -> Option<u8> {
self.nearest_one(
&[x.r as f32, x.g as f32, x.b as f32],
&kiddo::distance::squared_euclidean,
)
.map(|(_, i)| *i)
.ok()
}
}
impl<P: RTreeParams> NearestNeighbor for RTree<ColorPoint, P> {
fn nearest_neighbor(&self, x: &Rgb<u8>) -> Option<u8> {
self.nearest_neighbor(&[x.r as i32, x.g as i32, x.b as i32])
.map(|p| p.index)
}
}
#[derive(Debug, Clone, Copy)]
pub struct PaletteEncoding<'a, NN: NearestNeighbor, const N: u32>(&'a HashMap<BlockKind, NN>);
impl<'a, NN: NearestNeighbor, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, NN, N> {
#[allow(clippy::type_complexity)]
type Output = CompressedData<(Vec<u8>, [usize; 4])>;
#[allow(clippy::type_complexity)]
@ -539,13 +613,7 @@ impl<'a, const N: u32> VoxelImageEncoding for PaletteEncoding<'a, 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);
let i = self.0[&kind].nearest_neighbor(&rgb).unwrap_or(0);
ws.3.put_pixel(x / N, y / N, image::Luma([i]));
}
@ -877,7 +945,7 @@ fn main() {
if !SKIP_IMAGECHONK {
let jpegchonkgrid_pre = Instant::now();
let jpegchonkgrid =
image_terrain_chonk(JpegEncoding, GridLtrPacking, &chunk).unwrap();
image_terrain_chonk(&JpegEncoding, GridLtrPacking, &chunk).unwrap();
let jpegchonkgrid_post = Instant::now();
if false {
@ -892,19 +960,19 @@ fn main() {
let jpegchonktall_pre = Instant::now();
let jpegchonktall =
image_terrain_chonk(JpegEncoding, TallPacking { flip_y: false }, &chunk)
image_terrain_chonk(&JpegEncoding, TallPacking { flip_y: false }, &chunk)
.unwrap();
let jpegchonktall_post = Instant::now();
let jpegchonkflip_pre = Instant::now();
let jpegchonkflip =
image_terrain_chonk(JpegEncoding, TallPacking { flip_y: true }, &chunk)
image_terrain_chonk(&JpegEncoding, TallPacking { flip_y: true }, &chunk)
.unwrap();
let jpegchonkflip_post = Instant::now();
let pngchonk_pre = Instant::now();
let pngchonk =
image_terrain_chonk(PngEncoding, GridLtrPacking, &chunk).unwrap();
image_terrain_chonk(&PngEncoding, GridLtrPacking, &chunk).unwrap();
let pngchonk_post = Instant::now();
sizes.extend_from_slice(&[
@ -924,7 +992,7 @@ fn main() {
if !SKIP_MIXED {
let mixedchonk_pre = Instant::now();
let mixedchonk =
image_terrain_chonk(MixedEncoding, TallPacking { flip_y: true }, &chunk)
image_terrain_chonk(&MixedEncoding, TallPacking { flip_y: true }, &chunk)
.unwrap();
let mixedchonk_post = Instant::now();
@ -933,7 +1001,7 @@ fn main() {
let mixeddense_pre = Instant::now();
let mixeddense = image_terrain_chonk(
MixedEncodingDenseSprites,
&MixedEncodingDenseSprites,
TallPacking { flip_y: true },
&chunk,
)
@ -954,7 +1022,7 @@ fn main() {
let quadpngfull_pre = Instant::now();
let quadpngfull = image_terrain_chonk(
QuadPngEncoding::<1>(),
&QuadPngEncoding::<1>(),
TallPacking { flip_y: true },
&chunk,
)
@ -963,7 +1031,7 @@ fn main() {
let quadpnghalf_pre = Instant::now();
let quadpnghalf = image_terrain_chonk(
QuadPngEncoding::<2>(),
&QuadPngEncoding::<2>(),
TallPacking { flip_y: true },
&chunk,
)
@ -972,7 +1040,7 @@ fn main() {
let quadpngquarttall_pre = Instant::now();
let quadpngquarttall = image_terrain_chonk(
QuadPngEncoding::<4>(),
&QuadPngEncoding::<4>(),
TallPacking { flip_y: true },
&chunk,
)
@ -981,30 +1049,39 @@ fn main() {
let quadpngquartwide_pre = Instant::now();
let quadpngquartwide =
image_terrain_chonk(QuadPngEncoding::<4>(), WidePacking::<true>(), &chunk)
image_terrain_chonk(&QuadPngEncoding::<4>(), WidePacking::<true>(), &chunk)
.unwrap();
let quadpngquartwide_post = Instant::now();
let tripngaverage_pre = Instant::now();
let tripngaverage =
image_terrain_chonk(TriPngEncoding::<true>(), WidePacking::<true>(), &chunk)
image_terrain_chonk(&TriPngEncoding::<true>(), WidePacking::<true>(), &chunk)
.unwrap();
let tripngaverage_post = Instant::now();
let tripngconst_pre = Instant::now();
let tripngconst =
image_terrain_chonk(TriPngEncoding::<false>(), WidePacking::<true>(), &chunk)
image_terrain_chonk(&TriPngEncoding::<false>(), WidePacking::<true>(), &chunk)
.unwrap();
let tripngconst_post = Instant::now();
let palettepng_pre = Instant::now();
let palettepng = image_terrain_chonk(
PaletteEncoding::<4>(&PALETTE),
let palette_kdtree_pre = Instant::now();
let palette_kdtree = image_terrain_chonk(
&PaletteEncoding::<_, 4>(&PALETTE_KDTREE),
WidePacking::<true>(),
&chunk,
)
.unwrap();
let palettepng_post = Instant::now();
let palette_kdtree_post = Instant::now();
let palette_rtree_pre = Instant::now();
let palette_rtree = image_terrain_chonk(
&PaletteEncoding::<_, 4>(&PALETTE_RTREE),
WidePacking::<true>(),
&chunk,
)
.unwrap();
let palette_rtree_post = Instant::now();
#[rustfmt::skip]
sizes.extend_from_slice(&[
@ -1014,7 +1091,8 @@ 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),
("palette_kdtree", palette_kdtree.data.len() as f32 / n as f32),
("palette_rtree", palette_rtree.data.len() as f32 / n as f32),
]);
let best_idx = sizes
.iter()
@ -1035,7 +1113,8 @@ 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()),
("palette_kdtree", (palette_kdtree_post - palette_kdtree_pre).subsec_nanos()),
("palette_rtree", (palette_rtree_post - palette_rtree_pre).subsec_nanos()),
]);
if false {
let bucket = z_buckets
@ -1077,12 +1156,21 @@ fn main() {
}
if true {
let bucket = z_buckets
.entry("palettepng")
.entry("palette_kdtree")
.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;
bucket.1 += (palette_kdtree_post - palette_kdtree_pre).subsec_nanos() as f32;
}
if true {
let bucket = z_buckets
.entry("palette_rtree")
.or_default()
.entry(chunk.get_max_z() - chunk.get_min_z())
.or_insert((0, 0.0));
bucket.0 += 1;
bucket.1 += (palette_rtree_post - palette_rtree_pre).subsec_nanos() as f32;
}
trace!(
"{} {}: uncompressed: {}, {:?} {} {:?}",
@ -1111,12 +1199,13 @@ fn main() {
let mut f = File::create(&format!("chonkjpegs/{}_{}.jpg", sitename, count))
.unwrap();
let jpeg_volgrid =
image_terrain_volgrid(JpegEncoding, GridLtrPacking, &volgrid).unwrap();
image_terrain_volgrid(&JpegEncoding, GridLtrPacking, &volgrid).unwrap();
f.write_all(&*jpeg_volgrid).unwrap();
let mixedgrid_pre = Instant::now();
let (mixed_volgrid, indices) =
image_terrain_volgrid(MixedEncoding, GridLtrPacking, &volgrid).unwrap();
image_terrain_volgrid(&MixedEncoding, GridLtrPacking, &volgrid)
.unwrap();
let mixedgrid_post = Instant::now();
let seconds = (mixedgrid_post - mixedgrid_pre).as_secs_f64();
println!(