add ability to export to USDZ

Author: sukolsak

README.md

@@ -2,7 +2,7 @@
 
 [![License](http://img.shields.io/badge/license-MIT-blue.svg?style=flat)](./LICENSE)
 
-A tool for creating macromolecular images and 3D models. It is based on [molrender](https://github.com/molstar/molrender), which is based on [Mol\*](https://github.com/molstar/molstar). It can create PNG, JPEG, OBJ, and GLB files.
+A tool for creating macromolecular images and 3D models. It is based on [molrender](https://github.com/molstar/molrender), which is based on [Mol\*](https://github.com/molstar/molstar). It can create PNG, JPEG, OBJ, GLB, and USDZ files.
 
 The tool is being used on [MolAR](https://stanford.edu/~sukolsak/ar/) to create 3D models of proteins for augmented reality.
 
@@ -19,7 +19,7 @@ The `options` are
 
     --width WIDTH    image height
     --height HEIGHT  image width
-    --format FORMAT  image format (png, jpeg, obj, or glb)
+    --format FORMAT  image format (png, jpeg, obj, glb, or usdz)
 
 For example, the following command will create `6vxx.obj` in the `out` folder. The `6vxx.cif` file (structure data of SARS-CoV-2 spike glycoprotein) can be downloaded from https://files.rcsb.org/download/6vxx.cif.
 

src/3d-exporter.ts

@@ -6,6 +6,7 @@
 
 import { Vec3 } from 'molstar/lib/mol-math/linear-algebra';
 import { Color } from 'molstar/lib/mol-util/color/color';
+import { ValueType, UsdAttribute, UsdData, CrateFile } from './usdz';
 
 function computeBounding(points: Vec3[]) {
     const min = Vec3.create(Number.MAX_VALUE, Number.MAX_VALUE, Number.MAX_VALUE);
@@ -207,3 +208,66 @@ export function exportGlb(meshByColor: Map<Color, Mesh>) {
     ]);
     return glb;
 }
+
+export function exportUsdz(meshByColor: Map<Color, Mesh>) {
+    const root = new UsdData('', '');
+    const usdObj = root.createChild('ar', 'Xform');
+    usdObj.metadata['assetInfo'] = {'name': 'ar'};
+    usdObj.metadata['kind'] = 'component';
+
+    const materials_scope = usdObj.createChild('Materials', 'Scope');
+
+    let refId = 0;
+
+    const vecsToFloatArray = (v: number[][]) => {
+        const a = new Float32Array(v.length * 3);
+        for (let i = 0; i < v.length; ++i) {
+            a[i * 3] = v[i][0];
+            a[i * 3 + 1] = v[i][1];
+            a[i * 3 + 2] = v[i][2];
+        }
+        return a;
+    };
+
+    meshByColor.forEach((mesh, color) => {
+        const { positions, normals, faces } = mesh;
+        const positions2 = vecsToFloatArray(positions);
+        const normals2 = vecsToFloatArray(normals);
+        const faceVertexCounts = new Uint32Array(faces.length / 3);
+        faceVertexCounts.fill(3);
+
+        const rgb = new Float32Array(Color.toRgbNormalized(color));
+        const usdMaterial = materials_scope.createChild('k' + String(refId), 'Material');
+
+        const usdShader = usdMaterial.createChild('surfaceShader', 'Shader');
+        const infoIdAtt = new UsdAttribute('info:id', 'UsdPreviewSurface', ValueType.token, 'token');
+        infoIdAtt.addQualifier('uniform');
+        usdShader.addAttribute(infoIdAtt);
+        usdShader.addAttribute(new UsdAttribute('inputs:diffuseColor', rgb, ValueType.vec3f, 'color3f'));
+        usdShader.addAttribute(new UsdAttribute('inputs:roughness', 0.2, ValueType.float, 'float'));
+        const surface = new UsdAttribute('outputs:surface', null, ValueType.token, 'token');
+        usdShader.addAttribute(surface);
+
+        usdMaterial.addAttribute(new UsdAttribute('outputs:surface', surface, ValueType.Invalid, 'token'));
+
+        const usdMesh = usdObj.createChild('m' + String(refId), 'Mesh');
+        usdMesh.addAttribute(new UsdAttribute('material:binding', usdMaterial, ValueType.Invalid, 'rel'));
+        usdMesh.addAttribute(new UsdAttribute('doubleSided', false, ValueType.bool, 'bool'));
+        usdMesh.addAttribute(new UsdAttribute('faceVertexCounts', faceVertexCounts, ValueType.int, 'int[]', true));
+        usdMesh.addAttribute(new UsdAttribute('faceVertexIndices', faces, ValueType.int, 'int[]', true));
+        usdMesh.addAttribute(new UsdAttribute('points', positions2, ValueType.vec3f, 'point3f[]', true));
+        const normalsAtt = new UsdAttribute('primvars:normals', normals2, ValueType.vec3f, 'normal3f[]', true);
+        normalsAtt.metadata['interpolation'] = 'vertex';
+        usdMesh.addAttribute(normalsAtt);
+        const subdivAtt = new UsdAttribute('subdivisionScheme', 'none', ValueType.token, 'token');
+        subdivAtt.addQualifier('uniform');
+        usdMesh.addAttribute(subdivAtt);
+
+        refId++;
+    });
+
+    const crateFile = new CrateFile();
+    crateFile.writeUsd(root);
+    const usdz = crateFile.getUsdz();
+    return Buffer.from(usdz.buffer);
+}

src/index.ts

@@ -41,7 +41,7 @@ function addBasicArgs(currParser: argparse.ArgumentParser, isDir: boolean) {
     });
     currParser.addArgument([ '--format' ], {
         action: 'store',
-        help: 'image format (png, jpeg, obj, or glb)',
+        help: 'image format (png, jpeg, obj, glb, or usdz)',
         defaultValue: 'png'
     });
 }

src/render.ts

@@ -38,7 +38,7 @@ import { ColorNames } from 'molstar/lib/mol-util/color/names';
 import { Camera } from 'molstar/lib/mol-canvas3d/camera';
 import { SyncRuntimeContext } from 'molstar/lib/mol-task/execution/synchronous';
 import { AssetManager } from 'molstar/lib/mol-util/assets';
-import { Mesh, exportObj, exportGlb } from './3d-exporter';
+import { Mesh, exportObj, exportGlb, exportUsdz } from './3d-exporter';
 
 /**
  * Helper method to create PNG with given PNG data
@@ -65,9 +65,9 @@ async function writeObjFile(obj: string, objOutPath: string, mtl: string, mtlOut
         })
     ]);
 }
-async function writeGlbFile(glb: Buffer, outPath: string) {
+async function writeFile(data: Buffer, outPath: string) {
     await new Promise<void>(resolve => {
-        fs.writeFile(outPath, glb, () => resolve());
+        fs.writeFile(outPath, data, () => resolve());
     });
 }
 
@@ -157,7 +157,7 @@ export class ImageRenderer {
     imagePass: ImagePass
     assetManager = new AssetManager()
 
-    constructor(private width: number, private height: number, private format: 'png' | 'jpeg' | 'obj' | 'glb') {
+    constructor(private width: number, private height: number, private format: 'png' | 'jpeg' | 'obj' | 'glb' | 'usdz') {
         this.gl = getGLContext(this.width, this.height, {
             alpha: false,
             antialias: true,
@@ -261,7 +261,7 @@ export class ImageRenderer {
     }
 
     /**
-     * Creates OBJ/GLB with the current 3dcanvas data
+     * Creates OBJ/GLB/USDZ with the current 3dcanvas data
      */
     async create3DModel(outPath: string) {
         this.canvas3d.commit(true);
@@ -394,7 +394,7 @@ export class ImageRenderer {
         });
         const translate = Vec3.create(-(box.min[0] + box.max[0]) / 2, -box.min[1], -(box.min[2] + box.max[2]) / 2);
         const size = Math.max(box.max[0] - box.min[0], box.max[1] - box.min[1], box.max[2] - box.min[2]);
-        const scale = (this.format === 'glb') ? Math.min(0.4 / size, 0.01) : 1;
+        const scale = (this.format === 'glb') ? Math.min(0.4 / size, 0.01) : (this.format === 'usdz') ? Math.min(40 / size, 1) : 1;
         meshByColor.forEach(mesh => {
             for (const position of mesh.positions) {
                 Vec3.add(position, position, translate);
@@ -409,12 +409,15 @@ export class ImageRenderer {
             await writeObjFile(obj, `${outPath}.obj`, mtl, `${outPath}.mtl`);
         } else if (this.format === 'glb') {
             const glb = exportGlb(meshByColor);
-            await writeGlbFile(glb, `${outPath}.glb`);
+            await writeFile(glb, `${outPath}.glb`);
+        } else if (this.format === 'usdz') {
+            const usdz = exportUsdz(meshByColor);
+            await writeFile(usdz, `${outPath}.usdz`);
         }
     }
 
     async createImage(outPath: string, size: StructureSize) {
-        if (this.format === 'obj' || this.format === 'glb') {
+        if (this.format === 'obj' || this.format === 'glb' || this.format === 'usdz') {
             await this.create3DModel(outPath);
             return;
         }