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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::marker::PhantomData;
use std::ops::Deref;
use std::slice;

use libavif_sys as sys;

use crate::{AvifImage, Error, YuvFormat};

pub struct RgbPixels<'a> {
    owned: bool,
    inner: sys::avifRGBImage,

    phantom: PhantomData<&'a [u8]>,
}

impl<'a> RgbPixels<'a> {
    pub fn new(width: u32, height: u32, rgb: &'a [u8]) -> Result<Self, Error> {
        let (stride, format) = if (width * height * 3) as usize == rgb.len() {
            // RGB
            (3, sys::AVIF_RGB_FORMAT_RGB)
        } else if (width * height * 4) as usize == rgb.len() {
            // RGBA
            (4, sys::AVIF_RGB_FORMAT_RGBA)
        } else {
            return Err(Error::UnsupportedImageType);
        };

        Ok(Self {
            owned: true,
            inner: sys::avifRGBImage {
                width,
                height,
                depth: 8,
                format,
                chromaUpsampling: sys::AVIF_CHROMA_UPSAMPLING_BILINEAR,
                ignoreAlpha: 0,
                pixels: rgb.as_ptr() as *mut u8,
                rowBytes: stride * width,
            },
            phantom: PhantomData,
        })
    }

    /// width of the image in pixels
    pub fn width(&self) -> u32 {
        self.inner.width
    }

    /// height of the image in pixels
    pub fn height(&self) -> u32 {
        self.inner.height
    }

    pub fn pixel(&self, x: u32, y: u32) -> (u8, u8, u8, u8) {
        let stride = if self.inner.format == sys::AVIF_RGB_FORMAT_RGBA {
            4
        } else {
            3
        };

        let row_bytes = self.inner.rowBytes as usize;
        let i = (stride * x as usize) + (row_bytes * y as usize);

        let slice = self.as_slice();
        let slice = &slice[i..][..stride];
        (
            slice[0],
            slice[1],
            slice[2],
            if stride == 4 { slice[3] } else { 255 },
        )
    }

    /// Extracts a slice containg all of the pixels without doing clones or allocation.
    pub fn as_slice(&'a self) -> &'a [u8] {
        let size = self.inner.rowBytes * self.height();
        unsafe { slice::from_raw_parts(self.inner.pixels, size as usize) }
    }

    /// Converts `self` into a new vector by cloning all of the pixels.
    pub fn to_vec(&self) -> Vec<u8> {
        self.as_slice().to_vec()
    }

    pub fn to_image(&self, yuv_format: YuvFormat) -> AvifImage {
        unsafe {
            let image =
                sys::avifImageCreate(self.width() as _, self.height() as _, 8, yuv_format as i32);
            sys::avifImageAllocatePlanes(image, sys::AVIF_PLANES_YUV as _);

            sys::avifImageRGBToYUV(image, &self.inner as *const sys::avifRGBImage);
            AvifImage::from_raw(image)
        }
    }
}

impl<'a> From<sys::avifRGBImage> for RgbPixels<'a> {
    fn from(rgb: sys::avifRGBImage) -> Self {
        Self {
            owned: false,
            inner: rgb,
            phantom: PhantomData,
        }
    }
}

impl<'a> From<AvifImage> for RgbPixels<'a> {
    fn from(image: AvifImage) -> Self {
        Self::from(&image)
    }
}

impl<'a> From<&AvifImage> for RgbPixels<'a> {
    fn from(image: &AvifImage) -> Self {
        unsafe {
            let mut rgb = sys::avifRGBImage::default();
            let raw_rgb = &mut rgb as *mut sys::avifRGBImage;
            sys::avifRGBImageSetDefaults(raw_rgb, image.inner());
            rgb.format = sys::AVIF_RGB_FORMAT_RGBA;
            rgb.depth = 8;

            sys::avifRGBImageAllocatePixels(raw_rgb);
            sys::avifImageYUVToRGB(image.inner(), raw_rgb);

            RgbPixels::from(rgb)
        }
    }
}

impl<'a> Deref for RgbPixels<'a> {
    type Target = [u8];
    fn deref(&self) -> &Self::Target {
        self.as_slice()
    }
}

impl Drop for RgbPixels<'_> {
    fn drop(&mut self) {
        if !self.owned {
            // pixels were allocated by libavif
            unsafe {
                sys::avifRGBImageFreePixels(&mut self.inner as *mut sys::avifRGBImage);
            }
        }
    }
}