From da998ccbf13f20750f8d1e24f99848a56ffd7bd5 Mon Sep 17 00:00:00 2001 From: Dawid Pietrykowski Date: Tue, 8 Apr 2025 23:54:13 +0200 Subject: [PATCH] Add thumbnail cache --- Cargo.lock | 66 ++++++++++++++++++++ Cargo.toml | 1 + src/image.rs | 172 ++++++++++++++++++++++++++++++++++++++++++++------- src/store.rs | 1 + 4 files changed, 219 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 26ece92..25cf968 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -300,6 +300,15 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "block2" version = "0.5.1" @@ -581,6 +590,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.4.2" @@ -657,6 +675,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "43da5946c66ffcc7745f48db692ffbb10a83bfe0afd96235c5c2a4fb23994929" +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "csv" version = "1.3.1" @@ -750,6 +778,16 @@ dependencies = [ "syn", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -1021,6 +1059,16 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "795cbfc56d419a7ce47ccbb7504dd9a5b7c484c083c356e797de08bd988d9629" +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "gethostname" version = "0.4.3" @@ -1394,6 +1442,7 @@ dependencies = [ "libheif-rs", "pollster", "rexiv2", + "sha2", "threadpool", "winit", "zune-image", @@ -2804,6 +2853,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha2" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "1.3.0" @@ -3181,6 +3241,12 @@ dependencies = [ "rustc-hash", ] +[[package]] +name = "typenum" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" + [[package]] name = "unicode-ident" version = "1.0.18" diff --git a/Cargo.toml b/Cargo.toml index 109a619..868f0f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ itertools = "0.12" rexiv2 = "0.10.0" threadpool = "1.8.1" bytemuck = "1.22.0" +sha2 = "0.10.8" [profile.release] opt-level = 3 diff --git a/src/image.rs b/src/image.rs index 1ac26e7..e291ffa 100644 --- a/src/image.rs +++ b/src/image.rs @@ -8,17 +8,26 @@ use jpegxl_rs::decode::PixelFormat; use jpegxl_rs::decoder_builder; use libheif_rs::{HeifContext, LibHeif, RgbChroma}; use rexiv2::Metadata; +use sha2::Digest; +use sha2::Sha256; +use sha2::digest::consts::U32; +use sha2::digest::generic_array::GenericArray; use zune_image::codecs::jpeg::JpegDecoder; use zune_image::codecs::qoi::zune_core::colorspace::ColorSpace; use zune_image::codecs::qoi::zune_core::options::DecoderOptions; +use std::env; use std::fs; use std::fs::File; use std::fs::read; +use std::hash::Hash; use std::io::BufReader; use std::io::Cursor; +use std::io::Read; +use std::io::Write; use std::mem; use std::path::PathBuf; +use std::str::FromStr; use std::time::Instant; #[derive(Clone, Eq, Hash, PartialEq, PartialOrd)] @@ -28,14 +37,43 @@ pub enum ImageFormat { Heif, } -#[derive(Clone, Eq, Hash, PartialEq)] +#[derive(Clone)] pub struct ImageData { pub path: PathBuf, pub format: ImageFormat, pub embedded_thumbnail: bool, pub orientation: Orientation, + pub hash: GenericArray, } +impl ImageData { + pub fn get_cache_path(&self) -> PathBuf { + let home_dir = PathBuf::from_str(&env::var("HOME").unwrap()).unwrap(); + let cache_dir = home_dir.join(".cache/imflow"); + if !cache_dir.exists() { + fs::create_dir(&cache_dir).unwrap(); + } + let hash_hex = format!("{:x}", self.hash); + return cache_dir.join(hash_hex).to_path_buf(); + } +} + +impl Hash for ImageData { + fn hash(&self, state: &mut H) { + state.write(self.path.to_str().unwrap().as_bytes()); + state.write(self.hash.as_slice()); + } +} + +impl PartialEq for ImageData { + fn eq(&self, other: &Self) -> bool { + self.hash.eq(&other.hash) + } +} + +impl Eq for ImageData {} + +#[derive(Clone)] pub struct ImflowImageBuffer { pub width: usize, pub height: usize, @@ -213,14 +251,23 @@ pub fn load_available_images(dir: PathBuf) -> Vec { if let Some(format) = get_format(&path) { let meta = Metadata::new_from_path(&path) .expect(&format!("Image has no metadata: {:?}", path).to_string()); - let embedded_thumbnail = meta.get_preview_images().is_some(); + let embedded_thumbnail = if format == ImageFormat::Heif { + let ctx = HeifContext::read_from_file(path.to_str().unwrap()).unwrap(); + let binding = ctx.top_level_image_handles(); + let handle = binding.get(0).unwrap(); + handle.number_of_thumbnails() > 0 + } else { + meta.get_preview_images().is_some() + }; let orientation = Orientation::from_exif(meta.get_orientation() as u8) .unwrap_or(Orientation::NoTransforms); + let hash = get_file_hash(&path); Some(ImageData { path, format, embedded_thumbnail, orientation, + hash, }) } else { None @@ -253,13 +300,39 @@ pub fn get_embedded_thumbnail(image: &ImageData) -> Option> { } pub fn load_thumbnail(path: &ImageData) -> ImflowImageBuffer { - if path.format == ImageFormat::Heif { - return load_heif(path, true); - } - match load_thumbnail_exif(path) { - Some(thumbnail) => return thumbnail, - None => load_thumbnail_full(path), + let cache_path = path.get_cache_path(); + if cache_path.exists() { + let bytes = fs::read(cache_path).unwrap(); + let width = u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize; + let height = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize; + let buffer: &[u8] = &bytes[8..]; + let mut buffer_u8 = buffer.to_vec(); + let buffer_u32 = unsafe { + Vec::from_raw_parts( + buffer_u8.as_mut_ptr() as *mut u32, + buffer_u8.len() / 4, + buffer_u8.len() / 4, + ) + }; + std::mem::forget(buffer_u8); + + assert_eq!(width * height, buffer_u32.len()); + + return ImflowImageBuffer { + width, + height, + rgba_buffer: buffer_u32, + rating: 0, + }; } + let thumbnail = if path.format == ImageFormat::Heif { + load_heif(path, true) + } else { + load_thumbnail_exif(path).unwrap_or(load_thumbnail_full(path)) + }; + + save_thumbnail(&cache_path, thumbnail.clone()); + thumbnail } pub fn load_thumbnail_exif(path: &ImageData) -> Option { @@ -268,8 +341,9 @@ pub fn load_thumbnail_exif(path: &ImageData) -> Option { let decoder = image::ImageReader::new(Cursor::new(thumbnail)) .with_guessed_format() .unwrap(); - let image = decoder.decode().unwrap(); + let mut image = decoder.decode().unwrap(); + image.apply_orientation(path.orientation); let width: usize = image.width() as usize; let height: usize = image.height() as usize; let flat = image.into_rgba8().into_raw(); @@ -281,6 +355,7 @@ pub fn load_thumbnail_exif(path: &ImageData) -> Option { buffer.len() / 4, ) }; + std::mem::forget(buffer); let rating = get_rating(path.into()); @@ -303,7 +378,7 @@ pub fn load_thumbnail_full(path: &ImageData) -> ImflowImageBuffer { .unwrap() .decode() .unwrap() - .resize(640, 480, FilterType::Nearest); + .resize_to_fill(1920, 1920, FilterType::Lanczos3); let width = image.width() as usize; let height = image.height() as usize; let buffer = image_to_rgba_buffer(image); @@ -320,25 +395,47 @@ pub fn load_thumbnail_full(path: &ImageData) -> ImflowImageBuffer { pub fn load_heif(path: &ImageData, resize: bool) -> ImflowImageBuffer { let lib_heif = LibHeif::new(); let ctx = HeifContext::read_from_file(path.path.to_str().unwrap()).unwrap(); - let handle = ctx.primary_image_handle().unwrap(); - let mut image = lib_heif - .decode(&handle, libheif_rs::ColorSpace::Rgb(RgbChroma::Rgba), None) - .unwrap(); + + let image = if resize { + let binding = ctx.top_level_image_handles(); + let handle = binding.get(0).unwrap(); + let thumbnail_count = handle.number_of_thumbnails() as u32; + let mut thumbnail_ids = vec![0u32, thumbnail_count]; + handle.thumbnail_ids(&mut thumbnail_ids); + let handle = &handle.thumbnail(thumbnail_ids[0]).unwrap(); + + lib_heif + .decode(handle, libheif_rs::ColorSpace::Rgb(RgbChroma::Rgba), None) + .unwrap() + } else { + let binding = ctx.top_level_image_handles(); + let handle = binding.get(0).unwrap(); + + lib_heif + .decode(handle, libheif_rs::ColorSpace::Rgb(RgbChroma::Rgba), None) + .unwrap() + }; assert_eq!( image.color_space(), Some(libheif_rs::ColorSpace::Rgb(RgbChroma::Rgba)), ); - // Scale the image - if resize { - image = image.scale(640, 480, None).unwrap(); - assert_eq!(image.width(), 640); - assert_eq!(image.height(), 480); - } - let width = image.width() as usize; let height = image.height() as usize; + + // Scale the image + // if resize { + // const MAX: usize = 3000; + // let scale = max(width, height) as f32 / MAX as f32; + // width = (width as f32 / scale) as usize; + // height = (height as f32 / scale) as usize; + // // image = image.scale(width as u32, height as u32, None).unwrap(); + // image = image.scale(599, 300, None).unwrap(); + // width = image.width() as usize; + // height = image.height() as usize; + // } + let rating = get_rating(path); // Get "pixels" @@ -346,6 +443,7 @@ pub fn load_heif(path: &ImageData, resize: bool) -> ImflowImageBuffer { let interleaved_plane = planes.interleaved.unwrap(); assert!(!interleaved_plane.data.is_empty()); assert!(interleaved_plane.stride > 0); + assert_eq!(interleaved_plane.storage_bits_per_pixel, 32); let rgba_buffer = interleaved_plane.data; let u32_slice = unsafe { @@ -359,3 +457,35 @@ pub fn load_heif(path: &ImageData, resize: bool) -> ImflowImageBuffer { rating, } } + +pub fn get_file_hash(path: &PathBuf) -> GenericArray { + let mut file = File::open(path).unwrap(); + let mut buf = [0u8; 16 * 1024]; + file.read(&mut buf).unwrap(); + let mut hasher = Sha256::new(); + hasher.update(&buf); + hasher.update(file.metadata().unwrap().len().to_le_bytes()); + hasher.finalize() +} + +// TODO: optimize +pub fn save_thumbnail(path: &PathBuf, image: ImflowImageBuffer) { + let cache_dir = path.parent().unwrap(); + if !cache_dir.exists() { + fs::create_dir(cache_dir).unwrap(); + } + println!("path: {:?}", path); + let mut file = File::create(path).unwrap(); + let buffer = image.rgba_buffer; + let u8_buffer = unsafe { + Vec::from_raw_parts( + buffer.as_ptr() as *mut u8, + buffer.len() * 4, + buffer.len() * 4, + ) + }; + std::mem::forget(buffer); + file.write(&(image.width as u32).to_le_bytes()).unwrap(); + file.write(&(image.height as u32).to_le_bytes()).unwrap(); + file.write(&u8_buffer).unwrap(); +} diff --git a/src/store.rs b/src/store.rs index 681b382..407c001 100644 --- a/src/store.rs +++ b/src/store.rs @@ -171,6 +171,7 @@ impl ImageStore { .get(&self.current_image_path) .unwrap(); } + // panic!(); let buf = load_thumbnail(&self.current_image_path); self.loaded_images_thumbnails