boreholeio 0.1.0

use super::{Error, MeasurementArray};
use crate::{
    schema::{
        array_uri_reference::ArrayUriReference,
        axes::{Base, ChannelAxis, LabelledAxis},
        Label,
    },
    strided_array_file::{ArrayProxy, Statistics},
};
use image::{DynamicImage, ImageReader};
use ndarray::{ArrayViewD, IxDyn, ShapeBuilder, SliceInfoElem};
use std::{path::PathBuf, sync::Arc};

/// An object which owns (in memory) the data from a (PNG, JPG, etc) image file
/// and can quickly return views to it.
#[derive(Debug)]
pub(crate) struct ImageData<'a> {
    /// Data axes which describe every image dimension.
    axes: Vec<ChannelAxis>,
    /// Optional label object defining the implicit measurement axis, if one
    /// exists. Is not `None` for e.g. grayscale images.
    label: Option<Label>,
    /// The data itself.
    #[expect(dead_code)]
    im: DynamicImage,
    /// Data view whose elements type is incapsulated into `ArrayProxy`.
    data: ArrayProxy<'a>,
}

impl<'a> ImageData<'a> {
    /// Tries to read the input image file, verify its consistency with the
    /// provided metadata and constructs `Arc<Self>` out of it.
    pub(crate) fn from_uri(
        uri: &ArrayUriReference,
        mut axes: Vec<ChannelAxis>,
        label: Option<Label>,
    ) -> Result<Arc<dyn MeasurementArray<'a> + 'a>, Error> {
        // Verify the metadata.
        if axes.len() != 2 {
            return Err("ImageData expects to get exactly 2 axes".to_string());
        }
        // Read the image from file.
        let path: PathBuf = uri
            .try_into()
            .map_err(|_| format!("Failed to extract file path from URI: {uri:?}"))?;
        let dyn_im = ImageReader::open(path)
            .map_err(|e| e.to_string())?
            .decode()
            .map_err(|e| e.to_string())?;
        // Verify image data and metadata consistency.
        if dyn_im.width() as usize != axes[1].size()?
            || dyn_im.height() as usize != axes[0].size()?
        {
            return Err("Inconsistent image resolution".to_string());
        }
        // Parse image format and create a view to the data.
        let (data, colour_channels_names) = match &dyn_im {
            DynamicImage::ImageRgb8(im) => {
                let data = create_proxy::<u8>(&im.as_flat_samples())?;
                Ok((data, &["R", "G", "B"]))
            }
            _ => Err("Unsupported image type".to_string()),
        }?;
        // Update the metadata with colourspace information, which is known only
        // to the image file itself.
        axes.push(create_colour_channels_axis(&data, colour_channels_names)?);

        Ok(Arc::new(Self {
            axes,
            label,
            im: dyn_im,
            data,
        }))
    }
}

impl<'a> MeasurementArray<'a> for ImageData<'a> {
    fn data(&self) -> ArrayProxy<'a> {
        self.data.clone()
    }

    fn axes(&self) -> &Vec<ChannelAxis> {
        &self.axes
    }

    fn label(&self) -> &Option<Label> {
        &self.label
    }
}

/// Creates a view of the image data (stored in a data structure from the
/// `image` crate) with the type `T` incapsulated into `ArrayProxy`.
/// Note that the `image` crate assumes the row-major data order.
fn create_proxy<'a, T: 'static>(im: &image::FlatSamples<&[u8]>) -> Result<ArrayProxy<'a>, Error> {
    let (data, layout) = (im.samples, im.layout);
    let strided_shape = IxDyn(&[
        layout.height as usize,
        layout.width as usize,
        layout.channels as usize,
    ])
    .strides(IxDyn(&[
        layout.height_stride,
        layout.width_stride,
        layout.channel_stride,
    ]));
    let view = unsafe { ArrayViewD::from_shape_ptr(strided_shape, data.as_ptr() as *const T) };
    ArrayProxy::from_view(view)
}

/// When channel data is stored in a standard image file (like PNG, JPG, etc),
/// the metadata contains only 2 axes which describe only vertical and
/// horizontal image dimensions and the colour information is expected to be
/// stored in the image file itself. This method restores this information.
fn create_colour_channels_axis(
    image: &ArrayProxy<'_>,
    colour_channels_names: &[&str],
) -> Result<ChannelAxis, Error> {
    if image.dimensions() != 3 {
        return Err("Image is expected to have exactly 3 dimensions".to_string());
    }
    if image.shape()[2] != colour_channels_names.len() {
        return Err("Inconsistent size of colour channels names".to_string());
    }
    let mut slice = vec![
        SliceInfoElem::Slice {
            start: 0,
            end: None,
            step: 1,
        };
        image.dimensions()
    ];
    let mut labels = Vec::<Label>::with_capacity(image.dimensions());
    for (channel_index, channel_name) in colour_channels_names.iter().enumerate() {
        slice[2] = SliceInfoElem::Index(channel_index as isize);
        let values = image.slice(slice.clone())?;
        labels.push(Label {
            name: channel_name.to_string(),
            bounds: [
                values.calc_stats(Statistics::Min)?,
                values.calc_stats(Statistics::Max)?,
            ],
            units: None,
            transform: None,
        });
    }
    Ok(LabelledAxis { labels }.into())
}