docs: Background segmentation - custom models (#1951)

Author: lursz

apps/typegpu-docs/src/examples/image-processing/background-segmentation/index.ts

@@ -7,15 +7,16 @@ import tgpu, {
 } from 'typegpu';
 import { fullScreenTriangle } from 'typegpu/common';
 import * as d from 'typegpu/data';
-import { MODEL_HEIGHT, MODEL_WIDTH, prepareSession } from './model.ts';
+import { MODEL_HEIGHT, MODEL_WIDTH, MODELS, prepareSession } from './model.ts';
 import {
   blockDim,
   blurLayout,
   drawWithMaskLayout,
+  flipSlot,
   generateMaskLayout,
+  Params,
+  paramsAccessor,
   prepareModelInputLayout,
-  sampleBiasSlot,
-  useGaussianSlot,
 } from './schemas.ts';
 import {
   computeFn,
@@ -62,7 +63,7 @@ const oldRequestAdapter = navigator.gpu.requestAdapter;
 const oldRequestDevice = adapter.requestDevice;
 navigator.gpu.requestAdapter = async () => adapter;
 adapter.requestDevice = async () => device;
-const root = await tgpu.initFromDevice({ device });
+const root = tgpu.initFromDevice({ device });
 const context = canvas.getContext('webgpu') as GPUCanvasContext;
 const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
 
@@ -76,11 +77,13 @@ context.configure({
 
 let blurStrength = 5;
 let useGaussianBlur = false;
+let useSquareCrop = false;
 
-const zeroBuffer = root.createBuffer(d.u32, 0).$usage('uniform');
-const oneBuffer = root.createBuffer(d.u32, 1).$usage('uniform');
-const useGaussianUniform = root.createUniform(d.u32, 0);
-const sampleBiasUniform = root.createUniform(d.f32, 0);
+const paramsUniform = root.createUniform(Params, {
+  cropBounds: d.vec4f(0, 0, 1, 1),
+  useGaussian: 0,
+  sampleBias: blurStrength,
+});
 
 const sampler = root['~unstable'].createSampler({
   magFilter: 'linear',
@@ -123,27 +126,48 @@ let blurBindGroups: TgpuBindGroup<typeof blurLayout.entries>[];
 // pipelines
 
 const prepareModelInputPipeline = root['~unstable']
+  .with(paramsAccessor, paramsUniform)
   .createGuardedComputePipeline(
     prepareModelInput,
   );
 
-const session = await prepareSession(
+let currentModelIndex = 0;
+let session = await prepareSession(
   root.unwrap(modelInputBuffer),
   root.unwrap(modelOutputBuffer),
+  MODELS[currentModelIndex],
 );
+let isLoadingModel = false;
+
+async function switchModel(modelIndex: number) {
+  if (isLoadingModel || modelIndex === currentModelIndex) return;
+  isLoadingModel = true;
+
+  const oldSession = session;
+  currentModelIndex = modelIndex;
+  session = await prepareSession(
+    root.unwrap(modelInputBuffer),
+    root.unwrap(modelOutputBuffer),
+    MODELS[currentModelIndex],
+  );
+  oldSession.release();
+  isLoadingModel = false;
+}
 
 const generateMaskFromOutputPipeline = root['~unstable']
   .createGuardedComputePipeline(
     generateMaskFromOutput,
   );
 
-const blurPipeline = root['~unstable']
-  .withCompute(computeFn)
-  .createPipeline();
+const blurPipelines = [false, true].map((flip) =>
+  root['~unstable']
+    .with(flipSlot, flip)
+    .withCompute(computeFn)
+    .createPipeline()
+);
 
 const drawWithMaskPipeline = root['~unstable']
-  .with(useGaussianSlot, useGaussianUniform)
-  .with(sampleBiasSlot, sampleBiasUniform)
+  .with(paramsAccessor, paramsUniform)
   .withVertex(fullScreenTriangle, {})
   .withFragment(drawWithMaskFragment, { format: presentationFormat })
   .createPipeline();
@@ -153,7 +177,7 @@ const drawWithMaskPipeline = root['~unstable']
 let calculateMaskCallbackId: number | undefined;
 
 async function processCalculateMask() {
-  if (video.readyState < 2) {
+  if (video.readyState < 2 || isLoadingModel) {
     calculateMaskCallbackId = video.requestVideoFrameCallback(
       processCalculateMask,
     );
@@ -181,6 +205,29 @@ async function processCalculateMask() {
 calculateMaskCallbackId = video.requestVideoFrameCallback(processCalculateMask);
 
 // frame
+function updateCropBounds(aspectRatio: number) {
+  let uvMinX = 0;
+  let uvMinY = 0;
+  let uvMaxX = 1;
+  let uvMaxY = 1;
+
+  if (useSquareCrop) {
+    if (aspectRatio > 1) {
+      // wide (e.g. 16:9) - crop horizontally
+      const cropWidth = 1 / aspectRatio; // width of square in UV space
+      uvMinX = (1 - cropWidth) / 2;
+      uvMaxX = uvMinX + cropWidth;
+    } else if (aspectRatio < 1) {
+      // tall - crop vertically
+      const cropHeight = aspectRatio; // height of square in UV space
+      uvMinY = (1 - cropHeight) / 2;
+      uvMaxY = uvMinY + cropHeight;
+    }
+  }
+  paramsUniform.writePartial({
+    cropBounds: d.vec4f(uvMinX, uvMinY, uvMaxX, uvMaxY),
+  });
+}
 
 function onVideoChange(size: { width: number; height: number }) {
   const aspectRatio = size.width / size.height;
@@ -190,17 +237,18 @@ function onVideoChange(size: { width: number; height: number }) {
     canvas.parentElement.style.height =
       `min(100cqh, calc(100cqw/(${aspectRatio})))`;
   }
+
+  updateCropBounds(aspectRatio);
+
   blurredTextures = [0, 1].map(() =>
     root['~unstable'].createTexture({
       size: [size.width, size.height],
       format: 'rgba8unorm',
-      dimension: '2d',
       mipLevelCount: 10,
     }).$usage('sampled', 'render', 'storage')
   );
   blurBindGroups = [
     root.createBindGroup(blurLayout, {
-      flip: zeroBuffer,
       inTexture: blurredTextures[0],
       outTexture: blurredTextures[1].createView(
         d.textureStorage2d('rgba8unorm', 'read-only'),
@@ -209,7 +257,6 @@ function onVideoChange(size: { width: number; height: number }) {
       sampler,
     }),
     root.createBindGroup(blurLayout, {
-      flip: oneBuffer,
       inTexture: blurredTextures[1],
       outTexture: blurredTextures[0].createView(
         d.textureStorage2d('rgba8unorm', 'read-only'),
@@ -248,13 +295,13 @@ async function processVideoFrame(
 
   if (useGaussianBlur) {
     for (const _ of Array(blurStrength * 2)) {
-      blurPipeline
+      blurPipelines[0]
         .with(blurBindGroups[0])
         .dispatchWorkgroups(
           Math.ceil(frameWidth / blockDim),
           Math.ceil(frameHeight / 4),
         );
-      blurPipeline
+      blurPipelines[1]
         .with(blurBindGroups[1])
         .dispatchWorkgroups(
           Math.ceil(frameHeight / blockDim),
@@ -287,12 +334,22 @@ videoFrameCallbackId = video.requestVideoFrameCallback(processVideoFrame);
 // #region Example controls & Cleanup
 
 export const controls = {
+  model: {
+    initial: MODELS[0].name,
+    options: MODELS.map((m) => m.name),
+    async onSelectChange(value: string) {
+      const index = MODELS.findIndex((m) => m.name === value);
+      if (index !== -1) {
+        await switchModel(index);
+      }
+    },
+  },
   'blur type': {
     initial: 'mipmaps',
     options: ['mipmaps', 'gaussian'],
     async onSelectChange(value: string) {
       useGaussianBlur = value === 'gaussian';
-      useGaussianUniform.write(useGaussianBlur ? 1 : 0);
+      paramsUniform.writePartial({ useGaussian: useGaussianBlur ? 1 : 0 });
     },
   },
   'blur strength': {
@@ -302,7 +359,16 @@ export const controls = {
     step: 1,
     onSliderChange(newValue: number) {
       blurStrength = newValue;
-      sampleBiasUniform.write(blurStrength);
+      paramsUniform.writePartial({ sampleBias: blurStrength });
+    },
+  },
+  'square crop': {
+    initial: useSquareCrop,
+    onToggleChange(value: boolean) {
+      useSquareCrop = value;
+      if (lastFrameSize) {
+        updateCropBounds(lastFrameSize.width / lastFrameSize.height);
+      }
     },
   },
 };