Add toon3, toon5 material options (#125)

* Add toon3, toon5 material options * Fix types
nerfstudio-project · Nov 6, 2023 · a5b646c · a5b646c
1 parent 93e2699
commit a5b646c
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 4 deletions.
diff --git a/src/viser/_message_api.py b/src/viser/_message_api.py
@@ -452,6 +452,7 @@ def add_mesh_simple(
         color: RgbTupleOrArray = (90, 200, 255),
         wireframe: bool = False,
         opacity: Optional[float] = None,
+        material: Literal["standard", "toon3", "toon5"] = "standard",
         side: Literal["front", "back", "double"] = "front",
         wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),
         position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),
@@ -469,6 +470,7 @@ def add_mesh_simple(
                 wireframe=wireframe,
                 opacity=opacity,
                 side=side,
+                material=material,
             )
         )
         node_handle = MeshHandle._make(self, name, wxyz, position, visible)

diff --git a/src/viser/_messages.py b/src/viser/_messages.py
@@ -185,7 +185,8 @@ class MeshMessage(Message):
 
     wireframe: bool
     opacity: Optional[float]
-    side: Literal["front", "back", "double"] = "front"
+    side: Literal["front", "back", "double"]
+    material: Literal["standard", "toon3", "toon5"]
 
     def __post_init__(self):
         # Check shapes.

diff --git a/src/viser/client/src/App.tsx b/src/viser/client/src/App.tsx
@@ -268,6 +268,12 @@ function ViewerCanvas({ children }: { children: React.ReactNode }) {
         </EffectComposer>
       </Selection>
       <Environment path="/hdri/" files="potsdamer_platz_1k.hdr" />
+      <directionalLight color={0xffffff} intensity={1.0} position={[0, 1, 0]} />
+      <directionalLight
+        color={0xffffff}
+        intensity={0.2}
+        position={[0, -1, 0]}
+      />
     </Canvas>
   );
 }

diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx
@@ -236,18 +236,54 @@ function useMessageHandler() {
       // Add mesh
       case "MeshMessage": {
         const geometry = new THREE.BufferGeometry();
-        const material = new THREE.MeshStandardMaterial({
+
+        const generateGradientMap = (shades: 3 | 5) => {
+          const texture = new THREE.DataTexture(
+            Uint8Array.from(
+              shades == 3
+                ? [0, 0, 0, 255, 128, 128, 128, 255, 255, 255, 255, 255]
+                : [
+                    0, 0, 0, 255, 64, 64, 64, 255, 128, 128, 128, 255, 192, 192,
+                    192, 255, 255, 255, 255, 255,
+                  ],
+            ),
+            shades,
+            1,
+            THREE.RGBAFormat,
+          );
+
+          texture.needsUpdate = true;
+          return texture;
+        };
+        const standardArgs = {
           color: message.color || undefined,
           vertexColors: message.vertex_colors !== null,
           wireframe: message.wireframe,
           transparent: message.opacity !== null,
-          opacity: message.opacity ?? undefined,
+          opacity: message.opacity ?? 1.0,
           side: {
             front: THREE.FrontSide,
             back: THREE.BackSide,
             double: THREE.DoubleSide,
           }[message.side],
-        });
+        };
+        const assertUnreachable = (x: never): never => {
+          throw new Error(`Should never get here! ${x}`);
+        };
+        const material =
+          message.material == "standard"
+            ? new THREE.MeshStandardMaterial(standardArgs)
+            : message.material == "toon3"
+            ? new THREE.MeshToonMaterial({
+                gradientMap: generateGradientMap(3),
+                ...standardArgs,
+              })
+            : message.material == "toon5"
+            ? new THREE.MeshToonMaterial({
+                gradientMap: generateGradientMap(5),
+                ...standardArgs,
+              })
+            : assertUnreachable(message.material);
         geometry.setAttribute(
           "position",
           new THREE.Float32BufferAttribute(

diff --git a/src/viser/client/src/WebsocketMessages.tsx b/src/viser/client/src/WebsocketMessages.tsx
@@ -145,6 +145,7 @@ export interface MeshMessage {
   wireframe: boolean;
   opacity: number | null;
   side: "front" | "back" | "double";
+  material: "standard" | "toon3" | "toon5";
 }
 /** Message for transform gizmos.
  *