use crate::error::ImageError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DType {
F32,
U8,
}
#[derive(Debug, Clone)]
pub struct ImageBuf {
data: Vec<f32>,
width: usize,
height: usize,
channels: usize,
dtype: DType,
}
impl ImageBuf {
pub fn new(
data: Vec<f32>,
width: usize,
height: usize,
channels: usize,
) -> Result<Self, ImageError> {
let expected = width * height * channels;
if data.len() != expected {
return Err(ImageError::DimensionMismatch {
expected,
got: data.len(),
});
}
Ok(Self {
data,
width,
height,
channels,
dtype: DType::F32,
})
}
pub fn zeros(width: usize, height: usize, channels: usize) -> Self {
Self {
data: vec![0.0; width * height * channels],
width,
height,
channels,
dtype: DType::F32,
}
}
pub fn width(&self) -> usize {
self.width
}
pub fn height(&self) -> usize {
self.height
}
pub fn channels(&self) -> usize {
self.channels
}
pub fn dtype(&self) -> DType {
self.dtype
}
pub fn data(&self) -> &[f32] {
&self.data
}
pub fn data_mut(&mut self) -> &mut [f32] {
&mut self.data
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn channel(&self, channel: usize) -> Result<Self, ImageError> {
if channel >= self.channels {
return Err(ImageError::InvalidChannel {
channel,
max: self.channels,
});
}
let mut out = Vec::with_capacity(self.width * self.height);
for i in 0..self.width * self.height {
out.push(self.data[i * self.channels + channel]);
}
Ok(Self {
data: out,
width: self.width,
height: self.height,
channels: 1,
dtype: self.dtype,
})
}
}