scra_mirach_model/
data.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
use std::{fmt::{Debug, Display}, mem::MaybeUninit};

use burn::{
    prelude::Backend,
    tensor::{Device, Tensor, TensorData},
};

use crate::{IMAGE_CHANNELS, IMAGE_HEIGHT, IMAGE_WIDTH};

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
pub struct CaptchaCode(u16);

impl CaptchaCode {
    #[inline]
    pub fn new(code: u16) -> Self {
        assert!(code < 10000);
        Self(code)
    }

    #[inline]
    pub fn new_unchecked(code: u16) -> Self {
        Self(code)
    }

    #[inline]
    pub fn as_u16(&self) -> u16 {
        self.0
    }

    #[inline]
    pub fn as_digits(&self) -> [u8; 4] {
        [
            (self.0 / 1000) as u8,
            ((self.0 / 100) % 10) as u8,
            ((self.0 / 10) % 10) as u8,
            (self.0 % 10) as u8,
        ]
    }

    #[inline]
    pub fn from_digits(digits: [u8; 4]) -> Self {
        Self::new_unchecked(
            (digits[0] as u16 * 1000)
                + (digits[1] as u16 * 100)
                + (digits[2] as u16 * 10)
                + digits[3] as u16,
        )
    }
}

impl Debug for CaptchaCode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("CaptchaCode({:04})", self.0))
    }
}

impl Display for CaptchaCode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{:04}", self.0))
    }
}

#[derive(Clone, PartialEq, Eq)]
pub struct MirachImageData(Vec<u8>);

impl MirachImageData {
    pub fn new(pixels: Vec<u8>) -> Self {
        Self(pixels)
    }

    pub fn unwrap(self) -> Vec<u8> {
        self.0
    }
}

#[cfg(feature = "image")]
impl From<image::RgbImage> for MirachImageData {
    fn from(image: image::RgbImage) -> Self {
        // SAFETY: The array will be filled are allocation immediately,
        //         and no read will be performed before all data are fully written.
        //         Therefore no data can be leaked.
        #[allow(invalid_value)]
        let mut data: Box<[u8; IMAGE_HEIGHT * IMAGE_WIDTH * IMAGE_CHANNELS]> =
            unsafe { Box::new(MaybeUninit::uninit().assume_init()) };
        for x in 0..IMAGE_WIDTH {
            let offset = x * IMAGE_HEIGHT;
            for y in 0..IMAGE_HEIGHT {
                let offset = offset + y;
                let pixel = image.get_pixel(x as u32, y as u32);
                data[(0 * (IMAGE_WIDTH * IMAGE_HEIGHT)) + offset] = pixel[0];
                data[(1 * (IMAGE_WIDTH * IMAGE_HEIGHT)) + offset] = pixel[1];
                data[(2 * (IMAGE_WIDTH * IMAGE_HEIGHT)) + offset] = pixel[2];
            }
        }
        Self(Vec::from(data as Box<[u8]>))
    }
}

#[cfg(feature = "image")]
impl From<image::DynamicImage> for MirachImageData {
    fn from(image: image::DynamicImage) -> Self {
        image.into_rgb8().into()
    }
}

impl MirachImageData {
    pub fn into_tensor<B: Backend>(self, device: &Device<B>) -> Tensor<B, 4> {
        let image = TensorData::new::<u8, _>(
            self.unwrap(),
            [1, IMAGE_CHANNELS, IMAGE_WIDTH, IMAGE_HEIGHT],
        );
        let image = Tensor::<B, 4>::from(image.convert::<B::FloatElem>()).to_device(device);
        let image = image / 255;
        debug_assert_eq!(image.dims(), [1, IMAGE_CHANNELS, IMAGE_WIDTH, IMAGE_HEIGHT]);
        image
    }
}