Merge pull request #398 from salmanmkc/feat/world-segmentation

world: add a segmentation subsystem (person/background masks)

Author: ruofeidu

demos/magic_window/MagicWindow.js

@@ -1,8 +1,6 @@
 import * as THREE from 'three';
 import * as xb from 'xrblocks';
 
-import {SegmenterController} from './SegmenterController.js';
-
 // Backdrop modes for the cut-out background.
 //   0 = passthrough: background pixels are discarded so you see straight
 //       through the window to the real world behind it.
@@ -77,9 +75,10 @@ const FRAGMENT_SHADER = /* glsl */ `
     // The camera texture uses GL flipY; the mask DataTexture does not, so
     // flip the mask's v to line the two up.
     float id = texture2D(uMask, vec2(vUv.x, 1.0 - vUv.y)).r * 255.0;
-    // Until the first mask arrives, show the raw feed. Category 0 is the
-    // background; everything else is a person.
-    bool isPerson = (uHasMask < 0.5) || (id >= 0.5);
+    // Until the first mask arrives, treat the whole frame as background so we
+    // show the backdrop rather than flashing the raw camera (the room) during
+    // model warm-up. Category 0 is background; everything else is a person.
+    bool isPerson = (uHasMask > 0.5) && (id > 0.5);
     if (isPerson) {
       gl_FragColor = vec4(cam, 1.0);
       return;
@@ -95,7 +94,6 @@ const FRAGMENT_SHADER = /* glsl */ `
 export class MagicWindow extends xb.Script {
   constructor() {
     super();
-    this.segmenter = new SegmenterController();
     this.frameCanvas = document.createElement('canvas');
     this.frameCtx = this.frameCanvas.getContext('2d', {
       willReadFrequently: true,
@@ -145,8 +143,6 @@ export class MagicWindow extends xb.Script {
     this.plane.position.set(0, 1.5, -1.2);
     this.add(this.plane);
 
-    this.segmenter.load();
-
     // Quick keyboard control until the spatial panel lands: B cycles backdrop.
     this.onKeyDown_ = (event) => {
       if (event.key === 'b' || event.key === 'B') {
@@ -250,10 +246,14 @@ export class MagicWindow extends xb.Script {
   }
 
   updateMask_() {
-    if (!this.segmenter.isReady) {
+    // The Segmenter's own update() loop keeps latestMask fresh at its
+    // configured cadence. Reading it here is a synchronous poll — no
+    // MediaPipe inference is triggered by this call.
+    const segmentation = xb.core.world?.segmentation;
+    if (!segmentation) {
       return;
     }
-    const mask = this.segmenter.segment(this.frameCanvas);
+    const mask = segmentation.latestMask;
     if (!mask) {
       return;
     }
@@ -283,7 +283,6 @@ export class MagicWindow extends xb.Script {
 
   dispose() {
     window.removeEventListener('keydown', this.onKeyDown_);
-    this.segmenter.dispose();
     this.cameraTexture?.dispose();
     this.maskTexture?.dispose();
     this.plane.geometry.dispose();

demos/magic_window/SegmenterController.js

@@ -4,10 +4,11 @@ import {ControlPanel} from './ControlPanel.js';
 import {MagicWindow} from './MagicWindow.js';
 
 const options = new xb.Options();
-// The device camera feeds the segmenter, so we want a person in frame. The
-// user-facing mode pulls the real webcam through the camera module (on desktop
-// `facingMode: 'user'` skips the simulator camera and goes straight to
-// getUserMedia), giving an actual person to cut out.
+// Enable the SDK's segmentation primitive (MediaPipe person/background masks).
+// It turns the camera on for us; we then override to the user-facing camera so
+// there's a person to cut out (on desktop `facingMode: 'user'` skips the
+// simulator camera and goes straight to getUserMedia).
+options.enableSegmentation();
 options.enableCamera('user');
 options.setAppTitle('Magic Window');
 options.setAppDescription(

demos/magic_window/main.js

@@ -345,6 +345,19 @@ export class Options {
     return this;
   }
 
+  /**
+   * Enables semantic segmentation. Produces per-pixel person / background
+   * category masks from the device camera (MediaPipe, on-device). Unlike face
+   * and human detection it does not require depth.
+   * @returns The instance for chaining.
+   */
+  enableSegmentation() {
+    this.permissions.camera = true;
+    this.enableCamera();
+    this.world.enableSegmentation();
+    return this;
+  }
+
   /**
    * Enables device camera (passthrough) with a specific facing mode.
    * @param facingMode - The desired camera facing mode, either 'environment' or

src/core/Options.ts

@@ -15,6 +15,7 @@ import {placeOnHorizontalSurface} from './HorizontalPlacement';
 // import { LightEstimation } from '/lighting/LightEstimation.js';
 import {HumanRecognizer} from './humans/HumanRecognizer';
 import {FaceRecognizer} from './faces/FaceRecognizer';
+import {Segmenter} from './segmentation/Segmenter';
 
 /**
  * Manages all interactions with the real-world environment perceived by the XR
@@ -78,6 +79,11 @@ export class World extends Script {
    */
   faces?: FaceRecognizer;
 
+  /**
+   * The semantic segmentation module instance. Null if not enabled.
+   */
+  segmentation?: Segmenter;
+
   /**
    * A Three.js Raycaster for performing intersection tests.
    */
@@ -153,6 +159,11 @@ export class World extends Script {
       this.add(this.faces);
     }
 
+    if (this.options.segmentation.enabled) {
+      this.segmentation = new Segmenter();
+      this.add(this.segmentation);
+    }
+
     // TODO: Initialize other modules as they are available & implemented.
     /*
 

src/world/World.ts

@@ -7,6 +7,7 @@ import {PlanesOptions} from './planes/PlanesOptions';
 import {SoundsOptions} from './sounds/SoundsOptions';
 import {HumansOptions} from './humans/HumansOptions';
 import {FacesOptions} from './faces/FacesOptions';
+import {SegmentationOptions} from './segmentation/SegmentationOptions';
 
 export class WorldOptions {
   debugging = false;
@@ -18,6 +19,7 @@ export class WorldOptions {
   sounds = new SoundsOptions();
   humans = new HumansOptions();
   faces = new FacesOptions();
+  segmentation = new SegmentationOptions();
 
   constructor(options?: DeepPartial<WorldOptions>) {
     if (options) {
@@ -78,4 +80,13 @@ export class WorldOptions {
     this.faces.enable();
     return this;
   }
+
+  /**
+   * Enables semantic segmentation (person / background category masks).
+   */
+  enableSegmentation() {
+    this.enabled = true;
+    this.segmentation.enable();
+    return this;
+  }
 }

src/world/WorldOptions.ts

@@ -0,0 +1,22 @@
+import {describe, it, expect} from 'vitest';
+
+import {SegmentCategory} from './SegmentationMask';
+
+describe('SegmentCategory', () => {
+  it('matches the selfie multiclass model index ordering', () => {
+    // These indices are dictated by the model output, do not renumber.
+    expect(SegmentCategory.Background).toBe(0);
+    expect(SegmentCategory.Hair).toBe(1);
+    expect(SegmentCategory.BodySkin).toBe(2);
+    expect(SegmentCategory.FaceSkin).toBe(3);
+    expect(SegmentCategory.Clothes).toBe(4);
+    expect(SegmentCategory.Others).toBe(5);
+  });
+
+  it('treats only background as non-person (index 0)', () => {
+    // Convention used by consumers: anything >= 1 is part of a person.
+    expect(SegmentCategory.Background).toBe(0);
+    expect(SegmentCategory.Hair).toBeGreaterThanOrEqual(1);
+    expect(SegmentCategory.Others).toBeGreaterThanOrEqual(1);
+  });
+});

src/world/segmentation/SegmentationMask.test.ts

@@ -0,0 +1,26 @@
+/**
+ * Per-pixel semantic categories emitted by the selfie multiclass segmentation
+ * model. Index `0` is the background; every other index is part of a person,
+ * so anything `>= 1` can be treated as foreground.
+ */
+export enum SegmentCategory {
+  Background = 0,
+  Hair = 1,
+  BodySkin = 2,
+  FaceSkin = 3,
+  Clothes = 4,
+  Others = 5,
+}
+
+/**
+ * A single-frame segmentation result: a tightly packed, row-major map of
+ * per-pixel {@link SegmentCategory} indices at the given resolution.
+ */
+export interface SegmentationMask {
+  /** Row-major per-pixel category indices, length `width * height`. */
+  data: Uint8Array;
+  /** Mask width in pixels. */
+  width: number;
+  /** Mask height in pixels. */
+  height: number;
+}

src/world/segmentation/SegmentationMask.ts

@@ -0,0 +1,63 @@
+import {describe, it, expect} from 'vitest';
+
+import {SegmentationOptions} from './SegmentationOptions';
+
+describe('SegmentationOptions', () => {
+  it('disables segmentation by default', () => {
+    const opts = new SegmentationOptions();
+    expect(opts.enabled).toBe(false);
+  });
+
+  it('defaults to the mediapipe backend with the selfie multiclass model', () => {
+    const opts = new SegmentationOptions();
+    expect(opts.backendConfig.activeBackend).toBe('mediapipe');
+    expect(opts.backendConfig.mediapipe.modelAssetPath).toContain(
+      'selfie_multiclass'
+    );
+    expect(opts.backendConfig.mediapipe.wasmFilesUrl).toContain(
+      '@mediapipe/tasks-vision'
+    );
+  });
+
+  it('outputs the category mask by default', () => {
+    // The category mask is what produces a SegmentationMask, so it must be on
+    // out of the box.
+    const opts = new SegmentationOptions();
+    expect(opts.backendConfig.mediapipe.outputCategoryMask).toBe(true);
+  });
+
+  it('enable() turns segmentation on and returns the instance', () => {
+    const opts = new SegmentationOptions();
+    const returned = opts.enable();
+    expect(opts.enabled).toBe(true);
+    expect(returned).toBe(opts);
+  });
+
+  it('defaults pollingIntervalMs to 66 (~15 fps)', () => {
+    const opts = new SegmentationOptions();
+    expect(opts.pollingIntervalMs).toBe(66);
+  });
+
+  it('deep-merges constructor overrides while keeping unspecified defaults', () => {
+    const opts = new SegmentationOptions({
+      backendConfig: {
+        mediapipe: {
+          modelAssetPath: 'model://custom',
+        },
+      },
+    });
+    // Overridden field takes the new value.
+    expect(opts.backendConfig.mediapipe.modelAssetPath).toBe('model://custom');
+    // Untouched fields keep their defaults.
+    expect(opts.backendConfig.mediapipe.outputCategoryMask).toBe(true);
+    expect(opts.backendConfig.activeBackend).toBe('mediapipe');
+  });
+
+  it('allows overriding pollingIntervalMs via constructor', () => {
+    const opts = new SegmentationOptions({pollingIntervalMs: 100});
+    expect(opts.pollingIntervalMs).toBe(100);
+    // Other defaults are untouched.
+    expect(opts.backendConfig.activeBackend).toBe('mediapipe');
+    expect(opts.enabled).toBe(false);
+  });
+});