index.html 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. <html>
  2. <head>
  3. <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@3.7.0/dist/tf.min.js"></script>
  4. <script>
  5. async function main() {
  6. const video = document.querySelector('video');
  7. const canvas = document.querySelector('canvas');
  8. const select = document.querySelector('select');
  9. video.width = 640;
  10. video.height = 480;
  11. const webcam = await tf.data.webcam(video);
  12. const model = await tf.loadGraphModel('model/model.json');
  13. // Set initial recurrent state
  14. let [r1i, r2i, r3i, r4i] = [tf.tensor(0.), tf.tensor(0.), tf.tensor(0.), tf.tensor(0.)];
  15. // Set downsample ratio
  16. const downsample_ratio = tf.tensor(0.5);
  17. // Inference loop
  18. while (true) {
  19. await tf.nextFrame();
  20. const img = await webcam.capture();
  21. const src = tf.tidy(() => img.expandDims(0).div(255)); // normalize input
  22. const [fgr, pha, r1o, r2o, r3o, r4o] = await model.executeAsync(
  23. {src, r1i, r2i, r3i, r4i, downsample_ratio}, // provide inputs
  24. ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o'] // select outputs
  25. );
  26. // Draw the result based on selected view
  27. const viewOption = select.value;
  28. if (viewOption === 'recurrent1') {
  29. drawHidden(r1o, canvas);
  30. } else if (viewOption === 'recurrent2') {
  31. drawHidden(r2o, canvas);
  32. } else if (viewOption === 'recurrent3') {
  33. drawHidden(r3o, canvas);
  34. } else if (viewOption === 'recurrent4') {
  35. drawHidden(r4o, canvas);
  36. } else if (viewOption === 'white') {
  37. drawMatte(fgr.clone(), pha.clone(), canvas);
  38. canvas.style.background = 'rgb(255, 255, 255)';
  39. } else if (viewOption === 'green') {
  40. drawMatte(fgr.clone(), pha.clone(), canvas);
  41. canvas.style.background = 'rgb(120, 255, 155)';
  42. } else if (viewOption === 'alpha') {
  43. drawMatte(null, pha.clone(), canvas);
  44. canvas.style.background = 'rgb(0, 0, 0)';
  45. } else if (viewOption === 'foreground') {
  46. drawMatte(fgr.clone(), null, canvas);
  47. }
  48. // Dispose old tensors.
  49. tf.dispose([img, src, fgr, pha, r1i, r2i, r3i, r4i]);
  50. // Update recurrent states.
  51. [r1i, r2i, r3i, r4i] = [r1o, r2o, r3o, r4o];
  52. }
  53. }
  54. async function drawMatte(fgr, pha, canvas){
  55. const rgba = tf.tidy(() => {
  56. const rgb = (fgr !== null) ?
  57. fgr.squeeze(0).mul(255).cast('int32') :
  58. tf.fill([pha.shape[1], pha.shape[2], 3], 255, 'int32');
  59. const a = (pha !== null) ?
  60. pha.squeeze(0).mul(255).cast('int32') :
  61. tf.fill([fgr.shape[1], fgr.shape[2], 1], 255, 'int32');
  62. return tf.concat([rgb, a], -1);
  63. });
  64. fgr && fgr.dispose();
  65. pha && pha.dispose();
  66. const [height, width] = rgba.shape.slice(0, 2);
  67. const pixelData = new Uint8ClampedArray(await rgba.data());
  68. const imageData = new ImageData(pixelData, width, height);
  69. canvas.width = width;
  70. canvas.height = height;
  71. canvas.getContext('2d').putImageData(imageData, 0, 0);
  72. rgba.dispose();
  73. }
  74. async function drawHidden(r, canvas) {
  75. const rgba = tf.tidy(() => {
  76. r = r.unstack(-1)
  77. r = tf.concat(r, 1)
  78. r = r.split(4, 1)
  79. r = tf.concat(r, 2)
  80. r = r.squeeze(0)
  81. r = r.expandDims(-1)
  82. r = r.add(1).mul(127.5).cast('int32')
  83. r = r.tile([1, 1, 3])
  84. r = tf.concat([r, tf.fill([r.shape[0], r.shape[1], 1], 255, 'int32')], -1)
  85. return r;
  86. });
  87. const [height, width] = rgba.shape.slice(0, 2);
  88. const pixelData = new Uint8ClampedArray(await rgba.data());
  89. const imageData = new ImageData(pixelData, width, height);
  90. canvas.width = width;
  91. canvas.height = height;
  92. canvas.getContext('2d').putImageData(imageData, 0, 0);
  93. rgba.dispose();
  94. }
  95. window.addEventListener('load', main);
  96. </script>
  97. </head>
  98. <body>
  99. <select>
  100. <option value="white">White Background</option>
  101. <option value="green">Green Background</option>
  102. <option value="alpha">Alpha</option>
  103. <option value="foreground">Foreground</option>
  104. <option value="recurrent1">Recurrent State 1</option>
  105. <option value="recurrent2">Recurrent State 2</option>
  106. <option value="recurrent3">Recurrent State 3</option>
  107. <option value="recurrent4">Recurrent State 4</option>
  108. </select>
  109. <br>
  110. <video></video>
  111. <canvas></canvas>
  112. </body>
  113. </html>