[WIP] Segmentation based highlights recovery for bayer sensors

src/common/distance_transform.h

@@ -53,19 +53,14 @@
 
 #include "common/imagebuf.h"
 
-// We don't want to use the SIMD version as we might access unaligned memory
-static inline float sqrf(float a)
-{
-  return a * a;
-}
-
 typedef enum dt_distance_transform_t
 {
   DT_DISTANCE_TRANSFORM_NONE = 0,
   DT_DISTANCE_TRANSFORM_MASK = 1
 } dt_distance_transform_t;
 
 #define DT_DISTANCE_TRANSFORM_MAX (1e20)
+#define SQR(x) ((x) * (x))
 
 static void _image_distance_transform(const float *f, float *z, float *d, int *v, const int n)
 {
@@ -75,11 +70,11 @@ static void _image_distance_transform(const float *f, float *z, float *d, int *v
   z[1] = DT_DISTANCE_TRANSFORM_MAX;
   for(int q = 1; q <= n-1; q++)
   {
-    float s = (f[q] + sqrf((float)q)) - (f[v[k]] + sqrf((float)v[k]));
+    float s = (f[q] + SQR((float)q)) - (f[v[k]] + SQR((float)v[k]));
     while(s <= z[k] * (float)(2*q - 2*v[k]))
     {
       k--;
-      s = (f[q] + sqrf((float)q)) - (f[v[k]] + sqrf((float)v[k]));
+      s = (f[q] + SQR((float)q)) - (f[v[k]] + SQR((float)v[k]));
     }
     s /= (float)(2*q - 2*v[k]);
     k++;
@@ -93,7 +88,7 @@ static void _image_distance_transform(const float *f, float *z, float *d, int *v
   {
     while(z[k+1] < (float)q)
       k++;
-    d[q] = sqrf((float)(q-v[k])) + f[v[k]];
+    d[q] = SQR((float)(q-v[k])) + f[v[k]];
   }
 }
 
@@ -148,7 +143,7 @@ float dt_image_distance_transform(float *const restrict src, float *const restri
     // implicit barrier :-)
     // transform along rows
 #ifdef _OPENMP
-  #pragma omp for schedule(simd:static) nowait
+  #pragma omp for schedule(simd:static)
 #endif
     for(size_t y = 0; y < height; y++)
     {
@@ -167,4 +162,5 @@ float dt_image_distance_transform(float *const restrict src, float *const restri
   }
   return max_distance;
 }
+#undef SQR
 

src/iop/highlights.c

@@ -23,12 +23,17 @@
 #include <stdlib.h>
 #include <string.h>
 #include "bauhaus/bauhaus.h"
+#include "common/box_filters.h"
+#include "common/bspline.h"
 #include "common/opencl.h"
+#include "common/imagebuf.h"
 #include "control/control.h"
 #include "develop/develop.h"
 #include "develop/imageop.h"
+#include "develop/masks.h"
 #include "develop/imageop_math.h"
 #include "develop/imageop_gui.h"
+#include "develop/noise_generator.h"
 #include "develop/tiling.h"
 #include "gui/accelerators.h"
 #include "gui/gtk.h"
@@ -37,32 +42,49 @@
 #include <gtk/gtk.h>
 #include <inttypes.h>
 
-
-DT_MODULE_INTROSPECTION(2, dt_iop_highlights_params_t)
+DT_MODULE_INTROSPECTION(3, dt_iop_highlights_params_t)
 
 typedef enum dt_iop_highlights_mode_t
 {
-  DT_IOP_HIGHLIGHTS_CLIP = 0,    // $DESCRIPTION: "clip highlights"
-  DT_IOP_HIGHLIGHTS_LCH = 1,     // $DESCRIPTION: "reconstruct in LCh"
-  DT_IOP_HIGHLIGHTS_INPAINT = 2, // $DESCRIPTION: "reconstruct color"
+  DT_IOP_HIGHLIGHTS_CLIP = 0,      // $DESCRIPTION: "clip highlights"
+  DT_IOP_HIGHLIGHTS_LCH = 1,       // $DESCRIPTION: "reconstruct in LCh"
+  DT_IOP_HIGHLIGHTS_INPAINT = 2,   // $DESCRIPTION: "reconstruct color"
+  DT_IOP_HIGHLIGHTS_RECOVERY = 4,  // $DESCRIPTION: "highlights recovery"
 } dt_iop_highlights_mode_t;
 
 typedef struct dt_iop_highlights_params_t
 {
+  // params of v1
   dt_iop_highlights_mode_t mode; // $DEFAULT: DT_IOP_HIGHLIGHTS_CLIP $DESCRIPTION: "method"
   float blendL; // unused $DEFAULT: 1.0
   float blendC; // unused $DEFAULT: 0.0
   float blendh; // unused $DEFAULT: 0.0
-  float clip; // $MIN: 0.0 $MAX: 2.0 $DEFAULT: 1.0 $DESCRIPTION: "clipping threshold"
+  // params of v2
+  float clip; // $MIN: 0.0 $MAX: 2.0 $DEFAULT: 0.995 $DESCRIPTION: "clipping threshold"
+  // params of v3
+  float noise_level; // $MIN: 0. $MAX: 1.0 $DEFAULT: 0.05 $DESCRIPTION: "noise level"
+  float reconstructing;    // $MIN: 0.0 $MAX: 1.0  $DEFAULT: 0.4 $DESCRIPTION: "cast correction"
+  float combine;           // $MIN: 0.0 $MAX: 10.0 $DEFAULT: 2.0 $DESCRIPTION: "combine segments"
+  float synthesis;
 } dt_iop_highlights_params_t;
 
 typedef struct dt_iop_highlights_gui_data_t
 {
   GtkWidget *clip;
-  GtkWidget *mode;
+  GtkWidget *mode_bayer, *mode_xtrans, *nonraw;
+  GtkWidget *reconstructing, *combine;
+  GtkWidget *debug;
 } dt_iop_highlights_gui_data_t;
 
-typedef dt_iop_highlights_params_t dt_iop_highlights_data_t;
+typedef struct dt_iop_highlights_data_t
+{ 
+  dt_iop_highlights_mode_t mode;
+  float blendL, blendC, blendh;
+  float clip;
+  float noise_level;
+  float reconstructing, combine, synthesis;
+  gboolean debug;
+} dt_iop_highlights_data_t;
 
 typedef struct dt_iop_highlights_global_data_t
 {
@@ -105,13 +127,28 @@ int default_colorspace(dt_iop_module_t *self, dt_dev_pixelpipe_t *pipe, dt_dev_p
 int legacy_params(dt_iop_module_t *self, const void *const old_params, const int old_version,
                   void *new_params, const int new_version)
 {
-  if(old_version == 1 && new_version == 2)
+  if(old_version == 1 && new_version == 3)
   {
-    memcpy(new_params, old_params, sizeof(dt_iop_highlights_params_t) - sizeof(float));
+    memcpy(new_params, old_params, sizeof(dt_iop_highlights_params_t) - 5 * sizeof(float));
     dt_iop_highlights_params_t *n = (dt_iop_highlights_params_t *)new_params;
     n->clip = 1.0f;
+    n->noise_level = 0.05f;
+    n->reconstructing = 0.4f;
+    n->combine = 2.f;
+    n->synthesis = 0.f;
+    return 0;
+  }
+  if(old_version == 2 && new_version == 3)
+  {
+    memcpy(new_params, old_params, sizeof(dt_iop_highlights_params_t) - 4 * sizeof(float));
+    dt_iop_highlights_params_t *n = (dt_iop_highlights_params_t *)new_params;
+    n->noise_level = 0.05f;
+    n->reconstructing = 0.4f;
+    n->combine = 2.f;
+    n->synthesis = 0.f;
     return 0;
   }
+
   return 1;
 }
 
@@ -801,13 +838,55 @@ static void process_clip(dt_dev_pixelpipe_iop_t *piece, const void *const ivoid,
     }
   }
 }
+static void process_debug(dt_dev_pixelpipe_iop_t *piece, const void *const ivoid, void *const ovoid,
+                         const dt_iop_roi_t *const roi_in, const dt_iop_roi_t *const roi_out,
+                         const uint32_t filters, dt_iop_highlights_data_t *data)
+{
+  const float *const in = (const float *const)ivoid;
+  float *const out = (float *const)ovoid;
+  const int width = roi_out->width;
+  const int height = roi_out->height;
+  const float clip = data->clip;
+  float icoeffs[4] = { piece->pipe->dsc.temperature.coeffs[0], piece->pipe->dsc.temperature.coeffs[1], piece->pipe->dsc.temperature.coeffs[2], 1.0f};  
+  // make sure we have so wb coeffs
+  if((icoeffs[0] < 0.1f) || (icoeffs[0] < 0.1f) || (icoeffs[0] < 0.1f))
+  {
+    fprintf(stderr, "[highlights reconstruction in recovery mode] no white balance coeffs found, choosing stupid defaults\n");
+    icoeffs[0] = 2.0f;
+    icoeffs[1] = 1.0f;
+    icoeffs[2] = 1.5f;
+  }
+#ifdef _OPENMP
+  #pragma omp parallel for simd default(none) \
+  dt_omp_firstprivate(in, out, icoeffs) \
+  dt_omp_sharedconst(height, width, filters, clip) \
+  schedule(simd:static) aligned(in, out : 64)
+#endif
+  for(size_t row = 0; row < height; row++)
+  {
+    for(size_t col = 0, i = row*width; col < width; col++, i++)
+    {
+      const int c = FC(row, col, filters);
+      const float val = in[i] / icoeffs[c];
+      out[i] = (val < clip) ? 0.0f : 1.0f;
+    }
+  }
+}
+
+#include "iop/highlights_algos/hlrecovery.c"
 
 void process(struct dt_iop_module_t *self, dt_dev_pixelpipe_iop_t *piece, const void *const ivoid,
              void *const ovoid, const dt_iop_roi_t *const roi_in, const dt_iop_roi_t *const roi_out)
 {
   const uint32_t filters = piece->pipe->dsc.filters;
   dt_iop_highlights_data_t *data = (dt_iop_highlights_data_t *)piece->data;
 
+  if(data->debug)
+  {
+    process_debug(piece, ivoid, ovoid, roi_in, roi_out, filters, data);
+    return;
+  }
+
   const float clip
       = data->clip * fminf(piece->pipe->dsc.processed_maximum[0],
                            fminf(piece->pipe->dsc.processed_maximum[1], piece->pipe->dsc.processed_maximum[2]));
@@ -885,12 +964,18 @@ void process(struct dt_iop_module_t *self, dt_dev_pixelpipe_iop_t *piece, const
       }
       break;
     }
+
     case DT_IOP_HIGHLIGHTS_LCH:
       if(filters == 9u)
         process_lch_xtrans(self, piece, ivoid, ovoid, roi_in, roi_out, clip);
       else
         process_lch_bayer(self, piece, ivoid, ovoid, roi_in, roi_out, clip);
       break;
+
+    case DT_IOP_HIGHLIGHTS_RECOVERY:
+      process_recovery(piece, ivoid, ovoid, roi_in, roi_out, filters, data);
+      break;
+
     default:
     case DT_IOP_HIGHLIGHTS_CLIP:
       process_clip(piece, ivoid, ovoid, roi_in, roi_out, clip);
@@ -910,12 +995,24 @@ void commit_params(struct dt_iop_module_t *self, dt_iop_params_t *p1, dt_dev_pix
 {
   dt_iop_highlights_params_t *p = (dt_iop_highlights_params_t *)p1;
   dt_iop_highlights_data_t *d = (dt_iop_highlights_data_t *)piece->data;
+  dt_iop_highlights_gui_data_t *g = (dt_iop_highlights_gui_data_t *)self->gui_data;
+
   memcpy(d, p, sizeof(*p));
 
+  d->debug = (g != NULL) ? dt_bauhaus_widget_get_quad_active(g->debug) : FALSE;
+
+  const uint32_t filters = piece->pipe->dsc.filters;
+  // for safety
+  if((filters == 0 || filters == 9u) && (d->mode == DT_IOP_HIGHLIGHTS_RECOVERY))  
+    d->mode = DT_IOP_HIGHLIGHTS_CLIP;
+
   piece->process_cl_ready = 1;
 
-  // no OpenCL for DT_IOP_HIGHLIGHTS_INPAINT yet.
-  if(d->mode == DT_IOP_HIGHLIGHTS_INPAINT) piece->process_cl_ready = 0;
+  if(d->mode == DT_IOP_HIGHLIGHTS_INPAINT || d->mode == DT_IOP_HIGHLIGHTS_RECOVERY || d->debug)
+    piece->process_cl_ready = 0;
+
+  if(d->mode == DT_IOP_HIGHLIGHTS_RECOVERY)
+    piece->process_tiling_ready = 0;
 }
 
 void init_global(dt_iop_module_so_t *module)
@@ -952,12 +1049,43 @@ void cleanup_pipe(struct dt_iop_module_t *self, dt_dev_pixelpipe_t *pipe, dt_dev
   piece->data = NULL;
 }
 
-void gui_update(struct dt_iop_module_t *self)
+void gui_changed(dt_iop_module_t *self, GtkWidget *w, void *previous)
 {
   dt_iop_highlights_gui_data_t *g = (dt_iop_highlights_gui_data_t *)self->gui_data;
   dt_iop_highlights_params_t *p = (dt_iop_highlights_params_t *)self->params;
+
+  const gboolean bayer =     self->dev->image_storage.buf_dsc.filters != 9u;
+  const gboolean israw =     self->dev->image_storage.buf_dsc.filters != 0;
+  const gboolean recover =   bayer && (p->mode == DT_IOP_HIGHLIGHTS_RECOVERY);
+  dt_iop_highlights_mode_t mode = p->mode;
+  if((israw && !bayer && (mode > DT_IOP_HIGHLIGHTS_INPAINT)) || !israw)
+  {
+    mode = DT_IOP_HIGHLIGHTS_CLIP;
+    p->mode = mode;  
+  }
+
+  gtk_widget_set_visible(g->mode_bayer, bayer && israw);
+  gtk_widget_set_visible(g->mode_xtrans, !bayer && israw);
+  gtk_widget_set_visible(g->nonraw, !israw);
+
+  gtk_widget_set_visible(g->combine, recover);
+  gtk_widget_set_visible(g->reconstructing, recover);
+  gtk_widget_set_visible(g->debug, israw);
+
+  if(bayer && israw)  dt_bauhaus_combobox_set_from_value(g->mode_bayer, mode);  
+  if(!bayer && israw) dt_bauhaus_combobox_set_from_value(g->mode_xtrans, mode);  
+
   dt_bauhaus_slider_set(g->clip, p->clip);
-  dt_bauhaus_combobox_set(g->mode, p->mode);
+  dt_bauhaus_slider_set(g->reconstructing, p->reconstructing);
+  dt_bauhaus_slider_set(g->combine, p->combine);
+}
+
+void gui_update(struct dt_iop_module_t *self)
+{
+   dt_iop_highlights_gui_data_t *g = (dt_iop_highlights_gui_data_t *)self->gui_data;
+   dt_bauhaus_widget_set_quad_active(g->debug, FALSE);
+
+   gui_changed(self, NULL, NULL);
 }
 
 void reload_defaults(dt_iop_module_t *module)
@@ -966,17 +1094,68 @@ void reload_defaults(dt_iop_module_t *module)
   module->default_enabled = dt_image_is_rawprepare_supported(&(module->dev->image_storage));
 }
 
+static void debug_callback(GtkWidget *togglebutton, dt_iop_module_t *self)
+{
+  if(darktable.gui->reset) return;
+  dt_iop_request_focus(self);
+
+  gtk_toggle_button_set_active(GTK_TOGGLE_BUTTON(self->off), TRUE);
+  dt_iop_highlights_gui_data_t *g = (dt_iop_highlights_gui_data_t *)self->gui_data;
+
+  const gboolean state = dt_bauhaus_widget_get_quad_active(g->debug);
+  dt_bauhaus_widget_set_quad_active(g->debug, state);
+
+  dt_dev_reprocess_center(self->dev);
+}
+
+void gui_focus(struct dt_iop_module_t *self, gboolean in)
+{
+  dt_iop_highlights_gui_data_t *g = (dt_iop_highlights_gui_data_t *)self->gui_data;
+  if(!in)
+  {
+    const gboolean was_debug = dt_bauhaus_widget_get_quad_active(g->debug);
+    dt_bauhaus_widget_set_quad_active(GTK_WIDGET(g->debug), FALSE);
+    if(was_debug) dt_dev_reprocess_center(self->dev);
+  }
+}
+
 void gui_init(struct dt_iop_module_t *self)
 {
   dt_iop_highlights_gui_data_t *g = IOP_GUI_ALLOC(highlights);
+  g->nonraw = dt_bauhaus_combobox_from_params(self, "mode");
+  gtk_widget_set_tooltip_text(g->nonraw, _("highlight reconstruction method"));
+  for(int i=0;i<3;i++) dt_bauhaus_combobox_remove_at(g->nonraw, 1);
 
-  g->mode = dt_bauhaus_combobox_from_params(self, "mode");
-  gtk_widget_set_tooltip_text(g->mode, _("highlight reconstruction method"));
+  g->mode_bayer = dt_bauhaus_combobox_from_params(self, "mode");
+  gtk_widget_set_tooltip_text(g->mode_bayer, _("highlight reconstruction method"));
+
+  g->mode_xtrans = dt_bauhaus_combobox_from_params(self, "mode");
+  gtk_widget_set_tooltip_text(g->mode_xtrans, _("highlight reconstruction method"));
+  for(int i=0;i<1;i++) dt_bauhaus_combobox_remove_at(g->mode_xtrans, 3);
 
   g->clip = dt_bauhaus_slider_from_params(self, "clip");
   dt_bauhaus_slider_set_digits(g->clip, 3);
-  gtk_widget_set_tooltip_text(g->clip, _("manually adjust the clipping threshold against "
-                                         "magenta highlights (you shouldn't ever need to touch this)"));
+  gtk_widget_set_tooltip_text(g->clip, _("manually adjust the clipping threshold.\n"
+                                         " - helps against magenta highlights in images with incorrect white point settings.\n"
+                                         " - some algorithms might need a refinement."));
+
+  g->debug = dt_bauhaus_combobox_new(self);
+  dt_bauhaus_widget_set_label(g->debug, NULL, N_("above clip threshold"));
+  dt_bauhaus_widget_set_quad_paint(g->debug, dtgtk_cairo_paint_showmask, CPF_STYLE_FLAT | CPF_DO_NOT_USE_BORDER, NULL);
+  dt_bauhaus_widget_set_quad_toggle(g->debug, TRUE);
+  gtk_widget_set_tooltip_text(g->debug, _("visualize photosizes if above clip threshold")); 
+  g_signal_connect(G_OBJECT(g->debug), "quad-pressed", G_CALLBACK(debug_callback), self);
+  gtk_box_pack_start(GTK_BOX(self->widget), g->debug, FALSE, FALSE, 0);
+
+  g->reconstructing = dt_bauhaus_slider_from_params(self, "reconstructing");
+  gtk_widget_set_tooltip_text(g->reconstructing, _("reduces an existing color cast in regions where color planes are clipped")); 
+  dt_bauhaus_slider_set_factor(g->reconstructing, 100.0f);
+  dt_bauhaus_slider_set_format(g->reconstructing, "%.0f%%");
+
+  g->combine = dt_bauhaus_slider_from_params(self, "combine");
+  dt_bauhaus_slider_set_digits(g->combine, 0);
+  gtk_widget_set_tooltip_text(g->combine, _("combine closely related clipped areas.")); 
+
 }
 
 // modelines: These editor modelines have been set for all relevant files by tools/update_modelines.sh