diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index b23a587..2410923 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -61,6 +61,7 @@ dependencies = [ "csv", "image", "imageproc", + "img-parts", "opencv-yolov5", "pathdiff", "rayon", @@ -1425,6 +1426,17 @@ dependencies = [ "rusttype", ] +[[package]] +name = "img-parts" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b19358258d99a5fc34466fed27a5318f92ae636c3e36165cf9b1e87b5b6701f0" +dependencies = [ + "bytes", + "crc32fast", + "miniz_oxide 0.5.3", +] + [[package]] name = "indexmap" version = "1.9.1" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 873b7d5..69f7677 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -28,6 +28,7 @@ image = { version = "0.24.5", features = ["rgb"] } imageproc = "0.23.0" rayon = "1.6.1" chug = "1.1.0" +img-parts = "0.3.0" [features] # by default Tauri runs in production mode diff --git a/src-tauri/src/exports/image.rs b/src-tauri/src/exports/image.rs index cbba09d..99cb6d4 100644 --- a/src-tauri/src/exports/image.rs +++ b/src-tauri/src/exports/image.rs @@ -1,10 +1,10 @@ -use image; +use crate::structures::{CamTrapDetection, CamTrapImageDetections}; +use crate::util::magic_image::MagicImage; +use image::Rgba; use rayon::prelude::{IntoParallelRefIterator, ParallelIterator}; use serde::{Deserialize, Serialize}; use std::path::PathBuf; -use crate::structures::{CamTrapDetection, CamTrapImageDetections}; - #[derive(Serialize, Deserialize)] pub enum IncludeCriteria { Include, @@ -102,7 +102,7 @@ pub fn export_image( return; } - let img_result = image::open(&image.file); + let img_result = MagicImage::open(&image.file); if img_result.is_err() { return; @@ -112,23 +112,20 @@ pub fn export_image( for detection in &image.detections { if should_draw(detection, &draw_criteria) { - let rect = imageproc::rect::Rect::at( + let color = match detection.class_index { + 0 => Rgba([255, 255, 255, 255]), + 1 => Rgba([255, 0, 0, 255]), + 2 => Rgba([0, 0, 255, 255]), + _ => Rgba([0, 0, 0, 255]), + }; + + img.draw_bounding_box( (detection.x * img.width() as f32) as i32, (detection.y * img.height() as f32) as i32, - ) - .of_size( (detection.width * img.width() as f32) as u32, (detection.height * img.height() as f32) as u32, + color, ); - - let color = match detection.class_index { - 0 => image::Rgba([255, 255, 255, 255]), - 1 => image::Rgba([255, 0, 0, 255]), - 2 => image::Rgba([0, 0, 255, 255]), - _ => image::Rgba([0, 0, 0, 255]), - }; - - imageproc::drawing::draw_hollow_rect_mut(&mut img, rect, color); } } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 62060db..ccad6d1 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1,6 +1,7 @@ pub mod megadetector; pub mod structures; pub mod exports; +pub mod util; #[cfg(test)] mod tests { diff --git a/src-tauri/src/util/magic_image.rs b/src-tauri/src/util/magic_image.rs new file mode 100644 index 0000000..cf49264 --- /dev/null +++ b/src-tauri/src/util/magic_image.rs @@ -0,0 +1,110 @@ +//! Magic Image module +//! +//! This module contains the logic to load, draw bounding boxes and save images while persisting the original EXIF data. +//! +//! The module is based on the [image](https://crates.io/crates/image) and [img_parts](https://crates.io/crates/img_parts) crates. +//! +//! # Example +//! +//! ``` +//! use magic_image::MagicImage; +//! +//! let mut image = MagicImage::open("path/to/image.jpg").unwrap(); +//! +//! image.draw_bounding_box(0.1, 0.1, 0.2, 0.2, [255, 0, 0, 255]); +//! image.draw_bounding_box(0.3, 0.3, 0.2, 0.2, [0, 255, 0, 255]); +//! +//! image.save("path/to/output.jpg").unwrap(); +//! ``` + +use img_parts::ImageEXIF; +use std::fs; +use std::io::Cursor; +use std::path::Path; + +pub(crate) struct MagicImage { + image: image::DynamicImage, + exif: Option, + original_format: image::ImageFormat, +} + +impl MagicImage { + /// Open an image from a path + pub fn open(path: impl AsRef) -> Result> { + let image_bytes = fs::read(path.as_ref())?; + + let original_format = image::guess_format(&image_bytes)?; + let image = image::load_from_memory_with_format(&image_bytes, original_format)?; + + let exif = match original_format { + image::ImageFormat::Jpeg => { + img_parts::jpeg::Jpeg::from_bytes(image_bytes.into())?.exif() + } + image::ImageFormat::Png => img_parts::png::Png::from_bytes(image_bytes.into())?.exif(), + _ => None, + }; + + Ok(Self { + image, + exif, + original_format, + }) + } + + /// Image width + pub fn width(&self) -> u32 { + self.image.width() + } + + /// Image height + pub fn height(&self) -> u32 { + self.image.height() + } + + /// Draw a bounding box on the image + pub fn draw_bounding_box( + &mut self, + x: i32, + y: i32, + width: u32, + height: u32, + color: image::Rgba, + ) { + let rect = imageproc::rect::Rect::at(x, y).of_size(width, height); + imageproc::drawing::draw_hollow_rect_mut(&mut self.image, rect, color); + } + + /// Save the image to a path (preserving EXIF data if it is a JPEG or PNG) + /// + /// Preserves the original image format no matter what the path extension is. + pub fn save(&self, path: impl AsRef) -> Result<(), Box> { + // First, we need to check if the image is a JPEG or PNG. If it is, we need to use img_parts to + // preserve the EXIF data. Otherwise, we can use the image crate to save the image. + let mut output_file = fs::File::create(path)?; + + let mut buffer = Vec::new(); + let mut cursored_buffer = Cursor::new(&mut buffer); + + self.image + .write_to(&mut cursored_buffer, self.original_format)?; + + match self.original_format { + image::ImageFormat::Jpeg => { + let mut jpeg = img_parts::jpeg::Jpeg::from_bytes(buffer.into())?; + jpeg.set_exif(self.exif.clone()); + jpeg.encoder().write_to(&mut output_file)?; + } + image::ImageFormat::Png => { + let mut png = img_parts::png::Png::from_bytes(buffer.into())?; + png.set_exif(self.exif.clone()); + png.encoder().write_to(&mut output_file)?; + } + _ => { + self.image + .write_to(&mut output_file, self.original_format)?; + } + } + + Ok(()) + } +} diff --git a/src-tauri/src/util/mod.rs b/src-tauri/src/util/mod.rs new file mode 100644 index 0000000..09ded20 --- /dev/null +++ b/src-tauri/src/util/mod.rs @@ -0,0 +1 @@ +pub(crate) mod magic_image;