1
Fork 0
shimmeringmoon/src/arcaea/jacket.rs
prescientmoon ce18db3d14
Automatically run jacket processing
This commit makes it so jacket processing is automatically run if any of
it's outputs are missing from the filesystem, or if the hash of the raw
jacket directory has changes since the last recorded value.

Moreover, most assets and all fonts are now embedded in the binary!
2024-11-09 12:22:35 +01:00

186 lines
5.3 KiB
Rust

// {{{ Imports
use std::fs;
use anyhow::Context;
use faer::{Mat, MatRef};
use image::{GenericImageView, Pixel};
use num::{Integer, ToPrimitive};
use serde::{Deserialize, Serialize};
use crate::arcaea::chart::{Difficulty, Jacket, SongCache};
use crate::context::paths::ShimmeringPaths;
use crate::context::Error;
// }}}
/// How many sub-segments to split each side into
pub const SPLIT_FACTOR: u32 = 8;
pub const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR * 3) as usize;
pub const BITMAP_IMAGE_SIZE: u32 = 174;
pub const JACKET_RECOGNITITION_DIMENSIONS: usize = 20;
// {{{ (Image => vector) encoding
#[allow(clippy::identity_op)]
pub fn image_to_vec(image: &impl GenericImageView) -> MVec<f32> {
let mut colors = MVec::zeros(IMAGE_VEC_DIM, 1);
let chunk_width = image.width() / SPLIT_FACTOR;
let chunk_height = image.height() / SPLIT_FACTOR;
for i in 0..(SPLIT_FACTOR * SPLIT_FACTOR) {
let (iy, ix) = i.div_rem(&SPLIT_FACTOR);
let cropped = image.view(
chunk_width * ix,
chunk_height * iy,
chunk_width,
chunk_height,
);
let mut r = 0;
let mut g = 0;
let mut b = 0;
let mut count = 0;
for (_, _, pixel) in cropped.pixels() {
let channels = pixel.channels();
// I'm not sure this does what it's supposed to do for non rgb(a) pixels...
r += channels[0].to_u64().unwrap().pow(2);
g += channels[1].to_u64().unwrap().pow(2);
b += channels[2].to_u64().unwrap().pow(2);
count += 1;
}
let count = count as f64;
let r = (r as f64 / count).sqrt();
let g = (g as f64 / count).sqrt();
let b = (b as f64 / count).sqrt();
colors[(i as usize * 3 + 0, 0)] = r as f32;
colors[(i as usize * 3 + 1, 0)] = g as f32;
colors[(i as usize * 3 + 2, 0)] = b as f32;
}
colors
}
// }}}
/// A column vector
pub type MVec<T> = Mat<T>;
/// This struct holds:
/// - a set of (song_id, vec) pairs of different images projected through the
/// aforementioned transform.
/// - an projection matrix for dimensionality reduction
#[derive(Clone, Serialize, Deserialize)]
pub struct JacketCache {
/// A matrix with each column corresponding to the result of passing a jacket
/// through [[image_to_vec]], and then projecting it through `transform_matrix`
pub jacket_matrix: Mat<f32>,
/// Assigns each column of `jacket_matrix` a song id.
pub jacket_ids: Vec<u32>,
/// A projection matrix for dimensionality reduction.
pub transform_matrix: Mat<f32>,
}
// {{{ Read jackets
pub fn read_jackets(paths: &ShimmeringPaths, song_cache: &mut SongCache) -> Result<(), Error> {
let suffix = format!("_{BITMAP_IMAGE_SIZE}.jpg");
let songs_dir = paths.jackets_path();
let entries = fs::read_dir(songs_dir).with_context(|| "Couldn't read songs directory")?;
for entry in entries {
let dir = entry?;
let raw_dir_name = dir.file_name();
let dir_name = raw_dir_name.to_str().unwrap();
let song_id = dir_name
.parse()
.with_context(|| format!("Dir name {dir_name} could not be parsed as `u32` song id"))?;
let entries = fs::read_dir(dir.path()).with_context(|| "Couldn't read song directory")?;
for entry in entries {
let file = entry?;
let raw_name = file.file_name();
let name = raw_name.to_str().unwrap();
if !name.ends_with(&suffix) {
continue;
}
let name = name.strip_suffix(&suffix).unwrap();
let difficulty = Difficulty::DIFFICULTY_SHORTHANDS
.iter()
.zip(Difficulty::DIFFICULTIES)
.find_map(|(s, d)| Some(d).filter(|_| name == s.to_lowercase()));
let contents: &'static _ = fs::read(file.path())
.with_context(|| "Coult not read prepared jacket image")?
.leak();
let image = image::load_from_memory(contents)
.with_context(|| "Could not load jacket image from prepared bytes")?;
let bitmap: &'static _ = Box::leak(Box::new(image.into_rgb8()));
if let Some(difficulty) = difficulty {
let chart = song_cache
.lookup_by_difficulty_mut(song_id, difficulty)
.unwrap();
chart.jacket_source = Some(difficulty);
chart.cached_jacket = Some(Jacket {
raw: contents,
bitmap,
});
} else {
for (_, chart_id) in song_cache.lookup_song(song_id)?.charts() {
let chart = song_cache.lookup_chart_mut(chart_id)?;
if chart.jacket_source.is_none() {
chart.cached_jacket = Some(Jacket {
raw: contents,
bitmap,
});
chart.jacket_source = None;
}
}
}
}
}
Ok(())
}
// }}}
impl JacketCache {
// {{{ Generate
pub fn new(paths: &ShimmeringPaths) -> Result<Self, Error> {
let bytes = fs::read(paths.recognition_matrix_path())
.with_context(|| "Could not read jacket recognition matrix")?;
let result = postcard::from_bytes(&bytes)?;
// .with_context(|| "Could not decode jacket recognition matrix")?;
Ok(result)
}
// }}}
// {{{ Recognise
/// Transforms a vector from image space to recognition space.
#[inline]
pub fn transform_vec(&self, vec: MatRef<f32>) -> MVec<f32> {
&self.transform_matrix * vec
}
#[inline]
pub fn recognise(&self, image: &impl GenericImageView) -> Option<(f32, u32)> {
let vec = self.transform_vec(image_to_vec(image).as_ref());
self.jacket_ids
.iter()
.enumerate()
.map(|(idx, id)| {
(id, {
(self.jacket_matrix.subcols(idx, 1) - &vec).squared_norm_l2()
})
})
.min_by(|(_, d1), (_, d2)| d1.partial_cmp(d2).expect("NaN distance encountered"))
.map(|(i, d)| (d.sqrt(), *i))
}
// }}}
}