fix: enable export of scenes with KTX2 compressed textures

packages/model-viewer/src/features/scene-graph.ts

@@ -14,7 +14,7 @@
  */
 
 import {property} from 'lit/decorators.js';
-import {CanvasTexture, Object3D, RepeatWrapping, SRGBColorSpace, Texture, VideoTexture} from 'three';
+import {CanvasTexture, Mesh, NoToneMapping, Object3D, PerspectiveCamera, PlaneGeometry, RepeatWrapping, Scene, ShaderMaterial, SRGBColorSpace, Texture, VideoTexture, WebGLRenderer, WebGLRenderTarget} from 'three';
 import {GLTFExporter, GLTFExporterOptions} from 'three/examples/jsm/exporters/GLTFExporter.js';
 
 import ModelViewerElementBase, {$needsRender, $onModelLoad, $progressTracker, $renderer, $scene} from '../model-viewer-base.js';
@@ -67,6 +67,89 @@ export interface SceneGraphInterface {
   materialFromPoint(pixelX: number, pixelY: number): Material|null;
 }
 
+/**
+ * Decompresses a CompressedTexture into a CanvasTexture using a
+ * WebGLRenderTarget to avoid corrupting the renderer's main canvas.
+ */
+function decompressTexture(
+    texture: Texture, maxTextureSize: number, renderer: WebGLRenderer):
+    CanvasTexture {
+  const image = texture.image as {width: number, height: number};
+  const width = Math.min(image.width, maxTextureSize);
+  const height = Math.min(image.height, maxTextureSize);
+
+  const isSRGB = texture.colorSpace === SRGBColorSpace;
+  const material = new ShaderMaterial({
+    uniforms: {blitTexture: {value: texture}},
+    vertexShader: `
+      varying vec2 vUv;
+      void main() {
+        vUv = uv;
+        gl_Position = vec4(position.xy, 0.0, 1.0);
+      }`,
+    fragmentShader: `
+      uniform sampler2D blitTexture;
+      varying vec2 vUv;
+
+      vec3 linearToSRGB(vec3 linear) {
+        vec3 cutoff = step(linear, vec3(0.0031308));
+        vec3 higher = 1.055 * pow(linear, vec3(1.0 / 2.4)) - 0.055;
+        vec3 lower = linear * 12.92;
+        return mix(higher, lower, cutoff);
+      }
+
+      void main() {
+        vec4 texel = texture2D(blitTexture, vUv);
+        ${
+        isSRGB ? 'gl_FragColor = vec4(linearToSRGB(texel.rgb), texel.a);' :
+                 'gl_FragColor = texel;'}
+      }`
+  });
+
+  const mesh = new Mesh(new PlaneGeometry(2, 2), material);
+  mesh.frustumCulled = false;
+  const scene = new Scene();
+  scene.add(mesh);
+  const camera = new PerspectiveCamera();
+
+  const renderTarget = new WebGLRenderTarget(width, height);
+  const prevRenderTarget = renderer.getRenderTarget();
+  const prevToneMapping = renderer.toneMapping;
+
+  renderer.toneMapping = NoToneMapping;
+  renderer.setRenderTarget(renderTarget);
+  renderer.clear();
+  renderer.render(scene, camera);
+
+  const pixels = new Uint8Array(width * height * 4);
+  renderer.readRenderTargetPixels(renderTarget, 0, 0, width, height, pixels);
+
+  renderer.setRenderTarget(prevRenderTarget);
+  renderer.toneMapping = prevToneMapping;
+  renderTarget.dispose();
+  mesh.geometry.dispose();
+  material.dispose();
+
+  const canvas = document.createElement('canvas');
+  canvas.width = width;
+  canvas.height = height;
+  const ctx = canvas.getContext('2d')!;
+  const imageData = ctx.createImageData(width, height);
+  imageData.data.set(pixels);
+  ctx.putImageData(imageData, 0, 0);
+
+  const result = new CanvasTexture(canvas);
+  result.minFilter = texture.minFilter;
+  result.magFilter = texture.magFilter;
+  result.wrapS = texture.wrapS;
+  result.wrapT = texture.wrapT;
+  result.colorSpace = texture.colorSpace;
+  result.name = texture.name;
+  result.flipY = texture.flipY;
+
+  return result;
+}
+
 /**
  * SceneGraphMixin manages exposes a model API in order to support operations on
  * the <model-viewer> scene graph.
@@ -289,6 +372,15 @@ export const SceneGraphMixin = <T extends Constructor<ModelViewerElementBase>>(
                 .register(
                     (writer: any) =>
                         new GLTFExporterMaterialsVariantsExtension(writer));
+
+        const threeRenderer = this[$renderer].threeRenderer;
+        exporter.setTextureUtils({
+          decompress: (texture: Texture, maxTextureSize?: number) => {
+            return decompressTexture(
+                texture, maxTextureSize ?? Infinity, threeRenderer);
+          }
+        });
+
         let exportTarget: Object3D;
         if (scene.models.length > 1) {
           exportTarget = new Object3D();

packages/model-viewer/src/test/features/scene-graph/texture-spec.ts

@@ -75,6 +75,37 @@ suite('scene-graph/texture', () => {
           .to.be.equal('image/png');
     });
 
+    test(
+        'exports and re-imports a model with KTX2 compressed texture',
+        async () => {
+          const ktx2Texture = await element.createTexture(KTX2_TEXTURE_PATH);
+          element.model!.materials[0]
+              .pbrMetallicRoughness.baseColorTexture!.setTexture(ktx2Texture);
+
+          const materialCount = element.model!.materials.length;
+          const dim = element.getDimensions();
+
+          const exported = await element.exportScene({binary: true});
+          expect(exported).to.be.not.undefined;
+          expect(exported.size).to.be.greaterThan(500);
+
+          // Verify scene is still intact after export (no renderer
+          // corruption).
+          expect(element.model).to.not.be.null;
+          expect(element.model!.materials.length).to.equal(materialCount);
+          const dimAfter = element.getDimensions();
+          expect(dimAfter.x).to.be.closeTo(dim.x, 0.001);
+          expect(dimAfter.y).to.be.closeTo(dim.y, 0.001);
+          expect(dimAfter.z).to.be.closeTo(dim.z, 0.001);
+
+          const url = URL.createObjectURL(exported);
+          element.src = url;
+          await waitForEvent(element, 'load');
+
+          expect(element.model).to.not.be.null;
+          expect(element.model!.materials.length).to.be.greaterThan(0);
+        });
+
     test('Verify legacy correlatedObjects are updated.', async () => {
       const newUUID: string|undefined = texture?.source[$threeTexture]?.uuid;