sentryshot_convert/

lib.rs

// SPDX-License-Identifier: MPL-2.0+

mod yub2rgb;

#[cfg(test)]
mod test;

pub use sentryshot_util::{pixfmt::PixelFormat, ColorRange, Frame};

use crate::yub2rgb::yuv2rgb_get_func;
use sentryshot_util::{pixfmt::PIX_FMT_FLAG_RGB, ResetBufferError};
use std::{num::NonZeroU16, panic, sync::Mutex};
use thiserror::Error;
use yub2rgb::{yuv2rgb_init_tables, Yuv2RgbInitTablesError, FF_YUV2RGB_COEFFS};

const SWS_CS_DEFAULT: usize = 5;

pub(crate) const YUVRGB_TABLE_HEADROOM: u16 = 512;
#[allow(clippy::as_conversions)]
pub(crate) const YUVRGB_TABLE_HEADROOM_USIZE: usize = YUVRGB_TABLE_HEADROOM as usize;
pub(crate) const YUVRGB_TABLE_SIZE: usize = 256 + 2 * YUVRGB_TABLE_HEADROOM_USIZE;
pub(crate) const YUVRGB_TABLE_LUMA_HEADROOM: u16 = 512;
#[allow(clippy::as_conversions)]
pub(crate) const YUV_TABLE_SIZE: usize = 1024 + 2 * YUVRGB_TABLE_LUMA_HEADROOM as usize;

pub(crate) type ConvertFunc = fn(
    context: &PixelFormatConverter,
    src: &[Vec<u8>],
    src_stride: &mut [usize],
    src_slice_h: i32,
    dst: &mut [Vec<u8>],
    dst_stride: &[usize],
);

// alignment ensures the offset can be added in a single
// instruction on e.g. ARM
// ALIGN = 16.
#[repr(C, align(16))]
pub(crate) struct YuvTables {
    pub(crate) yuv_table: [u8; YUV_TABLE_SIZE],
    pub(crate) gv: [i16; YUVRGB_TABLE_SIZE],
    pub(crate) rv: [usize; YUVRGB_TABLE_SIZE],
    pub(crate) gu: [usize; YUVRGB_TABLE_SIZE],
    pub(crate) bu: [usize; YUVRGB_TABLE_SIZE],
}

// This struct should be aligned on at least a 32-byte boundary.
#[repr(C, align(32))]
pub struct PixelFormatConverter {
    pub(crate) convert: ConvertFunc,

    pub(crate) width: NonZeroU16,
    pub(crate) height: NonZeroU16,
    color_range: ColorRange,

    // Destination pixel format.
    pub(crate) dst_format: PixelFormat,
    // Source pixel format.
    pub(crate) src_format: PixelFormat,
    //
    // Number of bits per pixel o fthe destiantion pixel format.
    //dst_format_bpp: u8,
    //int dstFormatBpp;
    //
    //frame_src: &'a AVFrame,
    //AVFrame *frame_src;
    //frame_dst: &'a AVFrame,
    //AVFrame *frame_dst;
    //
    pub(crate) yuv_tables: Option<YuvTables>,
    //pub(crate) yuv_table: [u8; YUV_TABLE_SIZE],
    // alignment ensures the offset can be added in a single
    // instruction on e.g. ARM
    // ALIGN = 16.
    //pub(crate) table_gv: [i16; YUVRGB_TABLE_SIZE],
    //pub(crate) table_rv: [usize; YUVRGB_TABLE_SIZE],
    //pub(crate) table_gu: [usize; YUVRGB_TABLE_SIZE],
    //pub(crate) table_bu: [usize; YUVRGB_TABLE_SIZE],
    //DECLARE_ALIGNED(16, int32_t, input_rgb2yuv_table)
    //[16 + 40 * 4]; // This table can contain both C and SIMD formatted values, the
    // C vales are always at the XY_IDX points
    //#define RY_IDX 0
    //#define GY_IDX 1
    //#define BY_IDX 2
    //#define RU_IDX 3
    //#define GU_IDX 4
    //#define BU_IDX 5
    //#define RV_IDX 6
    //#define GV_IDX 7
    //#define BV_IDX 8
    //#define RGB2YUV_SHIFT 15

