diff options
| -rw-r--r-- | CMakeLists.txt | 1 | ||||
| -rw-r--r-- | src/kernels/tensor.hip | 106 | ||||
| -rw-r--r-- | src/main.cpp | 30 | ||||
| -rw-r--r-- | src/pipeline.h | 5 |
4 files changed, 138 insertions, 4 deletions
diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ac5561..e6aa46e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,6 +10,7 @@ find_package(hip REQUIRED) add_executable(kuwahara src/main.cpp src/kernels/pyramid.hip + src/kernels/tensor.hip ) target_include_directories(kuwahara PRIVATE diff --git a/src/kernels/tensor.hip b/src/kernels/tensor.hip new file mode 100644 index 0000000..d97bb5e --- /dev/null +++ b/src/kernels/tensor.hip @@ -0,0 +1,106 @@ +#include "pipeline.h" + +/* Structure tensor from Jähne's rotationally optimized Sobel derivatives */ + +__global__ void structure_tensor( + const float4* src, float4* dst, int w, int h) +{ + int x = blockIdx.x * blockDim.x + threadIdx.x; + int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= w || y >= h) return; + + const float p1 = 0.183f; + const float p2 = 1.0f - 2.0f * p1; + + int xm = max(x - 1, 0), xp = min(x + 1, w - 1); + int ym = max(y - 1, 0), yp = min(y + 1, h - 1); + + float4 c_mm = src[ym * w + xm], c_0m = src[ym * w + x], c_pm = src[ym * w + xp]; + float4 c_m0 = src[y * w + xm], c_p0 = src[y * w + xp]; + float4 c_mp = src[yp * w + xm], c_0p = src[yp * w + x], c_pp = src[yp * w + xp]; + + float3 gx, gy; + gx.x = 0.5f * (p1 * (c_pm.x - c_mm.x) + p2 * (c_p0.x - c_m0.x) + p1 * (c_pp.x - c_mp.x)); + gx.y = 0.5f * (p1 * (c_pm.y - c_mm.y) + p2 * (c_p0.y - c_m0.y) + p1 * (c_pp.y - c_mp.y)); + gx.z = 0.5f * (p1 * (c_pm.z - c_mm.z) + p2 * (c_p0.z - c_m0.z) + p1 * (c_pp.z - c_mp.z)); + gy.x = 0.5f * (p1 * (c_mp.x - c_mm.x) + p2 * (c_0p.x - c_0m.x) + p1 * (c_pp.x - c_pm.x)); + gy.y = 0.5f * (p1 * (c_mp.y - c_mm.y) + p2 * (c_0p.y - c_0m.y) + p1 * (c_pp.y - c_pm.y)); + gy.z = 0.5f * (p1 * (c_mp.z - c_mm.z) + p2 * (c_0p.z - c_0m.z) + p1 * (c_pp.z - c_pm.z)); + + dst[y * w + x] = make_float4( + gx.x * gx.x + gx.y * gx.y + gx.z * gx.z, + gx.x * gy.x + gx.y * gy.y + gx.z * gy.z, + gy.x * gy.x + gy.y * gy.y + gy.z * gy.z, + 0.0f); +} + +__device__ float4 gaussian_sample_axis( + const float4* src, int src_len, int src_stride, int src_base, + int center, float sigma) +{ + int radius = static_cast<int>(ceilf(3.0f * sigma)); + float inv_two_sigma2 = 1.0f / (2.0f * sigma * sigma); + + float4 sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float wsum = 0.0f; + + for (int i = -radius; i <= radius; ++i) { + int c = min(max(center + i, 0), src_len - 1); + float wt = expf(-(i * i) * inv_two_sigma2); + float4 s = src[src_base + c * src_stride]; + sum.x += s.x * wt; sum.y += s.y * wt; + sum.z += s.z * wt; sum.w += s.w * wt; + wsum += wt; + } + + float inv = 1.0f / wsum; + sum.x *= inv; sum.y *= inv; sum.z *= inv; sum.w *= inv; + return sum; +} + +__global__ void gaussian_horizontal( + const float4* src, float4* dst, int w, int h, float sigma) +{ + int x = blockIdx.x * blockDim.x + threadIdx.x; + int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= w || y >= h) return; + + dst[y * w + x] = gaussian_sample_axis(src, w, 1, y * w, x, sigma); +} + +__global__ void gaussian_vertical( + const float4* src, float4* dst, int w, int h, float sigma) +{ + int x = blockIdx.x * blockDim.x + threadIdx.x; + int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= w || y >= h) return; + + dst[y * w + x] = gaussian_sample_axis(src, h, w, x, y, sigma); +} + +void compute_tensor(PyramidLevel& lvl) { + dim3 block(16, 16); + dim3 grid((lvl.width + 15) / 16, (lvl.height + 15) / 16); + + structure_tensor<<<grid, block>>>( + lvl.color_smax, lvl.tensor, lvl.width, lvl.height); + hip_check(hipGetLastError()); +} + +void smooth_tensor(PyramidLevel& lvl, const Params& p) { + float4* tmp; + hip_check(hipMalloc(&tmp, sizeof(float4) * lvl.width * lvl.height)); + + dim3 block(16, 16); + dim3 grid((lvl.width + 15) / 16, (lvl.height + 15) / 16); + + gaussian_horizontal<<<grid, block>>>( + lvl.tensor, tmp, lvl.width, lvl.height, p.rho); + hip_check(hipGetLastError()); + + gaussian_vertical<<<grid, block>>>( + tmp, lvl.tensor, lvl.width, lvl.height, p.rho); + hip_check(hipGetLastError()); + + hipFree(tmp); +} 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); diff --git a/src/pipeline.h b/src/pipeline.h index 93a7ea9..d533b73 100644 --- a/src/pipeline.h +++ b/src/pipeline.h @@ -31,6 +31,7 @@ struct Pyramid { struct Params { float radius = 5.0f; + float rho = 2.0f; // structure tensor smoothing float alpha = 1.0f; // eccentricity tuning int num_sectors = 8; float q = 8.0f; @@ -47,6 +48,6 @@ Pyramid make_pyramid(int base_width, int base_height, int num_levels); void destroy_pyramid(Pyramid& pyr); void build_pyramid (Pyramid& pyr); -void compute_tensor(Pyramid& pyr, const Params& p); -void smooth_tensor (Pyramid& pyr, const Params& p); +void compute_tensor(PyramidLevel& lvl); +void smooth_tensor (PyramidLevel& lvl, const Params& p); void run_pipeline (Pyramid& pyr, const Params& p); |