2024-08-16 23:24:11 +02:00
|
|
|
use std::{fs, io::Cursor};
|
2024-06-22 23:07:11 +02:00
|
|
|
|
2024-07-20 04:52:24 +02:00
|
|
|
use image::{imageops::FilterType, GenericImageView, Rgba};
|
2024-06-22 23:07:11 +02:00
|
|
|
use num::Integer;
|
|
|
|
|
2024-07-01 18:00:03 +02:00
|
|
|
use crate::{
|
2024-08-08 23:26:13 +02:00
|
|
|
arcaea::chart::{Difficulty, Jacket, SongCache},
|
2024-08-16 23:24:11 +02:00
|
|
|
assets::{get_asset_dir, should_blur_jacket_art, should_skip_jacket_art},
|
2024-07-01 18:00:03 +02:00
|
|
|
context::Error,
|
2024-08-08 23:26:13 +02:00
|
|
|
recognition::fuzzy_song_name::guess_chart_name,
|
2024-07-01 18:00:03 +02:00
|
|
|
};
|
2024-06-22 23:07:11 +02:00
|
|
|
|
|
|
|
/// How many sub-segments to split each side into
|
2024-07-01 18:00:03 +02:00
|
|
|
pub const SPLIT_FACTOR: u32 = 8;
|
|
|
|
pub const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR * 3) as usize;
|
2024-07-20 04:52:24 +02:00
|
|
|
pub const BITMAP_IMAGE_SIZE: u32 = 174;
|
2024-06-22 23:07:11 +02:00
|
|
|
|
2024-07-18 20:17:39 +02:00
|
|
|
#[derive(Debug, Clone)]
|
2024-06-22 23:07:11 +02:00
|
|
|
pub struct ImageVec {
|
|
|
|
pub colors: [f32; IMAGE_VEC_DIM],
|
|
|
|
}
|
|
|
|
|
|
|
|
impl ImageVec {
|
|
|
|
// {{{ (Image => vector) encoding
|
2024-08-11 03:14:02 +02:00
|
|
|
fn from_image(image: &impl GenericImageView<Pixel = Rgba<u8>>) -> Self {
|
2024-06-22 23:07:11 +02:00
|
|
|
let mut colors = [0.0; IMAGE_VEC_DIM];
|
|
|
|
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() {
|
2024-07-18 20:17:39 +02:00
|
|
|
r += (pixel.0[0] as u64).pow(2);
|
|
|
|
g += (pixel.0[1] as u64).pow(2);
|
|
|
|
b += (pixel.0[2] as u64).pow(2);
|
2024-06-22 23:07:11 +02:00
|
|
|
count += 1;
|
|
|
|
}
|
|
|
|
|
|
|
|
let count = count as f64;
|
2024-07-18 20:17:39 +02:00
|
|
|
let r = (r as f64 / count).sqrt();
|
|
|
|
let g = (g as f64 / count).sqrt();
|
|
|
|
let b = (b as f64 / count).sqrt();
|
2024-06-22 23:07:11 +02:00
|
|
|
colors[i as usize * 3 + 0] = r as f32;
|
|
|
|
colors[i as usize * 3 + 1] = g as f32;
|
|
|
|
colors[i as usize * 3 + 2] = b as f32;
|
|
|
|
}
|
|
|
|
|
|
|
|
Self { colors }
|
|
|
|
}
|
|
|
|
|
2024-08-08 15:59:36 +02:00
|
|
|
#[inline]
|
|
|
|
pub fn distance_squared_to(&self, other: &Self) -> f32 {
|
|
|
|
let mut total = 0.0;
|
2024-06-22 23:07:11 +02:00
|
|
|
|
2024-08-08 15:59:36 +02:00
|
|
|
for i in 0..IMAGE_VEC_DIM {
|
|
|
|
let d = self.colors[i] - other.colors[i];
|
|
|
|
total += d * d;
|
|
|
|
}
|
2024-06-22 23:07:11 +02:00
|
|
|
|
2024-08-08 15:59:36 +02:00
|
|
|
total
|
2024-06-22 23:07:11 +02:00
|
|
|
}
|
2024-08-08 15:59:36 +02:00
|
|
|
// }}}
|
2024-06-22 23:07:11 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
pub struct JacketCache {
|
2024-08-08 15:59:36 +02:00
|
|
|
jackets: Vec<(u32, ImageVec)>,
|
2024-06-22 23:07:11 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
impl JacketCache {
|
2024-08-08 23:26:13 +02:00
|
|
|
// {{{ Generate
|
2024-06-23 02:51:50 +02:00
|
|
|
// This is a bit inefficient (using a hash set), but only runs once
|
2024-08-16 23:24:11 +02:00
|
|
|
pub fn new(song_cache: &mut SongCache) -> Result<Self, Error> {
|
2024-08-08 15:59:36 +02:00
|
|
|
let jacket_vectors = if should_skip_jacket_art() {
|
2024-08-16 23:24:11 +02:00
|
|
|
let path = get_asset_dir().join("placeholder_jacket.jpg");
|
2024-07-19 00:02:17 +02:00
|
|
|
let contents: &'static _ = fs::read(path)?.leak();
|
|
|
|
let image = image::load_from_memory(contents)?;
|
|
|
|
let bitmap: &'static _ = Box::leak(Box::new(
|
|
|
|
image
|
|
|
|
.resize(BITMAP_IMAGE_SIZE, BITMAP_IMAGE_SIZE, FilterType::Nearest)
|
|
|
|
.into_rgb8(),
|
|
|
|
));
|
|
|
|
|
2024-08-10 03:08:38 +02:00
|
|
|
for chart in song_cache.charts_mut() {
|
|
|
|
chart.cached_jacket = Some(Jacket {
|
|
|
|
raw: contents,
|
|
|
|
bitmap,
|
|
|
|
});
|
2024-06-23 02:51:50 +02:00
|
|
|
}
|
2024-07-01 18:00:03 +02:00
|
|
|
|
2024-07-19 00:02:17 +02:00
|
|
|
Vec::new()
|
|
|
|
} else {
|
|
|
|
let entries =
|
2024-08-16 23:24:11 +02:00
|
|
|
fs::read_dir(get_asset_dir().join("songs")).expect("Couldn't read songs directory");
|
2024-08-08 15:59:36 +02:00
|
|
|
let mut jacket_vectors = vec![];
|
2024-07-19 00:02:17 +02:00
|
|
|
|
|
|
|
for entry in entries {
|
|
|
|
let dir = entry?;
|
|
|
|
let raw_dir_name = dir.file_name();
|
|
|
|
let dir_name = raw_dir_name.to_str().unwrap();
|
|
|
|
for entry in fs::read_dir(dir.path()).expect("Couldn't read song directory") {
|
|
|
|
let file = entry?;
|
|
|
|
let raw_name = file.file_name();
|
|
|
|
let name = raw_name.to_str().unwrap().strip_suffix(".jpg").unwrap();
|
|
|
|
|
|
|
|
if !name.ends_with("_256") {
|
|
|
|
continue;
|
|
|
|
}
|
2024-08-10 03:08:38 +02:00
|
|
|
|
2024-07-19 00:02:17 +02:00
|
|
|
let name = name.strip_suffix("_256").unwrap();
|
|
|
|
|
|
|
|
let difficulty = match name {
|
|
|
|
"0" => Some(Difficulty::PST),
|
|
|
|
"1" => Some(Difficulty::PRS),
|
|
|
|
"2" => Some(Difficulty::FTR),
|
|
|
|
"3" => Some(Difficulty::BYD),
|
|
|
|
"4" => Some(Difficulty::ETR),
|
|
|
|
"base" => None,
|
|
|
|
"base_night" => None,
|
|
|
|
"base_ja" => None,
|
|
|
|
_ => Err(format!("Unknown jacket suffix {}", name))?,
|
|
|
|
};
|
|
|
|
|
2024-08-10 03:08:38 +02:00
|
|
|
let (song_id, chart_id) = {
|
|
|
|
let (song, chart) =
|
|
|
|
guess_chart_name(dir_name, &song_cache, difficulty, true)?;
|
|
|
|
(song.id, chart.id)
|
|
|
|
};
|
2024-07-19 00:02:17 +02:00
|
|
|
|
|
|
|
let contents: &'static _ = fs::read(file.path())?.leak();
|
|
|
|
|
|
|
|
let image = image::load_from_memory(contents)?;
|
2024-08-10 03:08:38 +02:00
|
|
|
jacket_vectors.push((song_id, ImageVec::from_image(&image)));
|
2024-08-12 03:13:41 +02:00
|
|
|
let mut image =
|
|
|
|
image.resize(BITMAP_IMAGE_SIZE, BITMAP_IMAGE_SIZE, FilterType::Nearest);
|
2024-07-19 00:02:17 +02:00
|
|
|
|
2024-08-12 03:13:41 +02:00
|
|
|
if should_blur_jacket_art() {
|
2024-08-16 15:38:00 +02:00
|
|
|
image = image.blur(40.0);
|
2024-08-12 03:13:41 +02:00
|
|
|
}
|
|
|
|
|
2024-08-16 15:38:00 +02:00
|
|
|
let encoded_pic = {
|
|
|
|
let mut processed_pic = Vec::new();
|
|
|
|
image.write_to(
|
|
|
|
&mut Cursor::new(&mut processed_pic),
|
|
|
|
image::ImageFormat::Jpeg,
|
|
|
|
)?;
|
|
|
|
processed_pic.leak()
|
|
|
|
};
|
2024-08-12 03:13:41 +02:00
|
|
|
let bitmap: &'static _ = Box::leak(Box::new(image.into_rgb8()));
|
2024-07-19 00:02:17 +02:00
|
|
|
|
|
|
|
if name == "base" {
|
2024-08-10 03:08:38 +02:00
|
|
|
// Inefficiently iterates over everything, but it's fine for ~1k entries
|
|
|
|
for chart in song_cache.charts_mut() {
|
|
|
|
if chart.song_id == song_id && chart.cached_jacket.is_none() {
|
2024-07-19 00:02:17 +02:00
|
|
|
chart.cached_jacket = Some(Jacket {
|
2024-08-16 15:38:00 +02:00
|
|
|
raw: encoded_pic,
|
2024-07-19 00:02:17 +02:00
|
|
|
bitmap,
|
|
|
|
});
|
|
|
|
}
|
|
|
|
}
|
|
|
|
} else if difficulty.is_some() {
|
2024-08-10 03:08:38 +02:00
|
|
|
let chart = song_cache.lookup_chart_mut(chart_id).unwrap();
|
2024-07-19 00:02:17 +02:00
|
|
|
chart.cached_jacket = Some(Jacket {
|
2024-08-16 15:38:00 +02:00
|
|
|
raw: encoded_pic,
|
2024-07-19 00:02:17 +02:00
|
|
|
bitmap,
|
|
|
|
});
|
|
|
|
}
|
2024-07-01 18:00:03 +02:00
|
|
|
}
|
|
|
|
}
|
2024-06-22 23:07:11 +02:00
|
|
|
|
2024-08-11 03:14:02 +02:00
|
|
|
for chart in song_cache.charts() {
|
|
|
|
if chart.cached_jacket.is_none() {
|
|
|
|
println!(
|
|
|
|
"No jacket found for '{} [{:?}]'",
|
|
|
|
song_cache.lookup_song(chart.song_id)?.song.title,
|
|
|
|
chart.difficulty
|
|
|
|
)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2024-08-08 15:59:36 +02:00
|
|
|
jacket_vectors
|
2024-07-19 00:02:17 +02:00
|
|
|
};
|
2024-07-18 20:17:39 +02:00
|
|
|
|
2024-06-22 23:07:11 +02:00
|
|
|
let result = Self {
|
2024-08-08 15:59:36 +02:00
|
|
|
jackets: jacket_vectors,
|
2024-06-22 23:07:11 +02:00
|
|
|
};
|
|
|
|
|
|
|
|
Ok(result)
|
|
|
|
}
|
|
|
|
// }}}
|
|
|
|
// {{{ Recognise
|
|
|
|
#[inline]
|
|
|
|
pub fn recognise(
|
|
|
|
&self,
|
|
|
|
image: &impl GenericImageView<Pixel = Rgba<u8>>,
|
2024-06-23 02:51:50 +02:00
|
|
|
) -> Option<(f32, &u32)> {
|
2024-08-08 15:59:36 +02:00
|
|
|
let vec = ImageVec::from_image(image);
|
|
|
|
self.jackets
|
|
|
|
.iter()
|
|
|
|
.map(|(i, v)| (i, v, v.distance_squared_to(&vec)))
|
|
|
|
.min_by(|(_, _, d1), (_, _, d2)| d1.partial_cmp(d2).expect("NaN distance encountered"))
|
2024-08-08 17:37:46 +02:00
|
|
|
.map(|(i, _, d)| (d.sqrt(), i))
|
2024-06-22 23:07:11 +02:00
|
|
|
}
|
|
|
|
// }}}
|
|
|
|
}
|