sophus_image/
mut_image.rs

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/// Mutable image of static tensors
16#[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    /// underlying mutable tensor
26    pub mut_tensor: MutTensor<TOTAL_RANK, 2, SRANK, Scalar, STensor, ROWS, COLS>,
27}
28
29/// is a mutable image
30pub 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    /// creates a mutable image view from image size
42    fn from_image_size(size: ImageSize) -> Self;
43    /// creates a mutable image view from image size and value
44    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            /// creates a mutable image view from image size
95            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            /// creates a mutable image from image size and value
111            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            /// creates a mutable image from image view
126            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            /// creates a mutable image from image size and slice
143            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            /// creates a mutable image from image size and byte slice
156            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            /// creates a mutable image from image size and byte slice
186            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            /// creates a mutable image from unary operator applied to image view
194            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            /// creates shared image from mutable image
242            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);