rten_imageio/
lib.rs

1//! Provides utilities for loading, saving and preprocessing images for use with
2//! [RTen](https://github.com/robertknight/rten).
3//!
4//! The APIs are limited to keep them simple for the most common use cases.
5//! If you need more flexibility from a function, copy and adjust the
6//! implementation.
7
8use std::error::Error;
9use std::path::Path;
10
11use rten_base::num::AsUsize;
12use rten_tensor::errors::FromDataError;
13use rten_tensor::prelude::*;
14use rten_tensor::{NdTensor, NdTensorView};
15
16/// Errors reported when creating a tensor from an image.
17#[derive(Debug)]
18pub enum ReadImageError {
19    /// The image could not be loaded.
20    ImageError(image::ImageError),
21    /// The loaded image could not be converted to a tensor.
22    ConvertError(FromDataError),
23}
24
25impl std::fmt::Display for ReadImageError {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            ReadImageError::ImageError(e) => write!(f, "failed to read image: {}", e),
29            ReadImageError::ConvertError(e) => write!(f, "failed to create tensor: {}", e),
30        }
31    }
32}
33
34impl Error for ReadImageError {}
35
36/// Convert an image into a CHW tensor with 3 channels and values in the range
37/// [0, 1].
38pub fn image_to_tensor(image: image::DynamicImage) -> Result<NdTensor<f32, 3>, ReadImageError> {
39    let image = image.into_rgb8();
40    let (width, height) = image.dimensions();
41    let layout = image.sample_layout();
42
43    let chw_tensor = NdTensorView::from_data_with_strides(
44        [height.as_usize(), width.as_usize(), 3],
45        image.as_raw().as_slice(),
46        [
47            layout.height_stride,
48            layout.width_stride,
49            layout.channel_stride,
50        ],
51    )
52    .map_err(ReadImageError::ConvertError)?
53    .permuted([2, 0, 1]) // HWC => CHW
54    .map(|x| *x as f32 / 255.); // Rescale from [0, 255] to [0, 1]
55
56    Ok(chw_tensor)
57}
58
59/// Read an image from a file into a CHW tensor.
60///
61/// To load an image from a byte buffer or other source, use [`image::open`]
62/// and pass the result to [`image_to_tensor`].
63pub fn read_image<P: AsRef<Path>>(path: P) -> Result<NdTensor<f32, 3>, ReadImageError> {
64    image::open(path)
65        .map_err(ReadImageError::ImageError)
66        .and_then(image_to_tensor)
67}
68
69/// Errors returned when writing a tensor to an image.
70#[derive(Debug)]
71pub enum WriteImageError {
72    /// The number of channels in the image tensor is unsupported.
73    UnsupportedChannelCount,
74    /// The image could not be written.
75    ImageError(image::ImageError),
76}
77
78impl std::fmt::Display for WriteImageError {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        match self {
81            Self::ImageError(e) => write!(f, "failed to write image: {}", e),
82            Self::UnsupportedChannelCount => write!(f, "image has unsupported number of channels"),
83        }
84    }
85}
86
87impl Error for WriteImageError {}
88
89/// Convert a CHW tensor to an image and write it to a PNG file.
90pub fn write_image(path: &str, img: NdTensorView<f32, 3>) -> Result<(), WriteImageError> {
91    let [channels, height, width] = img.shape();
92    let color_type = match channels {
93        1 => image::ColorType::L8,
94        3 => image::ColorType::Rgb8,
95        4 => image::ColorType::Rgba8,
96        _ => return Err(WriteImageError::UnsupportedChannelCount),
97    };
98
99    let hwc_img = img
100        .permuted([1, 2, 0]) // CHW => HWC
101        .map(|x| (x.clamp(0., 1.) * 255.0) as u8);
102
103    image::save_buffer(
104        path,
105        hwc_img.data().unwrap(),
106        width as u32,
107        height as u32,
108        color_type,
109    )
110    .map_err(WriteImageError::ImageError)?;
111
112    Ok(())
113}
rten_imageio/lib.rs

rten_imageio/
lib.rs
