1use sophus_tensor::{
2 MutTensor,
3 TensorView,
4};
5
6use crate::{
7 ImageSize,
8 arc_image::GenArcImage,
9 image_view::GenImageView,
10 prelude::*,
11};
12
13extern crate alloc;
14
15#[derive(Debug, Clone, Default)]
17pub struct GenMutImage<
18 const TOTAL_RANK: usize,
19 const SRANK: usize,
20 Scalar: IsCoreScalar + 'static,
21 STensor: IsStaticTensor<Scalar, SRANK, ROWS, COLS> + 'static,
22 const ROWS: usize,
23 const COLS: usize,
24> {
25 pub mut_tensor: MutTensor<TOTAL_RANK, 2, SRANK, Scalar, STensor, ROWS, COLS>,
27}
28
29pub trait IsMutImage<
31 'a,
32 const TOTAL_RANK: usize,
33 const SRANK: usize,
34 Scalar: IsCoreScalar + 'static,
35 STensor: IsStaticTensor<Scalar, SRANK, ROWS, COLS> + 'static,
36 const ROWS: usize,
37 const COLS: usize,
38> where
39 ndarray::Dim<[ndarray::Ix; TOTAL_RANK]>: ndarray::Dimension,
40{
41 fn from_image_size(size: ImageSize) -> Self;
43 fn from_image_size_and_val(size: ImageSize, val: STensor) -> Self;
45}
46
47macro_rules! mut_image {
48 ($scalar_rank:literal, $srank:literal) => {
49 impl<
50 'a,
51 Scalar: IsCoreScalar + 'static,
52 STensor: IsStaticTensor<Scalar, $srank, ROWS, COLS> + 'static,
53 const ROWS: usize,
54 const COLS: usize,
55
56 > IsImageView<'a, $scalar_rank, $srank, Scalar, STensor, ROWS, COLS>
57 for GenMutImage<$scalar_rank, $srank, Scalar, STensor, ROWS, COLS>
58 where
59 ndarray::Dim<[ndarray::Ix; $scalar_rank]>: ndarray::Dimension,
60 {
61 fn image_view(
62 &'a self,
63 ) -> crate::image_view::GenImageView<
64 'a,
65 $scalar_rank,
66 $srank,
67 Scalar,
68 STensor,
69 ROWS,
70 COLS
71 > {
72 let v = self.mut_tensor.view();
73 GenImageView { tensor_view: v }
74 }
75
76 fn pixel(&'a self, u: usize, v: usize) -> STensor {
77 self.mut_tensor.mut_array[[v, u]].clone()
78 }
79
80 fn image_size(&self) -> ImageSize {
81 self.image_view().image_size()
82 }
83 }
84
85 impl<
86 'a,
87 Scalar: IsCoreScalar + 'static,
88 STensor: IsStaticTensor<Scalar, $srank, ROWS, COLS> + 'static,
89 const ROWS: usize,
90 const COLS: usize,
91
92 > GenMutImage<$scalar_rank, $srank, Scalar, STensor, ROWS, COLS>
93 {
94 pub fn from_image_size(size: ImageSize) -> Self {
96 Self {
97 mut_tensor: MutTensor::<
98 $scalar_rank,
99 2,
100 $srank,
101 Scalar,
102 STensor,
103 ROWS,
104 COLS
105 >::from_shape(size.into()),
106 }
107
108 }
109
110 pub fn from_image_size_and_val(size: ImageSize, val: STensor) -> Self {
112 Self {
113 mut_tensor: MutTensor::<
114 $scalar_rank,
115 2,
116 $srank,
117 Scalar,
118 STensor,
119 ROWS,
120 COLS
121 >::from_shape_and_val(size.into(), val),
122 }
123 }
124
125 pub fn make_copy_from(
127 v: &GenImageView<$scalar_rank, $srank, Scalar, STensor, ROWS, COLS>,
128 ) -> Self {
129 Self {
130 mut_tensor: MutTensor::<
131 $scalar_rank,
132 2,
133 $srank,
134 Scalar,
135 STensor,
136 ROWS,
137 COLS
138 >::make_copy_from(&v.tensor_view),
139 }
140 }
141
142 pub fn make_copy_from_size_and_slice(image_size: ImageSize, slice: &'a [STensor]) -> Self {
144 Self::make_copy_from(&GenImageView::<
145 $scalar_rank,
146 $srank,
147 Scalar,
148 STensor,
149 ROWS,
150 COLS,
151 >::from_size_and_slice(image_size, slice))
152
153 }
154
155 pub fn try_make_copy_from_size_and_bytes(
157 image_size: ImageSize,
158 bytes: &'a [u8],
159 ) -> Result<Self, alloc::string::String> {
160 let num_scalars_in_pixel = STensor::num_scalars();
161 let num_scalars_in_image = num_scalars_in_pixel * image_size.width * image_size.height;
162 let size_in_bytes = num_scalars_in_image * core::mem::size_of::<Scalar>();
163 if(bytes.len() != size_in_bytes){
164 return Err(alloc::format!("bytes.len() = {} != size_in_bytes = {}",
165 bytes.len(), size_in_bytes));
166 }
167 let stensor_slice = unsafe {
168 core::slice::from_raw_parts(
169 bytes.as_ptr() as *const Scalar,
170 num_scalars_in_image)
171 };
172
173 let mut img = Self::from_image_size(image_size);
174
175 for v in 0..image_size.height {
176 for u in 0..image_size.width {
177 let idx = (v * image_size.width + u) * STensor::num_scalars() ;
178 let pixel = &stensor_slice[idx..idx + STensor::num_scalars()];
179 *img.mut_pixel(u, v) = STensor::from_slice(pixel);
180 }
181 }
182 Ok(img)
183 }
184
185 pub fn make_copy_from_make_from_size_and_bytes(
187 image_size: ImageSize,
188 bytes: &'a [u8],
189 ) -> Self {
190 Self::try_make_copy_from_size_and_bytes(image_size, bytes).unwrap()
191 }
192
193 pub fn from_map<
195 'b,
196 const OTHER_HRANK: usize,
197 const OTHER_SRANK: usize,
198 OtherScalar: IsCoreScalar + 'static,
199 OtherSTensor: IsStaticTensor<
200 OtherScalar,
201 OTHER_SRANK,
202 OTHER_ROWS,
203 OTHER_COLS,
204 > + 'static,
205 const OTHER_ROWS: usize,
206 const OTHER_COLS: usize,
207 F: FnMut(&OtherSTensor)-> STensor
208 >(
209 v: &'b GenImageView::<
210 'b,
211 OTHER_HRANK,
212 OTHER_SRANK,
213 OtherScalar,
214 OtherSTensor,
215 OTHER_ROWS,
216 OTHER_COLS,
217 >,
218 op: F,
219
220 ) -> Self
221 where ndarray::Dim<[ndarray::Ix; OTHER_HRANK]>: ndarray::Dimension,
222 TensorView<'b, OTHER_HRANK, 2, OTHER_SRANK, OtherScalar, OtherSTensor,
223 OTHER_ROWS, OTHER_COLS>:
224 IsTensorView<'b, OTHER_HRANK, 2, OTHER_SRANK, OtherScalar, OtherSTensor,
225 OTHER_ROWS, OTHER_COLS>,
226
227 {
228 Self {
229 mut_tensor: MutTensor::<
230 $scalar_rank,
231 2,
232 $srank,
233 Scalar,
234 STensor,
235 ROWS,
236 COLS
237 >::from_map(&v.tensor_view, op),
238 }
239 }
240
241 pub fn to_shared(
243 self,
244 ) -> GenArcImage<
245 $scalar_rank,
246 $srank,
247 Scalar,
248 STensor,
249 ROWS,
250 COLS
251 > {
252 GenArcImage {
253 tensor: self.mut_tensor.to_shared(),
254 }
255 }
256 }
257
258 impl<
259 'a,
260 Scalar: IsCoreScalar + 'static,
261 STensor: IsStaticTensor<Scalar, $srank, ROWS, COLS> + 'static,
262 const ROWS: usize,
263 const COLS: usize,
264
265 > IsMutImageView<'a, $scalar_rank, $srank, Scalar, STensor, ROWS, COLS>
266 for GenMutImage<$scalar_rank, $srank, Scalar, STensor, ROWS, COLS>
267 where
268 ndarray::Dim<[ndarray::Ix; $scalar_rank]>: ndarray::Dimension,
269 {
270 fn mut_image_view<'b: 'a>(
271 &'b mut self,
272 ) -> crate::mut_image_view::GenMutImageView<
273 'a,
274 $scalar_rank,
275 $srank,
276 Scalar,
277 STensor,
278 ROWS,
279 COLS
280 > {
281 crate::mut_image_view::GenMutImageView {
282 mut_tensor_view: self.mut_tensor.mut_view(),
283 }
284 }
285
286 fn mut_pixel(&'a mut self, u: usize, v: usize) -> &'a mut STensor {
287 self.mut_tensor.get_mut([v, u])
288 }
289 }
290 };
291}
292
293mut_image!(2, 0);
294mut_image!(3, 1);
295mut_image!(4, 2);