NF: shaders and materials modifiable & added smooth/flat shading

fury/actor.py

@@ -15,6 +15,7 @@ def sphere(
     opacity=None,
     material="phong",
     enable_picking=True,
+    smooth=True,
 ):
     """
     Visualize one or many spheres with different colors and radii.
@@ -40,6 +41,8 @@ def sphere(
         The material type for the spheres. Options are 'phong' and 'basic'.
     enable_picking : bool, optional
         Whether the spheres should be pickable in a 3D scene.
+    smooth : bool, optional
+        Whether to create a smooth sphere or a faceted sphere.
 
     Returns
     -------
@@ -95,7 +98,9 @@ def sphere(
         colors=big_colors.astype("float32"),
     )
 
-    mat = _create_mesh_material(material=material, enable_picking=enable_picking)
+    mat = _create_mesh_material(
+        material=material, enable_picking=enable_picking, flat_shading=not smooth
+    )
     obj = create_mesh(geometry=geo, material=mat)
     obj.local.position = centers[0]
     obj.prim_count = prim_count
@@ -121,7 +126,7 @@ def box(
         Box positions.
     directions : ndarray, shape (N, 3), optional
         The orientation vector of the box.
-    colors : ndarray (N,3) or (N, 4) or tuple (3,) or tuple (4,), optional
+    colors : ndarray, shape (N,3) or (N, 4) or tuple (3,) or tuple (4,), optional
         RGB or RGBA (for opacity) R, G, B and A should be at the range [0, 1].
     scales : int or ndarray (N,3) or tuple (3,), optional
         The size of the box in each dimension.  If a single value is provided,

fury/io.py

@@ -1,4 +1,7 @@
+import functools
+import importlib.resources
 import os
+import sys
 
 # from tempfile import TemporaryDirectory as InTemporaryDirectory
 from urllib.request import urlretrieve
@@ -236,6 +239,19 @@ def save_image(
     im.save(filename, quality=compression_quality, dpi=dpi)
 
 
+@functools.lru_cache(maxsize=None)
+def load_wgsl(shader_name, *, package_name="fury.wgsl"):
+    """Load wgsl code from pygfx builtin shader snippets."""
+    if sys.version_info < (3, 9):
+        context = importlib.resources.path(package_name, shader_name)
+    else:
+        ref = importlib.resources.files(package_name) / shader_name
+        context = importlib.resources.as_file(ref)
+    with context as path:
+        with open(path, "rb") as f:
+            return f.read().decode()
+
+
 # def load_polydata(file_name):
 #     """Load a vtk polydata to a supported format file.
 

fury/material.py

@@ -1,8 +1,41 @@
 import pygfx as gfx
+from pygfx import Mesh
+from pygfx.renderers.wgpu import register_wgpu_render_function
+
+from fury.shader import MeshBasicShader, MeshPhongShader
+
+
+class MeshPhongMaterial(gfx.MeshPhongMaterial):
+    """
+    Phong material.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+class MeshBasicMaterial(gfx.MeshBasicMaterial):
+    """
+    Basic material.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+# Register the custom shaders for the mesh materials
+register_wgpu_render_function(Mesh, MeshPhongMaterial)(MeshPhongShader)
+register_wgpu_render_function(Mesh, MeshBasicMaterial)(MeshBasicShader)
 
 
 def _create_mesh_material(
-    *, material="phong", enable_picking=True, color=None, opacity=1.0, mode="vertex"
+    *,
+    material="phong",
+    enable_picking=True,
+    color=None,
+    opacity=1.0,
+    mode="vertex",
+    flat_shading=True,
 ):
     """
     Create a mesh material.
@@ -23,10 +56,12 @@ def _create_mesh_material(
         final_alpha = alpha_in_RGBA * opacity
     mode : str, optional
         The color mode of the material. Options are 'auto' and 'vertex'.
+    flat_shading : bool, optional
+        Whether to use flat shading (True) or smooth shading (False).
 
     Returns
     -------
-    gfx.MeshMaterial
+    MeshMaterial
         A mesh material object of the specified type with the given properties.
     """
 
