342 lines
9.9 KiB
Rust
342 lines
9.9 KiB
Rust
//! Image processing utilities for NMF color unmixing
|
||
//!
|
||
//! Provides image loading, preprocessing, and conversion utilities
|
||
//! specifically designed for color blind unmixing applications.
|
||
|
||
use crate::{NMFError, Result};
|
||
use image::{ImageReader, Rgb, RgbImage};
|
||
use nalgebra::DMatrix;
|
||
use std::path::Path;
|
||
|
||
/// RGB color representation
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct Color {
|
||
pub r: f32,
|
||
pub g: f32,
|
||
pub b: f32,
|
||
}
|
||
|
||
impl Color {
|
||
pub fn new(r: f32, g: f32, b: f32) -> Self {
|
||
Self { r, g, b }
|
||
}
|
||
|
||
pub fn from_u8(r: u8, g: u8, b: u8) -> Self {
|
||
Self {
|
||
r: r as f32 / 255.0,
|
||
g: g as f32 / 255.0,
|
||
b: b as f32 / 255.0,
|
||
}
|
||
}
|
||
|
||
pub fn to_u8(&self) -> (u8, u8, u8) {
|
||
(
|
||
(self.r * 255.0).clamp(0.0, 255.0) as u8,
|
||
(self.g * 255.0).clamp(0.0, 255.0) as u8,
|
||
(self.b * 255.0).clamp(0.0, 255.0) as u8,
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Image data for NMF processing
|
||
#[derive(Debug, Clone)]
|
||
pub struct ImageData {
|
||
pub width: usize,
|
||
pub height: usize,
|
||
pub pixels: Vec<Color>,
|
||
}
|
||
|
||
impl ImageData {
|
||
/// Load image from file path
|
||
pub fn load_from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
|
||
let img = ImageReader::open(path)
|
||
.map_err(|e| NMFError::configuration_error(format!("Failed to open image: {e}")))?
|
||
.decode()
|
||
.map_err(|e| NMFError::configuration_error(format!("Failed to decode image: {e}")))?
|
||
.to_rgb8();
|
||
|
||
let (width, height) = (img.width() as usize, img.height() as usize);
|
||
let mut pixels = Vec::with_capacity(width * height);
|
||
|
||
for pixel in img.pixels() {
|
||
pixels.push(Color::from_u8(pixel[0], pixel[1], pixel[2]));
|
||
}
|
||
|
||
Ok(Self {
|
||
width,
|
||
height,
|
||
pixels,
|
||
})
|
||
}
|
||
|
||
/// Convert to RGB matrix for NMF processing
|
||
/// Returns: (red_channel, green_channel, blue_channel) as separate matrices
|
||
pub fn to_rgb_matrices(&self) -> (DMatrix<f32>, DMatrix<f32>, DMatrix<f32>) {
|
||
let mut red_data = Vec::with_capacity(self.width * self.height);
|
||
let mut green_data = Vec::with_capacity(self.width * self.height);
|
||
let mut blue_data = Vec::with_capacity(self.width * self.height);
|
||
|
||
for pixel in &self.pixels {
|
||
red_data.push(pixel.r);
|
||
green_data.push(pixel.g);
|
||
blue_data.push(pixel.b);
|
||
}
|
||
|
||
let red_matrix = DMatrix::from_vec(self.height, self.width, red_data);
|
||
let green_matrix = DMatrix::from_vec(self.height, self.width, green_data);
|
||
let blue_matrix = DMatrix::from_vec(self.height, self.width, blue_data);
|
||
|
||
(red_matrix, green_matrix, blue_matrix)
|
||
}
|
||
|
||
/// Convert to single matrix for processing (stacked RGB channels)
|
||
/// Format: [R1, G1, B1, R2, G2, B2, ...] as height×(width*3) matrix
|
||
pub fn to_stacked_matrix(&self) -> DMatrix<f32> {
|
||
let mut data = Vec::with_capacity(self.width * self.height * 3);
|
||
|
||
for y in 0..self.height {
|
||
for x in 0..self.width {
|
||
let idx = y * self.width + x;
|
||
let pixel = &self.pixels[idx];
|
||
data.push(pixel.r);
|
||
data.push(pixel.g);
|
||
data.push(pixel.b);
|
||
}
|
||
}
|
||
|
||
DMatrix::from_vec(self.height, self.width * 3, data)
|
||
}
|
||
|
||
/// Create from RGB matrices (inverse of to_rgb_matrices)
|
||
pub fn from_rgb_matrices(
|
||
red: &DMatrix<f32>,
|
||
green: &DMatrix<f32>,
|
||
blue: &DMatrix<f32>,
|
||
) -> Result<Self> {
|
||
if red.nrows() != green.nrows()
|
||
|| red.nrows() != blue.nrows()
|
||
|| red.ncols() != green.ncols()
|
||
|| red.ncols() != blue.ncols()
|
||
{
|
||
return Err(NMFError::configuration_error(
|
||
"RGB matrices must have same dimensions",
|
||
));
|
||
}
|
||
|
||
let height = red.nrows();
|
||
let width = red.ncols();
|
||
let mut pixels = Vec::with_capacity(width * height);
|
||
|
||
for y in 0..height {
|
||
for x in 0..width {
|
||
pixels.push(Color::new(
|
||
red[(y, x)].clamp(0.0, 1.0),
|
||
green[(y, x)].clamp(0.0, 1.0),
|
||
blue[(y, x)].clamp(0.0, 1.0),
|
||
));
|
||
}
|
||
}
|
||
|
||
Ok(Self {
|
||
width,
|
||
height,
|
||
pixels,
|
||
})
|
||
}
|
||
|
||
/// Convert to displayable RGB image
|
||
pub fn to_rgb_image(&self) -> RgbImage {
|
||
let mut img = RgbImage::new(self.width as u32, self.height as u32);
|
||
|
||
for (idx, pixel) in self.pixels.iter().enumerate() {
|
||
let x = idx % self.width;
|
||
let y = idx / self.width;
|
||
let (r, g, b) = pixel.to_u8();
|
||
img.put_pixel(x as u32, y as u32, Rgb([r, g, b]));
|
||
}
|
||
|
||
img
|
||
}
|
||
|
||
/// Get basic statistics about the image
|
||
pub fn statistics(&self) -> ImageStats {
|
||
let mut r_sum = 0.0f32;
|
||
let mut g_sum = 0.0f32;
|
||
let mut b_sum = 0.0f32;
|
||
let mut r_min = 1.0f32;
|
||
let mut g_min = 1.0f32;
|
||
let mut b_min = 1.0f32;
|
||
let mut r_max = 0.0f32;
|
||
let mut g_max = 0.0f32;
|
||
let mut b_max = 0.0f32;
|
||
|
||
for pixel in &self.pixels {
|
||
r_sum += pixel.r;
|
||
g_sum += pixel.g;
|
||
b_sum += pixel.b;
|
||
r_min = r_min.min(pixel.r);
|
||
g_min = g_min.min(pixel.g);
|
||
b_min = b_min.min(pixel.b);
|
||
r_max = r_max.max(pixel.r);
|
||
g_max = g_max.max(pixel.g);
|
||
b_max = b_max.max(pixel.b);
|
||
}
|
||
|
||
let count = self.pixels.len() as f32;
|
||
ImageStats {
|
||
width: self.width,
|
||
height: self.height,
|
||
pixel_count: self.pixels.len(),
|
||
red_stats: ChannelStats {
|
||
mean: r_sum / count,
|
||
min: r_min,
|
||
max: r_max,
|
||
},
|
||
green_stats: ChannelStats {
|
||
mean: g_sum / count,
|
||
min: g_min,
|
||
max: g_max,
|
||
},
|
||
blue_stats: ChannelStats {
|
||
mean: b_sum / count,
|
||
min: b_min,
|
||
max: b_max,
|
||
},
|
||
}
|
||
}
|
||
|
||
/// Convert image to PNG bytes for display
|
||
pub fn to_png_bytes(&self) -> Result<Vec<u8>> {
|
||
let rgb_image = self.to_rgb_image();
|
||
let mut bytes = Vec::new();
|
||
|
||
// Use a cursor to write PNG data to memory
|
||
use std::io::Cursor;
|
||
rgb_image
|
||
.write_to(&mut Cursor::new(&mut bytes), image::ImageFormat::Png)
|
||
.map_err(|e| std::io::Error::other(format!("Failed to encode PNG: {e}")))?;
|
||
|
||
Ok(bytes)
|
||
}
|
||
}
|
||
|
||
/// Statistics for a single color channel
|
||
#[derive(Debug, Clone)]
|
||
pub struct ChannelStats {
|
||
pub mean: f32,
|
||
pub min: f32,
|
||
pub max: f32,
|
||
}
|
||
|
||
/// Complete image statistics
|
||
#[derive(Debug, Clone)]
|
||
pub struct ImageStats {
|
||
pub width: usize,
|
||
pub height: usize,
|
||
pub pixel_count: usize,
|
||
pub red_stats: ChannelStats,
|
||
pub green_stats: ChannelStats,
|
||
pub blue_stats: ChannelStats,
|
||
}
|
||
|
||
impl ImageStats {
|
||
pub fn format_summary(&self) -> String {
|
||
format!(
|
||
"Image: {}×{} ({} pixels)\nChannels: R[{:.3}-{:.3}], G[{:.3}-{:.3}], B[{:.3}-{:.3}]",
|
||
self.width,
|
||
self.height,
|
||
self.pixel_count,
|
||
self.red_stats.min,
|
||
self.red_stats.max,
|
||
self.green_stats.min,
|
||
self.green_stats.max,
|
||
self.blue_stats.min,
|
||
self.blue_stats.max
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Preprocessing utilities for biomedical images
|
||
pub struct ImageProcessor;
|
||
|
||
impl ImageProcessor {
|
||
/// Normalize image to [0, 1] range
|
||
pub fn normalize(image: &mut ImageData) {
|
||
let stats = image.statistics();
|
||
let r_range = stats.red_stats.max - stats.red_stats.min;
|
||
let g_range = stats.green_stats.max - stats.green_stats.min;
|
||
let b_range = stats.blue_stats.max - stats.blue_stats.min;
|
||
|
||
for pixel in &mut image.pixels {
|
||
if r_range > 0.0 {
|
||
pixel.r = (pixel.r - stats.red_stats.min) / r_range;
|
||
}
|
||
if g_range > 0.0 {
|
||
pixel.g = (pixel.g - stats.green_stats.min) / g_range;
|
||
}
|
||
if b_range > 0.0 {
|
||
pixel.b = (pixel.b - stats.blue_stats.min) / b_range;
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Remove background by subtracting minimum value per channel
|
||
pub fn remove_background(image: &mut ImageData) {
|
||
let stats = image.statistics();
|
||
|
||
for pixel in &mut image.pixels {
|
||
pixel.r = (pixel.r - stats.red_stats.min).max(0.0);
|
||
pixel.g = (pixel.g - stats.green_stats.min).max(0.0);
|
||
pixel.b = (pixel.b - stats.blue_stats.min).max(0.0);
|
||
}
|
||
}
|
||
|
||
/// Apply gamma correction for better visualization
|
||
pub fn gamma_correction(image: &mut ImageData, gamma: f32) {
|
||
for pixel in &mut image.pixels {
|
||
pixel.r = pixel.r.powf(1.0 / gamma);
|
||
pixel.g = pixel.g.powf(1.0 / gamma);
|
||
pixel.b = pixel.b.powf(1.0 / gamma);
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_color_conversion() {
|
||
let color = Color::from_u8(128, 64, 192);
|
||
let (r, g, b) = color.to_u8();
|
||
assert_eq!((r, g, b), (128, 64, 192));
|
||
}
|
||
|
||
#[test]
|
||
fn test_matrix_conversion() {
|
||
let pixels = vec![
|
||
Color::new(1.0, 0.5, 0.0),
|
||
Color::new(0.0, 1.0, 0.5),
|
||
Color::new(0.5, 0.0, 1.0),
|
||
Color::new(0.25, 0.75, 0.25),
|
||
];
|
||
|
||
let image = ImageData {
|
||
width: 2,
|
||
height: 2,
|
||
pixels,
|
||
};
|
||
|
||
let (r, g, b) = image.to_rgb_matrices();
|
||
assert_eq!(r.shape(), (2, 2));
|
||
assert_eq!(g.shape(), (2, 2));
|
||
assert_eq!(b.shape(), (2, 2));
|
||
|
||
// Test round trip
|
||
let reconstructed = ImageData::from_rgb_matrices(&r, &g, &b).unwrap();
|
||
assert_eq!(reconstructed.width, 2);
|
||
assert_eq!(reconstructed.height, 2);
|
||
assert_eq!(reconstructed.pixels.len(), 4);
|
||
}
|
||
}
|