    // Colorspace stuff
    //int srcColorspaceTable[4];
    //int dstColorspaceTable[4];
    //src_range: i32, // 0 = MPG YUV range, 1 = JPG YUV range (source      image).

    //int dstRange; ///< 0 = MPG YUV range, 1 = JPG YUV range (destination image).
    //int yuv2rgb_y_offset;
    //int yuv2rgb_y_coeff;
    //int yuv2rgb_v2r_coeff;
    //int yuv2rgb_v2g_coeff;
    //int yuv2rgb_u2g_coeff;
    //int yuv2rgb_u2b_coeff;

    //DECLARE_ALIGNED(8, uint64_t, yCoeff);
    //DECLARE_ALIGNED(8, uint64_t, vrCoeff);
    //DECLARE_ALIGNED(8, uint64_t, ubCoeff);
    //DECLARE_ALIGNED(8, uint64_t, vgCoeff);
    //DECLARE_ALIGNED(8, uint64_t, ugCoeff);
    //DECLARE_ALIGNED(8, uint64_t, yOffset);
    //DECLARE_ALIGNED(8, uint64_t, uOffset);
    //DECLARE_ALIGNED(8, uint64_t, vOffset);
}

#[derive(Debug, Error)]
#[non_exhaustive]
pub enum NewConverterError {
    #[error("odd height is unsupported for yuv->rgb")]
    OddHeightUnsupported,

    #[error("unsupported pixel formats: {0} -> {1}")]
    UnsupportedPixelFormat(PixelFormat, PixelFormat),

    #[error("init tables: {0} {0}")]
    InitTables(Yuv2RgbInitTablesError, PixelFormat),
}

#[derive(Debug, Error)]
pub enum ConvertError {
    #[error("width changed: {0}vs{1}")]
    WidthChanged(NonZeroU16, NonZeroU16),

    #[error("height changed: {0}vs{1}")]
    HeightChanged(NonZeroU16, NonZeroU16),

    #[error("pixel format changed: {0}vs{1}")]
    PixelFormatChanged(PixelFormat, PixelFormat),

    #[error("color range changed: {0}vs{1}")]
    ColorRangeChanged(ColorRange, ColorRange),

    #[error("reset buffer: {0}")]
    ResetBuffer(#[from] ResetBufferError),

    #[error("converter paniced: {0}")]
    Panic(String),
}

impl PixelFormatConverter {
    pub fn new(
        width: NonZeroU16,
        height: NonZeroU16,
        color_range: ColorRange,
        mut src_format: PixelFormat,
        mut dst_format: PixelFormat,
    ) -> Result<Self, NewConverterError> {
        let dst_format2 = dst_format;
        src_format = handle_jpeg(src_format);
        dst_format = handle_jpeg(dst_format);

        if dst_format2 != dst_format {
            return Err(NewConverterError::UnsupportedPixelFormat(
                src_format,
                dst_format2,
            ));
        }

        //if (!c->dstFormatBpp)
        //  sws_setColorspaceDetails(c, ff_yuv2rgb_coeffs[SWS_CS_DEFAULT], c->srcRange,
        //                           ff_yuv2rgb_coeffs[SWS_CS_DEFAULT], c->dstRange, 0,
        //                           1 << 16, 1 << 16);

        let inv_table = FF_YUV2RGB_COEFFS[SWS_CS_DEFAULT];
        let brightness = 0;
        let contrast = 1 << 16;
        let saturation = 1 << 16;

        //const AVPixFmtDescriptor *desc_dst;

        //dstRange = 0;

        //int need_reinit = 0;
        /*if (c->dstRange != dstRange ||
              memcmp(c->dstColorspaceTable, table, sizeof(int) * 4)) {
            need_reinit = 1;
        }*/

        //memmove(c->srcColorspaceTable, inv_table, sizeof(int) * 4);
        //memmove(c->dstColorspaceTable, table, sizeof(int) * 4);

        //c->dstRange = dstRange;

        let dst_format_bpp = dst_format.bits_per_pixel();

        // if (need_reinit) {
        //    return 0;
        //}

        #[allow(clippy::if_then_some_else_none)]
        let yuv_tables = if is_yuv(src_format) && is_rgb(dst_format) {
            Some(
                yuv2rgb_init_tables(
                    dst_format_bpp,
                    &inv_table,
                    color_range,
                    brightness,
                    contrast,
                    saturation,
                )
                .map_err(|e| NewConverterError::InitTables(e, dst_format))?,
            )
        } else {
            None
        };

        Ok(PixelFormatConverter {
            convert: get_convert_func(src_format, dst_format, height)?,
            width,
            height,
            color_range,
            dst_format,
            src_format,
            yuv_tables,
        })
    }

