410 lines
14 KiB
Rust
410 lines
14 KiB
Rust
//! Demo utilities for image processing and visualization
|
|
//!
|
|
//! This module provides utilities for demonstrating NMF capabilities,
|
|
//! including image processing and interactive demonstrations.
|
|
|
|
use crate::{NMFConfig, NMFDecomposer, Result};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
/// Image processor for NMF demonstrations
|
|
pub struct ImageProcessor {
|
|
device: Device,
|
|
}
|
|
|
|
impl ImageProcessor {
|
|
/// Create new image processor
|
|
pub fn new(device: Device) -> Self {
|
|
Self { device }
|
|
}
|
|
|
|
/// Create mock image data for demonstration
|
|
pub fn create_mock_image(&self, width: usize, height: usize) -> Result<Tensor> {
|
|
// Create a synthetic image with some patterns
|
|
let mut image_data = Vec::with_capacity(height * width);
|
|
|
|
for y in 0..height {
|
|
for x in 0..width {
|
|
// Create a pattern with some geometric shapes
|
|
let center_x = width as f32 / 2.0;
|
|
let center_y = height as f32 / 2.0;
|
|
|
|
let dist_from_center =
|
|
((x as f32 - center_x).powi(2) + (y as f32 - center_y).powi(2)).sqrt();
|
|
let max_dist = (center_x.powi(2) + center_y.powi(2)).sqrt();
|
|
|
|
// Create concentric circles pattern
|
|
let circle_value = (dist_from_center / max_dist * 4.0 * std::f32::consts::PI)
|
|
.sin()
|
|
.abs();
|
|
|
|
// Add some diagonal stripes
|
|
let stripe_value = ((x + y) as f32 / 10.0).sin().abs();
|
|
|
|
// Combine patterns
|
|
let pixel_value = (circle_value * 0.7 + stripe_value * 0.3) * 255.0;
|
|
image_data.push(pixel_value.max(0.0));
|
|
}
|
|
}
|
|
|
|
Ok(Tensor::from_data(
|
|
image_data,
|
|
[height, width],
|
|
&self.device,
|
|
)?)
|
|
}
|
|
|
|
/// Create mock face-like image data
|
|
pub fn create_mock_faces(
|
|
&self,
|
|
num_faces: usize,
|
|
face_height: usize,
|
|
face_width: usize,
|
|
) -> Result<Tensor> {
|
|
let mut faces_data = Vec::with_capacity(num_faces * face_height * face_width);
|
|
|
|
for face_idx in 0..num_faces {
|
|
// Create different face-like patterns
|
|
let base_brightness = 50.0 + (face_idx as f32 * 30.0) % 150.0;
|
|
|
|
for y in 0..face_height {
|
|
for x in 0..face_width {
|
|
let center_x = face_width as f32 / 2.0;
|
|
let center_y = face_height as f32 / 2.0;
|
|
|
|
// Create oval face shape
|
|
let face_x = (x as f32 - center_x) / (face_width as f32 * 0.4);
|
|
let face_y = (y as f32 - center_y) / (face_height as f32 * 0.5);
|
|
let face_dist = (face_x.powi(2) + face_y.powi(2)).sqrt();
|
|
|
|
let mut pixel_value = base_brightness;
|
|
|
|
// Add face features if inside face boundary
|
|
if face_dist < 1.0 {
|
|
// Eyes (darker regions)
|
|
let eye_y = center_y - face_height as f32 * 0.15;
|
|
let left_eye_x = center_x - face_width as f32 * 0.15;
|
|
let right_eye_x = center_x + face_width as f32 * 0.15;
|
|
|
|
let left_eye_dist =
|
|
((x as f32 - left_eye_x).powi(2) + (y as f32 - eye_y).powi(2)).sqrt();
|
|
let right_eye_dist =
|
|
((x as f32 - right_eye_x).powi(2) + (y as f32 - eye_y).powi(2)).sqrt();
|
|
|
|
if left_eye_dist < face_width as f32 * 0.05
|
|
|| right_eye_dist < face_width as f32 * 0.05
|
|
{
|
|
pixel_value *= 0.3; // Dark eyes
|
|
}
|
|
|
|
// Mouth (darker region)
|
|
let mouth_y = center_y + face_height as f32 * 0.2;
|
|
let mouth_dist =
|
|
((x as f32 - center_x).powi(2) + (y as f32 - mouth_y).powi(2)).sqrt();
|
|
if mouth_dist < face_width as f32 * 0.08 {
|
|
pixel_value *= 0.5; // Dark mouth
|
|
}
|
|
|
|
// Add some noise for realism
|
|
let noise = ((x * 7 + y * 11 + face_idx * 13) as f32).sin() * 10.0;
|
|
pixel_value += noise;
|
|
} else {
|
|
pixel_value = 0.0; // Background
|
|
}
|
|
|
|
faces_data.push(pixel_value.max(0.0));
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(Tensor::from_data(
|
|
faces_data,
|
|
[num_faces, face_height * face_width],
|
|
&self.device,
|
|
)?)
|
|
}
|
|
|
|
/// Convert tensor back to image representation for display
|
|
pub fn tensor_to_image_info(&self, tensor: &Tensor) -> Result<ImageInfo> {
|
|
let shape = tensor.shape();
|
|
let data = tensor.to_cpu()?;
|
|
|
|
// Calculate real statistics
|
|
let min_value = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
let max_value = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
let mean_value = data.iter().sum::<f32>() / data.len() as f32;
|
|
|
|
Ok(ImageInfo {
|
|
shape: shape.dims().to_vec(),
|
|
min_value,
|
|
max_value,
|
|
mean_value,
|
|
data_size: shape.numel(),
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Information about an image tensor for display purposes
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ImageInfo {
|
|
pub shape: Vec<usize>,
|
|
pub min_value: f32,
|
|
pub max_value: f32,
|
|
pub mean_value: f32,
|
|
pub data_size: usize,
|
|
}
|
|
|
|
/// NMF demo runner for demonstrations
|
|
pub struct NMFDemo {
|
|
config: NMFConfig,
|
|
decomposer: NMFDecomposer,
|
|
image_processor: ImageProcessor,
|
|
}
|
|
|
|
impl NMFDemo {
|
|
/// Create new NMF demo
|
|
pub fn new(device: Device) -> Self {
|
|
let config = NMFConfig::new()
|
|
.with_components(10)
|
|
.with_max_iterations(50)
|
|
.with_tolerance(1e-4);
|
|
|
|
let decomposer = NMFDecomposer::new(config.clone());
|
|
let image_processor = ImageProcessor::new(device);
|
|
|
|
Self {
|
|
config,
|
|
decomposer,
|
|
image_processor,
|
|
}
|
|
}
|
|
|
|
/// Run image decomposition demo
|
|
pub fn run_image_demo(&mut self, width: usize, height: usize) -> Result<DemoResult> {
|
|
println!("🖼️ Creating mock image data ({width} x {height})...");
|
|
let image = self.image_processor.create_mock_image(width, height)?;
|
|
let image_info = self.image_processor.tensor_to_image_info(&image)?;
|
|
|
|
println!("📊 Image statistics:");
|
|
println!(" Shape: {:?}", image_info.shape);
|
|
println!(
|
|
" Range: [{:.2}, {:.2}]",
|
|
image_info.min_value, image_info.max_value
|
|
);
|
|
println!(" Mean: {:.2}", image_info.mean_value);
|
|
|
|
println!(
|
|
"\n⚡ Running NMF decomposition with {} components...",
|
|
self.config.components()
|
|
);
|
|
let result = self.decomposer.fit_transform_detailed(&image)?;
|
|
|
|
let w_info = self.image_processor.tensor_to_image_info(&result.w)?;
|
|
let h_info = self.image_processor.tensor_to_image_info(&result.h)?;
|
|
|
|
println!("\n✅ NMF decomposition completed!");
|
|
println!(" Iterations: {}", result.iterations);
|
|
println!(" Converged: {}", result.converged);
|
|
println!(" Error: {:.6}", result.reconstruction_error);
|
|
println!(" Time: {:.3}s", result.computation_time);
|
|
|
|
Ok(DemoResult {
|
|
original_image: image_info,
|
|
w_matrix: w_info,
|
|
h_matrix: h_info,
|
|
reconstruction_error: result.reconstruction_error,
|
|
iterations: result.iterations,
|
|
converged: result.converged,
|
|
computation_time: result.computation_time,
|
|
components: self.config.components(),
|
|
})
|
|
}
|
|
|
|
/// Run faces decomposition demo
|
|
pub fn run_faces_demo(
|
|
&mut self,
|
|
num_faces: usize,
|
|
face_height: usize,
|
|
face_width: usize,
|
|
) -> Result<DemoResult> {
|
|
println!(
|
|
"👥 Creating mock face dataset ({num_faces} faces, {face_height} x {face_width})..."
|
|
);
|
|
let faces = self
|
|
.image_processor
|
|
.create_mock_faces(num_faces, face_height, face_width)?;
|
|
let faces_info = self.image_processor.tensor_to_image_info(&faces)?;
|
|
|
|
println!("📊 Faces dataset statistics:");
|
|
println!(" Shape: {:?}", faces_info.shape);
|
|
println!(
|
|
" Range: [{:.2}, {:.2}]",
|
|
faces_info.min_value, faces_info.max_value
|
|
);
|
|
println!(" Mean: {:.2}", faces_info.mean_value);
|
|
|
|
// Update config for faces (usually need fewer components)
|
|
self.config = self.config.clone().with_components(5);
|
|
self.decomposer = NMFDecomposer::new(self.config.clone());
|
|
|
|
println!("\n⚡ Running NMF decomposition to find face basis...");
|
|
let result = self.decomposer.fit_transform_detailed(&faces)?;
|
|
|
|
let w_info = self.image_processor.tensor_to_image_info(&result.w)?;
|
|
let h_info = self.image_processor.tensor_to_image_info(&result.h)?;
|
|
|
|
println!("\n✅ Face decomposition completed!");
|
|
println!(" Found {} basis faces", self.config.components());
|
|
println!(" Iterations: {}", result.iterations);
|
|
println!(" Reconstruction error: {:.6}", result.reconstruction_error);
|
|
println!(" Time: {:.3}s", result.computation_time);
|
|
|
|
Ok(DemoResult {
|
|
original_image: faces_info,
|
|
w_matrix: w_info,
|
|
h_matrix: h_info,
|
|
reconstruction_error: result.reconstruction_error,
|
|
iterations: result.iterations,
|
|
converged: result.converged,
|
|
computation_time: result.computation_time,
|
|
components: self.config.components(),
|
|
})
|
|
}
|
|
|
|
/// Get performance metrics for display
|
|
pub fn get_performance_metrics(&self) -> HashMap<String, String> {
|
|
let mut metrics = HashMap::new();
|
|
|
|
metrics.insert(
|
|
"gpu_acceleration".to_string(),
|
|
if self.config.use_gpu() {
|
|
"Enabled"
|
|
} else {
|
|
"Disabled"
|
|
}
|
|
.to_string(),
|
|
);
|
|
metrics.insert(
|
|
"components".to_string(),
|
|
self.config.components().to_string(),
|
|
);
|
|
metrics.insert(
|
|
"max_iterations".to_string(),
|
|
self.config.max_iterations().to_string(),
|
|
);
|
|
metrics.insert(
|
|
"tolerance".to_string(),
|
|
format!("{:.0e}", self.config.tolerance()),
|
|
);
|
|
metrics.insert(
|
|
"device".to_string(),
|
|
format!("{:?}", self.image_processor.device),
|
|
);
|
|
|
|
// Add system-specific performance indicators
|
|
match &self.image_processor.device {
|
|
Device::Cuda(_) => {
|
|
metrics.insert("cuda_version".to_string(), "13.0+".to_string());
|
|
metrics.insert("memory_type".to_string(), "GPU VRAM".to_string());
|
|
metrics.insert("acceleration_type".to_string(), "CUDA Kernels".to_string());
|
|
}
|
|
Device::Rocm(_) => {
|
|
metrics.insert("acceleration_type".to_string(), "ROCm Kernels".to_string());
|
|
metrics.insert("memory_type".to_string(), "GPU VRAM".to_string());
|
|
}
|
|
Device::Metal(_) => {
|
|
metrics.insert("acceleration_type".to_string(), "Metal Kernels".to_string());
|
|
metrics.insert("memory_type".to_string(), "GPU Memory".to_string());
|
|
}
|
|
Device::Cpu => {
|
|
metrics.insert("acceleration_type".to_string(), "CPU".to_string());
|
|
metrics.insert("memory_type".to_string(), "System RAM".to_string());
|
|
}
|
|
}
|
|
|
|
metrics
|
|
}
|
|
}
|
|
|
|
/// Result of a demo run
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DemoResult {
|
|
pub original_image: ImageInfo,
|
|
pub w_matrix: ImageInfo,
|
|
pub h_matrix: ImageInfo,
|
|
pub reconstruction_error: f32,
|
|
pub iterations: usize,
|
|
pub converged: bool,
|
|
pub computation_time: f32,
|
|
pub components: usize,
|
|
}
|
|
|
|
impl DemoResult {
|
|
/// Format result for display
|
|
pub fn format_summary(&self) -> String {
|
|
format!(
|
|
"NMF Result: {} components, {:.6} error, {} iterations ({:.3}s)",
|
|
self.components, self.reconstruction_error, self.iterations, self.computation_time
|
|
)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_tensor::Device;
|
|
|
|
#[test]
|
|
fn test_image_processor() -> Result<()> {
|
|
let device = Device::Cuda(0);
|
|
let processor = ImageProcessor::new(device);
|
|
|
|
let image = processor.create_mock_image(32, 24)?;
|
|
assert_eq!(image.shape(), &[24, 32]);
|
|
|
|
let info = processor.tensor_to_image_info(&image)?;
|
|
assert_eq!(info.shape, vec![24, 32]);
|
|
assert!(info.min_value >= 0.0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_mock_faces() -> Result<()> {
|
|
let device = Device::Cuda(0);
|
|
let processor = ImageProcessor::new(device);
|
|
|
|
let faces = processor.create_mock_faces(5, 20, 20)?;
|
|
assert_eq!(faces.shape(), &[5, 400]); // 5 faces, 20x20 = 400 pixels each
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_nmf_demo() -> Result<()> {
|
|
let device = Device::Cuda(0);
|
|
let mut demo = NMFDemo::new(device);
|
|
|
|
// Small demo to avoid long test times
|
|
let result = demo.run_image_demo(8, 6)?;
|
|
|
|
assert_eq!(result.original_image.shape, vec![6, 8]);
|
|
assert!(result.reconstruction_error >= 0.0);
|
|
assert!(result.computation_time > 0.0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_performance_metrics() {
|
|
let device = Device::Cuda(0);
|
|
let demo = NMFDemo::new(device);
|
|
let metrics = demo.get_performance_metrics();
|
|
|
|
assert!(metrics.contains_key("gpu_acceleration"));
|
|
assert!(metrics.contains_key("components"));
|
|
assert!(metrics.contains_key("device"));
|
|
}
|
|
}
|