Add UV selector texture

include/mitsuba/render/mesh.h

@@ -282,6 +282,10 @@ class MI_EXPORT_LIB Mesh : public Shape<Float, Spectrum> {
                              const SurfaceInteraction3f &si,
                              Mask active = true) const override;
 
+    Vector2f eval_attribute_2(std::string_view name,
+                              const SurfaceInteraction3f &si,
+                              Mask active = true) const override;
+
     SurfaceInteraction3f eval_parameterization(const Point2f &uv,
                                                uint32_t ray_flags = +RayFlags::All,
                                                Mask active = true) const override;
@@ -531,8 +535,11 @@ class MI_EXPORT_LIB Mesh : public Shape<Float, Spectrum> {
         using StorageType =
             std::conditional_t<Size == 1,
                                dr::replace_scalar_t<Float, InputFloat>,
-                               dr::replace_scalar_t<Color3f, InputFloat>>;
-        using ReturnType = std::conditional_t<Size == 1, Float, Color3f>;
+                               std::conditional_t<Size == 2,
+                                                  dr::replace_scalar_t<Vector2f, InputFloat>,
+                                                  dr::replace_scalar_t<Color3f, InputFloat>>>;
+        using ReturnType = std::conditional_t<Size == 1, Float,
+                                              std::conditional_t<Size == 2, Vector2f, Color3f>>;
 
         if (type == MeshAttributeType::Vertex) {
             auto fi = face_indices(si.prim_index, active);

include/mitsuba/render/shape.h

@@ -808,6 +808,22 @@ class MI_EXPORT_LIB Shape : public JitObject<Shape<Float, Spectrum>> {
                                      const SurfaceInteraction3f &si,
                                      Mask active = true) const;
 
+    /**
+     * \brief 2D evaluation of a shape attribute at the given surface interaction
+     *
+     * \param name
+     *     Name of the attribute to evaluate
+     *
+     * \param si
+     *     Surface interaction associated with the query
+     *
+     * \return
+     *     A 2D vector value
+     */
+    virtual Vector2f eval_attribute_2(std::string_view name,
+                                      const SurfaceInteraction3f &si,
+                                      Mask active = true) const;
+
     /**
      * \brief Evaluate a dynamically sized shape attribute at the given surface interaction.
      *
@@ -1171,6 +1187,7 @@ DRJIT_CALL_TEMPLATE_BEGIN(mitsuba::Shape)
     DRJIT_CALL_METHOD(has_attribute)
     DRJIT_CALL_METHOD(eval_attribute)
     DRJIT_CALL_METHOD(eval_attribute_1)
+    DRJIT_CALL_METHOD(eval_attribute_2)
     DRJIT_CALL_METHOD(eval_attribute_3)
     DRJIT_CALL_METHOD(eval_attribute_x)
     DRJIT_CALL_METHOD(eval_parameterization)

src/render/mesh.cpp

@@ -1787,6 +1787,25 @@ Mesh<Float, Spectrum>::eval_attribute_3(std::string_view name,
     }
 }
 
+MI_VARIANT typename Mesh<Float, Spectrum>::Vector2f
+Mesh<Float, Spectrum>::eval_attribute_2(std::string_view name,
+                                        const SurfaceInteraction3f &si,
+                                        Mask active) const {
+    const auto& it = m_mesh_attributes.find(name);
+    if (it == m_mesh_attributes.end())
+        return Base::eval_attribute_2(name, si, active);
+
+    const auto& attr = it->second;
+    if (attr.size == 2) {
+        return interpolate_attribute<2, true>(attr.type, attr.buf, si, active);
+    } else {
+        if constexpr (dr::is_jit_v<Float>)
+            return 0.f;
+        else
+            Throw("eval_attribute_2(): Attribute \"%s\" requested but had size %u.", name, attr.size);
+    }
+}
+
 //! @}
 // =============================================================
 

src/render/python/shape_v.cpp

@@ -132,6 +132,12 @@ template <typename Ptr, typename Cls> void bind_shape_generic(Cls &cls) {
                 return shape->eval_attribute_1(name, si, active);
             },
             "name"_a, "si"_a, "active"_a = true, D(Shape, eval_attribute_1))
+       .def("eval_attribute_2",
+            [](Ptr shape, const std::string &name,
+               const SurfaceInteraction3f &si, const Mask &active) {
+                return shape->eval_attribute_2(name, si, active);
+            },
+            "name"_a, "si"_a, "active"_a = true)
        .def("eval_attribute_3",
             [](Ptr shape, const std::string &name,
                const SurfaceInteraction3f &si, const Mask &active) {

src/render/shape.cpp

@@ -571,6 +571,16 @@ Shape<Float, Spectrum>::eval_attribute_3(std::string_view name,
     return texture->eval_3(si, active);
 }
 
+MI_VARIANT typename Shape<Float, Spectrum>::Vector2f
+Shape<Float, Spectrum>::eval_attribute_2(std::string_view name,
+                                         const SurfaceInteraction3f & /*si*/,
+                                         Mask /*active*/) const {
+    if constexpr (dr::is_jit_v<Float>)
+        return 0.f;
+    else
+        Throw("Invalid attribute requested %s.", name);
+}
+
 MI_VARIANT typename dr::DynamicArray<Float>
 Shape<Float, Spectrum>::eval_attribute_x(std::string_view /*name*/,
                                          const SurfaceInteraction3f & /*si*/,

src/render/tests/test_mesh.py

@@ -1440,7 +1440,7 @@ def test38_ray_intersect_triangle(variants_all_rgb):
     mesh = mi.Mesh(name='', vertex_count=3, face_count=1)
     params = mi.traverse(mesh)
     params['vertex_positions'] = [
-        0.0, 0.0, 0.0, 
+        0.0, 0.0, 0.0,
         1.0, 0.0, 0.0,
         0.5, 1.0, 0.0,
     ]
@@ -1467,3 +1467,17 @@ def test39_custom_vertex_normals(variants_vec_rgb):
 
     # The custom vertex normals should not have been modified.
     assert dr.allclose(params['vertex_normals'], normals)
+
+
+def test40_write_and_load_2d_ply_attribute(variants_all_rgb, tmp_path):
+    mesh = mi.load_dict({"type": "rectangle"})
+    mesh.add_attribute(
+        "vertex_extra_uv", 2, [0.5, 0.5, 1.5, 0.5, 0.5, 1.5, 1.5, 1.5]
+    )
+    ply_file = str(tmp_path / "test_2d_attr.ply")
+    mesh.write_ply(ply_file)
+    mesh_loaded = mi.load_dict({"type": "ply", "filename": ply_file})
+    assert mesh_loaded.has_attribute("vertex_extra_uv")
+    si = mesh_loaded.eval_parameterization([0.5, 0.5])
+    uv = mesh_loaded.eval_attribute_2("vertex_extra_uv", si)
+    assert dr.allclose(uv, [1.0, 1.0])

src/render/tests/test_shape.py

@@ -112,3 +112,16 @@ def set_bsdf(self, bsdf):
     custom_shape.set_bsdf(new_bsdf)
     assert mi.has_flag(custom_shape.bsdf().flags(), mi.BSDFFlags.DeltaReflection)
     assert custom_shape.called
+
+
+def test04_eval_attribute_2(variants_all_backends_once):
+    # Test eval_attribute_2 directly on a mesh with 2D attributes
+    mesh = mi.load_dict({"type": "rectangle"})
+    mesh.add_attribute(
+        "vertex_uv2d", 2, [0.5, 0.5, 1.5, 0.5, 0.5, 1.5, 1.5, 1.5]
+    )
+
+    si = mesh.eval_parameterization([0.5, 0.5])
+    uv2d = mesh.eval_attribute_2("vertex_uv2d", si)
+    assert dr.allclose(uv2d, [1.0, 1.0])
+

src/shapes/ply.h

@@ -329,9 +329,9 @@ void find_other_fields(const std::string& type, std::vector<PLYAttributeDescript
         reading_attribute = false;
     };
     auto flush_attribute = [&]() {
-        if (current_postfix_level_index != 1 && current_postfix_level_index != 3) {
-            Log(Warn, "\"%s\": attribute must have either 1 or 3 fields (had %d) : attribute \"%s\" ignored",
-                name, current_postfix_level_index, (type + current_prefix).c_str());
+        if (current_postfix_level_index != 1 && current_postfix_level_index != 2 && current_postfix_level_index != 3) {
+            Log(Warn, "\"%s\": attribute must have 1, 2 or 3 fields (had %d) : attribute \"%s\" ignored",
+                name, (int) current_postfix_level_index, (type + current_prefix).c_str());
             ignore_attribute();
             return;
         }

src/textures/CMakeLists.txt

@@ -1,8 +1,9 @@
 set(MI_PLUGIN_PREFIX "textures")
 
-add_plugin(bitmap         bitmap.cpp)
-add_plugin(checkerboard   checkerboard.cpp)
-add_plugin(mesh_attribute mesh_attribute.cpp)
-add_plugin(volume         volume.cpp)
+add_plugin(bitmap          bitmap.cpp)
+add_plugin(checkerboard    checkerboard.cpp)
+add_plugin(mesh_attribute  mesh_attribute.cpp)
+add_plugin(uv_set_selector uv_set_selector.cpp)
+add_plugin(volume          volume.cpp)
 
 set(MI_PLUGIN_TARGETS "${MI_PLUGIN_TARGETS}" PARENT_SCOPE)

src/textures/tests/test_uv_set_selector.py

@@ -0,0 +1,45 @@
+import pytest
+import drjit as dr
+import mitsuba as mi
+import numpy as np
+
+
+def create_rectangle_with_extra_uvs():
+    mesh = mi.load_dict({"type": "rectangle"})
+    mesh.add_attribute(
+        "vertex_uv1", 2, [0.5, 0.5, 1.5, 0.5, 0.5, 1.5, 1.5, 1.5]
+    )
+    return mesh
+
+
+def test01_uv_selector_checkerboard(variant_scalar_rgb):
+    mesh = create_rectangle_with_extra_uvs()
+    tensor = mi.TensorXf(
+        np.array(
+            [[[0.1, 0.1, 0.1], [0.4, 0.9, 0.9]], [[0.5, 0.9, 0.1], [0.1, 0.3, 0.1]]]
+        )
+    )
+    bitmap = mi.load_dict({"type": "bitmap", "bitmap": mi.Bitmap(tensor)})
+
+    texture = mi.load_dict(
+        {"type": "uv_set_selector", "uv_set": "vertex_uv1", "nested": bitmap}
+    )
+
+    for u, v in [(0.0, 0.0), (0.4, 0.4), (0.5, 0.5), (0.9, 0.2)]:
+        si = mesh.eval_parameterization([u, v])
+        actual = texture.eval(si)
+        si.uv += mi.Point2f(0.5, 0.5)
+        expected = bitmap.eval(si)
+        assert dr.allclose(actual, expected)
+
+
+def test02_invalid_uv_set(variant_scalar_rgb):
+    with pytest.raises(Exception):
+        mi.load_dict(
+            {
+                "type": "uv_set_selector",
+                "uv_set": "invalid_name",
+                "nested": {"type": "checkerboard"},
+            }
+        )
+

src/textures/uv_set_selector.cpp

@@ -0,0 +1,118 @@
+#include <mitsuba/core/properties.h>
+#include <mitsuba/core/transform.h>
+#include <mitsuba/render/interaction.h>
+#include <mitsuba/render/texture.h>
+
+NAMESPACE_BEGIN(mitsuba)
+
+/**!
+
+.. _texture-uv_set_selector:
+
+UV Set Selector texture (:monosp:`uv_set_selector`)
+-------------------------------------------------
+
+.. pluginparameters::
+
+ * - uv_set
+   - |string|
+   - Name of the mesh attribute representing the alternative UV coordinates. It
+should start with ``"vertex_"`` or ``"face_"``.
+ * - nested
+   - |texture|
+   - The nested texture that should evaluate colors using the selected UV set.
+
+This plugin provides a mechanism to redirect the UV coordinates evaluated by a
+nested texture to a custom UV set (loaded as a mesh attribute).
+
+ */
+
+template <typename Float, typename Spectrum>
+class UvSetSelector final : public Texture<Float, Spectrum> {
+public:
+    MI_IMPORT_TYPES(Texture)
+
+    UvSetSelector(const Properties &props) : Texture(props) {
+        m_uv_set = props.get<std::string>("uv_set");
+        if (m_uv_set.find("vertex_") == std::string::npos &&
+            m_uv_set.find("face_") == std::string::npos)
+            Throw("Invalid mesh attribute name: must start with either "
+                  "\"vertex_\" or \"face_\" but was \"%s\".",
+                  m_uv_set.c_str());
+
+        for (auto &prop : props.objects()) {
+            if (Texture *tex = prop.try_get<Texture>()) {
+                if (m_nested)
+                    Throw("Only a single nested Texture can be specified.");
+                m_nested = tex;
+            }
+        }
+
+        if (!m_nested) {
+            if (props.has_property("nested")) {
+                m_nested = props.get_texture<Texture>("nested");
+            } else {
+                Throw("Exactly one nested Texture must be specified.");
+            }
+        }
+    }
+
+    void traverse(TraversalCallback *cb) override {
+        cb->put("nested", m_nested, ParamFlags::Differentiable);
+    }
+
+    SurfaceInteraction3f select_uv(const SurfaceInteraction3f &si,
+                                   Mask active) const {
+        SurfaceInteraction3f local_si(si);
+        Vector2f attr = si.shape->eval_attribute_2(m_uv_set, si, active);
+        local_si.uv  = Point2f(attr.x(), attr.y());
+        return local_si;
+    }
+
+    UnpolarizedSpectrum eval(const SurfaceInteraction3f &si,
+                             Mask active) const override {
+        MI_MASKED_FUNCTION(ProfilerPhase::TextureEvaluate, active);
+        return m_nested->eval(select_uv(si, active), active);
+    }
+
+    Float eval_1(const SurfaceInteraction3f &si,
+                 Mask active = true) const override {
+        MI_MASKED_FUNCTION(ProfilerPhase::TextureEvaluate, active);
+        return m_nested->eval_1(select_uv(si, active), active);
+    }
+
+    Vector2f eval_1_grad(const SurfaceInteraction3f &si,
+                         Mask active = true) const override {
+        MI_MASKED_FUNCTION(ProfilerPhase::TextureEvaluate, active);
+        return m_nested->eval_1_grad(select_uv(si, active), active);
+    }
+
+    Color3f eval_3(const SurfaceInteraction3f &si,
+                   Mask active = true) const override {
+        MI_MASKED_FUNCTION(ProfilerPhase::TextureEvaluate, active);
+        return m_nested->eval_3(select_uv(si, active), active);
+    }
+
+    Float mean() const override { return m_nested->mean(); }
+
+    bool is_spatially_varying() const override {
+        return m_nested->is_spatially_varying();
+    }
+
+    std::string to_string() const override {
+        std::ostringstream oss;
+        oss << "UvSetSelector[" << std::endl
+            << "  uv_set = \"" << m_uv_set << "\"," << std::endl
+            << "  nested = " << string::indent(m_nested) << std::endl
+            << "]";
+        return oss.str();
+    }
+
+    MI_DECLARE_CLASS(UvSetSelector)
+protected:
+    std::string m_uv_set;
+    ref<Texture> m_nested;
+};
+
+MI_EXPORT_PLUGIN(UvSetSelector)
+NAMESPACE_END(mitsuba)