    /// Converts `src_frame` to another pixel format and stores it in `dst_frame`.
    ///
    /// Will check that `src_frame` matches converter configuration. `dst_frame` will be resized if necessary.
    //
    // Assumes planar YUV to be in YUV order instead of YVU.
    #[allow(clippy::missing_panics_doc)]
    pub fn convert(
        &mut self,
        src_frame: &Frame,
        dst_frame: &mut Frame,
    ) -> Result<(), ConvertError> {
        use ConvertError::*;

        let width = self.width;
        let height = self.height;
        let color_range = self.color_range;
        let src_format = self.src_format;
        let dst_format = self.dst_format;

        if src_frame.width() != width {
            return Err(WidthChanged(src_frame.width(), width));
        }
        if src_frame.height() != height {
            return Err(HeightChanged(src_frame.height(), height));
        }
        if handle_jpeg(src_frame.pix_fmt()) != src_format {
            return Err(PixelFormatChanged(
                handle_jpeg(src_frame.pix_fmt()),
                src_format,
            ));
        }
        if src_frame.color_range() != color_range {
            return Err(ColorRangeChanged(src_frame.color_range(), color_range));
        }

        dst_frame.reset_buffer(width, height, dst_format, 1)?;
        dst_frame.set_width(width);
        dst_frame.set_height(height);
        dst_frame.set_pix_fmt(dst_format);
        dst_frame.set_color_range(color_range);

        let dst_frame = Mutex::new(dst_frame);
        let converter = Mutex::new(self);

        #[allow(clippy::unwrap_used)]
        let result = panic::catch_unwind(|| {
            let mut dst_frame = dst_frame.lock().unwrap();
            converter.lock().unwrap().convert_internal(
                src_frame.data(),
                src_frame.linesize(),
                src_frame.height(),
                &mut dst_frame,
            );
        });
        if result.is_err() {
            return Err(ConvertError::Panic(format!(
                "width={width} height={height} src_format={src_format} dst_format={dst_format}",
            )));
        }
        Ok(())
    }

