react-native-webrtc-kaleidoscope 1.0.0 → 2.0.0

package/android/src/main/java/com/simiancraft/kaleidoscope/segmentation/Mask.kt

@@ -1,35 +1,43 @@
-// Mask production: async MLKit Selfie Segmentation on a worker thread,
-// last-known-mask cache, mask uploaded to a 2D GL texture for the composite
-// shader to sample.
+// Mask production: GL-side adapter around the process-wide SegmentationEngine.
+//
+// This class owns the per-processor GL state (downsample FBO/program, mask
+// texture) and the per-stream temporal-smoothing (EMA) state. The actual
+// segmentation (the worker thread + the MediaPipe ImageSegmenter) lives in the
+// shared SegmentationEngine, so constructing a new Mask per effect switch does
+// NOT spin up a new thread or segmenter; see SegmentationEngine for why.
 //
 // Per-frame flow on the GL thread:
 //   1. If the worker produced a new mask bitmap since the last frame,
 //      upload it to the cached mask GL texture.
 //   2. If no segmentation is currently in flight, render a small downsample
-//      snapshot of the input, post it to the worker, set isProcessing=true.
+//      snapshot of the input and submit it to the SegmentationEngine.
 //   3. Return the current mask GL texture handle (or -1 if no mask has
 //      completed yet — caller falls through to the original frame).
 //
-// The worker thread is the bottleneck (~20-50 ms per MLKit call); decoupling
-// it from the frame thread keeps render at the camera's frame rate while
-// the mask updates ~10-20 Hz. One frame of latency on mask updates is
-// acceptable for this use case.
+// The engine hands the raw foreground-confidence mask back (on its worker
+// thread) via packMask(), which applies EMA smoothing, flips the mask back into
+// the bottom-up orientation the upload/composite expects (the engine segmented
+// an upright copy), quantizes to 8-bit, and stages the result for upload.
+//
+// RESIDUAL LEAK (bounded, documented): because upstream rebuilds the processor
+// per effect switch with no teardown hook, the dropped processor's GL resources
+// are not freed until the EGL context is destroyed (camera stop). That includes
+// this Mask's texture/FBO/program AND the dropped processor's own state (the
+// blur ping-pong FBOs/programs, the YuvConverter). All of it is bounded by the
+// number of switches in a session and small per item; the unbounded thread/
+// segmenter accumulation that would actually OOM is gone (it moved to the
+// process-lived SegmentationEngine).
 //
 // All failure paths log under Kaleidoscope.Mask and return -1 (or a stale
 // mask if one was previously computed).
 
 package com.simiancraft.kaleidoscope.segmentation
 
+import android.content.Context
 import android.graphics.Bitmap
 import android.opengl.GLES30
-import android.os.Handler
-import android.os.HandlerThread
 import android.util.Log
-import com.google.android.gms.tasks.Tasks
-import com.google.mlkit.vision.common.InputImage
-import com.google.mlkit.vision.segmentation.Segmentation
-import com.google.mlkit.vision.segmentation.Segmenter
-import com.google.mlkit.vision.segmentation.selfie.SelfieSegmenterOptions
+import com.simiancraft.kaleidoscope.EffectTuning
 import com.simiancraft.kaleidoscope.gpu.Fbo
 import com.simiancraft.kaleidoscope.gpu.GlDebug
 import com.simiancraft.kaleidoscope.gpu.GlProgram
@@ -39,9 +47,7 @@ import java.nio.ByteOrder
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicReference
 
