Skip to content

Commit

Permalink
mirror set_color_standard between rust and shader code
Browse files Browse the repository at this point in the history
  • Loading branch information
@YukumoHunter
YukumoHunter committed Jan 24, 2024
1 parent 5585f43 commit e9b75dc
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 45 deletions.
46 changes: 23 additions & 23 deletions crates/re_renderer/shader/decodings.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@ fn decode_nv12(texture: texture_2d<u32>, coords: vec2i) -> vec4f {
let uv_row = u32(coords.y / 2);
var uv_col = u32(coords.x / 2) * 2u;

let y = max(0.0, (f32(textureLoad(texture, vec2u(coords), 0).r) - 16.0)) / 219.0;
let u = (f32(textureLoad(texture, vec2u(u32(uv_col), uv_offset + uv_row), 0).r) - 128.0) / 224.0;
let v = (f32(textureLoad(texture, vec2u((u32(uv_col) + 1u), uv_offset + uv_row), 0).r) - 128.0) / 224.0;
let y = f32(textureLoad(texture, vec2u(coords), 0).r);
let u = f32(textureLoad(texture, vec2u(u32(uv_col), uv_offset + uv_row), 0).r);
let v = f32(textureLoad(texture, vec2u((u32(uv_col) + 1u), uv_offset + uv_row), 0).r);

let rgb = set_color_standard(vec3f(y, u, v));

Expand All @@ -19,18 +19,12 @@ fn decode_nv12(texture: texture_2d<u32>, coords: vec2i) -> vec4f {

/// Loads an RGBA texel from a texture holding an YUY2 encoded image at the given screen space coordinates.
fn decode_yuy2(texture: texture_2d<u32>, coords: vec2i) -> vec4f {
// texture is 2*width x height
// texture is 2 * width * height
// every 4 bytes is 2 pixels
let uv_row = u32(coords.y);
// multiply by 2 because the width is multiplied by 2
var uv_col = u32(coords.x) * 2u;

// YUYV
// y1 = 0
// u = 2
// y2 = 4
// v = 6

var y = 0.0;
if coords.x % 2 == 0 {
// we're on an even pixel, so we can sample the first y value
Expand All @@ -42,32 +36,38 @@ fn decode_yuy2(texture: texture_2d<u32>, coords: vec2i) -> vec4f {
// We subtract 2 from the column so that we can sample the u and v values
uv_col -= 2u;
}
let u = f32(textureLoad(texture, vec2u(uv_col + 1u, uv_row), 0).r);
let v = f32(textureLoad(texture, vec2u(uv_col + 3u, uv_row), 0).r);

y = (y - 16.0) / 219.0;
let u = (f32(textureLoad(texture, vec2u(uv_col + 1u, uv_row), 0).r) - 128.0) / 224.0;
let v = (f32(textureLoad(texture, vec2u(uv_col + 3u, uv_row), 0).r) - 128.0) / 224.0;
let rgb = set_color_standard(vec3f(y, u, v));

return vec4f(rgb, 1.0);
}

/// Sets the color standard for the given YUV color.
///
/// This conversion mirrors the function in `crates/re_types/src/datatypes/tensor_data_ext.rs`
///
/// Specifying the color standard should be exposed in the future (https://github.com/rerun-io/rerun/pull/3541)
fn set_color_standard(yuv: vec3f) -> vec3f {
let y = yuv.x;
let u = yuv.y;
let v = yuv.z;
// rescale YUV values
let y = (yuv.x - 16.0) / 219.0;
let u = (yuv.y - 128.0) / 224.0;
let v = (yuv.z - 128.0) / 224.0;

// BT.601 (aka. SDTV, aka. Rec.601). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion
let r = clamp(y + 1.402 * v, 0.0, 1.0);
let g = clamp(y - (0.344 * u + 0.714 * v), 0.0, 1.0);
let b = clamp(y + 1.772 * u, 0.0, 1.0);
let r = y + 1.402 * v;
let g = y - 0.344 * u - 0.714 * v;
let b = y + 1.772 * u;

// BT.709 (aka. HDTV, aka. Rec.709). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion
// let r = clamp(y + 1.5748 * v, 0.0, 1.0);
// let g = clamp(y + u * -0.1873 + v * -0.4681, 0.0, 1.0);
// let b = clamp(y + u * 1.8556, 0.0 , 1.0);
// let r = y + 1.575 * v;
// let g = y - 0.187 * u - 0.468 * v;
// let b = y + 1.856 * u;

return vec3f(r, g, b);
return vec3f(
clamp(r, 0.0, 1.0),
clamp(g, 0.0, 1.0),
clamp(b, 0.0, 1.0)
);
}
59 changes: 37 additions & 22 deletions crates/re_types/src/datatypes/tensor_data_ext.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -238,18 +238,16 @@ impl TensorData {
match self.image_height_width_channels() {
Some([h, w, _]) => {
let uv_offset = w * h;
let luma = ((buf[(y * w + x) as usize] as f64) - 16.0) / 216.0;
let u = ((buf[(uv_offset + (y / 2) * w + x) as usize] as f64) - 128.0) / 224.0;
let v =
((buf[((uv_offset + (y / 2) * w + x) as usize) + 1] as f64) - 128.0) / 224.0;
let r = luma + 1.402 * v;
let g = luma - 0.344 * u + 0.714 * v;
let b = luma + 1.772 * u;
let luma = buf[(y * w + x) as usize];
let u = buf[(uv_offset + (y / 2) * w + x) as usize];
let v = buf[(uv_offset + (y / 2) * w + x) as usize + 1];

let (r, g, b) = Self::set_color_standard(luma, u, v);

Some([
TensorElement::U8(f64::clamp(r * 255.0, 0.0, 255.0) as u8),
TensorElement::U8(f64::clamp(g * 255.0, 0.0, 255.0) as u8),
TensorElement::U8(f64::clamp(b * 255.0, 0.0, 255.0) as u8),
TensorElement::U8(r),
TensorElement::U8(g),
TensorElement::U8(b),
])
}
_ => None,
Expand All @@ -271,7 +269,7 @@ impl TensorData {
(buf[index], buf[index - 1], buf[index + 1])
};

let (r, g, b) = Self::yuv_to_rgb(luma, u, v);
let (r, g, b) = Self::set_color_standard(luma, u, v);

Some([
TensorElement::U8(r),
Expand All @@ -283,17 +281,34 @@ impl TensorData {
}
}

/// Convert YUV to RGB
fn yuv_to_rgb(y: u8, u: u8, v: u8) -> (u8, u8, u8) {
let y = y as f32 - 16.0;
let u = u as f32 - 128.0;
let v = v as f32 - 128.0;

let r = (y + 1.403 * v).clamp(0.0, 255.0) as u8;
let g = (y - 0.344 * u - 0.714 * v).clamp(0.0, 255.0) as u8;
let b = (y + 1.770 * u).clamp(0.0, 255.0) as u8;

(r, g, b)
/// Sets the color standard for the given YUV color.
///
/// This conversion mirrors the function of the same name in `crates/re_renderer/shader/decodings.wgsl`
///
/// Specifying the color standard should be exposed in the future (https://github.com/rerun-io/rerun/pull/3541)
fn set_color_standard(y: u8, u: u8, v: u8) -> (u8, u8, u8) {
let (y, u, v) = (y as f32, u as f32, v as f32);

// rescale YUV values
let y = (y - 16.0) / 219.0;
let u = (u - 128.0) / 224.0;
let v = (v - 128.0) / 224.0;

// BT.601 (aka. SDTV, aka. Rec.601). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion
let r = y + 1.402 * v;
let g = y - 0.344 * u - 0.714 * v;
let b = y + 1.772 * u;

// BT.709 (aka. HDTV, aka. Rec.709). wiki: https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.709_conversion
// let r = y + 1.575 * v;
// let g = y - 0.187 * u - 0.468 * v;
// let b = y + 1.856 * u;

(
(255.0 * r).clamp(0.0, 255.0) as u8,
(255.0 * g).clamp(0.0, 255.0) as u8,
(255.0 * b).clamp(0.0, 255.0) as u8,
)
}

#[inline]
Expand Down

0 comments on commit e9b75dc

Please sign in to comment.