Skip to content

Commit

Permalink
feat(client): ✨ Add RGB-based chroma keying
Browse files Browse the repository at this point in the history
  • Loading branch information
@zmerp
zmerp committed Feb 13, 2025
1 parent 539b977 commit 2289874
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 102 deletions.
6 changes: 5 additions & 1 deletion alvr/client_openxr/src/stream.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -423,7 +423,11 @@ impl StreamContext {
.map(|mode| ProjectionLayerAlphaConfig {
premultiplied: matches!(
mode,
PassthroughMode::AugmentedReality { .. } | PassthroughMode::ChromaKey(_)
PassthroughMode::Blend {
premultiplied_alpha: true,
..
} | PassthroughMode::RgbChromaKey(_)
| PassthroughMode::HsvChromaKey(_)
),
}),
);
Expand Down
36 changes: 21 additions & 15 deletions alvr/graphics/resources/stream.wgsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,12 @@ override C_RIGHT_Y: f32 = 0.0;
struct PushConstant {
reprojection_transform: mat4x4f,
view_idx: u32,
alpha: f32,
enable_chroma_key: u32,
passthrough_mode: u32, // 0: No passthrough, 1: Blend, 2: RGB chroma key, 3: HSV chroma key
blend_alpha: f32,
_align: u32,
ck_hue: vec4f,
ck_saturation: vec4f,
ck_value: vec4f,
ck_channel0: vec4f,
ck_channel1: vec4f,
ck_channel2: vec4f,
}
var<push_constant> pc: PushConstant;

Expand Down Expand Up @@ -131,9 +131,15 @@ fn fragment_main(@location(0) uv: vec2f) -> @location(0) vec4f {
color = enc_condition * enc_lowValues + (1.0 - enc_condition) * enc_highValues;
}

