use std::path::PathBuf; use image::{GenericImageView, Rgba}; use kd_tree::{KdMap, KdPoint}; use num::Integer; use crate::context::Error; /// How many sub-segments to split each side into const SPLIT_FACTOR: u32 = 5; const IMAGE_VEC_DIM: usize = (SPLIT_FACTOR * SPLIT_FACTOR * 3) as usize; #[derive(Debug, Clone)] pub struct ImageVec { pub colors: [f32; IMAGE_VEC_DIM], } #[derive(Debug, Clone)] pub struct Jacket { pub song_id: u32, pub path: PathBuf, } impl ImageVec { // {{{ (Image => vector) encoding fn from_image(image: &impl GenericImageView>) -> ImageVec { 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() { r += pixel.0[0] as u64; g += pixel.0[1] as u64; b += pixel.0[2] as u64; count += 1; } let count = count as f64; let r = r as f64 / count; let g = g as f64 / count; let b = b as f64 / count; 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 } } // }}} } impl KdPoint for ImageVec { type Dim = typenum::U75; type Scalar = f32; fn dim() -> usize { IMAGE_VEC_DIM } fn at(&self, i: usize) -> Self::Scalar { self.colors[i] } } pub struct JacketCache { tree: KdMap, } impl JacketCache { // {{{ Generate tree pub fn new(data_dir: &PathBuf) -> Result { let jacket_csv_path = data_dir.join("jackets.csv"); let mut reader = csv::Reader::from_path(jacket_csv_path)?; let mut entries = vec![]; for record in reader.records() { let record = record?; let filename = &record[0]; let song_id = u32::from_str_radix(&record[1], 10)?; let image_path = data_dir.join(format!("jackets/{}.png", filename)); let image = image::io::Reader::open(&image_path)?.decode()?; let jacket = Jacket { song_id, path: image_path, }; entries.push((ImageVec::from_image(&image), jacket)) } let result = Self { tree: KdMap::build_by_ordered_float(entries), }; Ok(result) } // }}} // {{{ Recognise #[inline] pub fn recognise( &self, image: &impl GenericImageView>, ) -> Option<(f32, &Jacket)> { self.tree .nearest(&ImageVec::from_image(image)) .map(|p| (p.squared_distance.sqrt(), &p.item.1)) } // }}} }