1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use crate::core::*;
use ndarray::prelude::*;
use ndarray_stats::{HistogramExt, histogram::Grid};
use num_traits::{Num, NumAssignOps};
use num_traits::cast::{FromPrimitive, ToPrimitive};

/// Extension trait to implement histogram equalisation on other types
pub trait HistogramEqExt<A> where A: Ord {
    /// Equalises an image histogram returning a new image.
    /// Grids should be for a 1xN image as the image is flattened during processing
    fn equalise_hist(&self, grid: Grid<A>) -> Self;

    /// Equalises an image histogram inplace
    /// Grids should be for a 1xN image as the image is flattened during processing
    fn equalise_hist_inplace(&mut self, grid: Grid<A>);
}


impl<T> HistogramEqExt<T> for Array3<T>
where
    T: Copy + Clone + Ord + Num + NumAssignOps + ToPrimitive + FromPrimitive + PixelBound
{
    fn equalise_hist(&self, grid: Grid<T>) -> Self {
        let mut result = self.clone();
        result.equalise_hist_inplace(grid);
        result
    }

    fn equalise_hist_inplace(&mut self, grid: Grid<T>) {
        for mut c in self.axis_iter_mut(Axis(2)) {
            // get the histogram
            let flat = Array::from_iter(c.iter())
                .mapv(|x| *x)
                .insert_axis(Axis(1));
            let hist = flat.histogram(grid.clone());
            // get cdf 
            let mut running_total = 0;
            let mut min = 0.0;
            let cdf = hist.counts().mapv(|x| {
                running_total += x;
                if min == 0.0 && running_total > 0 {
                    min = running_total as f32;
                }
                running_total as f32
            });
             
            // Rescale cdf writing back new values
            let scale = (T::max_pixel() - T::min_pixel()).to_f32().unwrap_or_default();
            let denominator = flat.len() as f32 - min;
            c.mapv_inplace(|x| {
                let index = match grid.index_of(&arr1(&[x])) {
                    Some(i) => {
                        if i.is_empty() {
                            0
                        } else {
                            i[0]
                        }
                    },
                    None => 0,
                };
                let mut f_res = ((cdf[index] - min)/denominator)*scale;
                if T::discrete_levels().is_some() {
                    f_res = f_res.round();
                }
                T::from_f32(f_res).unwrap_or_else(T::zero) + T::min_pixel()
            });
        }
    }
}


impl<T, C> HistogramEqExt<T> for Image<T, C> 
where 
    Image<T, C>: Clone,
    T: Copy + Clone + Ord + Num + NumAssignOps + ToPrimitive + FromPrimitive + PixelBound,
    C: ColourModel,
{
    fn equalise_hist(&self, grid: Grid<T>) -> Self {
        let mut result = self.clone();
        result.equalise_hist_inplace(grid);
        result
    }

    fn equalise_hist_inplace(&mut self, grid: Grid<T>) {
        self.data.equalise_hist_inplace(grid);
    }
}


#[cfg(test)]
mod tests {
    use crate::core::Gray;
    use ndarray_stats::histogram::{Bins, Edges};
    use super::*;

    #[test]
    fn hist_eq_test() {
        // test data from wikipedia
        let input_pixels = vec![52, 55, 61, 59, 70, 61, 76, 61, 
                                62, 59, 55, 104, 94, 85, 59, 71,
                                63, 65, 66, 113, 144, 104, 63, 72,
                                64, 70, 70, 126, 154, 109, 71, 69,
                                67, 73, 68, 106, 122, 88, 68, 68,
                                68, 79, 60, 79, 77, 66, 58, 75,
                                69, 85, 64, 58, 55, 61, 65, 83,
                                70, 87, 69, 68, 65, 73, 78, 90,];

        let output_pixels = vec![0, 12, 53, 32, 146, 53, 174, 53,
                                 57, 32, 12, 227, 219, 202, 32, 154,
                                 65, 85, 93, 239, 251, 227, 65, 158,
                                 73, 146, 146, 247, 255, 235, 154, 130,
                                 97, 166, 117, 231, 243, 210, 117, 117,
                                 117, 190, 36, 190, 178, 93, 20, 170,
                                 130, 202, 73, 20, 12, 53, 85, 194,
                                 146, 206, 130, 117, 85, 166, 182, 215];

        let input = Image::<u8, Gray>::from_shape_data(8, 8, input_pixels);
        
        let expected = Image::<u8, Gray>::from_shape_data(8, 8, output_pixels);
    
        let edges_vec: Vec<u8> = (0..255).collect();
        let grid = Grid::from(vec![Bins::new(Edges::from(edges_vec))]);

        let equalised = input.equalise_hist(grid);

        assert_eq!(expected, equalised);
    }
}