cv_convert_fork/
with_tch_image.rs

1#[cfg(feature = "image_0-23")]
2mod with_image_0_23 {
3    use crate::image;
4    use crate::tch;
5    use crate::{common::*, FromCv, IntoCv, TchTensorAsImage, TchTensorImageShape, TryFromCv};
6    use std::ops::Deref;
7
8    impl<P, Container> FromCv<&image::ImageBuffer<P, Container>> for TchTensorAsImage
9    where
10        P: image::Pixel + 'static,
11        P::Subpixel: 'static + tch::kind::Element,
12        Container: Deref<Target = [P::Subpixel]>,
13    {
14        fn from_cv(from: &image::ImageBuffer<P, Container>) -> Self {
15            let (width, height) = from.dimensions();
16            let channels = P::CHANNEL_COUNT;
17            let tensor =
18                tch::Tensor::from_slice(&*from).view([width as i64, height as i64, channels as i64]);
19            TchTensorAsImage {
20                tensor,
21                kind: TchTensorImageShape::Whc,
22            }
23        }
24    }
25
26    impl<P, Container> FromCv<image::ImageBuffer<P, Container>> for TchTensorAsImage
27    where
28        P: image::Pixel + 'static,
29        P::Subpixel: 'static + tch::kind::Element,
30        Container: Deref<Target = [P::Subpixel]>,
31    {
32        fn from_cv(from: image::ImageBuffer<P, Container>) -> Self {
33            Self::from_cv(&from)
34        }
35    }
36
37    impl TryFromCv<&image::DynamicImage> for TchTensorAsImage {
38        type Error = Error;
39
40        fn try_from_cv(from: &image::DynamicImage) -> Result<Self, Self::Error> {
41            use image::DynamicImage as D;
42
43            let tensor = match from {
44                D::ImageLuma8(image) => image.into_cv(),
45                D::ImageLumaA8(image) => image.into_cv(),
46                D::ImageRgb8(image) => image.into_cv(),
47                D::ImageRgba8(image) => image.into_cv(),
48                D::ImageBgr8(image) => image.into_cv(),
49                D::ImageBgra8(image) => image.into_cv(),
50                _ => bail!("cannot convert an image with u16 components to a tensor"),
51            };
52            Ok(tensor)
53        }
54    }
55
56    impl TryFromCv<image::DynamicImage> for TchTensorAsImage {
57        type Error = Error;
58
59        fn try_from_cv(from: image::DynamicImage) -> Result<Self, Self::Error> {
60            Self::try_from_cv(&from)
61        }
62    }
63}
64
65#[cfg(feature = "image_0-24")]
66mod with_image_0_24 {
67    use crate::image;
68    use crate::tch;
69    use crate::{common::*, FromCv, IntoCv, TchTensorAsImage, TchTensorImageShape, TryFromCv};
70    use std::ops::Deref;
71
72    impl<P, Container> FromCv<&image::ImageBuffer<P, Container>> for TchTensorAsImage
73    where
74        P: image::Pixel + 'static,
75        P::Subpixel: 'static + tch::kind::Element,
76        Container: Deref<Target = [P::Subpixel]>,
77    {
78        fn from_cv(from: &image::ImageBuffer<P, Container>) -> Self {
79            let (width, height) = from.dimensions();
80            let channels = P::CHANNEL_COUNT;
81            let tensor =
82                tch::Tensor::from_slice(&*from).view([width as i64, height as i64, channels as i64]);
83            TchTensorAsImage {
84                tensor,
85                kind: TchTensorImageShape::Whc,
86            }
87        }
88    }
89
90    impl<P, Container> FromCv<image::ImageBuffer<P, Container>> for TchTensorAsImage
91    where
92        P: image::Pixel + 'static,
93        P::Subpixel: 'static + tch::kind::Element,
94        Container: Deref<Target = [P::Subpixel]>,
95    {
96        fn from_cv(from: image::ImageBuffer<P, Container>) -> Self {
97            Self::from_cv(&from)
98        }
99    }
100
101    impl TryFromCv<&image::DynamicImage> for TchTensorAsImage {
102        type Error = Error;
103
104        fn try_from_cv(from: &image::DynamicImage) -> Result<Self, Self::Error> {
105            use image::DynamicImage as D;
106
107            let tensor = match from {
108                D::ImageLuma8(image) => image.into_cv(),
109                D::ImageLumaA8(image) => image.into_cv(),
110                D::ImageRgb8(image) => image.into_cv(),
111                D::ImageRgba8(image) => image.into_cv(),
112                D::ImageRgb32F(image) => image.into_cv(),
113                D::ImageRgba32F(image) => image.into_cv(),
114                _ => bail!("the color type {:?} is not supported", from.color()),
115            };
116            Ok(tensor)
117        }
118    }
119
120    impl TryFromCv<image::DynamicImage> for TchTensorAsImage {
121        type Error = Error;
122
123        fn try_from_cv(from: image::DynamicImage) -> Result<Self, Self::Error> {
124            Self::try_from_cv(&from)
125        }
126    }
127}
cv_convert_fork/with_tch_image.rs

cv_convert_fork/
with_tch_image.rs
