YUV <-> RGB

2025-09-27
四川
本文字数：3178 字
阅读完需：约 10 分钟

前面犯了一个小小的错误，我实现了了一个 HSV 到 RGB 算法，但是这是个基础算法，在实际的图像处理中，我们面临的是各种 YUV 的格式和 RGB 格式的转化，也就是 24bit 数据之间的转化，所以我们的转换算法应该是

    fn yuv2rgb(yuv) {        let y = ((yuv[0] - 16.0) * 255.0 / 219.0).clamp(0.0, 255.0);        let u = yuv[1] - 128.0;        let v = yuv[2] - 128.0;        let r = (y + 1.402 * v).clamp(0.0, 255.0).round();        let g = (y - 0.344136 * u - 0.714136 * v).clamp(0.0, 255.0).round();        let b = (y + 1.772 * u).clamp(0.0, 255.0).round();        let rgb = (r as u32) << 16 | (g as u32) << 8 | (b as u32);        return rgb;    }

复制代码

这是一个使用浮点数的 yuv 转换到 RGB 的算法

我们使用 spirv 的代码来实现几个内建的函数 clamp 和 round

            "clamp"=> {                let (ty, ty_id, s1, s2) = self.binary(BinaryOp::Null, &args[0], &args[1])?;//-> Result<(Type, u32, u32, u32)> {                let v3 = self.get(&args[2])?;                let s3 = self.get_var(v3);                let ids: Vec<rspirv::dr::Operand> = [s1, s2, s3.0].iter().map(|id| rspirv::dr::Operand::IdRef(*id) ).collect();                let id = self.builder.ext_inst(ty_id, None, self.glsl_std_450, spirv::GLOp::FClamp as u32, ids)?;                self.set(ret, Var{id, ty: SpirvType::Var(ty)});            }            "round"=> {                let v = self.get(&args[0])?;                let s = self.get_var(v);                let id = self.builder.ext_inst(s.2, None, self.glsl_std_450, spirv::GLOp::Round as u32, [rspirv::dr::Operand::IdRef(s.0)])?;                self.set(ret, Var{id, ty: SpirvType::Var(s.1)});            }

复制代码

其中 self.glsl_std_450 是标准的扩展库，需要在创建 builder 的时候加载

        let glsl_std_450 = builder.ext_inst_import("GLSL.std.450");

复制代码

我们使用给定的测试值来测试 YUV 转换函数

设定 main 函数这样

    pub fn main(param: Param) {        let rgb = yuv2rgb([81u8, 90u8, 240u8]);        let id = spirv_local();        let z = id[1] * 32 + id[0];        param.id[z] = rgb;    }"#;

复制代码

测试 [81,90,240]这一组，结果应该是 RGB 的 [233，9,8]

使用下面的函数调用

let spirv_code = b.assemble();        let mut rt = spirv::Runtime::default();        let buf = rt.run::<[u32; 1024], _>(&spirv_code, |buf| {} , [1, 1, 1])?;        println!("{:x?}", buf.read().unwrap()[0]);

复制代码

结果是这样

完美！

现在我们开始实际的转换

我们先使用 yuv 库将 rgb 图像转换为 YUV

    let img = image::open("demo.jpg")?.to_rgb8();    let (w, h) = (img.width(), img.height());    let mut yuv = yuv::YuvPlanarImageMut::alloc(w, h, yuv::YuvChromaSubsampling::Yuv444);    let start = std::time::Instant::now();    for _ in 0..1 {        yuv::rgb_to_yuv444(&mut yuv, img.as_raw(), w * 3, yuv::YuvRange::Full, yuv::YuvStandardMatrix::Bt601, yuv::YuvConversionMode::Balanced)?;    }    let elapsed = start.elapsed();

复制代码

YUV 数据分三个平面存放我们的图像大小是 2560 * 1440

我们的 Zeta 转换函数是这样的没有做任何优化

    fn yuv2rgb(y, u, v) {        let r = (y + 1.402f32 * (v - 128.0)).clamp(0.0, 255.0);        let g = (y - 0.344136f32 * (u - 128.0) - 0.714136f32 * (v - 128.0)).clamp(0.0, 255.0);        let b = (y + 1.772f32 * (u - 128.0)).clamp(0.0, 255.0);        let rgb = (r as u32) | (g as u32) << 8 | (b as u32) << 16 | 0xFF000000u32;        return rgb;    }

复制代码

转换回来的结果应该是 RGBA Alpha 固定设置为 0xFF

我们使用固定大小的缓冲区传递数据

pub struct Param {    y: [u8; 3686400],    u: [u8; 3686400],    v: [u8; 3686400],    rgb: [u32; 3686400],}

复制代码

这个 Zeta 和 Rust 可以共用

整个 Zeta Shader 代码是这样的

    pub struct Param {        y: [u8; 3686400],        u: [u8; 3686400],        v: [u8; 3686400],        rgb: [u32; 3686400],    }
    fn yuv2rgb(y, u, v) {        let r = (y + 1.402f32 * (v - 128.0)).clamp(0.0, 255.0);        let g = (y - 0.344136f32 * (u - 128.0) - 0.714136f32 * (v - 128.0)).clamp(0.0, 255.0);        let b = (y + 1.772f32 * (u - 128.0)).clamp(0.0, 255.0);        let rgb = (r as u32) | (g as u32) << 8 | (b as u32) << 16 | 0xFF000000u32;        return rgb;    }        pub fn main(param: Param) {        let g_id = spirv_group();        let id = spirv_local();        let x = g_id[0] * 32 + id[0];   // x 方向的坐标 分成 32 x 32 的小块        let y = g_id[1] * 32 + id[1];   // y 方向的坐标 分成 32 x 32 的小块        let offset = y * 2560 + x;        let rgb = yuv2rgb(param.y[offset], param.u[offset], param.v[offset]);        param.rgb[offset] = rgb;    }"#;

复制代码

每个组 32 x 32 一个组同时启动 1024 个线程我们

运行的代码是这样的

    let mut rt = spirv::Runtime::default();    let start = std::time::Instant::now();    let buf = rt.run::<Param, _>(spirv_code, |param| {        param.y.copy_from_slice(yuv.y_plane.borrow());        param.u.copy_from_slice(yuv.u_plane.borrow());        param.v.copy_from_slice(yuv.v_plane.borrow());        //param.rgb.copy_from_slice(&yuv.rgb());    } , [80, 44, 1])?;    let elapsed = start.elapsed();    println!("spirv use {}", elapsed.as_micros());

复制代码

我们启动 80 个 x 44 个 y 正好处理了 2560 x 1440 的数据

验证的记过完美！！

yuv 是高度优化的代码，我们执行 1000 次看看

    let start = std::time::Instant::now();    for _ in 0..1000 {        yuv::rgb_to_yuv444(&mut yuv, img.as_raw(), w * 3, yuv::YuvRange::Full, yuv::YuvStandardMatrix::Bt601, yuv::YuvConversionMode::Balanced)?;    }    let elapsed = start.elapsed();

复制代码

耗时

微秒，极度强大！！

但是我们的 spirv 版本更牛，不能用只能计算执行时间如果算上拷贝时间那就太不公平了

所以我们在 spirv Runtime 里面计算时间

        let start = std::time::Instant::now();        for _ in 0..1000 {            let mut builder = AutoCommandBufferBuilder::primary(command_buffer_allocator.clone(), self.queue.queue_family_index(), CommandBufferUsage::OneTimeSubmit).unwrap();
            builder.bind_pipeline_compute(pipeline.clone())?;            builder.bind_descriptor_sets(PipelineBindPoint::Compute, pipeline_layout.clone(), 0, set.clone())?;
            unsafe { builder.dispatch(dim) }.unwrap();            let command_buffer = builder.build().unwrap();            let future = sync::now(self.device.clone()).then_execute(self.queue.clone(), command_buffer).unwrap().then_signal_fence_and_flush().unwrap();            future.wait(None).unwrap();        }        let elapsed = start.elapsed();        println!("spirv use {}", elapsed.as_micros());