Files
rustytorch/demos/rtx-slidescope/src/lib.rs
T
2026-03-04 00:08:42 +00:00

273 lines
8.7 KiB
Rust

//! rtx-slidescope - GPU-accelerated pathology slide processing
//!
//! This crate provides NMF-based stain unmixing and tile pyramid generation
//! for digital pathology slide analysis.
//!
//! # Features
//!
//! - **NMF Stain Unmixing**: Separate stain components (H&E, IHC-DAB) using
//! non-negative matrix factorization
//! - **Tile Pyramid Generation**: Create deep-zoom compatible tile pyramids
//! for fast viewing of large images
//! - **Optical Density Conversion**: Convert between RGB and optical density
//! color spaces for accurate stain analysis
//!
//! # Example
//!
//! ```no_run
//! use rtx_slidescope::{
//! SlidescopeProcessor, ProcessingConfig,
//! image_io::load_image,
//! };
//! use std::path::Path;
//!
//! # fn main() -> anyhow::Result<()> {
//! // Load an image
//! let image = load_image(Path::new("slide.tiff"))?;
//!
//! // Create processor with default config
//! let config = ProcessingConfig::default();
//! let processor = SlidescopeProcessor::new(config)?;
//!
//! // Generate tile pyramid
//! let pyramid = processor.generate_pyramid(&image)?;
//!
//! // Get DZI metadata for OpenSeadragon
//! let dzi = pyramid.get_dzi_metadata();
//! println!("Image: {}x{}, {} levels", dzi.width, dzi.height, dzi.max_level);
//! # Ok(())
//! # }
//! ```
pub mod config;
pub mod gpu;
pub mod gpu_nmf;
pub mod image_io;
pub mod nmf;
pub mod optical_density;
pub mod pyramid;
pub mod stain_vectors;
// Re-exports
pub use config::{
EstimationMethod, NmfProcessingConfig, ProcessingConfig, PyramidProcessingConfig,
StainEstimationConfig,
};
pub use gpu::{BackendType, DeviceInfo, GpuBackend, gpu_available, select_backend};
pub use gpu_nmf::GpuNmfProcessor;
pub use nmf::{NmfOutput, NmfProcessor};
pub use optical_density::{od_to_rgb, rgb_to_od};
pub use pyramid::{MultiLayerPyramid, TilePyramid};
pub use stain_vectors::{
StainEstimationResult, StainEstimator, he_stain_matrix, ihc_dab_stain_matrix,
};
use anyhow::Result;
use image::DynamicImage;
use ndarray::{Array2, Array3};
use slidescope_shared::TileLayer;
use std::time::Instant;
use tracing::info;
/// Main processor for SlideScope operations
pub struct SlidescopeProcessor {
config: ProcessingConfig,
}
impl SlidescopeProcessor {
/// Create a new processor with the given configuration
pub fn new(config: ProcessingConfig) -> Result<Self> {
Ok(Self { config })
}
/// Create a processor with default configuration
pub fn with_defaults() -> Result<Self> {
Self::new(ProcessingConfig::default())
}
/// Generate a tile pyramid from an image
pub fn generate_pyramid(&self, image: &DynamicImage) -> Result<TilePyramid> {
TilePyramid::generate(image, &self.config.pyramid)
}
/// Process an image with NMF stain unmixing
///
/// Returns a multi-layer pyramid containing the original image and
/// individual stain concentration maps.
pub fn process_nmf(&self, image: &DynamicImage) -> Result<NmfProcessingResult> {
let start_time = Instant::now();
// Convert to array
let rgb_array = image_io::image_to_array(image);
let (height, width, _) = rgb_array.dim();
info!("Processing {}x{} image with NMF", width, height);
// Convert to optical density
let od_array = rgb_to_od_array(&rgb_array);
// Estimate stain vectors
let estimator = StainEstimator::new(
self.config.stain_estimation.method,
self.config.stain_estimation.background_threshold,
self.config.stain_estimation.sample_size,
);
let stain_result = estimator.estimate(&od_array, self.config.nmf.n_components)?;
// Flatten OD to pixel matrix
let od_pixels = optical_density::flatten_to_pixels(&od_array);
// Perform NMF unmixing
let concentrations =
nmf::unmix_stains(&od_pixels, &stain_result.stain_matrix, &self.config.nmf)?;
// Reshape concentrations back to images
let mut stain_images = Vec::new();
for i in 0..self.config.nmf.n_components {
let conc_col = concentrations.column(i);
let conc_2d = Array2::from_shape_vec((height, width), conc_col.to_vec())?;
stain_images.push(conc_2d);
}
// Create multi-layer pyramid
let mut multi_pyramid = MultiLayerPyramid::new(width as u32, height as u32);
// Add original image pyramid
let original_pyramid = self.generate_pyramid(image)?;
multi_pyramid.add_layer(TileLayer::Original, original_pyramid);
// Add stain concentration pyramids
let stain_colors = [
[0.4, 0.2, 0.7], // Hematoxylin-ish (purple)
[0.9, 0.5, 0.5], // Eosin-ish (pink)
[0.6, 0.4, 0.2], // DAB-ish (brown)
];
for (i, (conc, color)) in stain_images.iter().zip(stain_colors.iter()).enumerate() {
let colored = image_io::concentration_to_colored(conc, *color)?;
let dynamic = DynamicImage::ImageRgb8(colored);
let pyramid = self.generate_pyramid(&dynamic)?;
let layer = match i {
0 => TileLayer::Stain1,
1 => TileLayer::Stain2,
2 => TileLayer::Stain3,
_ => continue,
};
multi_pyramid.add_layer(layer, pyramid);
}
let processing_time_ms = start_time.elapsed().as_millis() as u64;
info!(
"NMF processing complete in {}ms, {} stain components",
processing_time_ms,
stain_images.len()
);
Ok(NmfProcessingResult {
multi_pyramid,
stain_matrix: stain_result.stain_matrix,
concentrations: stain_images,
processing_time_ms,
})
}
/// Get the current configuration
pub fn config(&self) -> &ProcessingConfig {
&self.config
}
}
/// Result of NMF processing
pub struct NmfProcessingResult {
/// Multi-layer tile pyramid with original and stain layers
pub multi_pyramid: MultiLayerPyramid,
/// Estimated stain matrix
pub stain_matrix: Array2<f32>,
/// Stain concentration maps
pub concentrations: Vec<Array2<f32>>,
/// Processing time in milliseconds
pub processing_time_ms: u64,
}
/// Convert RGB array to OD array (helper)
fn rgb_to_od_array(rgb: &Array3<u8>) -> Array3<f32> {
optical_density::rgb_to_od(rgb)
}
/// Get GPU device name if available
pub fn gpu_device_name() -> Option<String> {
gpu::select_backend()
.ok()
.filter(|b| b.device_info().backend != gpu::BackendType::Cpu)
.map(|b| b.device_info().name)
}
#[cfg(test)]
mod tests {
use super::*;
use image::{Rgb, RgbImage};
fn create_test_image(width: u32, height: u32) -> DynamicImage {
let mut img = RgbImage::new(width, height);
for y in 0..height {
for x in 0..width {
// Create a gradient with some color variation
let r = ((x as f32 / width as f32) * 200.0 + 50.0) as u8;
let g = ((y as f32 / height as f32) * 150.0 + 100.0) as u8;
let b = 180;
img.put_pixel(x, y, Rgb([r, g, b]));
}
}
DynamicImage::ImageRgb8(img)
}
#[test]
fn test_processor_creation() {
let processor = SlidescopeProcessor::with_defaults().unwrap();
assert_eq!(processor.config().nmf.n_components, 2);
}
#[test]
fn test_pyramid_generation() {
let processor = SlidescopeProcessor::with_defaults().unwrap();
let image = create_test_image(200, 150);
let pyramid = processor.generate_pyramid(&image).unwrap();
assert_eq!(pyramid.width, 200);
assert_eq!(pyramid.height, 150);
assert!(pyramid.total_tiles() > 0);
}
#[test]
fn test_nmf_processing() {
// Use a smaller config for faster testing
let mut config = ProcessingConfig::default();
config.nmf = NmfProcessingConfig::fast();
config.pyramid.tile_size = 64;
let processor = SlidescopeProcessor::new(config).unwrap();
let image = create_test_image(100, 100);
let result = processor.process_nmf(&image).unwrap();
// Should have original + 2 stain layers
assert!(result.multi_pyramid.has_layer(TileLayer::Original));
assert!(result.multi_pyramid.has_layer(TileLayer::Stain1));
assert!(result.multi_pyramid.has_layer(TileLayer::Stain2));
// Should have 2 concentration maps
assert_eq!(result.concentrations.len(), 2);
assert_eq!(result.concentrations[0].shape(), &[100, 100]);
}
#[test]
fn test_gpu_availability() {
// This should not panic regardless of GPU presence
let _ = gpu_available();
let _ = gpu_device_name();
}
}