Skip to main content

fits_io/image/
image.rs

1use crate::header::{BayerPattern, Bitpix, Header};
2use crate::image::{ImageData, Normalizer};
3use image::{ImageBuffer, Luma, Primitive, Rgb, RgbImage};
4use std::error::Error;
5
6/// One image, in whatever type its BITPIX card called for.
7#[derive(Debug, Clone)]
8pub enum Image {
9    /// A double precision image, from BITPIX -64.
10    F64(ImageData<f64>),
11    /// A single precision image, from BITPIX -32.
12    F32(ImageData<f32>),
13    /// A signed 32-bit image, from BITPIX 32.
14    I32(ImageData<i32>),
15    /// A signed 16-bit image, from BITPIX 16.
16    I16(ImageData<i16>),
17    /// An 8-bit image, from BITPIX 8.
18    U8(ImageData<u8>),
19}
20
21impl Image {
22    /// Image width in pixels
23    pub fn width(&self) -> u32 {
24        match self {
25            Self::F64(image) => image.width(),
26            Self::F32(image) => image.width(),
27            Self::I32(image) => image.width(),
28            Self::I16(image) => image.width(),
29            Self::U8(image) => image.width(),
30        }
31    }
32
33    /// Image height in pixels
34    pub fn height(&self) -> u32 {
35        match self {
36            Self::F64(image) => image.height(),
37            Self::F32(image) => image.height(),
38            Self::I32(image) => image.height(),
39            Self::I16(image) => image.height(),
40            Self::U8(image) => image.height(),
41        }
42    }
43
44    /// The camera bayer pattern or None if the camera is monochrome
45    pub fn bayer_pattern(&self) -> &Option<BayerPattern> {
46        match self {
47            Self::F64(image) => image.bayer_pattern(),
48            Self::F32(image) => image.bayer_pattern(),
49            Self::I32(image) => image.bayer_pattern(),
50            Self::I16(image) => image.bayer_pattern(),
51            Self::U8(image) => image.bayer_pattern(),
52        }
53    }
54
55    /// Returns a normalised version of the image, where all values are converted into f64 in the range of 0.0 - 1.0
56    pub fn normalized(&self) -> ImageBuffer<Luma<f64>, Vec<f64>> {
57        match self {
58            Self::F64(image) => image.normalized(),
59            Self::F32(image) => image.normalized(),
60            Self::I32(image) => image.normalized(),
61            Self::I16(image) => image.normalized(),
62            Self::U8(image) => image.normalized(),
63        }
64    }
65
66    /// Performs a superpixel demosaic and returns a normalised version. The superpixel algorithm is fast, but essentially cutting the resolution in half.
67    pub fn normalized_superpixel(
68        &self,
69    ) -> Result<ImageBuffer<Rgb<f64>, Vec<f64>>, Box<dyn Error + Send + Sync>> {
70        match self {
71            Self::F64(image) => image.normalized_superpixel(),
72            Self::F32(image) => image.normalized_superpixel(),
73            Self::I32(image) => image.normalized_superpixel(),
74            Self::I16(image) => image.normalized_superpixel(),
75            Self::U8(image) => image.normalized_superpixel(),
76        }
77    }
78
79    /// Converts this image into a RgbImage from image-rs
80    pub fn rgb_image(&self) -> Result<RgbImage, Box<dyn Error + Send + Sync>> {
81        if self.bayer_pattern().is_some() {
82            let normalized = self.normalized_superpixel()?;
83            let mut buffer = RgbImage::new(normalized.width(), normalized.height());
84            for (x, y, pixel) in buffer.enumerate_pixels_mut() {
85                let r_pixel = (u8::MAX as f64 * normalized.get_pixel(x, y)[0]) as u8;
86                let g_pixel = (u8::MAX as f64 * normalized.get_pixel(x, y)[1]) as u8;
87                let b_pixel = (u8::MAX as f64 * normalized.get_pixel(x, y)[2]) as u8;
88
89                pixel[0] = r_pixel;
90                pixel[1] = g_pixel;
91                pixel[2] = b_pixel;
92            }
93
94            Ok(buffer)
95        } else {
96            let normalized = self.normalized();
97            let mut buffer = RgbImage::new(self.width(), self.height());
98            for (x, y, pixel) in buffer.enumerate_pixels_mut() {
99                let gray_pixel = (u8::MAX as f64 * normalized.get_pixel(x, y)[0]) as u8;
100                pixel[0] = gray_pixel;
101                pixel[1] = gray_pixel;
102                pixel[2] = gray_pixel;
103            }
104            Ok(buffer)
105        }
106    }
107
108    pub(crate) fn from_data_and_header(
109        data: Vec<u8>,
110        header: &Header,
111    ) -> Result<Self, Box<dyn Error + Send + Sync>> {
112        let bitpix = header
113            .bitpix()
114            .ok_or("Cannot read an image from a header without a BITPIX card")?;
115        let width = header
116            .naxis_n(0)
117            .ok_or("Cannot read an image from a header without a NAXIS1 card")?;
118        let height = header
119            .naxis_n(1)
120            .ok_or("Cannot read an image from a header without a NAXIS2 card")?;
121        let width = usize::try_from(width)
122            .map_err(|_| format!("NAXIS1 must not be negative, but was {}", width))?;
123        let height = usize::try_from(height)
124            .map_err(|_| format!("NAXIS2 must not be negative, but was {}", height))?;
125
126        let bayer_pattern = header.bayer_pattern();
127
128        let expected = width
129            .checked_mul(height)
130            .and_then(|pixels| pixels.checked_mul(bitpix.byte_size()))
131            .ok_or("Image dimensions overflow the address space")?;
132        if data.len() < expected {
133            return Err(format!(
134                "Image data is too short, expected {} bytes for a {}x{} {:?} image, but got {}",
135                expected,
136                width,
137                height,
138                bitpix,
139                data.len()
140            )
141            .into());
142        }
143
144        match bitpix {
145            Bitpix::F64 => {
146                let image_data = data
147                    .as_chunks::<8>()
148                    .0
149                    .iter()
150                    .map(|i| f64::from_be_bytes(*i))
151                    .collect::<Vec<_>>();
152                Ok(Image::F64(ImageData::<f64>::from_data(
153                    width,
154                    height,
155                    normalizer_for(header, &image_data),
156                    bayer_pattern,
157                    image_data,
158                )?))
159            }
160            Bitpix::F32 => {
161                let image_data = data
162                    .as_chunks::<4>()
163                    .0
164                    .iter()
165                    .map(|i| f32::from_be_bytes(*i))
166                    .collect::<Vec<_>>();
167                Ok(Image::F32(ImageData::<f32>::from_data(
168                    width,
169                    height,
170                    normalizer_for(header, &image_data),
171                    bayer_pattern,
172                    image_data,
173                )?))
174            }
175            Bitpix::U8 => Ok(Image::U8(ImageData::<u8>::from_data(
176                width,
177                height,
178                normalizer_for(header, &data),
179                bayer_pattern,
180                data,
181            )?)),
182            Bitpix::I16 => {
183                let image_data = data
184                    .as_chunks::<2>()
185                    .0
186                    .iter()
187                    .map(|i| i16::from_be_bytes(*i))
188                    .collect::<Vec<_>>();
189                Ok(Image::I16(ImageData::<i16>::from_data(
190                    width,
191                    height,
192                    normalizer_for(header, &image_data),
193                    bayer_pattern,
194                    image_data,
195                )?))
196            }
197            Bitpix::I32 => {
198                let image_data = data
199                    .as_chunks::<4>()
200                    .0
201                    .iter()
202                    .map(|i| i32::from_be_bytes(*i))
203                    .collect::<Vec<_>>();
204                Ok(Image::I32(ImageData::<i32>::from_data(
205                    width,
206                    height,
207                    normalizer_for(header, &image_data),
208                    bayer_pattern,
209                    image_data,
210                )?))
211            }
212        }
213    }
214}
215
216/// Chooses the scaling for an image whose samples are already in memory.
217///
218/// Follows [`Normalizer::from_header`] — DATAMIN and DATAMAX first, then the
219/// representable range of BITPIX — but adds a fallback the streaming path cannot
220/// use: a floating point image has no representable range, and here the whole
221/// array is in hand, so its actual extent can be measured.
222fn normalizer_for<T: Primitive>(header: &Header, data: &[T]) -> Normalizer {
223    Normalizer::from_header(header).unwrap_or_else(|_| {
224        // `from_header` only fails for a floating point image with neither
225        // DATAMIN nor DATAMAX, and BLANK does not apply to those.
226        Normalizer::from_samples(
227            header.bzero_or_default(),
228            header.bscale_or_default(),
229            data.iter().filter_map(|sample| sample.to_f64()),
230        )
231    })
232}
233
234impl From<ImageBuffer<Luma<f64>, Vec<f64>>> for Image {
235    fn from(image: ImageBuffer<Luma<f64>, Vec<f64>>) -> Self {
236        Image::F64(ImageData::from_buffer(image))
237    }
238}