blob: 787d4b4c5e2416da8ee897e3b49b08b529e99080 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
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 landscape from './vendor/models/selfie_segmentation_landscape.tflite'
const models = {
modelLandscape: landscape.split('/').pop(),
}
self.compiled = false
self.onmessage = (e) => {
const message = e.data.message
switch (message) {
case 'makeTFLite':
self.segmentationPixelCount = e.data.segmentationPixelCount
makeTFLite(e.data.simd)
break
case 'resizeSource':
if (!self.compiled) return
resizeSource(e.data.imageData, e.data.frameId)
break
case 'runInference':
runInference()
break
default:
console.error('JitsiStreamBackgroundEffect.worker: Message unknown.')
console.error(message)
break
}
}
/**
* @param {boolean} isSimd whether WebAssembly SIMD is available or not
*/
async function makeTFLite(isSimd) {
try {
switch (isSimd) {
case true:
self.wasmUrl = withSIMD.split('/').pop()
self.tflite = await createTFLiteSIMDModule({ locateFile: (path) => { return self.wasmUrl } })
break
case false:
self.wasmUrl = withoutSIMD.split('/').pop()
self.tflite = await createTFLiteModule({ locateFile: (path) => { return self.wasmUrl } })
break
default:
return
}
self.modelBufferOffset = self.tflite._getModelBufferMemoryOffset()
self.modelResponse = await fetch(models.modelLandscape)
if (!self.modelResponse.ok) {
throw new Error('Failed to download tflite model!')
}
self.model = await self.modelResponse.arrayBuffer()
self.tflite.HEAPU8.set(new Uint8Array(self.model), self.modelBufferOffset)
await self.tflite._loadModel(self.model.byteLength)
// Even if the wrong tflite file is downloaded (for example, if an HTML
// error is downloaded instead of the file) loading the model will
// succeed. However, if the model does not have certain values it could
// be assumed that the model failed to load.
if (!self.tflite._getInputWidth() || !self.tflite._getInputHeight()
|| !self.tflite._getOutputWidth() || !self.tflite._getOutputHeight()) {
throw new Error('Failed to load tflite model!')
}
self.compiled = true
self.postMessage({ message: 'loaded' })
} catch (error) {
console.error(error)
console.error('JitsiStreamBackgroundEffect.worker: tflite compilation failed. The web server may not be properly configured to send wasm and/or tflite files.')
self.postMessage({ message: 'loadFailed' })
}
}
/**
* @param {ImageData} imageData the image data from the canvas
* @param {number} frameId the ID of the frame that the image data belongs to
*/
function resizeSource(imageData, frameId) {
const inputMemoryOffset = self.tflite._getInputMemoryOffset() / 4
for (let i = 0; i < self.segmentationPixelCount; i++) {
self.tflite.HEAPF32[inputMemoryOffset + (i * 3)] = imageData.data[i * 4] / 255
self.tflite.HEAPF32[inputMemoryOffset + (i * 3) + 1] = imageData.data[(i * 4) + 1] / 255
self.tflite.HEAPF32[inputMemoryOffset + (i * 3) + 2] = imageData.data[(i * 4) + 2] / 255
}
runInference(frameId)
}
/**
* @param {number} frameId the ID of the frame that the image data belongs to
*/
function runInference(frameId) {
self.tflite._runInference()
const outputMemoryOffset = self.tflite._getOutputMemoryOffset() / 4
const segmentationMaskData = []
// All consts in Worker in obj array.
for (let i = 0; i < self.segmentationPixelCount; i++) {
const person = self.tflite.HEAPF32[outputMemoryOffset + i]
segmentationMaskData.push({
person,
})
}
self.postMessage({ message: 'inferenceRun', segmentationResult: segmentationMaskData, frameId })
}
// This is needed to make the linter happy, but even if nothing is actually
// exported the worker is loaded as expected.
export default null
|
