NLM: add smoothing control parameter *h*

Which is what *sigma* used to do, but in fact it should have been used
to improve weights computation.

3 files changed, 39 insertions, 14 deletions

diff --git a/docs/filters.rst b/docs/filters.rst
index 530ca67..9d6fed0 100644
--- a/docs/filters.rst
+++ b/docs/filters.rst
@@ -538,9 +538,15 @@ Non-local-means denoising
 
         Radius of patches.
 
+    .. gobj:prop:: h:float
+
+        Smoothing control parameter, should be around noise standard deviation
+        or slightly less. Higher h results in a smoother image but with blurred
+        features.
+
     .. gobj:prop:: sigma:float
 
-        Sigma influencing the Gaussian weighting.
+        Noise standard deviation, improves weights computation.
 
     .. gobj:prop:: addressing-mode:enum
 
diff --git a/src/kernels/nlm.cl b/src/kernels/nlm.cl
index c3ae3f3..df96fbd 100644
--- a/src/kernels/nlm.cl
+++ b/src/kernels/nlm.cl
@@ -24,7 +24,8 @@ dist (read_only image2d_t input,
       float2 q,
       int radius,
       int width,
-      int height)
+      int height,
+      float variance)
 {
     float dist = 0.0f, tmp;
     float wsize = (2.0f * radius + 1.0f);
@@ -34,7 +35,7 @@ dist (read_only image2d_t input,
         for (int j = -radius; j < radius + 1; j++) {
             tmp = read_imagef (input, sampler, (float2) ((p.x + i) / width, (p.y + j) / height)).x -
                     read_imagef (input, sampler, (float2) ((q.x + i) / width, (q.y + j) / height)).x;
-            dist += tmp * tmp;
+            dist += fmax (0.0f, tmp * tmp - 2 * variance);
         }
     }
 
@@ -47,13 +48,13 @@ nlm_noise_reduction (read_only image2d_t input,
                      sampler_t sampler,
                      const int search_radius,
                      const int patch_radius,
-                     const float sigma)
+                     const float h_2,
+                     const float variance)
 {
     const int x = get_global_id (0);
     const int y = get_global_id (1);
     const int width = get_global_size (0);
     const int height = get_global_size (1);
-    const float sigma_2 = sigma * sigma;
     float d, weight;
 
     float total_weight = 0.0f;
@@ -62,8 +63,8 @@ nlm_noise_reduction (read_only image2d_t input,
     for (int i = x - search_radius; i < x + search_radius + 1; i++) {
         for (int j = y - search_radius; j < y + search_radius + 1; j++) {
             d = dist (input, sampler, (float2) (x + 0.5f, y + 0.5f), (float2) (i + 0.5f, j + 0.5f),
-                      patch_radius, width, height);
-            weight = exp (- sigma_2 * d);
+                      patch_radius, width, height, variance);
+            weight = exp (- h_2 * d);
             pixel_value += weight * read_imagef (input, sampler, (float2) ((i + 0.5f) / width, (j + 0.5f) / height)).x;
             total_weight += weight;
         }
diff --git a/src/ufo-non-local-means-task.c b/src/ufo-non-local-means-task.c
index b953b1c..764778d 100644
--- a/src/ufo-non-local-means-task.c
+++ b/src/ufo-non-local-means-task.c
@@ -30,6 +30,7 @@
 struct _UfoNonLocalMeansTaskPrivate {
     guint search_radius;
     guint patch_radius;
+    gfloat h;
     gfloat sigma;
     cl_kernel kernel;
     cl_sampler sampler;
@@ -49,6 +50,7 @@ enum {
     PROP_0,
     PROP_SEARCH_RADIUS,
     PROP_PATCH_RADIUS,
+    PROP_H,
     PROP_SIGMA,
     PROP_ADDRESSING_MODE,
     N_PROPERTIES
@@ -127,21 +129,23 @@ ufo_non_local_means_task_process (UfoTask *task,
     cl_command_queue cmd_queue;
     cl_mem in_mem;
     cl_mem out_mem;
-    gfloat sigma;
+    gfloat h, var;
 
     priv = UFO_NON_LOCAL_MEANS_TASK_GET_PRIVATE (task);
     node = UFO_GPU_NODE (ufo_task_node_get_proc_node (UFO_TASK_NODE (task)));
     cmd_queue = ufo_gpu_node_get_cmd_queue (node);
     in_mem = ufo_buffer_get_device_image (inputs[0], cmd_queue);
     out_mem = ufo_buffer_get_device_array (output, cmd_queue);
-    sigma = 1 / priv->sigma;
+    h = 1 / priv->h / priv->h;
+    var = priv->sigma * priv->sigma;
 
     UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 0, sizeof (cl_mem), &in_mem));
     UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 1, sizeof (cl_mem), &out_mem));
     UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 2, sizeof (cl_sampler), &priv->sampler));
     UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 3, sizeof (guint), &priv->search_radius));
     UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 4, sizeof (guint), &priv->patch_radius));
-    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 5, sizeof (gfloat), &sigma));
+    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 5, sizeof (gfloat), &h));
+    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->kernel, 6, sizeof (gfloat), &var));
 
     profiler = ufo_task_node_get_profiler (UFO_TASK_NODE (task));
     ufo_profiler_call (profiler, cmd_queue, priv->kernel, 2, requisition->dims, NULL);
@@ -165,6 +169,9 @@ ufo_non_local_means_task_set_property (GObject *object,
         case PROP_PATCH_RADIUS:
             priv->patch_radius = g_value_get_uint (value);
             break;
+        case PROP_H:
+            priv->h = g_value_get_float (value);
+            break;
         case PROP_SIGMA:
             priv->sigma = g_value_get_float (value);
             break;
@@ -192,6 +199,9 @@ ufo_non_local_means_task_get_property (GObject *object,
         case PROP_PATCH_RADIUS:
             g_value_set_uint (value, priv->patch_radius);
             break;
+        case PROP_H:
+            g_value_set_float (value, priv->h);
+            break;
         case PROP_SIGMA:
             g_value_set_float (value, priv->sigma);
             break;
@@ -261,11 +271,18 @@ ufo_non_local_means_task_class_init (UfoNonLocalMeansTaskClass *klass)
             1, 100, 3,
             G_PARAM_READWRITE);
 
+    properties[PROP_H] =
+        g_param_spec_float ("h",
+            "Smoothing control parameter, should be around noise standard deviation or slightly less",
+            "Smoothing control parameter, should be around noise standard deviation or slightly less",
+            0.0f, G_MAXFLOAT, 0.1f,
+            G_PARAM_READWRITE);
+
     properties[PROP_SIGMA] =
         g_param_spec_float ("sigma",
-            "Sigma",
-            "Sigma",
-            0.0f, G_MAXFLOAT, 0.1f,
+            "Noise standard deviation",
+            "Noise standard deviation",
+            0.0f, G_MAXFLOAT, 0.0f,
             G_PARAM_READWRITE);
 
     properties[PROP_ADDRESSING_MODE] =
@@ -289,6 +306,7 @@ ufo_non_local_means_task_init(UfoNonLocalMeansTask *self)
 
     self->priv->search_radius = 10;
     self->priv->patch_radius = 3;
-    self->priv->sigma = 0.1f;
+    self->priv->h = 0.1f;
+    self->priv->sigma = 0.0f;
     self->priv->addressing_mode = CL_ADDRESS_MIRRORED_REPEAT;
 }