Rendering improvements

simulator/model.cpp

@@ -18,8 +18,9 @@ namespace mrover {
             Assimp::Importer importer;
             importer.SetPropertyInteger(AI_CONFIG_PP_SBP_REMOVE, aiPrimitiveType_POINT | aiPrimitiveType_LINE); // Drop points and lines
 
-            // aiScene const* scene = importer.ReadFile(uri.data(),aiProcessPreset_TargetRealtime_MaxQuality);
-            aiScene const* scene = importer.ReadFile(uriToPath(uri), aiProcessPreset_TargetRealtime_Fast | aiProcess_FlipUVs);
+            // we switch to GenSmoothNormals because this seems to fix normals on some meshes
+            // if you shade flat in blender it shouldn't smooth it
+            aiScene const* scene = importer.ReadFile(uriToPath(uri), (aiProcessPreset_TargetRealtime_Fast ^ aiProcess_GenNormals) | aiProcess_GenSmoothNormals | aiProcess_FlipUVs);
             if (!scene) throw std::runtime_error{std::format("Scene import error: {} on path: {}", importer.GetErrorString(), uri)};
 
             RCLCPP_INFO_STREAM(logger, std::format("Loaded scene: {} with mesh count: {}", uri, scene->mNumMeshes));
@@ -34,7 +35,7 @@ namespace mrover {
                 if (!mesh->HasNormals()) throw std::invalid_argument{std::format("Mesh #{} has no normals", meshIndex)};
                 if (!mesh->HasTextureCoords(0)) throw std::invalid_argument{std::format("Mesh #{} has no texture coordinates", meshIndex)};
 
-                auto& [vertices, normals, uvs, indices, texture] = meshes.emplace_back();
+                auto& [vertices, normals, tangents, bitangents, uvs, indices, texture, normal_map, roughness, metallic] = meshes.emplace_back();
 
                 assert(mesh->HasPositions());
                 vertices.data.resize(mesh->mNumVertices);
@@ -67,6 +68,17 @@ namespace mrover {
                     uvs.data[uvIndex] = {uv.x, uv.y};
                 }
 
+                assert(mesh->HasTangentsAndBitangents());
+                tangents.data.resize(mesh->mNumVertices);
+                bitangents.data.resize(mesh->mNumVertices);
+                for (std::size_t vertexIndex = 0; vertexIndex < mesh->mNumVertices; ++vertexIndex) {
+                    aiVector3D const& tangent = mesh->mTangents[vertexIndex];
+                    aiVector3D const& bitangent = mesh->mBitangents[vertexIndex];
+                    tangents.data[vertexIndex] = Eigen::Vector3f{tangent.x, tangent.y, tangent.z};
+                    // from testing it seems like the bitangents blender exports are backwards, but this seems sus
+                    bitangents.data[vertexIndex] = Eigen::Vector3f{-bitangent.x, -bitangent.y, -bitangent.z};
+                }
+
                 if (aiMaterial const* material = scene->mMaterials[mesh->mMaterialIndex]) {
                     if (aiString path; material->GetTextureCount(aiTextureType_DIFFUSE) > 0 && material->GetTexture(aiTextureType_DIFFUSE, 0, &path) == AI_SUCCESS) {
                         texture.data = readTexture(path.C_Str());
@@ -82,6 +94,23 @@ namespace mrover {
                     }
                     aiString name;
                     material->Get(AI_MATKEY_NAME, name);
+
+                    if (aiString path; material->GetTextureCount(aiTextureType_NORMALS) > 0 && material->GetTexture(aiTextureType_NORMALS, 0, &path) == AI_SUCCESS) {
+                        normal_map.data = readTexture(path.C_Str());
+                    } else {
+                        // create a 1x1 texture with normal pointing straight out
+                        // this is simply the vector (0, 0, 1) (points up in the Z direction)
+                        cv::Scalar color(255, 128, 128);
+                        normal_map.data = cv::Mat{1, 1, CV_8UC4, color};
+                    }
+
+                    material->Get(AI_MATKEY_ROUGHNESS_FACTOR, roughness);
+                    RCLCPP_INFO_STREAM(logger, std::format("\t\troughness: {}", roughness));
+
+                    // don't ask me why metallic is this and not AI_MATKEY_METALLIC_FACTOR
+                    material->Get(AI_MATKEY_REFLECTIVITY, metallic);
+                    RCLCPP_INFO_STREAM(logger, std::format("\t\tmetallic: {}", metallic));
+
                     RCLCPP_INFO_STREAM(logger, std::format("\tLoaded material: {}", name.C_Str()));
                 }
 

simulator/shaders/shaders.wgsl

@@ -14,28 +14,38 @@ struct SceneUniforms {
 struct MeshUniforms {
     modelToWorld: mat4x4f,
     modelToWorldForNormals: mat4x4f,
+    roughness: f32,
+    metallic: f32,
 }
 @group(0) @binding(0) var<uniform> mu: MeshUniforms;
 @group(0) @binding(1) var texture: texture_2d<f32>;
 @group(0) @binding(2) var textureSampler: sampler;
+@group(0) @binding(3) var normalTexture: texture_2d<f32>;
 
 struct InVertex {
     @location(0) positionInModel: vec3f,
     @location(1) normalInModel: vec3f,
-    @location(2) uv: vec2f,
+    @location(2) tangentInModel: vec3f,
+    @location(3) bitangentInModel: vec3f,
+    @location(4) uv: vec2f,
 }
 struct OutVertex {
    @builtin(position) positionInClip: vec4f,
    @location(0) positionInWorld: vec4f,
    @location(1) normalInWorld: vec4f,
-   @location(2) uv: vec2f,
+   @location(2) tangentInWorld: vec4f,
+   @location(3) bitangentInWorld: vec4f,
+   @location(4) uv: vec2f,
 }
 @vertex fn vs_main(in: InVertex) -> OutVertex {
     let positionInWorld = mu.modelToWorld * vec4(in.positionInModel, 1);
     var out : OutVertex;
     out.positionInClip = su.cameraToClip * su.worldToCamera * positionInWorld;
     out.positionInWorld = positionInWorld;
     out.normalInWorld = mu.modelToWorldForNormals * vec4(in.normalInModel, 0);
+    // TODO: should I just use the regular modelToWorld matrix here?
+    out.tangentInWorld = mu.modelToWorldForNormals * vec4(in.tangentInModel, 0);
+    out.bitangentInWorld = mu.modelToWorldForNormals * vec4(in.bitangentInModel, 0);
     out.uv = in.uv;
     return out;
 }
@@ -44,27 +54,151 @@ struct OutFragment {
     @location(0) color: vec4f,
     @location(1) normalInCamera: vec4f,
 }
+// Blinn-Phong shader
 @fragment fn fs_main(in: OutVertex) -> OutFragment {
     let baseColor = textureSample(texture, textureSampler, in.uv);
 
+    let normalMapStrength = 1.0;
+    // each component is in the range [0,1]
+    let encodedNormal = textureSample(normalTexture, textureSampler, in.uv).rgb;
+    // shift to get negative values, then normalize because we only care about direction
+    // this normal is in a frame local to the surface of this face
+    let normalInLocal = normalize(encodedNormal - 0.5);
+
+    let localToWorld = mat3x3f(
+        normalize(in.tangentInWorld).xyz,
+        normalize(in.bitangentInWorld).xyz,
+        normalize(in.normalInWorld).xyz,
+    );
+
+    let normalInWorld = normalize(localToWorld * normalInLocal);
+    let mixedNormal = normalize(mix(in.normalInWorld, vec4(normalInWorld, 0), normalMapStrength));
+
     var out : OutFragment;
-    out.normalInCamera = (su.worldToCamera * in.normalInWorld + vec4(1, 1, 1, 0)) / 2;
+    out.normalInCamera = (su.worldToCamera * mixedNormal + vec4(1, 1, 1, 0)) / 2;
     out.normalInCamera.a = 1;
     // Ambient
-    let ambientStrength = 0.6;
+    let ambientStrength = 0.4;
     let ambient = ambientStrength * su.lightColor;
     // Diffuse
     let lightDirInWorld = normalize(su.lightInWorld - in.positionInWorld);
-    let diff = max(dot(in.normalInWorld, lightDirInWorld), 0.0);
+    let diff = max(dot(mixedNormal, lightDirInWorld), 0.0);
     let diffuse = diff * su.lightColor;
     // Specular
-    let specularStrength = 0.5;
-    let viewDirInWolrd = normalize(su.cameraInWorld - in.positionInWorld);
-    let reflectDir = reflect(-lightDirInWorld, in.normalInWorld);
-    let spec = pow(max(dot(viewDirInWolrd, reflectDir), 0.0), 32);
+    let specularStrength = 0.08;
+    let viewDirInWorld = normalize(su.cameraInWorld - in.positionInWorld);
+    // blinn-phong
+    let halfwayDir = normalize(lightDirInWorld + viewDirInWorld);
+    let spec = pow(max(dot(mixedNormal, halfwayDir), 0.0), 32);
+
     let specular = specularStrength * spec * su.lightColor;
     // Combination
-    out.color = vec4(((ambient + diffuse + specular) * baseColor).rgb, 1);
+    out.color = vec4(((ambient + diffuse) * baseColor + specular).rgb, 1);
+    return out;
+}
+
+// PBR FS
+// see https://learnopengl.com/PBR/Lighting
+fn fresnelSchlick(cosTheta: f32, F0: vec3f) -> vec3f {
+    return F0 + (1.0 - F0) * pow(clamp(1.0 - cosTheta, 0.0, 1.0), 5.0);
+}
+
+fn DistributionGGX(normal: vec3f, halfway: vec3f, roughness: f32) -> f32 {
+    let a = roughness * roughness;
+    let a2 = a * a;
+    let NdotH = max(dot(normal, halfway), 0.0);
+    let NdotH2 = NdotH * NdotH;
+
+    let num = a2;
+    let x = (NdotH2 * (a2 - 1.0) + 1.0);
+    let denom = 3.14159 * x * x;
+
+    return num / denom;
+}
+
+fn GeometrySchlickGGX(NdotV: f32, roughness: f32) -> f32 {
+    let r = roughness + 1.0;
+    let k = r * r / 8.0;
+
+    let num = NdotV;
+    let denom = NdotV * (1.0 - k) + k;
+
+    return num / denom;
+}
+
+fn GeometrySmith(normal: vec3f, viewDir: vec3f, lightDir: vec3f, roughness: f32) -> f32 {
+    let NdotV = max(dot(normal, viewDir), 0.0);
+    let NdotL = max(dot(normal, lightDir), 0.0);
+    let ggx1 = GeometrySchlickGGX(NdotL, roughness);
+    let ggx2 = GeometrySchlickGGX(NdotV, roughness);
+
+    return ggx1 * ggx2;
+}
+
+@fragment fn fs_main_pbr(in: OutVertex) -> OutFragment {
+    // let metallic = 0.0;
+    let metallic = mu.metallic;
+    // let roughness = 0.93;
+    let roughness = mu.roughness;
+    // crank up the ambient lighting effect otherwise everything is really dark
+    let ao = 0.8;
+
+    // get the normal vector
+    let normalMapStrength = 1.0;
+    // each component is in the range [0,1]
+    let encodedNormal = textureSample(normalTexture, textureSampler, in.uv).rgb;
+    // shift to get negative values, then normalize because we only care about direction
+    // this normal is in a frame local to the surface of this face
+    let normalInLocal = normalize(encodedNormal - 0.5);
+
+    let localToWorld = mat3x3f(
+        normalize(in.tangentInWorld).xyz,
+        normalize(in.bitangentInWorld).xyz,
+        normalize(in.normalInWorld).xyz,
+    );
+
+    let normalInWorld = normalize(localToWorld * normalInLocal);
+    let mixedNormal = normalize(mix(in.normalInWorld.xyz, normalInWorld, normalMapStrength));
+
+    let viewDirInWorld = normalize((su.cameraInWorld - in.positionInWorld).xyz);
+    let lightDirInWorld = normalize((su.lightInWorld - in.positionInWorld).xyz);
+    let halfwayDir = normalize((lightDirInWorld + viewDirInWorld).xyz);
+
+    // now PBR stuff
+    let distanceToLight = length(su.lightInWorld - in.positionInWorld);
+    // let attenuation = 1 / (distanceToLight * distanceToLight);
+    let attenuation = 1.0;
+    let radiance = su.lightColor.rgb * attenuation * 4;
+
+    let albedo = textureSample(texture, textureSampler, in.uv).rgb;
+    let F0 = mix(vec3(0.04), albedo, metallic);
+    let F = fresnelSchlick(max(dot(halfwayDir, viewDirInWorld), 0), F0);
+
+    let NDF = DistributionGGX(mixedNormal, halfwayDir, roughness);
+    let G = GeometrySmith(mixedNormal, viewDirInWorld, lightDirInWorld, roughness);
+    let numerator = NDF * G * F;
+    let denominator = 4.0 * max(dot(mixedNormal, viewDirInWorld), 0.0) * max(dot(mixedNormal, lightDirInWorld), 0.0) + 0.0001;
+    let specular = numerator / denominator;
+
+    let kS = F;
+    let kD = (vec3(1.0) - kS) * (1.0 - metallic);
+
+    let PI = 3.14159265;
+    let NdotL = max(dot(mixedNormal, lightDirInWorld), 0.0);
+    let Lo = (kD * albedo / PI + specular) * radiance * NdotL;
+    let ambient = 0.4 * albedo * ao;
+
+    var color = ambient + Lo;
+    // gamma correct
+    // TODO: why does this make it look bad...
+    // color = color / (color + vec3(1.0));
+    // color = pow(color, vec3(1.0/2.2));
+
+    var out : OutFragment;
+    out.normalInCamera = (su.worldToCamera * vec4(mixedNormal, 0) + vec4(1, 1, 1, 0)) / 2;
+    out.normalInCamera.a = 1;
+    // TODO: think about alpha??
+    out.color = vec4(color, 1.0);
     return out;
 }
 

simulator/simulator.gui.cpp

@@ -68,6 +68,8 @@ namespace mrover {
             ImGui::Checkbox("Enable Physics (P)", &mEnablePhysics);
             ImGui::Checkbox("Render Models (M)", &mRenderModels);
             ImGui::Checkbox("Render Wireframe Colliders (C)", &mRenderWireframeColliders);
+            ImGui::Checkbox("Use PBR Rendering", &mPbrEnabled);
+            ImGui::Checkbox("Use Normal Mapping", &mNormalMapEnabled);
             ImGui::Text("Camera Locked: %s", mCameraInRoverTarget ? "True" : "False");
             ImGui::SliderFloat("Camera Lock Lerp", &mCameraLockSlerp, 0.0f, 1.0f);
 

simulator/simulator.hpp

@@ -37,6 +37,8 @@ namespace mrover {
     struct ModelUniforms {
         Eigen::Matrix4f modelToWorld{};
         Eigen::Matrix4f modelToWorldForNormals{};
+        float roughness;
+        float metallic;
 
         EIGEN_MAKE_ALIGNED_OPERATOR_NEW
     };
@@ -63,9 +65,15 @@ namespace mrover {
         struct Mesh {
             SharedBuffer<Eigen::Vector3f> vertices;
             SharedBuffer<Eigen::Vector3f> normals;
+            SharedBuffer<Eigen::Vector3f> tangents;
+            SharedBuffer<Eigen::Vector3f> bitangents;
             SharedBuffer<Eigen::Vector2f> uvs;
             SharedBuffer<std::uint32_t> indices;
             MeshTexture texture;
+            MeshTexture normal_map;
+            // could make these also textures, but that might be doing too much
+            float roughness;
+            float metallic;
         };
 
         // DO NOT access the mesh unless you are certain it has been set from the async loader
@@ -242,6 +250,8 @@ namespace mrover {
         bool mEnablePhysics{};
         bool mRenderModels = true;
         bool mRenderWireframeColliders = false;
+        bool mPbrEnabled = false;
+        bool mNormalMapEnabled = false;
         double mPublishHammerDistanceThreshold = 3;
         double mPublishBottleDistanceThreshold = 3;
         float mCameraLockSlerp = 0.02;
@@ -313,6 +323,7 @@ namespace mrover {
         wgpu::TextureView mNormalTextureView;
 
         wgpu::ShaderModule mShaderModule;
+        wgpu::RenderPipeline mBlinnPhongPipeline;
         wgpu::RenderPipeline mPbrPipeline;
         wgpu::RenderPipeline mWireframePipeline;
 
@@ -327,6 +338,9 @@ namespace mrover {
 
         Eigen::Vector4f mSkyColor{0.05f, 0.8f, 0.92f, 1.0f};
 
+        // 1x1 normal map texture that's just flat
+        MeshTexture mDefaultNormalMap = {.data = cv::Mat{1, 1, CV_8UC4, {255, 128, 128}}};
+
         // Physics
 
         std::unique_ptr<btDefaultCollisionConfiguration> mCollisionConfig;