    fn convert_internal(
        &mut self,
        src_slice: &[Vec<u8>],
        src_stride: &[usize; 8],
        src_slice_h: NonZeroU16,
        dst_frame: &mut Frame,
    ) {
        // Copy strides, so they can safely be modified.
        let mut src2: [Vec<u8>; 4] = Default::default();
        for i in 0..4 {
            src2[i].extend_from_slice(&src_slice[i]);
        }
        //memcpy(src2, srcSlice, sizeof(src2));

        let mut src_stride2: [usize; 8] = src_stride.to_owned();

        let dst_stride = dst_frame.linesize_mut();
        let dst_stride2: [usize; 8] = *dst_stride;

        (self.convert)(
            self,
            &src2,
            &mut src_stride2,
            src_slice_h.get().into(),
            dst_frame.data_mut(),
            &dst_stride2[..],
        );
    }
}

fn is_yuv(pix_fmt: PixelFormat) -> bool {
    (pix_fmt.flags() & PIX_FMT_FLAG_RGB == 0) && pix_fmt.comps().len() >= 2
}

fn is_rgb(pix_fmt: PixelFormat) -> bool {
    pix_fmt.flags() & PIX_FMT_FLAG_RGB != 0
}

// Set c.convert_unscaled to an unscaled converter if one exists for the
// specific source and destination formats, bit depths, flags, etc.
pub(crate) fn get_convert_func(
    src_format: PixelFormat,
    dst_format: PixelFormat,
    height: NonZeroU16,
) -> Result<ConvertFunc, NewConverterError> {
    use NewConverterError::*;

    // yuv2rgb.
    if src_format == PixelFormat::YUV420P
        /*|| src_format == PixelFormat::YUV422P
        || src_format == PixelFormat::YUVA420P*/ && dst_format == PixelFormat::RGB24
    {
        if height.get() & 1 != 0 {
            return Err(NewConverterError::OddHeightUnsupported);
        }
        return yuv2rgb_get_func(dst_format).ok_or(UnsupportedPixelFormat(src_format, dst_format));
    }
    // yuv2gray.
    if src_format == PixelFormat::YUV420P && dst_format == PixelFormat::GRAY8 {
        return yuv2gray_get_func(dst_format).ok_or(UnsupportedPixelFormat(src_format, dst_format));
    }

    Err(UnsupportedPixelFormat(src_format, dst_format))
}

fn yuv2gray_get_func(dst_format: PixelFormat) -> Option<ConvertFunc> {
    match dst_format {
        PixelFormat::GRAY8 => Some(yuv2gray_8),
        _ => None,
    }
}

#[allow(clippy::unnecessary_wraps)]
pub(crate) fn yuv2gray_8(
    _c: &PixelFormatConverter,
    src: &[Vec<u8>],
    src_stride: &mut [usize],
    src_slice_h: i32,
    dst: &mut [Vec<u8>],
    dst_stride: &[usize],
) {
    dst[0].clear();

    let src_width = src_stride[0];
    let dst_width = dst_stride[0];

    let mut src: &[u8] = &src[0];
    for _ in 0..src_slice_h {
        dst[0].extend_from_slice(&src[..dst_width]);
        src = &src[src_width..];
    }
}

pub(crate) fn handle_jpeg(pix_fmt: PixelFormat) -> PixelFormat {
    match pix_fmt {
        PixelFormat::YUVJ420P => PixelFormat::YUV420P,
        // YUVJ411P.
        PixelFormat::YUVJ422P => PixelFormat::YUV422P,
        PixelFormat::YUVJ444P => PixelFormat::YUV444P,
        // YUVJ440P.
        /*PixelFormat::GRAY8
        | PixelFormat::YA8
        | PixelFormat::GRAY9LE
        | PixelFormat::GRAY9BE
        | PixelFormat::GRAY10LE
        | PixelFormat::GRAY10BE
        | PixelFormat::GRAY12LE
        | PixelFormat::GRAY12BE
        | PixelFormat::GRAY14LE
        | PixelFormat::GRAY14BE
        | PixelFormat::GRAY16LE
        | PixelFormat::GRAY16BE
        | PixelFormat::YA16BE
        | PixelFormat::YA16LE => 1,*/
        _ => pix_fmt,
    }
    /*switch (*format) {
    case AV_PIX_FMT_YUVJ420P:
      *format = AV_PIX_FMT_YUV420P;
      return 1;
    case AV_PIX_FMT_YUVJ411P:
      *format = AV_PIX_FMT_YUV411P;
      return 1;
    case AV_PIX_FMT_YUVJ422P:
      *format = AV_PIX_FMT_YUV422P;
      return 1;
    case AV_PIX_FMT_YUVJ444P:
      *format = AV_PIX_FMT_YUV444P;
      return 1;
    case AV_PIX_FMT_YUVJ440P:
      *format = AV_PIX_FMT_YUV440P;
      return 1;
    case AV_PIX_FMT_GRAY8:
    case AV_PIX_FMT_YA8:
    case AV_PIX_FMT_GRAY9LE:
    case AV_PIX_FMT_GRAY9BE:
    case AV_PIX_FMT_GRAY10LE:
    case AV_PIX_FMT_GRAY10BE:
    case AV_PIX_FMT_GRAY12LE:
    case AV_PIX_FMT_GRAY12BE:
    case AV_PIX_FMT_GRAY14LE:
    case AV_PIX_FMT_GRAY14BE:
    case AV_PIX_FMT_GRAY16LE:
    case AV_PIX_FMT_GRAY16BE:
    case AV_PIX_FMT_YA16BE:
    case AV_PIX_FMT_YA16LE:
      return 1;
    default:
      return 0;
    }*/
}