1 /* This file is an image processing operation for GEGL 2 * 3 * GEGL is free software; you can redistribute it and/or 4 * modify it under the terms of the GNU Lesser General Public 5 * License as published by the Free Software Foundation; either 6 * version 3 of the License, or (at your option) any later version. 7 * 8 * GEGL is distributed in the hope that it will be useful, 9 * but WITHOUT ANY WARRANTY; without even the implied warranty of 10 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 11 * Lesser General Public License for more details. 12 * 13 * You should have received a copy of the GNU Lesser General Public 14 * License along with GEGL; if not, see <https://www.gnu.org/licenses/>. 15 * 16 * Segment colors using K-means clustering 17 * 18 * Copyright 2017 Thomas Manni <thomas.manni@free.fr> 19 * 20 */ 21 22 #include "config.h" 23 #include <glib/gi18n-lib.h> 24 25 #ifdef GEGL_PROPERTIES 26 27 property_int (n_clusters, _("Number of clusters"), 5) 28 description (_("Number of clusters")) 29 value_range (2, 255) 30 ui_range (2, 30) 31 32 property_int (max_iterations, _("Max. Iterations"), 5) 33 description (_("Maximum number of iterations")) 34 value_range (1, G_MAXINT) 35 ui_range (1, 30) 36 37 property_seed (seed, _("Random seed"), rand) 38 39 #else 40 41 #define GEGL_OP_FILTER 42 #define GEGL_OP_NAME segment_kmeans 43 #define GEGL_OP_C_SOURCE segment-kmeans.c 44 45 #include "gegl-op.h" 46 47 #define MAX_PIXELS 100000 48 49 #define POW2(x) ((x)*(x)) 50 51 typedef struct 52 { 53 gfloat center[3]; 54 gfloat sum[3]; 55 glong count; 56 } Cluster; 57 58 static void 59 downsample_buffer (GeglBuffer *input, 60 GeglBuffer **downsampled) 61 { 62 gint w, h; 63 glong n_pixels; 64 65 w = gegl_buffer_get_width (input); 66 h = gegl_buffer_get_height (input); 67 n_pixels = w * h; 68 69 if (n_pixels > MAX_PIXELS) 70 { 71 GeglNode *node; 72 GeglNode *source; 73 GeglNode *scale; 74 GeglNode *sink; 75 76 gdouble factor = sqrt (MAX_PIXELS / (gdouble) n_pixels); 77 78 node = gegl_node_new (); 79 80 source = gegl_node_new_child (node, "operation", "gegl:buffer-source", 81 "buffer", input, 82 NULL); 83 84 scale = gegl_node_new_child (node, "operation", "gegl:scale-ratio", 85 "x", factor, 86 "y", factor, 87 NULL); 88 89 sink = gegl_node_new_child (node, "operation", "gegl:buffer-sink", 90 "buffer", downsampled, 91 NULL); 92 93 gegl_node_link_many (source, scale, sink, NULL); 94 gegl_node_process (sink); 95 g_object_unref (node); 96 } 97 else 98 { 99 *downsampled = input; 100 } 101 } 102 103 static inline gfloat 104 get_distance (gfloat *c1, gfloat *c2) 105 { 106 return POW2(c2[0] - c1[0]) + 107 POW2(c2[1] - c1[1]) + 108 POW2(c2[2] - c1[2]); 109 } 110 111 static inline gint 112 find_nearest_cluster (gfloat *pixel, 113 Cluster *clusters, 114 gint n_clusters) 115 { 116 gfloat min_distance = G_MAXFLOAT; 117 gint min_cluster = 0; 118 gint i; 119 120 for (i = 0; i < n_clusters; i++) 121 { 122 gfloat distance = get_distance (clusters[i].center, pixel); 123 124 if (distance < min_distance) 125 { 126 min_distance = distance; 127 min_cluster = i; 128 } 129 } 130 131 return min_cluster; 132 } 133 134 static Cluster * 135 init_clusters (GeglBuffer *input, 136 GeglProperties *o) 137 { 138 Cluster *clusters = g_new0 (Cluster, o->n_clusters); 139 GRand *prg = g_rand_new_with_seed (o->seed); 140 141 gint width = gegl_buffer_get_width (input); 142 gint height = gegl_buffer_get_height (input); 143 gint i; 144 145 for (i = 0; i < o->n_clusters; i++) 146 { 147 gfloat color[3]; 148 GeglRectangle one_pixel = {0, 0, 1, 1}; 149 Cluster *c = clusters + i; 150 151 one_pixel.x = g_rand_int_range (prg, 0, width); 152 one_pixel.y = g_rand_int_range (prg, 0, height); 153 154 gegl_buffer_get (input, &one_pixel, 1.0, babl_format ("CIE Lab float"), 155 color, GEGL_AUTO_ROWSTRIDE, GEGL_ABYSS_NONE); 156 157 c->center[0] = color[0]; 158 c->center[1] = color[1]; 159 c->center[2] = color[2]; 160 c->sum[0] = 0.0f; 161 c->sum[1] = 0.0f; 162 c->sum[2] = 0.0f; 163 c->count = 0; 164 } 165 166 g_rand_free (prg); 167 168 return clusters; 169 } 170 171 static void 172 assign_pixels_to_clusters (GeglBuffer *input, 173 Cluster *clusters, 174 gint n_clusters) 175 { 176 GeglBufferIterator *iter; 177 178 iter = gegl_buffer_iterator_new (input, NULL, 0, babl_format ("CIE Lab float"), 179 GEGL_ACCESS_READ, GEGL_ABYSS_NONE, 1); 180 181 while (gegl_buffer_iterator_next (iter)) 182 { 183 gfloat *pixel = iter->items[0].data; 184 glong n_pixels = iter->length; 185 186 while (n_pixels--) 187 { 188 gint index = find_nearest_cluster (pixel, clusters, n_clusters); 189 190 clusters[index].sum[0] += pixel[0]; 191 clusters[index].sum[1] += pixel[1]; 192 clusters[index].sum[2] += pixel[2]; 193 clusters[index].count++; 194 195 pixel += 3; 196 } 197 } 198 } 199 200 static gboolean 201 update_clusters (Cluster *clusters, 202 gint n_clusters) 203 { 204 gboolean has_changed = FALSE; 205 gint i; 206 207 for (i = 0; i < n_clusters; i++) 208 { 209 gfloat new_center[3]; 210 211 if (!clusters[i].count) 212 continue; 213 214 new_center[0] = clusters[i].sum[0] / clusters[i].count; 215 new_center[1] = clusters[i].sum[1] / clusters[i].count; 216 new_center[2] = clusters[i].sum[2] / clusters[i].count; 217 218 if (new_center[0] != clusters[i].center[0] || 219 new_center[1] != clusters[i].center[1] || 220 new_center[2] != clusters[i].center[2]) 221 has_changed = TRUE; 222 223 clusters[i].center[0] = new_center[0]; 224 clusters[i].center[1] = new_center[1]; 225 clusters[i].center[2] = new_center[2]; 226 clusters[i].sum[0] = 0.0f; 227 clusters[i].sum[1] = 0.0f; 228 clusters[i].sum[2] = 0.0f; 229 clusters[i].count = 0; 230 } 231 232 return has_changed; 233 } 234 235 static void 236 set_output (GeglBuffer *input, 237 GeglBuffer *output, 238 Cluster *clusters, 239 gint n_clusters) 240 { 241 GeglBufferIterator *iter; 242 243 iter = gegl_buffer_iterator_new (output, NULL, 0, babl_format ("CIE Lab float"), 244 GEGL_ACCESS_WRITE, GEGL_ABYSS_NONE, 2); 245 246 gegl_buffer_iterator_add (iter, input, NULL, 0, babl_format ("CIE Lab float"), 247 GEGL_ACCESS_READ, GEGL_ABYSS_NONE); 248 249 while (gegl_buffer_iterator_next (iter)) 250 { 251 gfloat *out_pixel = iter->items[0].data; 252 gfloat *in_pixel = iter->items[1].data; 253 glong n_pixels = iter->length; 254 255 while (n_pixels--) 256 { 257 gint index = find_nearest_cluster (in_pixel, clusters, n_clusters); 258 259 out_pixel[0] = clusters[index].center[0]; 260 out_pixel[1] = clusters[index].center[1]; 261 out_pixel[2] = clusters[index].center[2]; 262 263 out_pixel += 3; 264 in_pixel += 3; 265 } 266 } 267 } 268 269 static void 270 prepare (GeglOperation *operation) 271 { 272 const Babl *format = babl_format ("CIE Lab float"); 273 274 gegl_operation_set_format (operation, "input", format); 275 gegl_operation_set_format (operation, "output", format); 276 } 277 278 static GeglRectangle 279 get_required_for_output (GeglOperation *operation, 280 const gchar *input_pad, 281 const GeglRectangle *roi) 282 { 283 GeglRectangle result = *gegl_operation_source_get_bounding_box (operation, "input"); 284 285 /* Don't request an infinite plane */ 286 if (gegl_rectangle_is_infinite_plane (&result)) 287 return *roi; 288 289 return result; 290 } 291 292 static GeglRectangle 293 get_cached_region (GeglOperation *operation, 294 const GeglRectangle *roi) 295 { 296 GeglRectangle result = *gegl_operation_source_get_bounding_box (operation, "input"); 297 298 if (gegl_rectangle_is_infinite_plane (&result)) 299 return *roi; 300 301 return result; 302 } 303 304 static gboolean 305 process (GeglOperation *operation, 306 GeglBuffer *input, 307 GeglBuffer *output, 308 const GeglRectangle *result, 309 gint level) 310 { 311 GeglProperties *o = GEGL_PROPERTIES (operation); 312 gint iterations = o->max_iterations; 313 Cluster *clusters; 314 GeglBuffer *source; 315 316 /* if pixels count of input buffer > MAX_PIXELS, compute a smaller buffer */ 317 318 downsample_buffer (input, &source); 319 320 /* clusters initialization */ 321 322 clusters = init_clusters (source, o); 323 324 /* perform segmentation */ 325 326 while (iterations--) 327 { 328 assign_pixels_to_clusters (source, clusters, o->n_clusters); 329 330 if (!update_clusters (clusters, o->n_clusters)) 331 break; 332 } 333 334 /* apply cluster colors to output */ 335 336 set_output (input, output, clusters, o->n_clusters); 337 338 g_free (clusters); 339 340 if (source != input) 341 g_object_unref (source); 342 343 return TRUE; 344 } 345 346 static gboolean 347 operation_process (GeglOperation *operation, 348 GeglOperationContext *context, 349 const gchar *output_prop, 350 const GeglRectangle *result, 351 gint level) 352 { 353 GeglOperationClass *operation_class; 354 355 const GeglRectangle *in_rect = 356 gegl_operation_source_get_bounding_box (operation, "input"); 357 358 operation_class = GEGL_OPERATION_CLASS (gegl_op_parent_class); 359 360 if (in_rect && gegl_rectangle_is_infinite_plane (in_rect)) 361 { 362 gpointer in = gegl_operation_context_get_object (context, "input"); 363 gegl_operation_context_take_object (context, "output", 364 g_object_ref (G_OBJECT (in))); 365 return TRUE; 366 } 367 368 return operation_class->process (operation, context, output_prop, result, 369 gegl_operation_context_get_level (context)); 370 } 371 372 static void 373 gegl_op_class_init (GeglOpClass *klass) 374 { 375 GeglOperationClass *operation_class; 376 GeglOperationFilterClass *filter_class; 377 378 operation_class = GEGL_OPERATION_CLASS (klass); 379 filter_class = GEGL_OPERATION_FILTER_CLASS (klass); 380 381 filter_class->process = process; 382 operation_class->prepare = prepare; 383 operation_class->process = operation_process; 384 operation_class->get_required_for_output = get_required_for_output; 385 operation_class->get_cached_region = get_cached_region; 386 operation_class->opencl_support = FALSE; 387 operation_class->threaded = FALSE; 388 389 gegl_operation_class_set_keys (operation_class, 390 "name", "gegl:segment-kmeans", 391 "title", _("K-means Segmentation"), 392 "categories", "color:segmentation", 393 "description", _("Segment colors using K-means clustering"), 394 NULL); 395 } 396 397 #endif 398