cg_vertex_denoiser.c (9242B)
1 /* Copyright (C) 2022 Université de Pau et des Pays de l'Adour UPPA 2 * Copyright (C) 2022 CNRS 3 * Copyright (C) 2022 Sorbonne Université 4 * Copyright (C) 2022 Université Paul Sabatier 5 * Copyright (C) 2022 |Meso|Star> (contact@meso-star.com) 6 * 7 * This program is free software: you can redistribute it and/or modify 8 * it under the terms of the GNU General Public License as published by 9 * the Free Software Foundation, either version 3 of the License, or 10 * (at your option) any later version. 11 * 12 * This program is distributed in the hope that it will be useful, 13 * but WITHOUT ANY WARRANTY; without even the implied warranty of 14 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 * GNU General Public License for more details. 16 * 17 * You should have received a copy of the GNU General Public License 18 * along with this program. If not, see <http://www.gnu.org/licenses/>. */ 19 20 #include "cg.h" 21 #include "cg_vertex_denoiser.h" 22 23 #include <rsys/rsys.h> 24 #include <rsys/dynamic_array_double.h> 25 #include <rsys/double3.h> 26 #include <rsys/mem_allocator.h> 27 #include <rsys/str.h> 28 29 #include <float.h> 30 31 #define NODE_COUNT_MAX 64 32 33 enum node_type { 34 LEAF, 35 NODE 36 }; 37 38 struct leaf { 39 double vertices[3*NODE_COUNT_MAX]; 40 unsigned count; 41 }; 42 43 struct node { 44 enum node_type type; 45 unsigned split_dim; 46 double split_pos; 47 union { 48 struct leaf* leaf; 49 struct node* children[2]; 50 } c; 51 }; 52 53 static res_T 54 alloc_leaf_node 55 (struct mem_allocator* allocator, 56 struct node** leaf) 57 { 58 res_T res = RES_OK; 59 struct node* node = NULL; 60 61 ASSERT(allocator && leaf); 62 63 node = MEM_CALLOC(allocator, 1, sizeof(*node)); 64 if(!node) { 65 res = RES_MEM_ERR; 66 goto error; 67 } 68 node->c.leaf = MEM_CALLOC(allocator, 1, sizeof(*node->c.leaf)); 69 if(!node->c.leaf) { 70 res = RES_MEM_ERR; 71 goto error; 72 } 73 node->type = LEAF; 74 75 exit: 76 *leaf = node; 77 return res; 78 error: 79 if(node && node->c.leaf) MEM_RM(allocator, node->c.leaf); 80 if(node) MEM_RM(allocator, node); 81 node = NULL; 82 goto exit; 83 } 84 85 static void 86 release_node 87 (struct mem_allocator* allocator, 88 struct node* node) 89 { 90 ASSERT(allocator && node); 91 if(node->type == LEAF) { 92 MEM_RM(allocator, node->c.leaf); 93 MEM_RM(allocator, node); 94 } else { 95 ASSERT(node->type == NODE); 96 release_node(allocator, node->c.children[0]); 97 release_node(allocator, node->c.children[1]); 98 MEM_RM(allocator, node); 99 } 100 } 101 102 static void 103 add_to_leaf 104 (struct node* node, 105 const double v[3]) 106 { 107 ASSERT(node && v && node->c.leaf->count < NODE_COUNT_MAX); 108 d3_set(node->c.leaf->vertices + (node->c.leaf->count++ * 3), v); 109 } 110 111 static res_T 112 search_and_add 113 (struct node* node, 114 struct vertex_denoiser* denoiser, 115 const double in[3], 116 const int can_add, 117 double out[3], 118 int* done, 119 size_t* added_count, 120 size_t* denoised_count) 121 { 122 res_T res = RES_OK; 123 struct mem_allocator* allocator; 124 125 ASSERT(node && denoiser && in && out && added_count && denoised_count); 126 allocator = denoiser->allocator; 127 128 if(node->type == LEAF) { 129 int same = 0; 130 unsigned i; 131 double* replacement = NULL; 132 struct leaf* leaf = node->c.leaf; 133 for(i = 0; i < leaf->count; i++) { 134 double dist2, tmp[3]; 135 if(d3_eq(leaf->vertices + i*3, in)) { 136 same = 1; 137 break; 138 } 139 d3_sub(tmp, leaf->vertices + i*3, in); 140 dist2 = d3_dot(tmp, tmp); 141 if(dist2 <= denoiser->eps2) { 142 replacement = leaf->vertices + i*3; 143 break; 144 } 145 } 146 if(same) { 147 if(in != out) d3_set(out, in); 148 if(done) *done = 1; 149 } 150 else if(replacement) { 151 d3_set(out, replacement); 152 (*denoised_count)++; 153 if(done) *done = 1; 154 } 155 else if(can_add) { 156 /* Vertex not registered yet: register and keep */ 157 if(leaf->count < NODE_COUNT_MAX) { 158 add_to_leaf(node, in); 159 (*added_count)++; 160 } else { 161 /* Need mode room: change leaf to node with 2 leaf children */ 162 double var[3], tmp1[3], tmp2[3], tmp3[3], split_pos; 163 unsigned split_dim; 164 double range[2] = { DBL_MAX, -DBL_MAX }; 165 double sum[3] = { 0, 0, 0 }, sum2[3] = { 0, 0, 0 }; 166 /* Chose the dim to split wrt variance */ 167 for(i = 0; i < NODE_COUNT_MAX; i++) { 168 double tmp[3]; 169 d3_mul(tmp, leaf->vertices + 3*i, leaf->vertices + 3*i); 170 d3_add(sum, sum, leaf->vertices + 3*i); 171 d3_add(sum2, sum2, tmp); 172 } 173 d3_divd(tmp1, sum2, NODE_COUNT_MAX); 174 d3_divd(tmp2, sum, NODE_COUNT_MAX); 175 d3_sub(var, tmp1, d3_mul(tmp3, tmp2, tmp2)); 176 split_dim = (var[0] > var[1]) 177 ? (var[0] > var[2] ? 0 : 2) : (var[1] > var[2] ? 1 : 2); 178 for(i = 0; i < NODE_COUNT_MAX; i++) { 179 range[0] = MMIN(range[0], leaf->vertices[3*i + split_dim]); 180 range[1] = MMAX(range[1], leaf->vertices[3*i + split_dim]); 181 } 182 split_pos = (range[0] + range[1]) * 0.5; 183 /* Slit dimension dim */ 184 ERR(alloc_leaf_node(allocator, &node->c.children[0])); 185 ERR(alloc_leaf_node(allocator, &node->c.children[1])); 186 node->split_dim = split_dim; 187 node->split_pos = split_pos; 188 node->type = NODE; 189 for(i = 0; i < NODE_COUNT_MAX; i++) { 190 double* v = leaf->vertices + 3*i; 191 double pos = v[split_dim]; 192 struct node* side 193 = (pos <= split_pos) ? node->c.children[0] : node->c.children[1]; 194 /* Copy content to children wrt split */ 195 add_to_leaf(side, v); 196 } 197 /* Free old leaf */ 198 MEM_RM(allocator, leaf); 199 /* Retry same search_and_add */ 200 ERR(search_and_add(node, denoiser, in, can_add, out, done, added_count, 201 denoised_count)); 202 } 203 if(in != out) d3_set(out, in); 204 if(done) *done = 1; 205 } 206 } else { 207 int ldone = 0; 208 ASSERT(node->type == NODE); 209 /* Due to epsilon, one can have to search in both children. 210 * If the vertex doesn't exist yet, it should be added only once though! */ 211 if(node->split_pos + denoiser->eps >= in[node->split_dim]) { 212 int in_range = (node->split_pos >= in[node->split_dim]); 213 ERR(search_and_add(node->c.children[0], denoiser, in, in_range, out, 214 &ldone, added_count, denoised_count)); 215 if(done && ldone) *done = 1; 216 } 217 if(!ldone && node->split_pos - denoiser->eps <= in[node->split_dim]) { 218 int in_range = (node->split_pos < in[node->split_dim]); 219 ERR(search_and_add(node->c.children[1], denoiser, in, in_range, out, done, 220 added_count, denoised_count)); 221 } 222 } 223 224 exit: 225 return res; 226 error: 227 goto exit; 228 } 229 230 res_T 231 vertex_denoiser_create 232 (struct mem_allocator* allocator, 233 const double eps, 234 struct vertex_denoiser** denoiser) 235 { 236 res_T res = RES_OK; 237 struct vertex_denoiser* den = NULL; 238 239 ASSERT(allocator && eps >= 0 && denoiser); 240 241 den = MEM_CALLOC(allocator, 1, sizeof(*den)); 242 if(!den) { 243 res = RES_MEM_ERR; 244 goto error; 245 } 246 den->allocator = allocator; 247 den->eps = eps; 248 den->eps2 = eps*eps; 249 ERR(alloc_leaf_node(allocator, &den->root)); 250 251 exit: 252 *denoiser = den; 253 return res; 254 error: 255 if(den) vertex_denoiser_release(den); 256 den = NULL; 257 goto exit; 258 } 259 260 void 261 vertex_denoiser_release 262 (struct vertex_denoiser* denoiser) 263 { 264 ASSERT(denoiser); 265 release_node(denoiser->allocator, denoiser->root); 266 MEM_RM(denoiser->allocator, denoiser); 267 } 268 269 size_t 270 vertex_denoiser_get_count 271 (struct vertex_denoiser* denoiser) 272 { 273 ASSERT(denoiser); 274 return denoiser->count; 275 } 276 277 size_t 278 vertex_denoiser_get_denoised_count 279 (struct vertex_denoiser* denoiser) 280 { 281 ASSERT(denoiser); 282 return denoiser->denoised_count; 283 } 284 285 res_T 286 vertex_denoiser_denoise 287 (struct vertex_denoiser* denoiser, 288 const double *in, 289 const size_t count, 290 double *out) 291 { 292 res_T res = RES_OK; 293 size_t n; 294 295 ASSERT(denoiser && in && out); 296 297 for(n = 0; n < count; n++) { 298 size_t added = 0, denoised = 0; 299 ERR(search_and_add(denoiser->root, denoiser, in + 3*n, 1, out + 3*n, NULL, 300 &added, &denoised)); 301 denoiser->count += added; 302 denoiser->denoised_count += denoised; 303 } 304 305 exit: 306 return res; 307 error: 308 goto exit; 309 } 310 311 res_T 312 denoise_array 313 (struct vertex_denoiser* denoiser, 314 struct darray_double* array) 315 { 316 res_T res = RES_OK; 317 318 ASSERT(denoiser && array); 319 320 ASSERT(darray_double_size_get(array) % 3 == 0); 321 ERR(vertex_denoiser_denoise(denoiser, 322 darray_double_cdata_get(array), darray_double_size_get(array) / 3, 323 darray_double_data_get(array))); 324 325 exit: 326 return res; 327 error: 328 goto exit; 329 } 330 331 res_T 332 stl_export_denoised_geometry 333 (struct vertex_denoiser* denoiser, 334 struct scad_geometry* geom, 335 const enum scad_normals_orientation orientation, 336 const int binary) 337 { 338 res_T res = RES_OK; 339 struct darray_double trg; 340 const char* name; 341 struct str filename; 342 343 ASSERT(denoiser && geom); 344 345 darray_double_init(denoiser->allocator, &trg); 346 str_init(denoiser->allocator, &filename); 347 ERR(scad_geometry_get_name(geom, &name)); 348 ERR(str_set(&filename, name)); 349 ERR(str_append(&filename, ".stl")); 350 ERR(scad_stl_get_data(geom, &trg)); 351 ERR(denoise_array(denoiser, &trg)); 352 ERR(scad_stl_data_write(&trg, str_cget(&filename), orientation, binary)); 353 354 exit: 355 darray_double_release(&trg); 356 str_release(&filename); 357 return res; 358 error: 359 goto exit; 360 } 361