Implement waveform envelope visualization with peak and RMS analysis

willibrandon · willibrandon · commit a2d6b083ae72 · 2025-09-16T20:40:04.000-07:00
diff --git a/examples/README.md b/examples/README.md
@@ -20,6 +20,7 @@ This creates a `samples/` directory with test WAV files including sine waves, ch
 - **onset_detection.rs** - Detect note onsets and transients in audio
 - **cached_analysis.rs** - Using the cache system for faster repeated analysis
 - **batch_processing.rs** - Process multiple files in parallel
+- **envelope_visualization.rs** - Generate waveform visualization with peak and RMS envelopes
 - **generate_samples.rs** - Generate test WAV files for the examples
 
 ## Running Examples
@@ -33,8 +34,11 @@ cargo run --example spectral_analysis
 cargo run --example onset_detection
 cargo run --example cached_analysis
 cargo run --example batch_processing
+cargo run --example envelope_visualization
 ```
 
+The `envelope_visualization` example creates a PNG image showing waveform with peak and RMS envelopes.
+
 ## MCP Server
 
 The `mcp_client.rs` example shows the JSON format for MCP tool calls. To use the actual MCP server:
diff --git a/examples/envelope_visualization.rs b/examples/envelope_visualization.rs
@@ -0,0 +1,58 @@
+use ferrous_waves::visualization::waveform::render_waveform_with_envelope;
+use std::fs;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // Make sure samples directory exists
+    fs::create_dir_all("samples")?;
+
+    println!("Generating envelope visualization example...");
+
+    // Create an interesting audio signal with varying dynamics
+    let sample_rate = 44100;
+    let duration = 3.0;
+    let num_samples = (sample_rate as f32 * duration) as usize;
+    let mut samples = Vec::with_capacity(num_samples);
+
+    for i in 0..num_samples {
+        let t = i as f32 / sample_rate as f32;
+
+        // Create an ADSR-like envelope
+        let envelope = if t < 0.1 {
+            // Attack: quick rise
+            t * 10.0
+        } else if t < 0.3 {
+            // Decay: fall to sustain level
+            1.0 - (t - 0.1) * 2.0
+        } else if t < 2.0 {
+            // Sustain: hold at 60% with some variation
+            0.6 + (t * 5.0).sin() * 0.1
+        } else {
+            // Release: fade out
+            (3.0 - t).max(0.0) * 0.6
+        };
+
+        // Mix multiple frequencies for richer sound
+        let carrier = (2.0 * std::f32::consts::PI * 440.0 * t).sin();
+        let modulator = (2.0 * std::f32::consts::PI * 110.0 * t).sin() * 0.3;
+        let high = (2.0 * std::f32::consts::PI * 880.0 * t).sin() * 0.2;
+
+        let sample = (carrier + modulator + high) * envelope * 0.5;
+        samples.push(sample);
+    }
+
+    // Generate the envelope visualization
+    let output_path = "samples/envelope_example.png";
+    render_waveform_with_envelope(&samples, output_path.as_ref(), 1200, 400)?;
+
+    println!("✓ Created envelope visualization at: {}", output_path);
+    println!();
+    println!("The visualization shows:");
+    println!("  - Blue filled area: Peak envelope (min/max values)");
+    println!("  - Blue lines: Peak envelope outline");
+    println!("  - Red lines: RMS envelope (average power)");
+    println!("  - Black line: Zero crossing");
+    println!();
+    println!("Open {} to view the result!", output_path);
+
+    Ok(())
+}
diff --git a/src/visualization/waveform.rs b/src/visualization/waveform.rs
@@ -13,7 +13,158 @@ pub fn render_waveform_with_envelope(
     width: u32,
     height: u32,
 ) -> Result<()> {
-    // For now, delegate to regular waveform rendering
-    // Envelope rendering can be added as enhancement
-    render_waveform(samples, output_path, width, height)
+    use plotters::prelude::*;
+
+    // Calculate envelope with a window size
+    let window_size = (samples.len() / width as usize).max(1);
+    let mut peak_envelope_upper = Vec::new();
+    let mut peak_envelope_lower = Vec::new();
+    let mut rms_envelope = Vec::new();
+
+    // Process samples in windows
+    for chunk in samples.chunks(window_size) {
+        let mut max_val = 0.0f32;
+        let mut min_val = 0.0f32;
+        let mut rms_sum = 0.0f32;
+
+        for &sample in chunk {
+            max_val = max_val.max(sample);
+            min_val = min_val.min(sample);
+            rms_sum += sample * sample;
+        }
+
+        peak_envelope_upper.push(max_val);
+        peak_envelope_lower.push(min_val);
+
+        let rms = (rms_sum / chunk.len() as f32).sqrt();
+        rms_envelope.push(rms);
+    }
+
+    // Create the plot
+    let root = BitMapBackend::new(output_path, (width, height)).into_drawing_area();
+
+    root.fill(&WHITE).map_err(|e| {
+        crate::utils::error::FerrousError::Visualization(format!(
+            "Failed to fill background: {}",
+            e
+        ))
+    })?;
+
+    let max_val = samples.iter().fold(0.0f32, |a, &b| a.max(b.abs()));
+    let y_range = if max_val > 0.0 { max_val * 1.1 } else { 1.0 };
+
+    let mut chart = ChartBuilder::on(&root)
+        .margin(10)
+        .build_cartesian_2d(0f32..peak_envelope_upper.len() as f32, -y_range..y_range)
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to build chart: {}",
+                e
+            ))
+        })?;
+
+    // Draw peak envelope as filled area
+    let upper_points: Vec<(f32, f32)> = peak_envelope_upper
+        .iter()
+        .enumerate()
+        .map(|(i, &v)| (i as f32, v))
+        .collect();
+
+    let lower_points: Vec<(f32, f32)> = peak_envelope_lower
+        .iter()
+        .enumerate()
+        .map(|(i, &v)| (i as f32, v))
+        .collect();
+
+    // Draw filled area between envelopes using polygons
+    let mut polygon_points = upper_points.clone();
+    polygon_points.extend(lower_points.iter().rev());
+    polygon_points.push(upper_points[0]); // Close the polygon
+
+    chart
+        .draw_series(std::iter::once(Polygon::new(
+            polygon_points,
+            BLUE.mix(0.2).filled(),
+        )))
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to draw envelope fill: {}",
+                e
+            ))
+        })?;
+
+    // Draw RMS envelope
+    let rms_points_upper: Vec<(f32, f32)> = rms_envelope
+        .iter()
+        .enumerate()
+        .map(|(i, &v)| (i as f32, v))
+        .collect();
+
+    let rms_points_lower: Vec<(f32, f32)> = rms_envelope
+        .iter()
+        .enumerate()
+        .map(|(i, &v)| (i as f32, -v))
+        .collect();
+
+    chart
+        .draw_series(LineSeries::new(rms_points_upper, &RED))
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to draw RMS upper: {}",
+                e
+            ))
+        })?;
+
+    chart
+        .draw_series(LineSeries::new(rms_points_lower, &RED))
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to draw RMS lower: {}",
+                e
+            ))
+        })?;
+
+    // Draw peak envelope outlines
+    chart
+        .draw_series(LineSeries::new(
+            upper_points,
+            ShapeStyle::from(&BLUE).stroke_width(2),
+        ))
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to draw peak upper: {}",
+                e
+            ))
+        })?;
+
+    chart
+        .draw_series(LineSeries::new(
+            lower_points,
+            ShapeStyle::from(&BLUE).stroke_width(2),
+        ))
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to draw peak lower: {}",
+                e
+            ))
+        })?;
+
+    // Draw zero line
+    chart
+        .draw_series(LineSeries::new(
+            vec![(0.0, 0.0), (peak_envelope_upper.len() as f32, 0.0)],
+            ShapeStyle::from(&BLACK).stroke_width(1),
+        ))
+        .map_err(|e| {
+            crate::utils::error::FerrousError::Visualization(format!(
+                "Failed to draw zero line: {}",
+                e
+            ))
+        })?;
+
+    root.present().map_err(|e| {
+        crate::utils::error::FerrousError::Visualization(format!("Failed to present: {}", e))
+    })?;
+
+    Ok(())
 }
diff --git a/tests/envelope_test.rs b/tests/envelope_test.rs
@@ -0,0 +1,77 @@
+use ferrous_waves::visualization::waveform::render_waveform_with_envelope;
+use std::fs;
+use tempfile::TempDir;
+
+#[test]
+fn test_envelope_rendering() {
+    // Create test signal with varying amplitude
+    let sample_rate = 44100;
+    let duration = 2.0;
+    let num_samples = (sample_rate as f32 * duration) as usize;
+    let mut samples = Vec::with_capacity(num_samples);
+
+    for i in 0..num_samples {
+        let t = i as f32 / sample_rate as f32;
+
+        // Create amplitude envelope that changes over time
+        let envelope = if t < 0.5 {
+            t * 2.0 // Ramp up
+        } else if t < 1.5 {
+            1.0 // Sustain
+        } else {
+            (2.0 - t) * 2.0 // Ramp down
+        };
+
+        // Modulate a 440Hz sine wave with the envelope
+        let sample = (2.0 * std::f32::consts::PI * 440.0 * t).sin() * envelope * 0.8;
+        samples.push(sample);
+    }
+
+    // Render to temp file
+    let temp_dir = TempDir::new().unwrap();
+    let output_path = temp_dir.path().join("envelope_test.png");
+
+    let result = render_waveform_with_envelope(&samples, &output_path, 800, 400);
+    assert!(result.is_ok());
+    assert!(output_path.exists());
+
+    // Verify the file is not empty
+    let metadata = fs::metadata(&output_path).unwrap();
+    assert!(metadata.len() > 0);
+}
+
+#[test]
+fn test_envelope_with_silence() {
+    // Create signal with silence periods
+    let mut samples = vec![0.0; 44100]; // 1 second of silence
+
+    // Add a burst in the middle
+    for (i, sample) in samples.iter_mut().enumerate().take(33075).skip(11025) {
+        let t = (i - 11025) as f32 / 44100.0;
+        *sample = (2.0 * std::f32::consts::PI * 1000.0 * t).sin() * 0.5;
+    }
+
+    let temp_dir = TempDir::new().unwrap();
+    let output_path = temp_dir.path().join("envelope_silence_test.png");
+
+    let result = render_waveform_with_envelope(&samples, &output_path, 600, 300);
+    assert!(result.is_ok());
+    assert!(output_path.exists());
+}
+
+#[test]
+fn test_envelope_with_clipping() {
+    // Create signal that clips
+    let mut samples = Vec::new();
+    for i in 0..22050 {
+        let t = i as f32 / 44100.0;
+        let sample = (2.0 * std::f32::consts::PI * 200.0 * t).sin() * 2.0; // Amplitude > 1
+        samples.push(sample.clamp(-1.0, 1.0)); // Clip to valid range
+    }
+
+    let temp_dir = TempDir::new().unwrap();
+    let output_path = temp_dir.path().join("envelope_clipping_test.png");
+
+    let result = render_waveform_with_envelope(&samples, &output_path, 600, 300);
+    assert!(result.is_ok());
+}
