使用Rust转换bayer和yuv图像

这是我学习rust之后写的第一个程序，所以记录下，rust和c代码的使用方式还有有非常大的不同，我们不应该用for index的方式来写rust代码。反之使用map/reduce这类函数式的写法会更好一些。不论是性能还是可读性上都更好一些

这里的raw图是2byte一个像素，低12bit有效。raw图前面有128byte的header，我暂时先直接将它略过了，没去解析。由于只是临时写一个工具。所以没有对error进行很好的处理，一路unwrap()

Rust的零开销iter抽象还是相当好用的，另外强推下vscode + rust-analyzer作为开发环境，目前看起来是最棒的组合了

use std::{
    fs::File,
    io::{Cursor, Read, Seek, SeekFrom},
    vec,
};

use bayer::{BayerDepth, Demosaic, RasterDepth, RasterMut};
use image::RgbImage;

fn main() {
    let (w, h) = (1920, 1080);
    let mut file = File::open("1.raw").unwrap();
    let mut buf = Vec::new();
    file.seek(SeekFrom::Start(128)).unwrap();
    file.read_to_end(&mut buf).unwrap();

    let raw_vec: Vec<_> = buf.chunks_exact(2).map(|v| v[0] >> 4 | v[1] << 4).collect();

    let mut rgb8_vec = vec![0u8; w * h * 3];
    let mut dst = RasterMut::new(w, h, RasterDepth::Depth8, &mut rgb8_vec);
    bayer::run_demosaic(
        &mut Cursor::new(raw_vec.as_slice()),
        BayerDepth::Depth8,
        bayer::CFA::GRBG,
        Demosaic::Cubic,
        &mut dst,
    )
    .unwrap();
    RgbImage::from_vec(w as u32, h as u32, rgb8_vec)
        .unwrap()
        .save("rgb8.png")
        .unwrap();
}
}

对于YCbCr的图像, 也可以采用类似的方式来处理。我实际测试下来，下面的写法还是比纯C代码要慢的(慢了50%左右了)。可能是iterator做了实际的line copy带来的开销吧, 不过相应的这个写法的可读性比c循环要高很多。不过这个速度比之前用numpy写的哪个也要快非常多了。不知道除了使用unsafe模式外，是否还有更快的实现方式

use std::{
    fs::File,
    io::{Read, Seek, SeekFrom},
};

use image::RgbImage;

fn read_ycbcr(filename: &str, capacity: usize) -> Vec<u8> {
    let mut buf = Vec::with_capacity(capacity);
    let mut file = File::open(filename).unwrap();
    file.seek(SeekFrom::Start(0)).unwrap();
    file.read_to_end(&mut buf).unwrap();
    buf
}

fn write_rgb(filename: &str, buf: Vec<u8>, w: usize, h: usize) {
    let rgb = RgbImage::from_vec(w as u32, h as u32, buf).unwrap();
    rgb.save(filename).unwrap();
}

fn csc(
    y: impl Iterator<Item = u8>,
    u: impl Iterator<Item = u8>,
    v: impl Iterator<Item = u8>,
) -> Vec<u8> {
    y.zip(u.zip(v))
        .flat_map(|(y, (u, v))| {
            [
                (y as f32 + 1.280 * (v as f32 - 128.)) as u8,
                (y as f32 - 0.215 * (u as f32 - 128.) - 0.380 * (v as f32 - 128.)) as u8,
                (y as f32 + 2.128 * (u as f32 - 128.)) as u8,
            ]
        })
        .collect()
}

pub fn nv12_to_rgb(nv12: &str, rgb: &str, w: usize, h: usize) {
    let buf = read_ycbcr(nv12, w * h * 3 / 2);
    let y = buf[..w * h].iter().copied();
    let u = buf[w * h..]
        .chunks_exact(w)
        .flat_map(|line| [line, line])
        .flatten()
        .step_by(2)
        .flat_map(|u| [u, u])
        .copied();
    let v = buf[w * h..]
        .chunks_exact(w)
        .flat_map(|line| [line, line])
        .flatten()
        .skip(1)
        .step_by(2)
        .flat_map(|v| [v, v])
        .copied();

    write_rgb(rgb, csc(y, u, v), w, h);
}

pub fn nv16_to_rgb(nv16: &str, rgb: &str, w: usize, h: usize) {
    let buf = read_ycbcr(nv16, w * h * 2);
    let y = buf[..w * h].iter().copied();
    let u = buf[w * h..].iter().step_by(2).flat_map(|u| [u, u]).copied();
    let v = buf[w * h + 1..]
        .iter()
        .step_by(2)
        .flat_map(|v| [v, v])
        .copied();

    write_rgb(rgb, csc(y, u, v), w, h);
}

pub fn yuyv_to_rgb(yuyv: &str, rgb: &str, w: usize, h: usize) {
    let buf = read_ycbcr(yuyv, w * h * 2);
    let y = buf.iter().step_by(2).copied();
    let u = buf[1..].iter().step_by(4).flat_map(|v| [v, v]).copied();
    let v = buf[3..].iter().step_by(4).flat_map(|v| [v, v]).copied();

    write_rgb(rgb, csc(y, u, v), w, h);
}

fn main() {
    let (w, h) = (2560, 1440);

    nv12_to_rgb("nv12.yuv", "nv12.jpg", w, h);
    nv16_to_rgb("nv16.yuv", "nv16.jpg", w, h);
    yuyv_to_rgb("yuyv.yuv", "yuyv.jpg", w, h);
}