var alpha = pc.alpha;
if pc.enable_chroma_key == 1 {
let mask = chroma_key_mask(rgb_to_hsv(color));
var alpha = 1.0;
if pc.passthrough_mode == 1 { // Blend mode
alpha = pc.blend_alpha;
} if pc.passthrough_mode >= 2 { // Chroma key
var current = color;
if pc.passthrough_mode == 3 { // HSV mode
current = rgb_to_hsv(color);
}
let mask = chroma_key_mask(current);

// Note: because of this calculation, we require premultiplied alpha option in the XR layer
color = max(color * mask, vec3f(0.0));
Expand All @@ -143,14 +149,14 @@ fn fragment_main(@location(0) uv: vec2f) -> @location(0) vec4f {
return vec4f(color, alpha);
}

fn chroma_key_mask(hsv: vec3f) -> f32 {
let start_max = vec3f(pc.ck_hue.x, pc.ck_saturation.x, pc.ck_value.x);
let start_min = vec3f(pc.ck_hue.y, pc.ck_saturation.y, pc.ck_value.y);
let end_min = vec3f(pc.ck_hue.z, pc.ck_saturation.z, pc.ck_value.z);
let end_max = vec3f(pc.ck_hue.w, pc.ck_saturation.w, pc.ck_value.w);
fn chroma_key_mask(color: vec3f) -> f32 {
let start_max = vec3f(pc.ck_channel0.x, pc.ck_channel1.x, pc.ck_channel2.x);
let start_min = vec3f(pc.ck_channel0.y, pc.ck_channel1.y, pc.ck_channel2.y);
let end_min = vec3f(pc.ck_channel0.z, pc.ck_channel1.z, pc.ck_channel2.z);
let end_max = vec3f(pc.ck_channel0.w, pc.ck_channel1.w, pc.ck_channel2.w);

let start_mask = smoothstep(start_min, start_max, hsv);
let end_mask = smoothstep(end_min, end_max, hsv);
let start_mask = smoothstep(start_min, start_max, color);
let end_mask = smoothstep(end_min, end_max, color);

return max(start_mask.x, max(start_mask.y, max(start_mask.z, max(end_mask.x, max(end_mask.y, end_mask.z)))));
}
Expand Down
157 changes: 85 additions & 72 deletions alvr/graphics/src/stream.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,12 @@ const TRANSFORM_SIZE: u32 = mem::size_of::<Mat4>() as u32;

const TRANSFORM_CONST_OFFSET: u32 = 0;
const VIEW_INDEX_CONST_OFFSET: u32 = TRANSFORM_SIZE;
const ALPHA_CONST_OFFSET: u32 = VIEW_INDEX_CONST_OFFSET + U32_SIZE;
const ENABLE_CHROMA_KEY_CONST_OFFSET: u32 = ALPHA_CONST_OFFSET + FLOAT_SIZE;
const CK_HUE_CONST_OFFSET: u32 = ENABLE_CHROMA_KEY_CONST_OFFSET + U32_SIZE + ALIGN4_SIZE;
const CK_SATURATION_CONST_OFFSET: u32 = CK_HUE_CONST_OFFSET + VEC4_SIZE;
const CK_VALUE_CONST_OFFSET: u32 = CK_SATURATION_CONST_OFFSET + VEC4_SIZE;
const PUSH_CONSTANTS_SIZE: u32 = CK_VALUE_CONST_OFFSET + VEC4_SIZE;
const PASSTHROUGH_MODE: u32 = VIEW_INDEX_CONST_OFFSET + U32_SIZE;
const ALPHA_CONST_OFFSET: u32 = PASSTHROUGH_MODE + U32_SIZE;
const CK_CHANNEL0_CONST_OFFSET: u32 = ALPHA_CONST_OFFSET + FLOAT_SIZE + ALIGN4_SIZE;
const CK_CHANNEL1_CONST_OFFSET: u32 = CK_CHANNEL0_CONST_OFFSET + VEC4_SIZE;
const CK_CHANNEL2_CONST_OFFSET: u32 = CK_CHANNEL1_CONST_OFFSET + VEC4_SIZE;
const PUSH_CONSTANTS_SIZE: u32 = CK_CHANNEL2_CONST_OFFSET + VEC4_SIZE;

const _: () = assert!(
PUSH_CONSTANTS_SIZE <= MAX_PUSH_CONSTANTS_SIZE,
Expand Down Expand Up @@ -304,89 +304,102 @@ impl StreamRenderer {
fn set_passthrough_push_constants(render_pass: &mut RenderPass, config: Option<&PassthroughMode>) {
const DEG_TO_NORM: f32 = 1. / 360.;

fn set_u32(render_pass: &mut RenderPass, offset: u32, value: u32) {
render_pass.set_push_constants(ShaderStages::VERTEX_FRAGMENT, offset, &value.to_le_bytes());
}

fn set_float(render_pass: &mut RenderPass, offset: u32, value: f32) {
render_pass.set_push_constants(ShaderStages::VERTEX_FRAGMENT, offset, &value.to_le_bytes());
}

fn set_vec4(render_pass: &mut RenderPass, offset: u32, value: Vec4) {
render_pass.set_push_constants(
ShaderStages::VERTEX_FRAGMENT,
offset,
&value.x.to_le_bytes(),
);
render_pass.set_push_constants(
ShaderStages::VERTEX_FRAGMENT,
offset + FLOAT_SIZE,
&value.y.to_le_bytes(),
);
render_pass.set_push_constants(
ShaderStages::VERTEX_FRAGMENT,
offset + 2 * FLOAT_SIZE,
&value.z.to_le_bytes(),
);
render_pass.set_push_constants(
ShaderStages::VERTEX_FRAGMENT,
offset + 3 * FLOAT_SIZE,
&value.w.to_le_bytes(),
);
}

match config {
Some(PassthroughMode::AugmentedReality { brightness }) => {
set_float(render_pass, ALPHA_CONST_OFFSET, 1. - brightness);
set_float(render_pass, ENABLE_CHROMA_KEY_CONST_OFFSET, 0.);
None => {
set_u32(render_pass, PASSTHROUGH_MODE, 0);
}
Some(PassthroughMode::Blend { opacity }) => {
set_float(render_pass, ALPHA_CONST_OFFSET, 1. - opacity);
set_float(render_pass, ENABLE_CHROMA_KEY_CONST_OFFSET, 0.);
Some(PassthroughMode::Blend { threshold, .. }) => {
set_u32(render_pass, PASSTHROUGH_MODE, 1);
set_float(render_pass, ALPHA_CONST_OFFSET, 1. - threshold);
}
Some(PassthroughMode::ChromaKey(config)) => {
render_pass.set_push_constants(
ShaderStages::VERTEX_FRAGMENT,
ENABLE_CHROMA_KEY_CONST_OFFSET,
&1_u32.to_le_bytes(),
);
Some(PassthroughMode::RgbChromaKey(config)) => {
set_u32(render_pass, PASSTHROUGH_MODE, 2);

set_float(
render_pass,
CK_HUE_CONST_OFFSET,
config.hue_start_max_deg * DEG_TO_NORM,
);
set_float(
render_pass,
CK_HUE_CONST_OFFSET + FLOAT_SIZE,
config.hue_start_min_deg * DEG_TO_NORM,
);
set_float(
render_pass,
CK_HUE_CONST_OFFSET + 2 * FLOAT_SIZE,
config.hue_end_min_deg * DEG_TO_NORM,
);
set_float(
render_pass,
CK_HUE_CONST_OFFSET + 3 * FLOAT_SIZE,
config.hue_end_max_deg * DEG_TO_NORM,
);
let norm = |v| v as f32 / 255.;

set_float(
render_pass,
CK_SATURATION_CONST_OFFSET,
config.saturation_start_max,
);
set_float(
render_pass,
CK_SATURATION_CONST_OFFSET + FLOAT_SIZE,
config.saturation_start_min,
);
set_float(
render_pass,
CK_SATURATION_CONST_OFFSET + 2 * FLOAT_SIZE,
config.saturation_end_min,
);
set_float(
render_pass,
CK_SATURATION_CONST_OFFSET + 3 * FLOAT_SIZE,
config.saturation_end_max,
);
let red = norm(config.red);
let green = norm(config.green);
let blue = norm(config.blue);

set_float(render_pass, CK_VALUE_CONST_OFFSET, config.value_start_max);
set_float(
let thresh = norm(config.distance_threshold);

let up_feather = 1. + config.feathering;
let down_feather = 1. - config.feathering;

let range_vec =
thresh * Vec4::new(-up_feather, -down_feather, down_feather, up_feather);

set_vec4(render_pass, CK_CHANNEL0_CONST_OFFSET, red + range_vec);
set_vec4(render_pass, CK_CHANNEL1_CONST_OFFSET, green + range_vec);
set_vec4(render_pass, CK_CHANNEL2_CONST_OFFSET, blue + range_vec);
}
Some(PassthroughMode::HsvChromaKey(config)) => {
set_u32(render_pass, PASSTHROUGH_MODE, 3);

set_vec4(
render_pass,
CK_VALUE_CONST_OFFSET + FLOAT_SIZE,
config.value_start_min,
CK_CHANNEL0_CONST_OFFSET,
Vec4::new(
config.hue_start_max_deg,
config.hue_start_min_deg,
config.hue_end_min_deg,
config.hue_end_max_deg,
) * DEG_TO_NORM,
);
set_float(

set_vec4(
render_pass,
CK_VALUE_CONST_OFFSET + 2 * FLOAT_SIZE,
config.value_end_min,
CK_CHANNEL1_CONST_OFFSET,
Vec4::new(
config.saturation_start_max,
config.saturation_start_min,
config.saturation_end_min,
config.saturation_end_max,
),
);
set_float(

set_vec4(
render_pass,
CK_VALUE_CONST_OFFSET + 3 * FLOAT_SIZE,
config.value_end_max,
CK_CHANNEL2_CONST_OFFSET,
Vec4::new(
config.value_start_max,
config.value_start_min,
config.value_end_min,
config.value_end_max,
),
);
}
None => {
set_float(render_pass, ALPHA_CONST_OFFSET, 1.0);
set_float(render_pass, ENABLE_CHROMA_KEY_CONST_OFFSET, 0.);
}
}
}

Expand Down
66 changes: 52 additions & 14 deletions alvr/session/src/settings.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -557,7 +557,31 @@ pub enum H264Profile {
}

#[derive(SettingsSchema, Serialize, Deserialize, Clone, PartialEq, Debug)]
pub struct ChromaKeyConfig {
pub struct RgbChromaKeyConfig {
#[schema(flag = "real-time")]
#[schema(gui(slider(min = 0, max = 255)))]
pub red: u8,

#[schema(flag = "real-time")]
#[schema(gui(slider(min = 0, max = 255)))]
pub green: u8,

#[schema(flag = "real-time")]
#[schema(gui(slider(min = 0, max = 255)))]
pub blue: u8,

#[schema(strings(help = "The threshold is applied per-channel"))]
#[schema(flag = "real-time")]
#[schema(gui(slider(min = 1, max = 255)))]
pub distance_threshold: u8,

#[schema(flag = "real-time")]
#[schema(gui(slider(min = 0.01, max = 1.0, step = 0.01)))]
pub feathering: f32,
}

#[derive(SettingsSchema, Serialize, Deserialize, Clone, PartialEq, Debug)]
pub struct HsvChromaKeyConfig {
#[schema(strings(display_name = "Hue start max"), suffix = "°")]
#[schema(flag = "real-time")]
#[schema(gui(slider(min = -179.0, max = 539.0, step = 1.0)))]
Expand Down Expand Up @@ -614,23 +638,28 @@ pub struct ChromaKeyConfig {
#[derive(SettingsSchema, Serialize, Deserialize, Clone, PartialEq, Debug)]
#[schema(gui = "button_group")]
pub enum PassthroughMode {
AugmentedReality {
#[schema(flag = "real-time")]
#[schema(gui(slider(min = 0.0, max = 1.0, step = 0.01)))]
brightness: f32,
},
Blend {
#[schema(strings(
help = "Enabling this will adapt transparency based on the brightness of each pixel.
This is a similar effect to AR glasses."
))]
#[schema(flag = "real-time")]
premultiplied_alpha: bool,

#[schema(flag = "real-time")]
#[schema(gui(slider(min = 0.0, max = 1.0, step = 0.01)))]
opacity: f32,
threshold: f32,
},
ChromaKey(#[schema(flag = "real-time")] ChromaKeyConfig),

#[schema(strings(display_name = "RGB Chroma Key"))]
RgbChromaKey(#[schema(flag = "real-time")] RgbChromaKeyConfig),

#[schema(strings(display_name = "HSV Chroma Key"))]
HsvChromaKey(#[schema(flag = "real-time")] HsvChromaKeyConfig),
}

#[derive(SettingsSchema, Serialize, Deserialize, Clone)]
pub struct VideoConfig {
#[schema(strings(help = r"Augmented reality: corresponds to premultiplied alpha
Blend: corresponds to un-premultiplied alpha"))]
#[schema(flag = "real-time")]
pub passthrough: Switch<PassthroughMode>,

Expand Down Expand Up @@ -1486,10 +1515,19 @@ pub fn session_settings_default() -> SettingsDefault {
passthrough: SwitchDefault {
enabled: false,
content: PassthroughModeDefault {
variant: PassthroughModeDefaultVariant::AugmentedReality,
AugmentedReality: PassthroughModeAugmentedRealityDefault { brightness: 0.4 },
Blend: PassthroughModeBlendDefault { opacity: 0.5 },
ChromaKey: ChromaKeyConfigDefault {
variant: PassthroughModeDefaultVariant::Blend,
Blend: PassthroughModeBlendDefault {
premultiplied_alpha: true,
threshold: 0.5,
},
RgbChromaKey: RgbChromaKeyConfigDefault {
red: 0,
green: 255,
blue: 0,
distance_threshold: 85,
feathering: 0.05,
},
HsvChromaKey: HsvChromaKeyConfigDefault {
hue_start_max_deg: 70.0,
hue_start_min_deg: 80.0,
hue_end_min_deg: 160.0,
Expand Down

0 comments on commit 2289874

Please sign in to comment.