quantize colors when exporting OBJ

src/extensions/geo-export/obj-exporter.ts

@@ -4,6 +4,7 @@
  * @author Sukolsak Sakshuwong <sukolsak@stanford.edu>
  */
 
+import { sort, arraySwap } from '../../mol-data/util';
 import { asciiWrite } from '../../mol-io/common/ascii';
 import { Vec3, Mat3, Mat4 } from '../../mol-math/linear-algebra';
 import { RuntimeContext } from '../../mol-task';
@@ -66,6 +67,70 @@ export class ObjExporter extends MeshExporter<ObjData> {
         }
     }
 
+    private static quantizeColors(colorArray: Uint8Array, vertexCount: number) {
+        if (vertexCount <= 1024) return;
+        const rgb = Vec3();
+        const min = Vec3();
+        const max = Vec3();
+        const sum = Vec3();
+        const colorMap = new Map<Color, Color>();
+        const colorComparers = [
+            (colors: Color[], i: number, j: number) => (Color.toVec3(rgb, colors[i])[0] - Color.toVec3(rgb, colors[j])[0]),
+            (colors: Color[], i: number, j: number) => (Color.toVec3(rgb, colors[i])[1] - Color.toVec3(rgb, colors[j])[1]),
+            (colors: Color[], i: number, j: number) => (Color.toVec3(rgb, colors[i])[2] - Color.toVec3(rgb, colors[j])[2]),
+        ];
+
+        const medianCut = (colors: Color[], l: number, r: number, depth: number) => {
+            if (l > r) return;
+            if (l === r || depth >= 10) {
+                // Find the average color.
+                Vec3.set(sum, 0, 0, 0);
+                for (let i = l; i <= r; ++i) {
+                    Color.toVec3(rgb, colors[i]);
+                    Vec3.add(sum, sum, rgb);
+                }
+                Vec3.round(rgb, Vec3.scale(rgb, sum, 1 / (r - l + 1)));
+                const averageColor = Color.fromArray(rgb, 0);
+                for (let i = l; i <= r; ++i) colorMap.set(colors[i], averageColor);
+                return;
+            }
+
+            // Find the color channel with the greatest range.
+            Vec3.set(min, 255, 255, 255);
+            Vec3.set(max, 0, 0, 0);
+            for (let i = l; i <= r; ++i) {
+                Color.toVec3(rgb, colors[i]);
+                for (let j = 0; j < 3; ++j) {
+                    Vec3.min(min, min, rgb);
+                    Vec3.max(max, max, rgb);
+                }
+            }
+            let k = 0;
+            if (max[1] - min[1] > max[k] - min[k]) k = 1;
+            if (max[2] - min[2] > max[k] - min[k]) k = 2;
+
+            sort(colors, l, r + 1, colorComparers[k], arraySwap);
+
+            const m = (l + r) >> 1;
+            medianCut(colors, l, m, depth + 1);
+            medianCut(colors, m + 1, r, depth + 1);
+        };
+
+        // Create an array of unique colors and use the median cut algorithm.
+        const colorSet = new Set<Color>();
+        for (let i = 0; i < vertexCount; ++i) {
+            colorSet.add(Color.fromArray(colorArray, i * 3));
+        }
+        const colors = Array.from(colorSet);
+        medianCut(colors, 0, colors.length - 1, 0);
+
+        // Map actual colors to quantized colors.
+        for (let i = 0; i < vertexCount; ++i) {
+            const color = colorMap.get(Color.fromArray(colorArray, i * 3));
+            Color.toArray(color!, colorArray, i * 3);
+        }
+    }
+
     protected async addMeshWithColors(input: AddMeshInput) {
         const { mesh, meshes, values, isGeoTexture, webgl, ctx } = input;
 
@@ -87,6 +152,7 @@ export class ObjExporter extends MeshExporter<ObjData> {
         let interpolatedColors: Uint8Array;
         if (colorType === 'volume' || colorType === 'volumeInstance') {
             interpolatedColors = ObjExporter.getInterpolatedColors(mesh!.vertices, mesh!.vertexCount, values, stride, colorType, webgl!);
+            ObjExporter.quantizeColors(interpolatedColors, mesh!.vertexCount);
         }
 
         await ctx.update({ isIndeterminate: false, current: 0, max: instanceCount });