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
//! Simple RGB to YUV420 converter
//! (full swing)
//! (only supports 8 bit RGB color depth)

/// Converts an RGB image to YUV420p (planar/3 planes)
///
/// # Arguments
///
/// * `img` - should contain the pixel data in the following format:
/// `[r, g, b, ... , r, g, b, ... , r, g, b, ...]`
///
/// * `bytes_per_pixel` - should contain the number of bytes used by one pixel
/// (eg.: RGB is 3 bytes and RGBA is 4 bytes)
///
/// # Return
///
/// `[y, y, y, ... , u, u, u, ... , v, v, v, ...]`
///
/// # Examples
///
/// ```
/// let rgb = vec![0u8; 12];
/// let yuv = rgb2yuv420::convert_rgb_to_yuv420p(&rgb, 2, 2, 3);
/// assert_eq!(yuv.len(), rgb.len() / 2);
/// ```
pub fn convert_rgb_to_yuv420p(img: &[u8], width: u32, height: u32, bytes_per_pixel: usize) -> Vec<u8> {
    convert_rgb_to_yuv420(img, width, height, bytes_per_pixel, |yuv, uv_index, chroma_size, u, v| {
        yuv[*uv_index] = u;
        yuv[*uv_index + (f32::ceil(chroma_size as f32 / 2.0) as usize)] = v;
        *uv_index += 1;
    })
}

/// Converts an RGB image to YUV420sp NV12 (semi-planar/2 planes)
///
/// # Arguments
///
/// * `img` - should contain the pixel data in the following format:
/// [r, g, b, ... , r, g, b, ... , r, g, b, ...]
///
/// * `bytes_per_pixel` - should contain the number of bytes used by one pixel
/// (eg.: RGB is 3 bytes and RGBA is 4 bytes)
///
/// # Return
///
/// `[y, y, y, ... , u, v, u, v, ...]`
///
/// # Examples
///
/// ```
/// let rgb = vec![0u8; 12];
/// let yuv = rgb2yuv420::convert_rgb_to_yuv420sp_nv12(&rgb, 2, 2, 3);
/// assert_eq!(yuv.len(), rgb.len() / 2);
/// ```
pub fn convert_rgb_to_yuv420sp_nv12(img: &[u8], width: u32, height: u32, bytes_per_pixel: usize) -> Vec<u8> {
    convert_rgb_to_yuv420(img, width, height, bytes_per_pixel, |yuv, uv_index, _cs, u, v| {
        yuv[*uv_index] = u;
        *uv_index += 1;
        yuv[*uv_index] = v;
        *uv_index += 1;
    })
}

fn convert_rgb_to_yuv420<T>(img: &[u8], width: u32, height: u32, bytes_per_pixel: usize, store_uv: T) -> Vec<u8>
    where T: Fn(&mut Vec<u8>, &mut usize, usize, u8, u8) -> () {
    let frame_size: usize = (width * height) as usize;
    let chroma_size: usize = frame_size / 4;
    let mut y_index: usize = 0;
    let mut uv_index = frame_size;
    let mut yuv = vec![0; (width * height * 3 / 2) as usize];
    let mut index: usize = 0;
    for j in 0..height {
        for _ in 0..width {
            let r = i32::from(img[index * bytes_per_pixel]);
            let g = i32::from(img[index * bytes_per_pixel + 1]);
            let b = i32::from(img[index * bytes_per_pixel + 2]);
            index += 1;
            yuv[y_index] = clamp((77 * r + 150 * g + 29 * b + 128) >> 8);
            y_index += 1;
            if j % 2 == 0 && index % 2 == 0 {
                store_uv(&mut yuv,
                         &mut uv_index,
                         chroma_size,
                         clamp(((-43 * r - 84 * g + 127 * b + 128) >> 8) + 128),
                         clamp(((127 * r - 106 * g - 21 * b + 128) >> 8) + 128));
            }
        }
    }
    yuv
}

fn clamp(val: i32) -> u8 {
    match val {
        ref v if *v < 0 => 0,
        ref v if *v > 255 => 255,
        v => v as u8,
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn rgb_to_yuv() {
        use super::convert_rgb_to_yuv420p;
        let rgb = vec![0u8; 12];
        let expected = vec![0u8, 0u8, 0u8, 0u8, 128u8, 128u8];
        let yuv = convert_rgb_to_yuv420p(&rgb, 2, 2, 3);
        assert_eq!(yuv.len(), rgb.len() / 2);
        for (val, exp) in yuv.iter().zip(expected.iter()) {
            assert_eq!(val, exp);
        }
    }

    #[test]
    fn rgba_to_yuv_from_file() {
        extern crate png;
        use std::fs::File;
        use super::convert_rgb_to_yuv420p;
        let decoder = png::Decoder::new(File::open("pic/ferris.png").unwrap());
        let (info, mut reader) = decoder.read_info().unwrap();
        let mut buf = vec![0; info.buffer_size()];
        reader.next_frame(&mut buf).unwrap();
        let yuv = convert_rgb_to_yuv420p(&buf, info.width, info.height, info.line_size / info.width as usize);
        assert_eq!(yuv.len(), buf.len() / 4 * 3 / 2);
    }
}