@@ -49,16 +84,18 @@ def _create_mesh_material(
         color = (1, 1, 1)
 
     if material == "phong":
-        return gfx.MeshPhongMaterial(
+        return MeshPhongMaterial(
             pick_write=enable_picking,
             color_mode=mode,
             color=color,
+            flat_shading=flat_shading,
         )
     elif material == "basic":
-        return gfx.MeshBasicMaterial(
+        return MeshBasicMaterial(
             pick_write=enable_picking,
             color_mode=mode,
             color=color,
+            flat_shading=flat_shading,
         )
     else:
         raise ValueError(f"Unsupported material type: {material}")

fury/shader.py

@@ -0,0 +1,22 @@
+from pygfx.renderers.wgpu.shaders.meshshader import MeshShader
+
+from fury.io import load_wgsl
+
+
+class MeshBasicShader(MeshShader):
+    """Base class for mesh shaders."""
+
+    def __init__(self, wobject):
+        super().__init__(wobject)
+
+    def get_code(self):
+        return load_wgsl("mesh.wgsl")
+
+
+class MeshPhongShader(MeshBasicShader):
+    """Phong shader for meshes."""
+
+    def __init__(self, wobject):
+        super().__init__(wobject)
+
+        self["lighting"] = "phong"

fury/tests/test_material.py

@@ -1,35 +1,78 @@
-from fury import material
+import numpy as np
+
+from fury import material, window
+from fury.geometry import buffer_to_geometry, create_mesh
+from fury.material import _create_mesh_material
+from fury.primitive import prim_sphere
 
 
 def test_create_mesh_material():
     color = (1, 0, 0)
     mat = material._create_mesh_material(
         material="phong", color=color, opacity=0.5, mode="auto"
     )
-    assert type(mat) == material.gfx.MeshPhongMaterial
+    assert isinstance(mat, material.MeshPhongMaterial)
     assert mat.color == color + (0.5,)
     assert mat.color_mode == "auto"
 
     color = (1, 0, 0, 0.5)
     mat = material._create_mesh_material(
-        material="phong", color=color, opacity=0.5, mode="auto"
+        material="phong", color=color, opacity=0.5, mode="auto", flat_shading=False
     )
-    assert type(mat) == material.gfx.MeshPhongMaterial
+    assert isinstance(mat, material.MeshPhongMaterial)
     assert mat.color == (1, 0, 0, 0.25)
     assert mat.color_mode == "auto"
+    assert mat.flat_shading is False
 
     color = (1, 0, 0)
     mat = material._create_mesh_material(
         material="phong", color=color, opacity=0.5, mode="vertex"
     )
-    assert type(mat) == material.gfx.MeshPhongMaterial
+    assert isinstance(mat, material.MeshPhongMaterial)
     assert mat.color == (1, 1, 1)
     assert mat.color_mode == "vertex"
 
     color = (1, 0, 0)
     mat = material._create_mesh_material(
-        material="basic", color=color, mode="vertex", enable_picking=False
+        material="basic",
+        color=color,
+        mode="vertex",
+        enable_picking=False,
+        flat_shading=True,
     )
-    assert type(mat) == material.gfx.MeshBasicMaterial
+    assert isinstance(mat, material.MeshBasicMaterial)
     assert mat.color == (1, 1, 1)
     assert mat.color_mode == "vertex"
+    assert mat.flat_shading is True
+
+    verts, faces = prim_sphere()
+
+    geo = buffer_to_geometry(
+        indices=faces.astype("int32"),
+        positions=verts.astype("float32"),
+        texcoords=verts.astype("float32"),
+        colors=np.ones_like(verts).astype("float32"),
+    )
+
+    mat = _create_mesh_material(
+        material="phong", enable_picking=False, flat_shading=False
+    )
+
+    obj = create_mesh(geometry=geo, material=mat)
+
+    scene = window.Scene()
+
+    scene.add(obj)
+
+    # window.snapshot(scene=scene, fname="mat_test_1.png")
+    #
+    # img = Image.open("mat_test_1.png")
+    # img_array = np.array(img)
+    #
+    # mean_r, mean_g, mean_b, _ = np.mean(
+    #     img_array.reshape(-1, img_array.shape[2]), axis=0
+    # )
+    #
+    # assert 0 <= mean_r <= 255 and 0 <= mean_g <= 255 and 0 <= mean_b <= 255
+    #
+    # assert sum([mean_r, mean_g, mean_b]) > 0

fury/tests/test_shader.py

@@ -0,0 +1,54 @@
+import numpy as np
+from pygfx import Mesh
+from pygfx.renderers.wgpu import register_wgpu_render_function
+
+from fury.actor import sphere
+from fury.io import load_wgsl
+from fury.material import MeshBasicMaterial
+from fury.shader import MeshBasicShader, MeshPhongShader
+
+
+def test_shader():
+    class CustomBasicMaterial(MeshBasicMaterial):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+
+    class CustomPhongMaterial(MeshBasicMaterial):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+
+    register_wgpu_render_function(Mesh, CustomBasicMaterial)(MeshBasicShader)
+    register_wgpu_render_function(Mesh, CustomPhongMaterial)(MeshPhongShader)
+
+    try:
+        register_wgpu_render_function(Mesh, CustomBasicMaterial)(MeshBasicShader)
+        register_wgpu_render_function(Mesh, CustomPhongMaterial)(MeshPhongShader)
+    except ValueError:
+        ...
+    else:
+        raise AssertionError("Shouldn't be able to register the same material twice.")
+
+
+def test_wgsl():
+    shader_code = load_wgsl("mesh.wgsl")
+
+    assert isinstance(shader_code, str)
+    assert "fn vs_main" in shader_code
+    assert "fn fs_main" in shader_code
+
+    actor = sphere(centers=np.array([[0, 0, 0]]), colors=np.array([[1, 0, 0]]))
+    kwargs = {
+        "blending_code": "placeholder",
+        "write_pick": True,
+        "indexer": None,
+        "used_uv": {"uv": None},
+    }
+
+    cs = MeshBasicShader(actor)
+
+    assert isinstance(cs.get_code(), str)
+
+    gen_sh_code = cs.generate_wgsl(**kwargs)
+    assert isinstance(gen_sh_code, str)
+
+    assert "placeholder" in gen_sh_code