Add filter per pyramid levelHEADmain

1 files changed, 28 insertions, 2 deletions

diff --git a/src/main.cpp b/src/main.cpp
index 6526f82..c0ec0e0 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -5,6 +5,7 @@
 
 #include "pipeline.h"
 #include <hip/hip_runtime.h>
+#include <cmath>
 #include <cstdio>
 #include <string>
 #include <vector>
@@ -52,6 +53,25 @@ static void dump_level(const PyramidLevel& lvl, const char* path) {
     printf("Wrote %s (%dx%d)\n", path, lvl.width, lvl.height);
 }
 
+static void dump_anisotropy(const PyramidLevel& lvl, const char* path) {
+    const int n = lvl.width * lvl.height;
+    std::vector<float4> host(n);
+    hip_check(hipMemcpy(host.data(), lvl.tensor,
+                        sizeof(float4) * n, hipMemcpyDeviceToHost));
+
+    std::vector<unsigned char> bytes(n);
+    for (int i = 0; i < n; ++i) {
+        const float E = host[i].x, F = host[i].y, G = host[i].z;
+        const float trace = E + G;
+        const float A = trace > 1e-8f
+            ? std::sqrt((E - G) * (E - G) + 4.0f * F * F) / trace
+            : 0.0f;
+        bytes[i] = static_cast<unsigned char>(std::min(A, 1.0f) * 255.0f);
+    }
+    stbi_write_png(path, lvl.width, lvl.height, 1, bytes.data(), lvl.width);
+    printf("Wrote %s (%dx%d)\n", path, lvl.width, lvl.height);
+}
+
 int main(int argc, char* argv[]) {
     if (argc < 2) {
         fprintf(stderr, "Usage: %s <image>\n", argv[0]);
@@ -71,9 +91,15 @@ int main(int argc, char* argv[]) {
 
         build_pyramid(pyr);
 
+        Params params;
         for (int i = 0; i < static_cast<int>(pyr.levels.size()); ++i) {
-            std::string path = "level_" + std::to_string(i) + ".png";
-            dump_level(pyr.levels[i], path.c_str());
+            PyramidLevel& lvl = pyr.levels[i];
+            compute_tensor(lvl);
+            smooth_tensor(lvl, params);
+
+            std::string idx = std::to_string(i);
+            dump_level(lvl, ("level_" + idx + ".png").c_str());
+            dump_anisotropy(lvl, ("anisotropy_" + idx + ".png").c_str());
         }
 
         destroy_pyramid(pyr);