1use crate::error::ImageError;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum DType {
10 F32,
12 U8,
14}
15
16#[derive(Debug, Clone)]
20pub struct ImageBuf {
21 data: Vec<f32>,
22 width: usize,
23 height: usize,
24 channels: usize,
25 dtype: DType,
26}
27
28impl ImageBuf {
29 pub fn new(
35 data: Vec<f32>,
36 width: usize,
37 height: usize,
38 channels: usize,
39 ) -> Result<Self, ImageError> {
40 let expected = width * height * channels;
41 if data.len() != expected {
42 return Err(ImageError::DimensionMismatch {
43 expected,
44 got: data.len(),
45 });
46 }
47 Ok(Self {
48 data,
49 width,
50 height,
51 channels,
52 dtype: DType::F32,
53 })
54 }
55
56 pub fn zeros(width: usize, height: usize, channels: usize) -> Self {
58 Self {
59 data: vec![0.0; width * height * channels],
60 width,
61 height,
62 channels,
63 dtype: DType::F32,
64 }
65 }
66
67 pub fn width(&self) -> usize {
69 self.width
70 }
71
72 pub fn height(&self) -> usize {
74 self.height
75 }
76
77 pub fn channels(&self) -> usize {
79 self.channels
80 }
81
82 pub fn dtype(&self) -> DType {
84 self.dtype
85 }
86
87 pub fn data(&self) -> &[f32] {
89 &self.data
90 }
91
92 pub fn data_mut(&mut self) -> &mut [f32] {
94 &mut self.data
95 }
96
97 pub fn len(&self) -> usize {
99 self.data.len()
100 }
101
102 pub fn is_empty(&self) -> bool {
104 self.data.is_empty()
105 }
106
107 pub fn channel(&self, channel: usize) -> Result<Self, ImageError> {
113 if channel >= self.channels {
114 return Err(ImageError::InvalidChannel {
115 channel,
116 max: self.channels,
117 });
118 }
119 let mut out = Vec::with_capacity(self.width * self.height);
120 for i in 0..self.width * self.height {
121 out.push(self.data[i * self.channels + channel]);
122 }
123 Ok(Self {
124 data: out,
125 width: self.width,
126 height: self.height,
127 channels: 1,
128 dtype: self.dtype,
129 })
130 }
131}