Modularize and cleanHEADmain

1 files changed, 87 insertions, 0 deletions

diff --git a/src/main.cpp b/src/main.cpp
new file mode 100644
index 0000000..6526f82
--- /dev/null
+++ b/src/main.cpp
@@ -0,0 +1,87 @@
+#define STB_IMAGE_IMPLEMENTATION
+#include "stb_image.h"
+#define STB_IMAGE_WRITE_IMPLEMENTATION
+#include "stb_image_write.h"
+
+#include "pipeline.h"
+#include <hip/hip_runtime.h>
+#include <cstdio>
+#include <string>
+#include <vector>
+
+static float4* upload_image(const char* path, int& out_width, int& out_height) {
+    int w, h, channels;
+    unsigned char* data = stbi_load(path, &w, &h, &channels, 4);
+    if (!data) throw std::runtime_error(std::string("Failed to load: ") + path);
+
+    out_width = w; out_height = h;
+    const int n = w * h;
+
+    std::vector<float4> host(n);
+    for (int i = 0; i < n; ++i) {
+        host[i] = make_float4(
+            data[4*i + 0] / 255.0f,
+            data[4*i + 1] / 255.0f,
+            data[4*i + 2] / 255.0f,
+            0.0f);
+    }
+    stbi_image_free(data);
+
+    float4* dev;
+    hip_check(hipMalloc(&dev, sizeof(float4) * n));
+    hip_check(hipMemcpy(dev, host.data(), sizeof(float4) * n, hipMemcpyHostToDevice));
+    return dev;
+}
+
+static void dump_level(const PyramidLevel& lvl, const char* path) {
+    const int n = lvl.width * lvl.height;
+    std::vector<float4> host(n);
+    hip_check(hipMemcpy(host.data(), lvl.color_smax,
+                        sizeof(float4) * n, hipMemcpyDeviceToHost));
+
+    std::vector<unsigned char> bytes(n * 3);
+    for (int i = 0; i < n; ++i) {
+        bytes[3*i + 0] = static_cast<unsigned char>(
+            std::min(host[i].x, 1.0f) * 255.0f);
+        bytes[3*i + 1] = static_cast<unsigned char>(
+            std::min(host[i].y, 1.0f) * 255.0f);
+        bytes[3*i + 2] = static_cast<unsigned char>(
+            std::min(host[i].z, 1.0f) * 255.0f);
+    }
+    stbi_write_png(path, lvl.width, lvl.height, 3, bytes.data(), lvl.width * 3);
+    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]);
+        return 1;
+    }
+
+    try {
+        int w, h;
+        float4* src = upload_image(argv[1], w, h);
+
+        const int num_levels = 5;
+        Pyramid pyr = make_pyramid(w, h, num_levels);
+
+        hip_check(hipMemcpy(pyr.levels[0].color_smax, src,
+                            sizeof(float4) * w * h, hipMemcpyDeviceToDevice));
+        hipFree(src);
+
+        build_pyramid(pyr);
+
+        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());
+        }
+
+        destroy_pyramid(pyr);
+
+    } catch (const std::exception& e) {
+        fprintf(stderr, "Error: %s\n", e.what());
+        return 1;
+    }
+
+    return 0;
+}