Merge branch 'haosulab:rt' into rt

Author: angli66

3rd_party/sapien-vulkan-2

@@ -65,6 +65,9 @@ class SVulkan2Material : public IPxrMaterial {
   void setNormalTexture(std::shared_ptr<IPxrTexture> texture) override;
   void setNormalTextureFromFilename(std::string_view filename) override;
   [[nodiscard]] std::shared_ptr<IPxrTexture> getNormalTexture() const override;
+  void setTransmissionTexture(std::shared_ptr<IPxrTexture> texture) override;
+  void setTransmissionTextureFromFilename(std::string_view filename) override;
+  [[nodiscard]] std::shared_ptr<IPxrTexture> getTransmissionTexture() const override;
 
   [[nodiscard]] std::shared_ptr<svulkan2::resource::SVMetallicMaterial> getMaterial() const {
     return mMaterial;

include/sapien/renderer/svulkan2_material.h

@@ -1341,7 +1341,7 @@ def build_ik_window(self):
             )
 
     def take_screenshot(self, _):
-        picture = self.window.get_float_texture("Color")
+        picture = self.window.get_float_texture(self.target_name)
         for i in range(100000000):
             n = f"sapien_screenshot_{i}.png"
             if os.path.exists(n):

python/py_package/utils/viewer.py

@@ -506,7 +506,7 @@ void buildRenderer(py::module &parent) {
       .def("set_textures", &resource::SVMetallicMaterial::setTextures,
            py::arg("base_color") = nullptr, py::arg("roughness") = nullptr,
            py::arg("normal") = nullptr, py::arg("metallic") = nullptr,
-           py::arg("emission") = nullptr)
+           py::arg("emission") = nullptr, py::arg("transmission") = nullptr)
       .def(
           "set_base_color",
           [](resource::SVMetallicMaterial &mat, py::array_t<float> color) {
@@ -515,7 +515,9 @@ void buildRenderer(py::module &parent) {
           py::arg("rgba"))
       .def("set_specular", &resource::SVMetallicMaterial::setFresnel, py::arg("specular"))
       .def("set_metallic", &resource::SVMetallicMaterial::setMetallic, py::arg("metallic"))
-      .def("set_roughness", &resource::SVMetallicMaterial::setRoughness, py::arg("roughness"));
+      .def("set_roughness", &resource::SVMetallicMaterial::setRoughness, py::arg("roughness"))
+      .def("set_emission", &resource::SVMetallicMaterial::setRoughness, py::arg("emission"))
+      .def("set_transmission", &resource::SVMetallicMaterial::setRoughness, py::arg("transmission"));
 
   PyScene
       .def(

python/pysapien_renderer.hpp

@@ -176,6 +176,18 @@ std::shared_ptr<IPxrTexture> SVulkan2Material::getMetallicTexture() const {
   return tex ? std::make_shared<SVulkan2Texture>(tex) : nullptr;
 }
 
+void SVulkan2Material::setTransmissionTexture(std::shared_ptr<IPxrTexture> texture) {
+  if (auto tex = std::dynamic_pointer_cast<SVulkan2Texture>(texture)) {
+    mMaterial->setTransmissionTexture(tex->getTexture());
+  } else {
+    mMaterial->setTransmissionTexture(nullptr);
+  }
+}
+std::shared_ptr<IPxrTexture> SVulkan2Material::getTransmissionTexture() const {
+  auto tex = mMaterial->getTransmissionTexture();
+  return tex ? std::make_shared<SVulkan2Texture>(tex) : nullptr;
+}
+
 void SVulkan2Material::setEmissionTextureFromFilename(std::string_view filename) {
   setEmissionTexture(mParentRenderer->createTexture(filename));
 }
@@ -191,6 +203,9 @@ void SVulkan2Material::setMetallicTextureFromFilename(std::string_view filename)
 void SVulkan2Material::setRoughnessTextureFromFilename(std::string_view filename) {
   setRoughnessTexture(mParentRenderer->createTexture(filename));
 }
+void SVulkan2Material::setTransmissionTextureFromFilename(std::string_view filename) {
+  setTransmissionTexture(mParentRenderer->createTexture(filename));
+}
 
 } // namespace Renderer
 } // namespace sapien

src/renderer/svulkan2_material.cpp

@@ -37,6 +37,7 @@ layout(set = 2, binding = 2) uniform sampler2D roughnessTexture;
 layout(set = 2, binding = 3) uniform sampler2D normalTexture;
 layout(set = 2, binding = 4) uniform sampler2D metallicTexture;
 layout(set = 2, binding = 5) uniform sampler2D emissionTexture;
+layout(set = 2, binding = 6) uniform sampler2D transmissionTexture;
 
 layout(location = 0) in vec4 inPosition;
 layout(location = 1) in vec4 inPrevPosition;

vulkan_shader/ibl/gbuffer.frag

@@ -45,6 +45,7 @@ layout(set = 2, binding = 2) uniform sampler2D roughnessTexture;
 layout(set = 2, binding = 3) uniform sampler2D normalTexture;
 layout(set = 2, binding = 4) uniform sampler2D metallicTexture;
 layout(set = 2, binding = 5) uniform sampler2D emissionTexture;
+layout(set = 2, binding = 6) uniform sampler2D transmissionTexture;
 
 #include "../common/lights.glsl"
 

vulkan_shader/point/gbuffer1.frag

@@ -19,6 +19,8 @@ layout(set = 2, binding = 1) uniform sampler2D colorTexture;
 layout(set = 2, binding = 2) uniform sampler2D roughnessTexture;
 layout(set = 2, binding = 3) uniform sampler2D normalTexture;
 layout(set = 2, binding = 4) uniform sampler2D metallicTexture;
+layout(set = 2, binding = 5) uniform sampler2D emissionTexture;
+layout(set = 2, binding = 6) uniform sampler2D transmissionTexture;
 
 layout(location = 0) in vec4 inPosition;
 layout(location = 1) in vec2 inUV;

vulkan_shader/point/gbuffer2.frag

@@ -238,7 +238,7 @@ vec3 evalGGXTransmission(vec3 N, vec3 L, vec3 V, float eta, float roughness) {
 
     float dotNH = clamp(dot(N, H), 1e-6, 1);
     float dotVH = clamp(dot(V, H), 1e-6, 1);
-    float dotNL = dot(N, L);
+    float dotNL = clamp(dot(N, L), 1e-6, 1);
     float dotNV = clamp(dot(N, V), 1e-6, 1);
 
     float F = dielectricFresnel(dotVH, 1.0 / eta);
@@ -512,7 +512,7 @@ void main() {
   }
 
   if (rnd(ray.seed) < obj.transparency) {
-    ray.origin = worldPosition + 0.001 * ray.direction;
+    ray.origin = worldPosition;
     ray.attenuation = vec3(1.0);
     ray.normal = worldShadingNormal;
     ray.albedo = baseColor;

vulkan_shader/rt/camera.rchit

@@ -15,6 +15,7 @@ layout(set = 0, binding = 1, rgba32f) uniform image2D outAlbedo;
 layout(set = 0, binding = 2, rgba32f) uniform image2D outNormal;
 layout(set = 0, binding = 3, rgba32ui) uniform uimage2D outSegmentation;
 layout(set = 0, binding = 4, rgba32f) uniform image2D outRadiance;
+layout(set = 0, binding = 5, rgba32f) uniform image2D outPosition;
 
 layout(set = 1, binding = 0) uniform accelerationStructureEXT tlas;
 
@@ -104,6 +105,10 @@ void main() {
         imageStore(outNormal, ivec2(gl_LaunchIDEXT.xy), vec4(mat3(cameraBuffer.viewMatrix) * ray.normal, 1.0));
         imageStore(outAlbedo, ivec2(gl_LaunchIDEXT.xy), vec4(ray.albedo, 1.0));
         imageStore(outSegmentation, ivec2(gl_LaunchIDEXT.xy), ray.segmentation);
+
+        vec4 position = cameraBuffer.viewMatrix * vec4(ray.origin, 1.0);
+        vec4 ndc = cameraBuffer.projectionMatrix * position;
+        imageStore(outPosition, ivec2(gl_LaunchIDEXT.xy), vec4(position.xyz, ndc.z / ndc.w));
       }
     }
     radiance += attenuation * ambientLight;  // HACK: implement unterminated ray properly