improve-heatmap

zenustech · Dec 17, 2024 · 5cd56e4 · 5cd56e4
1 parent b518b64
commit 5cd56e4
Showing 1 changed file with 20 additions and 20 deletions.
diff --git a/zeno/src/nodes/mtl/ShaderTexture.cpp b/zeno/src/nodes/mtl/ShaderTexture.cpp
@@ -190,6 +190,26 @@ struct SmartTexture2D : ShaderNodeClone<SmartTexture2D>
         auto texId = em->tex2Ds.size();
         auto tex = std::make_shared<zeno::Texture2DObject>();
         auto texture_path = get_input2<std::string>("path");
+        if (has_input("heatmap")) {
+            if (texture_path.empty()) {
+                std::srand(std::time(0));
+                texture_path = std::filesystem::temp_directory_path().string() + '/' + "heatmap-" + std::to_string(std::rand()) + ".png";
+            }
+            auto heatmap = get_input<zeno::HeatmapObject>("heatmap");
+            std::vector<uint8_t> col;
+            int width = heatmap->colors.size();
+            int height = width;
+            col.reserve(width * height * 3);
+            for (auto i = 0; i < height; i++) {
+                for (auto & color : heatmap->colors) {
+                    col.push_back(zeno::clamp(int(color[0] * 255.99), 0, 255));
+                    col.push_back(zeno::clamp(int(color[1] * 255.99), 0, 255));
+                    col.push_back(zeno::clamp(int(color[2] * 255.99), 0, 255));
+                }
+            }
+            stbi_flip_vertically_on_write(false);
+            stbi_write_png(texture_path.c_str(), width, height, 3, col.data(), 0);
+        }
         if(!std::filesystem::exists(std::filesystem::u8path(texture_path))){
             zeno::log_warn("texture file not found!");
             auto type = get_input2<std::string>("type");
@@ -242,26 +262,6 @@ struct SmartTexture2D : ShaderNodeClone<SmartTexture2D>
         }
 
         tex->path = texture_path;
-        if (has_input("heatmap")) {
-            if (tex->path.empty()) {
-                std::srand(std::time(0));
-                tex->path = std::filesystem::temp_directory_path().string() + '/' + "heatmap-" + std::to_string(std::rand()) + ".png";
-            }
-            auto heatmap = get_input<zeno::HeatmapObject>("heatmap");
-            std::vector<uint8_t> col;
-            int width = heatmap->colors.size();
-            int height = width;
-            col.reserve(width * height * 3);
-            for (auto i = 0; i < height; i++) {
-                for (auto & color : heatmap->colors) {
-                    col.push_back(zeno::clamp(int(color[0] * 255.99), 0, 255));
-                    col.push_back(zeno::clamp(int(color[1] * 255.99), 0, 255));
-                    col.push_back(zeno::clamp(int(color[2] * 255.99), 0, 255));
-                }
-            }
-            stbi_flip_vertically_on_write(false);
-            stbi_write_png(tex->path.c_str(), width, height, 3, col.data(), 0);
-        }
         tex->blockCompression = get_input2<bool>("blockCompression");
 
     #define SET_TEX_WRAP(TEX, WRAP)                                    \