Replace virtual background segmentation model

The original segmentation models used in Jitsi were relicensed by Google
from the Apache license to a proprietary one. While in theory the
previous versions should still be usable under the Apache license it
seems that they were never intended to be released as such. Therefore
the segmentation model is now replaced with the MediaPipe Selfie
Segmentation model.

Code changes are based on commit
"9b6b335c60223dc7615b308b8a25a263c7fc95eb" of repository
"https://github.com/jitsi/jitsi-meet".

"selfie_segmentation_landscape.tflite" was copied from
"mediapipe/modules/selfie_segmentation/selfie_segmentation.tflite" of
repository "https://github.com/google/mediapipe" at commit
"8b57bf879b419173b26277d220b643dac0402334".

"Model Card MediaPipe Selfie Segmentation.pdf" was downloaded from
"https://drive.google.com/file/d/1dCfozqknMa068vVsO2j_1FgZkW_e3VWv/view".

Signed-off-by: Daniel Calviño Sánchez <[email protected]>

src/utils/media/effects/virtual-background/JitsiStreamBackgroundEffect.js

@@ -235,7 +235,7 @@ export default class JitsiStreamBackgroundEffect {
 	runInference(data) {
 		// All consts in Worker in obj array.
 		for (let i = 0; i < this._segmentationPixelCount; i++) {
-			this._segmentationMask.data[(i * 4) + 3] = (255 * data[i].personExp) / (data[i].backgroundExp + data[i].personExp)
+			this._segmentationMask.data[(i * 4) + 3] = 255 * data[i].person
 		}
 		this._segmentationMaskCtx.putImageData(this._segmentationMask, 0, 0)
 	}

src/utils/media/effects/virtual-background/JitsiStreamBackgroundEffect.worker.js

@@ -2,12 +2,10 @@ import createTFLiteModule from './vendor/tflite/tflite'
 import createTFLiteSIMDModule from './vendor/tflite/tflite-simd'
 import withoutSIMD from './vendor/tflite/tflite.wasm'
 import withSIMD from './vendor/tflite/tflite-simd.wasm'
-import v681 from './vendor/models/segm_lite_v681.tflite'
-import v679 from './vendor/models/segm_full_v679.tflite'
+import landscape from './vendor/models/selfie_segmentation_landscape.tflite'
 
 const models = {
-	model96: v681.split('/').pop(),
-	model144: v679.split('/').pop(),
+	modelLandscape: landscape.split('/').pop(),
 }
 
 self.compiled = false
@@ -51,7 +49,7 @@ async function makeTFLite(isSimd) {
 			return
 		}
 		self.modelBufferOffset = self.tflite._getModelBufferMemoryOffset()
-		self.modelResponse = await fetch(isSimd === true ? models.model144 : models.model96)
+		self.modelResponse = await fetch(models.modelLandscape)
 
 		if (!self.modelResponse.ok) {
 			throw new Error('Failed to download tflite model!')
@@ -107,16 +105,10 @@ function runInference(frameId) {
 	// All consts in Worker in obj array.
 	for (let i = 0; i < self.segmentationPixelCount; i++) {
 
-		const background = self.tflite.HEAPF32[outputMemoryOffset + (i * 2)]
-		const person = self.tflite.HEAPF32[outputMemoryOffset + (i * 2) + 1]
-		const shift = Math.max(background, person)
+		const person = self.tflite.HEAPF32[outputMemoryOffset + i]
 
 		segmentationMaskData.push({
-			background,
 			person,
-			shift,
-			backgroundExp: Math.exp(background - shift),
-			personExp: Math.exp(person - shift),
 		})
 	}
 	self.postMessage({ message: 'inferenceRun', segmentationResult: segmentationMaskData, frameId })

src/utils/media/effects/virtual-background/index.js

@@ -4,16 +4,11 @@ import JitsiStreamBackgroundEffect from './JitsiStreamBackgroundEffect'
 import createTFLiteModule from './vendor/tflite/tflite'
 import createTFLiteSIMDModule from './vendor/tflite/tflite-simd'
 const models = {
-	model96: 'libs/segm_lite_v681.tflite',
-	model144: 'libs/segm_full_v679.tflite',
+	modelLandscape: 'libs/selfie_segmentation_landscape.tflite',
 }
 
 const segmentationDimensions = {
-	model96: {
-		height: 96,
-		width: 160,
-	},
-	model144: {
+	modelLandscape: {
 		height: 144,
 		width: 256,
 	},
@@ -53,7 +48,7 @@ export async function createVirtualBackgroundEffect(virtualBackground, dispatch)
 	}
 
 	const modelBufferOffset = tflite._getModelBufferMemoryOffset()
-	const modelResponse = await fetch(wasmCheck.feature.simd ? models.model144 : models.model96)
+	const modelResponse = await fetch(models.modelLandscape)
 
 	if (!modelResponse.ok) {
 		throw new Error('Failed to download tflite model!')
@@ -66,7 +61,7 @@ export async function createVirtualBackgroundEffect(virtualBackground, dispatch)
 	tflite._loadModel(model.byteLength)
 
 	const options = {
-		...wasmCheck.feature.simd ? segmentationDimensions.model144 : segmentationDimensions.model96,
+		...segmentationDimensions.modelLandscape,
 		virtualBackground,
 	}
 

src/utils/media/effects/virtual-background/vendor/README.md

@@ -1,6 +1,6 @@
 # Virtual Background on stream effects
 
-> Inspired from https://ai.googleblog.com/2020/10/background-features-in-google-meet.html and https://github.com/Volcomix/virtual-background.git
+> From https://google.github.io/mediapipe/solutions/models.html#selfie-segmentation
 
 #### Canvas 2D + CPU
 
@@ -22,15 +22,3 @@ More details:
 - [WebAssembly](https://webassembly.org/)
 - [WebAssembly SIMD](https://github.com/WebAssembly/simd)
 - [TFLite](https://blog.tensorflow.org/2020/07/accelerating-tensorflow-lite-xnnpack-integration.html)
-
-## LICENSE
-
-The mdoels vendored here were downloaded early January (they were available as early as the 4th), before Google switched the license away from Apache 2. Thus we understand they are not covered by the new license which according to the [model card](https://drive.google.com/file/d/1lnP1bRi9CSqQQXUHa13159vLELYDgDu0/view) dates from the 21st of January.
-
-We are not lawyers so do get legal advise if in doubt.
-
-References:
-
-- Model license discussion: https://github.com/tensorflow/tfjs/issues/4177
-- Current vendored model is discovered: https://github.com/tensorflow/tfjs/issues/4177#issuecomment-753934631
-- License change is noticed: https://github.com/tensorflow/tfjs/issues/4177#issuecomment-771536641

src/utils/media/pipeline/VirtualBackground.js

@@ -137,11 +137,7 @@ export default class VirtualBackground extends TrackSinkSource {
 
 	_initJitsiStreamBackgroundEffect() {
 		const segmentationDimensions = {
-			model96: {
-				height: 96,
-				width: 160,
-			},
-			model144: {
+			modelLandscape: {
 				height: 144,
 				width: 256,
 			},
@@ -158,7 +154,7 @@ export default class VirtualBackground extends TrackSinkSource {
 			blurValue: 10,
 		}
 		const options = {
-			...isSimd ? segmentationDimensions.model144 : segmentationDimensions.model96,
+			...segmentationDimensions.modelLandscape,
 			virtualBackground,
 			simd: isSimd,
 		}