Starting to test against numpy

proc_blocks/image-normalization/src/distribution.rs

@@ -11,7 +11,12 @@ pub struct Distribution {
 }
 
 impl Distribution {
-    pub const fn new(mean: f32, standard_deviation: f32) -> Self {
+    pub fn new(mean: f32, standard_deviation: f32) -> Self {
+        assert_ne!(
+            standard_deviation, 0.0,
+            "The standard deviation must be non-zero"
+        );
+
         Distribution {
             mean,
             standard_deviation,

proc_blocks/image-normalization/src/lib.rs

@@ -8,13 +8,34 @@ use core::{convert::TryInto, fmt::Display};
 use runic_types::{HasOutputs, Tensor, Transform, TensorViewMut};
 
 #[derive(Debug, Default, Clone, PartialEq)]
+#[non_exhaustive]
 pub struct ImageNormalization {
-    red: Distribution,
-    green: Distribution,
-    blue: Distribution,
+    pub red: Distribution,
+    pub green: Distribution,
+    pub blue: Distribution,
 }
 
 impl ImageNormalization {
+    /// A shortcut for initializing the red, green, and blue distributions in
+    /// one call.
+    pub fn with_rgb<D>(self, distribution: D) -> Self
+    where
+        D: TryInto<Distribution>,
+        D::Error: Display,
+    {
+        let d = match distribution.try_into() {
+            Ok(d) => d,
+            Err(e) => panic!("Invalid distribution: {}", e),
+        };
+
+        ImageNormalization {
+            red: d,
+            green: d,
+            blue: d,
+            ..self
+        }
+    }
+
     pub fn with_red<D>(self, distribution: D) -> Self
     where
         D: TryInto<Distribution>,
@@ -86,9 +107,9 @@ fn transform(
 ) {
     let [_, rows, columns] = view.dimensions();
 
-    for y in 0..rows {
-        for x in 0..columns {
-            let ix = [channel, x, y];
+    for row in 0..rows {
+        for column in 0..columns {
+            let ix = [channel, row, column];
             let current_value = view[ix];
             view[ix] = d.z_score(current_value);
         }
@@ -107,3 +128,21 @@ impl HasOutputs for ImageNormalization {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn normalizing_with_default_distribution_is_noop() {
+        let red = [[1.0], [4.0], [7.0], [10.0]];
+        let green = [[2.0], [5.0], [8.0], [11.0]];
+        let blue = [[3.0], [6.0], [9.0], [12.0]];
+        let image: Tensor<f32> = Tensor::from([red, green, blue]);
+        let mut norm = ImageNormalization::default();
+
+        let got = norm.transform(image.clone());
+
+        assert_eq!(got, image);
+    }
+}

python/src/proc_blocks/image_normalization.rs

@@ -1,6 +1,11 @@
+use std::convert::TryFrom;
+
 use pyo3::{
-    PyAny, PyObject, PyResult, Python, ToPyObject,
-    prelude::{pyclass, pymethods},
+    IntoPy, PyAny, PyErr, PyObject, PyObjectProtocol, PyRef, PyResult, Python,
+    ToPyObject,
+    basic::CompareOp,
+    exceptions::{PyAssertionError, PyTypeError},
+    prelude::{FromPyObject, pyclass, pymethods, pyproto},
 };
 use runic_types::{Tensor, Transform};
 
@@ -17,14 +22,154 @@ pub struct ImageNormalization {
 #[pymethods]
 impl ImageNormalization {
     #[new]
-    pub fn new() -> ImageNormalization { ImageNormalization::default() }
+    #[args("*", red = "None", green = "None", blue = "None")]
+    pub fn new(
+        red: Option<&PyAny>,
+        green: Option<&PyAny>,
+        blue: Option<&PyAny>,
+    ) -> PyResult<ImageNormalization> {
+        let mut inner = image_normalization::ImageNormalization::default();
+
+        update_distribution(&mut inner.red, red)?;
+        update_distribution(&mut inner.blue, blue)?;
+        update_distribution(&mut inner.green, green)?;
+
+        Ok(ImageNormalization { inner })
+    }
 
     #[call]
     pub fn call(&mut self, py: Python, iter: &PyAny) -> PyResult<PyObject> {
         let input: Tensor<f32> = utils::to_tensor(iter)?;
 
-        let spectrum = py.allow_threads(move || self.inner.transform(input));
+        let normalized = py.allow_threads(move || self.inner.transform(input));
+
+        utils::to_numpy(py, &normalized).map(|obj| obj.to_object(py))
+    }
+
+    #[getter]
+    pub fn red(&self) -> Distribution { self.inner.red.into() }
+
+    #[setter]
+    pub fn set_red(&mut self, d: Distribution) { self.inner.red = d.into(); }
+
+    #[getter]
+    pub fn blue(&self) -> Distribution { self.inner.blue.into() }
+
+    #[setter]
+    pub fn set_blue(&mut self, d: Distribution) { self.inner.blue = d.into(); }
+
+    #[getter]
+    pub fn green(&self) -> Distribution { self.inner.green.into() }
+
+    #[setter]
+    pub fn set_green(&mut self, d: Distribution) {
+        self.inner.green = d.into();
+    }
+}
+
+#[pyproto]
+impl PyObjectProtocol for ImageNormalization {
+    fn __repr__(&self) -> String { format!("{:?}", self) }
+}
+
+fn update_distribution(
+    dest: &mut image_normalization::Distribution,
+    value: Option<&PyAny>,
+) -> PyResult<()> {
+    if let Some(value) = value {
+        let distribution = Distribution::try_from(value)?;
+        *dest = distribution.into();
+    }
+
+    Ok(())
+}
+
+/// A normal distribution.
+#[pyclass]
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct Distribution {
+    #[pyo3(get, set)]
+    mean: f32,
+    #[pyo3(get, set)]
+    std_dev: f32,
+}
+
+#[pymethods]
+impl Distribution {
+    #[new]
+    pub fn new(mean: f32, std_dev: f32) -> PyResult<Self> {
+        if std_dev == 0.0 {
+            return Err(PyAssertionError::new_err(
+                "The standard deviation must be non-zero",
+            ));
+        }
+
+        Ok(Distribution { mean, std_dev })
+    }
+}
+
+#[pyproto]
+impl PyObjectProtocol for Distribution {
+    fn __repr__(&self) -> String { format!("{:?}", self) }
+
+    fn __richcmp__(
+        &'p self,
+        other: PyRef<'p, Distribution>,
+        op: CompareOp,
+    ) -> PyResult<PyObject> {
+        let py = other.py();
+
+        match op {
+            CompareOp::Eq => Ok((*self == *other).into_py(py)),
+            CompareOp::Ne => Ok((*self != *other).into_py(py)),
+            _ => Ok(py.NotImplemented()),
+        }
+    }
+}
+
+impl<'py> TryFrom<&'py PyAny> for Distribution {
+    type Error = PyErr;
+
+    fn try_from(ob: &'py PyAny) -> Result<Self, Self::Error> {
+        if let Ok(d) = ob.extract() {
+            return Ok(d);
+        }
+
+        if let Ok((mean, std_dev)) = ob.extract() {
+            return Ok(Distribution { mean, std_dev });
+        }
+
+        if let Ok([mean, std_dev]) = ob.extract::<[f32; 2]>() {
+            return Ok(Distribution { mean, std_dev });
+        }
+
+        #[derive(FromPyObject)]
+        struct DictLike {
+            #[pyo3(item)]
+            mean: f32,
+            #[pyo3(item)]
+            std_dev: f32,
+        }
+
+        if let Ok(DictLike { mean, std_dev }) = ob.extract() {
+            return Ok(Distribution { mean, std_dev });
+        }
+
+        Err(PyTypeError::new_err("Expected a 2-element tuple or a dict"))
+    }
+}
+
+impl From<Distribution> for image_normalization::Distribution {
+    fn from(d: Distribution) -> Self {
+        image_normalization::Distribution::new(d.mean, d.std_dev)
+    }
+}
 
-        utils::to_numpy(py, &spectrum).map(|obj| obj.to_object(py))
+impl From<image_normalization::Distribution> for Distribution {
+    fn from(d: image_normalization::Distribution) -> Self {
+        Distribution {
+            mean: d.mean,
+            std_dev: d.standard_deviation,
+        }
     }
 }

python/src/proc_blocks/mod.rs

@@ -7,8 +7,10 @@ mod noise_filtering;
 mod normalize;
 
 use self::{
-    fft::Fft, image_normalization::ImageNormalization,
-    noise_filtering::NoiseFiltering, normalize::Normalize,
+    fft::Fft,
+    image_normalization::{ImageNormalization, Distribution},
+    noise_filtering::NoiseFiltering,
+    normalize::Normalize,
 };
 use pyo3::{PyResult, Python, types::PyModule};
 
@@ -19,6 +21,7 @@ pub fn register(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<Normalize>()?;
     m.add_class::<NoiseFiltering>()?;
     m.add_class::<ImageNormalization>()?;
+    m.add_class::<Distribution>()?;
 
     Ok(())
 }

python/tests/integration_tests.py

@@ -1,6 +1,7 @@
 import unittest
 import math
-from rune_py import Normalize, Fft
+from rune_py import Normalize, Fft, ImageNormalization, Distribution
+import numpy as np
 
 
 class NormalizeTests(unittest.TestCase):
@@ -46,3 +47,53 @@ def test_calculate_spectrum(self):
         spectrum = fft(samples)
 
         self.assertEqual(1960, len(spectrum))
+
+
+class ImageNormalizationTest(unittest.TestCase):
+    def test_constructor_and_setters(self):
+        norm = ImageNormalization(red=(5.0, 1.5), blue=Distribution(10.0, 2.5))
+
+        self.assertEqual(norm.red, Distribution(5.0, 1.5))
+        self.assertEqual(norm.blue, Distribution(10.0, 2.5))
+        self.assertEqual(norm.green, Distribution(0.0, 1.0))
+
+    def test_normalizing(self):
+        image = np.array(
+            [
+                [[1], [4], [7], [10]],
+                [[2], [5], [8], [11]],
+                [[3], [6], [9], [12]],
+            ],
+            # [
+            #     [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
+            # ],
+            dtype="float32",
+        )
+        mean = image.mean(axis=(1, 2))
+        std = image.std(axis=(1, 2))
+        print("shape", image.shape)
+        print("mean", mean)
+        print("std", std)
+        norm = ImageNormalization(
+            red=(mean[0], std[0]), green=(mean[1], std[1]), blue=(mean[2], std[2])
+        )
+        should_be = normalize_with_numpy(image, mean, std)
+
+        got = norm(image)
+
+        print(Distribution(10, 0.75))
+
+        self.assertEqual(got, should_be)
+
+
+def normalize_with_numpy(image: np.ndarray, mean: np.ndarray, std: np.ndarray):
+    image = np.copy(image)
+
+    image[0, ...] -= mean[0]
+    image[1, ...] -= mean[1]
+    image[2, ...] -= mean[2]
+    image[0, ...] /= std[0]
+    image[1, ...] /= std[1]
+    image[2, ...] /= std[2]
+
+    return image

runic-types/src/tensor.rs

@@ -240,7 +240,13 @@ fn index_of(dimensions: &[usize], indices: &[usize]) -> Option<usize> {
         }
     }
 
-    Some(index)
+    let num_elements: usize = dimensions.iter().product();
+
+    if index < num_elements {
+        Some(index)
+    } else {
+        None
+    }
 }
 
 impl<'t, T, const RANK: usize> Index<[usize; RANK]>