-internal class Mask {
-  private var segmenter: Segmenter? = null
-
+internal class Mask(private val context: Context) {
   // Cached small-FBO state for the downsample pass.
   private var downsampleFbo: Fbo? = null
   private var downsampleProgram: GlProgram? = null
@@ -51,14 +57,8 @@ internal class Mask {
   private var maskTexWidth: Int = 0
   private var maskTexHeight: Int = 0
 
-  // Worker thread for blocking MLKit calls. Started lazily on first frame.
-  private val workerThread: HandlerThread = HandlerThread("Kaleidoscope.MaskWorker").apply {
-    start()
-  }
-  private val workerHandler: Handler = Handler(workerThread.looper)
-
   // Throttle to a single in-flight segmentation at a time. Set true on
-  // kickoff (GL thread), reset false when the worker finishes.
+  // kickoff (GL thread), reset false when the engine reports done.
   private val isProcessing = AtomicBoolean(false)
 
   // Worker -> GL thread handoff: a Bitmap ready to upload as the mask
@@ -70,17 +70,24 @@ internal class Mask {
   // Pre-allocated readback buffer (resized only on input-dim change).
   private var pixelByteBuffer: ByteBuffer? = null
 
+  // Temporal-smoothing (EMA) state: the previous smoothed confidence buffer and
+  // its dims. Touched only on the SegmentationEngine worker thread (packMask),
+  // which is single-threaded, so no locking. (If a future change ever reads
+  // these off that thread, e.g. on the GL thread, they would need @Volatile.)
+  private var smoothedMask: FloatArray? = null
+  private var smoothedMaskW: Int = 0
+  private var smoothedMaskH: Int = 0
+
   /**
-   * Per-frame mask production. Always returns immediately (no MLKit blocking
-   * on the GL thread). Returns the GL texture handle of the latest available
-   * mask, or -1 if no segmentation has completed yet. Callers must treat -1
-   * as "no mask this frame" and fall through to the original frame.
+   * Per-frame mask production. Always returns immediately (no segmentation
+   * blocking on the GL thread). Returns the GL texture handle of the latest
+   * available mask, or -1 if no segmentation has completed yet. Callers must
+   * treat -1 as "no mask this frame" and fall through to the original frame.
    */
   fun produce(
     source2D: Int,
     sourceWidth: Int,
     sourceHeight: Int,
-    downsampleSize: Int = 256,
   ): Int {
     // Step 1: drain any pending mask the worker has produced. getAndSet
     // claims the bitmap atomically; the GL thread is now its sole owner.
@@ -98,9 +105,15 @@ internal class Mask {
     // Step 2: kick off a new segmentation if the worker is idle.
     if (isProcessing.compareAndSet(false, true)) {
       try {
-        val downsampleBmp = renderAndReadback(source2D, sourceWidth, sourceHeight, downsampleSize)
+        val downsampleBmp =
+          renderAndReadback(source2D, sourceWidth, sourceHeight, EffectTuning.targetShortSide)
         if (downsampleBmp != null) {
-          workerHandler.post { runSegmentation(downsampleBmp) }
+          SegmentationEngine.submit(
+            downsampleBmp,
+            context,
+            onMask = { raw, w, h -> packMask(raw, w, h) },
+            onDone = { isProcessing.set(false) },
+          )
         } else {
           isProcessing.set(false)
         }
@@ -114,13 +127,14 @@ internal class Mask {
   }
 
   /**
-   * Release MLKit + worker thread + GL resources. Call from the GL thread.
-   * Not currently invoked by any caller because VideoFrameProcessor has no
-   * explicit teardown hook; worker thread leaks for the app's lifetime.
+   * Release this Mask's GL resources. Call from the GL thread. Does NOT touch
+   * the SegmentationEngine's worker thread or segmenter (those are process-
+   * lived and shared). Not currently invoked by any caller because
+   * VideoFrameProcessor has no explicit teardown hook; see the file header on
+   * the bounded GL leak this implies.
    */
   fun release() {
     try {
-      workerThread.quitSafely()
       if (maskTextureId != 0) {
         GLES30.glDeleteTextures(1, intArrayOf(maskTextureId), 0)
         maskTextureId = 0
@@ -129,14 +143,57 @@ internal class Mask {
       downsampleFbo = null
       downsampleProgram?.delete()
       downsampleProgram = null
-      segmenter?.close()
-      segmenter = null
       pendingMaskBitmap.getAndSet(null)?.recycle()
+      smoothedMask = null
     } catch (t: Throwable) {
       Log.w(TAG, "Mask.release encountered an error; resources may leak", t)
     }
   }
 
+  // --- Worker thread (invoked by SegmentationEngine) -----------------------
+
+  /**
+   * Apply EMA smoothing to the raw upright confidence mask, flip it back into
+   * the bottom-up orientation the upload/composite expects, quantize to 8-bit
+   * RGBA, and stage it for the GL thread to upload. Runs on the engine's single
+   * worker thread, so the EMA state below needs no locking.
+   */
+  private fun packMask(raw: FloatArray, maskW: Int, maskH: Int) {
+    val pixelCount = maskW * maskH
+
+    // Temporal smoothing (exponential moving average) across mask updates, to
+    // damp shoulder-popping / edge shimmer. History resets on dim change.
+    val prevSmoothed = smoothedMask
+    val blend = prevSmoothed != null && smoothedMaskW == maskW && smoothedMaskH == maskH
+    val smoothed = if (blend) prevSmoothed!! else FloatArray(pixelCount)
+
+    val outPixels = IntArray(pixelCount)
+    for (i in 0 until pixelCount) {
+      val r = raw[i].coerceIn(0f, 1f)
+      val s = if (blend) MASK_EMA_ALPHA * r + (1f - MASK_EMA_ALPHA) * smoothed[i] else r
+      smoothed[i] = s
+      val c = (s * 255f + 0.5f).toInt() and 0xFF
+      // Flip vertically back into the orientation the upload/composite expects
+      // (the engine segmented an upright copy). smoothedMask stays in upright
+      // space for frame-to-frame EMA consistency; only the output is flipped.
+      val row = i / maskW
+      val col = i - row * maskW
+      outPixels[(maskH - 1 - row) * maskW + col] = (0xFF shl 24) or (c shl 16) or (c shl 8) or c
+    }
+    smoothedMask = smoothed
+    smoothedMaskW = maskW
+    smoothedMaskH = maskH
+
+    val outBmp = Bitmap.createBitmap(maskW, maskH, Bitmap.Config.ARGB_8888)
+    outBmp.setPixels(outPixels, 0, maskW, 0, 0, maskW, maskH)
+
+    // Hand off to GL thread. getAndSet atomically claims any previously
+    // unconsumed bitmap as `prev` so we own the recycle; the GL thread
+    // can never observe the same reference we are about to free.
+    val prev = pendingMaskBitmap.getAndSet(outBmp)
+    prev?.recycle()
+  }
+
   // --- GL thread -----------------------------------------------------------
 
   private fun renderAndReadback(
@@ -184,59 +241,8 @@ internal class Mask {
     GlDebug.check("mask upload texImage2D")
   }
 
-  // --- Worker thread -------------------------------------------------------
-
-  private fun runSegmentation(inputBmp: Bitmap) {
-    try {
-      val seg = ensureSegmenter()
-      val inputImage = InputImage.fromBitmap(inputBmp, 0)
-      val rawMask = Tasks.await(seg.process(inputImage))
-
-      val maskBuffer = rawMask.buffer.order(ByteOrder.nativeOrder()).asFloatBuffer()
-      val maskW = rawMask.width
-      val maskH = rawMask.height
-
-      val outPixels = IntArray(maskW * maskH)
-      maskBuffer.rewind()
-      for (i in 0 until maskW * maskH) {
-        val c = (maskBuffer.get().coerceIn(0f, 1f) * 255f + 0.5f).toInt() and 0xFF
-        outPixels[i] = (0xFF shl 24) or (c shl 16) or (c shl 8) or c
-      }
-
-      val outBmp = Bitmap.createBitmap(maskW, maskH, Bitmap.Config.ARGB_8888)
-      outBmp.setPixels(outPixels, 0, maskW, 0, 0, maskW, maskH)
-
-      // Hand off to GL thread. getAndSet atomically claims any previously
-      // unconsumed bitmap as `prev` so we own the recycle; the GL thread
-      // can never observe the same reference we are about to free.
-      val prev = pendingMaskBitmap.getAndSet(outBmp)
-      prev?.recycle()
-    } catch (t: Throwable) {
-      Log.e(TAG, "runSegmentation failed on worker", t)
-    } finally {
-      inputBmp.recycle()
-      isProcessing.set(false)
-    }
-  }
-
   // --- Lazy init helpers ---------------------------------------------------
 
-  private fun ensureSegmenter(): Segmenter {
-    val existing = segmenter
-    if (existing != null) return existing
-    val seg = Segmentation.getClient(
-      SelfieSegmenterOptions.Builder()
-        .setDetectorMode(SelfieSegmenterOptions.STREAM_MODE)
-        // Raw model resolution; MLKit returns the mask at the segmenter's
-        // native size instead of upsampling internally. Faster per call;
-        // the composite shader's smoothstep softens the coarser edge.
-        .enableRawSizeMask()
-        .build(),
-    )
-    segmenter = seg
-    return seg
-  }
-
   private fun ensureDownsampleFbo(w: Int, h: Int): Fbo {
     val existing = downsampleFbo
     if (existing != null && existing.width == w && existing.height == h) return existing
@@ -285,6 +291,11 @@ internal class Mask {
   companion object {
     private const val TAG = "Kaleidoscope.Mask"
 
+    // EMA weight for the new mask vs history. Higher = more responsive, lower =
+    // smoother (more lag). 0.5 is ~a 1-2 update time constant at the ~10-20 Hz
+    // mask rate: damps flicker without obvious lag.
+    private const val MASK_EMA_ALPHA = 0.5f
+
     private const val TWO_D_PASSTHROUGH_FRAG = """#version 300 es
 precision mediump float;
 uniform sampler2D uTex;

package/android/src/main/java/com/simiancraft/kaleidoscope/segmentation/SegmentationEngine.kt

@@ -0,0 +1,148 @@
+// Process-wide segmentation worker. Owns the ONE MediaPipe Tasks
+// ImageSegmenter and the ONE worker thread for the whole process; every Mask
+// instance (one per active VideoFrameProcessor) submits through here instead of
+// spinning up its own thread + segmenter.
+//
+// WHY THIS IS A SINGLETON: react-native-webrtc calls VideoFrameProcessorFactory
+// .build() on EVERY effect switch from JS, constructing a fresh processor (and
+// thus a fresh Mask) each time, and upstream has no teardown hook to dispose
+// the old one. If each Mask owned a HandlerThread + ImageSegmenter, toggling
+// effects would leak a thread and a native segmenter handle per switch and
+// eventually OOM on a constrained device. Hoisting the thread + segmenter here
+// bounds them to one-per-process regardless of how many Masks come and go.
+//
+// The segmenter is CPU (BitmapImageBuilder input), so sharing it across Masks
+// is safe; VIDEO running mode only needs monotonically increasing timestamps,
+// which a single shared counter guarantees. Per-stream temporal state (the EMA
+// smoothing) stays on each Mask, so two simultaneous tracks do not blend masks.
+//
+// The thread and segmenter are created lazily on first submit and live for the
+// process lifetime (the intended steady state: one thread, one segmenter). They
+// are intentionally never torn down.
+
+package com.simiancraft.kaleidoscope.segmentation
+
+import android.content.Context
+import android.graphics.Bitmap
+import android.graphics.Matrix
+import android.os.Handler
+import android.os.HandlerThread
+import android.util.Log
+import com.google.mediapipe.framework.image.BitmapImageBuilder
+import com.google.mediapipe.framework.image.ByteBufferExtractor
+import com.google.mediapipe.framework.image.MPImage
+import com.google.mediapipe.tasks.core.BaseOptions
+import com.google.mediapipe.tasks.vision.core.RunningMode
+import com.google.mediapipe.tasks.vision.imagesegmenter.ImageSegmenter
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+
+internal object SegmentationEngine {
+  private const val TAG = "Kaleidoscope.SegEngine"
+
+  // Lazy, process-lived. Created on first submit, never torn down.
+  private val workerThread: HandlerThread by lazy {
+    HandlerThread("Kaleidoscope.MaskWorker").also { it.start() }
+  }
+  private val workerHandler: Handler by lazy { Handler(workerThread.looper) }
+
+  // Worker-thread-confined state.
+  private var segmenter: ImageSegmenter? = null
+  private var videoTimestamp: Long = 0
+
+  /**
+   * Run segmentation for [inputBmp] off the GL thread. [inputBmp] is the
+   * downsampled, bottom-up GL readback; this uprights it (the selfie model is
+   * trained on upright people and does its worst on an inverted one), segments
+   * it, and hands the raw foreground-confidence mask back via [onMask] in the
+   * SAME upright orientation the caller's EMA/flip-back expects.
+   *
+   * [onMask] is invoked on the worker thread only on success. [onDone] is
+   * invoked on the worker thread exactly once, success or failure, so the
+   * caller can reset its in-flight throttle. [inputBmp] is recycled here.
+   */
+  fun submit(
+    inputBmp: Bitmap,
+    context: Context,
+    onMask: (rawUpright: FloatArray, width: Int, height: Int) -> Unit,
+    onDone: () -> Unit,
+  ) {
+    workerHandler.post {
+      var upright: Bitmap? = null
+      var mpImage: MPImage? = null
+      try {
+        val seg = ensureSegmenter(context)
+        // The GL readback (glReadPixels reads bottom-to-top) is vertically
+        // flipped (head-down). Feed the segmenter an upright copy; the caller
+        // flips the mask back so the upload/composite alignment is untouched.
+        // Vertical flip only, not a 180 rotate (columns are already correct).
+        upright = flipVertical(inputBmp)
+        mpImage = BitmapImageBuilder(upright).build()
+        val result = seg.segmentForVideo(mpImage, videoTimestamp++)
+
+        val confidenceMasks = result.confidenceMasks()
+        if (!confidenceMasks.isPresent || confidenceMasks.get().isEmpty()) {
+          Log.w(TAG, "segmentation produced no confidence mask")
+          return@post
+        }
+        // General selfie segmenter: a single foreground-confidence mask
+        // (float [0,1], higher = person). Copy it out of MediaPipe's buffer
+        // into a plain array the caller owns.
+        val maskImage = confidenceMasks.get()[0]
+        val maskW = maskImage.width
+        val maskH = maskImage.height
+        val maskBuffer = ByteBufferExtractor.extract(maskImage)
+          .order(ByteOrder.nativeOrder())
+          .asFloatBuffer()
+        val raw = FloatArray(maskW * maskH)
+        maskBuffer.rewind()
+        maskBuffer.get(raw)
+        onMask(raw, maskW, maskH)
+      } catch (t: Throwable) {
+        Log.e(TAG, "segmentation failed on worker", t)
+      } finally {
+        // close() in finally: if segmentForVideo throws, the MPImage native
+        // handle would otherwise leak (unbounded if the error recurs per frame).
+        mpImage?.close()
+        upright?.recycle()
+        inputBmp.recycle()
+        onDone()
+      }
+    }
+  }
+
+  private fun ensureSegmenter(context: Context): ImageSegmenter {
+    val existing = segmenter
+    if (existing != null) return existing
+    // Load the model as a direct ByteBuffer (setModelAssetBuffer) rather than
+    // setModelAssetPath: the path variant memory-maps the asset and requires it
+    // to be stored uncompressed (aaptOptions noCompress), which we cannot
+    // guarantee in the CONSUMING app's build. Reading the asset into a direct
+    // buffer ourselves works regardless of apk compression.
+    val appContext = context.applicationContext
+    val modelBytes = appContext.assets.open("selfie_segmenter.tflite").use { it.readBytes() }
+    val modelBuffer = ByteBuffer.allocateDirect(modelBytes.size).order(ByteOrder.nativeOrder())
+    modelBuffer.put(modelBytes)
+    modelBuffer.rewind()
+
+    val baseOptions = BaseOptions.builder()
+      .setModelAssetBuffer(modelBuffer)
+      .build()
+    val options = ImageSegmenter.ImageSegmenterOptions.builder()
+      .setBaseOptions(baseOptions)
+      .setRunningMode(RunningMode.VIDEO)
+      .setOutputConfidenceMasks(true)
+      .setOutputCategoryMask(false)
+      .build()
+    val seg = ImageSegmenter.createFromOptions(appContext, options)
+    segmenter = seg
+    return seg
+  }
+
+  /** Vertical mirror (flip across the horizontal axis). Uprights the
+   * bottom-to-top glReadPixels frame before segmentation. */
+  private fun flipVertical(src: Bitmap): Bitmap {
+    val m = Matrix().apply { postScale(1f, -1f) }
+    return Bitmap.createBitmap(src, 0, 0, src.width, src.height, m, true)
+  }
+}