minor fixes
[colorize.git] / lib / colorize.c
1 #include <stdio.h>
2 #include <string.h>
3 #include <stdlib.h>
4 #include <math.h>
5 #include "darray.h"
6 #include "multigrid.h"
7 #include "colorize.h"
8
9 /* colorize image array using image_mark
10  *
11  * image:
12  *      first dimansion = width
13  *      second dimension = height
14  *      thrid dimension = components (first component is luminance, other components are features, like u, v)
15  *      fourth dimension = time
16  *
17  * image_mark:
18  *      first dimansion = width
19  *      second dimension = height
20  *      fourth dimension = time
21  */
22 int colorize_prepare(struct colorize *col)
23 {
24         int dimensions[4];
25         int ww, hh, kk, last_kk, d, temp, z;
26         darray_t *luminance = NULL;
27         int rc = -1;
28
29         ASSERT(darrayGetNumberOfDimensions(col->image) == 4, "image array must have 4 dimensions");
30         col->w = darrayGetDimensions(col->image)[0];
31         col->h = darrayGetDimensions(col->image)[1];
32         col->comp = darrayGetDimensions(col->image)[2];
33         col->k = darrayGetDimensions(col->image)[3];
34         ASSERT(col->comp > 0, "image array requires at least one component (Y-component)");
35
36         ASSERT(darrayGetNumberOfDimensions(col->image_mark) == 3, "image_mark array must have 3 dimensions");
37         ASSERT(col->w == darrayGetDimensions(col->image_mark)[0], "image_mark and image must both have same x-size");
38         ASSERT(col->h == darrayGetDimensions(col->image_mark)[1], "image_mark and image must both have same y-size");
39         ASSERT(col->k == darrayGetDimensions(col->image_mark)[2], "image_mark and image must both have same z-size");
40
41         /* Calculate maximum depth
42          * Scale down dimensions until at least 2x2 remains
43          */
44         temp = min(col->w, col->h);
45         for (col->max_depth = 0; temp > 2; col->max_depth++)
46                 temp=(temp+1)>>1;
47         ASSERT(col->max_depth >= 2, "image too small");
48
49         col->nb_list = calloc(col->max_depth, sizeof(struct neighbors *));
50         if (!col->nb_list) {
51                 rc = -1;
52                 goto error;
53         }
54         col->marks = calloc(col->max_depth, sizeof(darray_t *));
55         if (!col->marks) {
56                 rc = -1;
57                 goto error;
58         }
59         col->values = calloc(col->max_depth, sizeof(darray_t *));
60         if (!col->values) {
61                 rc = -1;
62                 goto error;
63         }
64         col->flows = calloc(col->max_depth, sizeof(darray_t *));
65         if (!col->flows) {
66                 rc = -1;
67                 goto error;
68         }
69         col->flows_i = calloc(col->max_depth, sizeof(darray_t *));
70         if (!col->flows_i) {
71                 rc = -1;
72                 goto error;
73         }
74         if (col->scalexyz) {
75                 /* the flows_xy are scaled in xy in the current level */
76                 col->flows_xy = calloc(col->max_depth, sizeof(darray_t *));
77                 if (!col->flows_xy) {
78                         rc = -1;
79                         goto error;
80                 }
81         }
82         dimensions[0] = col->w;
83         dimensions[1] = col->h;
84         dimensions[2] = col->k;
85         luminance = darrayCreate(3, dimensions);
86         if (!luminance) {
87                 rc = -1;
88                 goto error;
89         }
90         /* generate sparse multigrid */
91         last_kk = kk = col->k;
92         printf("-> prepare sparse grid");
93         for (d = 0, ww=col->w, hh=col->h; d < col->max_depth; d++, ww=(ww+1)>>1, hh=(hh+1)>>1) {
94                 /* create temporary arrays
95                  * copy fist array to level 0 or shrink previous level
96                  */
97                 printf(" #%d(w=%d h=%d k=%d)", d, ww, hh, kk); fflush(stdout);
98                 dimensions[0] = ww;
99                 dimensions[1] = hh;
100                 dimensions[2] = last_kk;
101                 col->marks[d] = darrayCreate(3, dimensions);
102                 if (!col->marks[d]) {
103                         rc = -1;
104                         goto error;
105                 }
106                 dimensions[0] = ww;
107                 dimensions[1] = hh;
108                 dimensions[2] = last_kk-1;
109                 dimensions[3] = 2;
110                 col->flows[d] = darrayCreate(4, dimensions);
111                 if (!col->flows[d]) {
112                         rc = -1;
113                         goto error;
114                 }
115                 col->flows_i[d] = darrayCreate(4, dimensions);
116                 if (!col->flows_i[d]) {
117                         rc = -1;
118                         goto error;
119                 }
120                 if (d < col->scalexyz) {
121                         dimensions[0] = (ww+1)>>1;
122                         dimensions[1] = (hh+1)>>1;
123                         dimensions[2] = kk-1;
124                         dimensions[3] = 2;
125                         col->flows_xy[d] = darrayCreate(4, dimensions);
126                         if (!col->flows_xy[d]) {
127                                 rc = -1;
128                                 goto error;
129                         }
130                 }
131                 if (d == 0) {
132                         for (z = 0; z < kk; z++) {
133                                 darrayCheckPr(luminance, darrayGetPr(luminance) + z*col->w*col->h, col->w*col->h);
134                                 darrayCheckPr(col->image, darrayGetPr(col->image) + z*col->w*col->h*col->comp, col->w*col->h);
135                                 memcpy(darrayGetPr(luminance) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*col->comp, col->w*col->h*sizeof(double));
136                         }
137                         if (col->flow && col->flow_i) {
138                                 darrayCheckPr(col->flows[d], darrayGetPr(col->flows[d]), col->w*col->h*(kk-1)*2);
139                                 darrayCheckPr(col->flow, darrayGetPr(col->flow), col->w*col->h*(kk-1)*2);
140                                 memcpy(darrayGetPr(col->flows[d]), darrayGetPr(col->flow), col->w*col->h*(kk-1)*2*sizeof(double));
141                                 darrayCheckPr(col->flows_i[d], darrayGetPr(col->flows_i[d]), col->w*col->h*(kk-1)*2);
142                                 darrayCheckPr(col->flow_i, darrayGetPr(col->flow_i), col->w*col->h*(kk-1)*2);
143                                 memcpy(darrayGetPr(col->flows_i[d]), darrayGetPr(col->flow_i), col->w*col->h*(kk-1)*2*sizeof(double));
144                         }
145                         darrayCheckPr(col->marks[d], darrayGetPr(col->marks[d]), col->w*col->h*kk);
146                         darrayCheckPr(col->image_mark, darrayGetPr(col->image_mark), col->w*col->h*kk);
147                         memcpy(darrayGetPr(col->marks[d]), darrayGetPr(col->image_mark), col->w*col->h*kk*sizeof(double));
148                         /* invert, so 0 = marked */
149                         quantize(col->marks[d], 0.5, 0, 1);
150                 } else {
151                         rc = shrink_xy(luminance, NULL);
152                         if (rc < 0)
153                                 goto error;
154                         rc = shrink_xy(col->marks[d-1], col->marks[d]);
155                         if (rc < 0)
156                                 goto error;
157                         rc = shrink_xy(col->flows[d-1], col->flows[d]);
158                         if (rc < 0)
159                                 goto error;
160                         rc = shrink_xy(col->flows_i[d-1], col->flows_i[d]);
161                         if (rc < 0)
162                                 goto error;
163                         /* shrink length of flow vector on smaller grid.
164                          * This is not performed when comparing with
165                          * the original algorithm. (it is a bug there)
166                          */
167                         multiply(col->flows[d], 0.5);
168                         multiply(col->flows_i[d], 0.5);
169                         if (d-1 < col->scalexyz) { /* process if previous d shall be shrinked in z direction */
170                                 /* shrink time */
171                                 rc = shrink_z(luminance, NULL, col->flows[d]);
172                                 if (rc < 0)
173                                         goto error;
174                                 rc = shrink_z(col->marks[d], NULL, col->flows[d]);
175                                 if (rc < 0)
176                                         goto error;
177                                 rc = shrink_flow_z(col->flows[d], NULL);
178                                 if (rc < 0)
179                                         goto error;
180                                 rc = shrink_flow_z(col->flows_i[d], NULL);
181                                 if (rc < 0)
182                                         goto error;
183                         }
184                         /* If any of the restricted pixles is 0 (marked),
185                          * the result must be 0 (marked) too.
186                          */
187                         quantize(col->marks[d], 0.9999, 1, 0);
188                 }
189                 if (d < col->scalexyz) {
190                         /* create x and y shrinked flow vectors for restriction and prolonging process */
191                         rc = shrink_xy(col->flows[d], col->flows_xy[d]);
192                         if (rc < 0)
193                                 goto error;
194                         multiply(col->flows_xy[d], 0.5);
195                 }
196
197                 /* generate values array */
198                 dimensions[0] = ww;
199                 dimensions[1] = hh;
200                 dimensions[2] = kk;
201                 col->values[d] = darrayCreate(3, dimensions);
202                 if (!col->values[d]) {
203                         rc = -1;
204                         goto error;
205                 }
206
207                 /* generate array of neighbors */
208 #ifndef TEST_NO_RENDER
209                 col->nb_list[d] = gen_neighbor(ww, hh, kk, col->flows[d], col->flows_i[d]);
210                 if (!col->nb_list[d]) {
211                         rc = -1;
212                         goto error;
213                 }
214                 /* weighten */
215                 weighten(luminance, col->nb_list[d]);
216 #endif
217
218                 last_kk = kk;
219                 if (d < col->scalexyz)
220                         kk = (kk+1)>>1;
221         }
222         darrayDestroy(luminance);
223         luminance = NULL;
224         printf("\n");
225
226         dimensions[0] = col->w;
227         dimensions[1] = col->h;
228         dimensions[2] = col->k;
229         col->init = darrayCreate(3, dimensions);
230         if (!col->init) {
231                 rc = -1;
232                 goto error;
233         }
234
235 #ifdef TEST_PYRAMID
236         dimensions[0] = col->w;
237         dimensions[1] = col->h;
238         dimensions[2] = col->k;
239         luminance = darrayCreate(3, dimensions);
240         for (z = 0; z < col->k; z++) {
241                 darrayCheckPr(luminance, darrayGetPr(luminance) + z*col->w*col->h, col->w*col->h);
242                 darrayCheckPr(col->image, darrayGetPr(col->image) + z*col->w*col->h*comp, col->w*col->h);
243                 memcpy(darrayGetPr(luminance) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*comp, col->w*col->h*sizeof(double));
244                 darrayCheckPr(col->values[0], darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h);
245                 darrayCheckPr(col->image, darrayGetPr(col->image) + z*col->w*col->h*comp, col->w*col->h);
246                 memcpy(darrayGetPr(col->values[0]) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*comp, col->w*col->h*sizeof(double));
247 #ifdef TEST_FLOW
248                 draw_flow(col->values[0], col->flows[0], col->flows_i[0]);
249 #endif
250 #ifdef TEST_MARKED
251                 paste(col->marks[0], col->values[0], 0, 0, 0);
252 #endif
253
254         }
255         for (d = 1; d < col->max_depth; d++) {
256                 rc = shrink_xy(luminance, NULL);
257                 if (rc < 0)
258                         goto error;
259                 if (d-1 < col->scalexyz)
260                         shrink_z(luminance, NULL, col->flows_xy[d-1]);
261                 paste(luminance, col->values[0], 0, 0, 0);
262 #ifdef TEST_FLOW
263                 draw_flow(col->values[0], col->flows_x[d], col->flows_ix[d]);
264 #endif
265 #ifdef TEST_MARKED
266                 paste(col->marks[d], col->values[0], 0, 0, 0);
267 #endif
268         }
269         darrayDestroy(luminance);
270         luminance = NULL;
271         /* copy test pyramid image as Y-component */
272         for (z = 0; z < col->k; z++) {
273                 darrayCheckPr(col->image, darrayGetPr(col->image) + z*col->w*col->h*col->comp, col->w*col->h);
274                 darrayCheckPr(col->values[0], darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h);
275                 memcpy(darrayGetPr(col->image) + z*col->w*col->h*col->comp, darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h*sizeof(double));
276         }
277 #endif
278
279         rc = 0;
280
281 error:
282         darrayDestroy(luminance);
283
284         return rc;
285 }
286
287 int colorize_solve(struct colorize *col, int quick)
288 {
289         int z, d, c, ww, hh, kk;
290         int rc = -1;
291
292         if (!quick)
293                 printf("-> apply component to grid and solve");
294         for (c = 1; c < col->comp; c++) {
295                 int pixles, j;
296                 double *value_ptr, *init_ptr, *mark_ptr;
297
298                 printf(" #%d", c); fflush(stdout);
299                 /* apply component
300                  * - copy component into value array (level 0), otherwise 0
301                  * - use maked colors for init array, otherwise use 0
302                  */
303                 for (z = 0; z < col->k; z++) {
304                         darrayCheckPr(col->values[0], darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h);
305                         darrayCheckPr(col->image, darrayGetPr(col->image) + c*col->w*col->h + z*col->w*col->h*col->comp, col->w*col->h);
306                         memcpy(darrayGetPr(col->values[0]) + z*col->w*col->h, darrayGetPr(col->image) + c*col->w*col->h + z*col->w*col->h*col->comp, col->w*col->h*sizeof(double));
307                 }
308                 for (d = 1, ww=(col->w+1)>>1, hh=(col->h+1)>>1, kk=col->k; d < col->max_depth; d++, ww=(ww+1)>>1, hh=(hh+1)>>1) {
309                         if (d-1 < col->scalexyz) /* scale kk, if last array was scaled */
310                                 kk=(kk+1)>>1;   
311                         value_ptr = darrayGetPr(col->values[d]);
312                         pixles = ww*hh*kk;
313                         for (j = 0; j < pixles; j++) {
314                                 darrayCheckPr(col->values[d], value_ptr, 1);
315                                 *value_ptr++ = 0.0;
316                         }
317                 }
318                 value_ptr = darrayGetPr(col->values[0]);
319                 init_ptr = darrayGetPr(col->init);
320                 mark_ptr = darrayGetPr(col->image_mark);
321                 pixles = col->w*col->h*col->k;
322                 for (j = 0; j < pixles; j++) {
323                         darrayCheckPr(col->image_mark, mark_ptr, 1);
324                         if (*mark_ptr++ > 0.5) {
325                                 darrayCheckPr(col->init, init_ptr, 1);
326                                 darrayCheckPr(col->values[0], value_ptr, 1);
327                                 *init_ptr++ = *value_ptr;
328                         } else {
329                                 darrayCheckPr(col->init, init_ptr, 1);
330                                 *init_ptr++ = 0.0;
331                         }
332                         value_ptr++;
333                 }
334
335 #ifndef TEST_NO_RENDER
336                 /* colorize component */
337                 if (quick) {
338                         rc = solve_mg(col->max_depth, col->inner_iter, 1, col->values, col->marks, col->init, col->nb_list, col->flows_xy, col->scalexyz);
339                         if (rc < 0)
340                                 goto error;
341                 } else {
342                         int i;
343                         for (i = 0; i < col->outer_iter; i++) {
344                                 rc = solve_mg(col->max_depth, col->inner_iter, 2, col->values, col->marks, col->init, col->nb_list, col->flows_xy, col->scalexyz);
345                                 if (rc < 0)
346                                         goto error;
347                                 rc = solve_mg((col->max_depth+1) / 2, col->inner_iter, 2, col->values, col->marks, col->init, col->nb_list, col->flows_xy, col->scalexyz);
348                                 if (rc < 0)
349                                         goto error;
350                                 rc = solve_mg(2, col->inner_iter, 2, col->values, col->marks, col->init, col->nb_list, col->flows_xy, col->scalexyz);
351                                 if (rc < 0)
352                                         goto error;
353                                 rc = solve_mg(1, col->inner_iter, 4, col->values, col->marks, col->init, col->nb_list, col->flows_xy, col->scalexyz);
354                                 if (rc < 0)
355                                         goto error;
356                         }
357                 }
358 #endif
359
360 #ifdef TEST_PYRAMID
361                 for (d = 1; d < col->max_depth; d++)
362                         paste(col->values[d], col->values[0], 0, 0, 0);
363 #endif
364
365                 /* apply colorized component */
366                 for (z = 0; z < col->k; z++) {
367                         darrayCheckPr(col->image, darrayGetPr(col->image) + c*col->w*col->h + z*col->w*col->h*col->comp, col->w*col->h);
368                         darrayCheckPr(col->values[0], darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h);
369                         memcpy(darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h*sizeof(double));
370                 }
371         }
372         if (!quick)
373                 printf("\n");
374
375 #ifdef TEST_SHRINK_XY
376         for (c = 0; c < col->comp; c++) {
377                 darray_t *temp;
378                 int dimensions[3] = {col->w, col->h, col->k};
379                 for (z = 0; z < col->k; z++)
380                         memcpy(darrayGetPr(col->values[0]) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, col->w*col->h*sizeof(double));
381                 dimensions[0] = (col->w+1)>>1;
382                 dimensions[1] = (col->h+1)>>1;
383                 temp = darrayCreate(3, dimensions);
384                 ASSERT(temp, "no memory");
385                 rc = shrink_xy(col->values[0], temp);
386                 ASSERT(rc >= 0, "no memory");
387                 paste(temp, col->values[0], 0, 0, 0);
388                 darrayDestroy(temp);
389                 for (z = 0; z < col->k; z++)
390                         memcpy(darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h*sizeof(double));
391         }
392 #endif
393
394 #ifdef TEST_SHRINK_Z
395         for (c = 0; c < col->comp; c++) {
396                 darray_t *temp;
397                 int dimensions[3] = {col->w, col->h, col->k};
398                 for (z = 0; z < col->k; z++)
399                         memcpy(darrayGetPr(col->values[0]) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, col->w*col->h*sizeof(double));
400                 dimensions[2] = (col->k+1)>>1;
401                 temp = darrayCreate(3, dimensions);
402                 ASSERT(temp, "no memory");
403                 rc = shrink_z(col->values[0], temp, col->flows[0]);
404                 ASSERT(rc >= 0, "no memory");
405                 paste(temp, col->values[0], 0, 0, 0);
406                 darrayDestroy(temp);
407                 for (z = 0; z < col->k; z++)
408                         memcpy(darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h*sizeof(double));
409         }
410 #endif
411
412 #ifdef TEST_EXPAND_XY
413         for (c = 0; c < col->comp; c++) {
414                 darray_t *temp;
415                 int dimensions[3] = {col->w, col->h, col->k};
416                 for (z = 0; z < col->k; z++)
417                         memcpy(darrayGetPr(col->values[0]) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, col->w*col->h*sizeof(double));
418                 dimensions[0] = (col->w+1)>>1;
419                 dimensions[1] = (col->h+1)>>1;
420                 temp = darrayCreate(3, dimensions);
421                 ASSERT(temp, "no memory");
422                 rc = shrink_xy(col->values[0], temp);
423                 ASSERT(rc >= 0, "no memory");
424                 expand_xy(temp, col->values[0]);
425                 darrayDestroy(temp);
426                 for (z = 0; z < col->k; z++)
427                         memcpy(darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h*sizeof(double));
428         }
429 #endif
430
431 #ifdef TEST_EXPAND_Z
432         for (c = 0; c < col->comp; c++) {
433                 darray_t *temp;
434                 int dimensions[3] = {col->w, col->h, col->k};
435                 for (z = 0; z < col->k; z++)
436                         memcpy(darrayGetPr(col->values[0]) + z*col->w*col->h, darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, col->w*col->h*sizeof(double));
437                 dimensions[2] = (col->k+1)>>1;
438                 temp = darrayCreate(3, dimensions);
439                 ASSERT(temp, "no memory");
440                 rc = shrink_z(col->values[0], temp, col->flows[0]);
441                 ASSERT(rc >= 0, "no memory");
442                 rc = expand_z(temp, col->values[0], col->flows[0]);
443                 ASSERT(rc >= 0, "no memory");
444                 darrayDestroy(temp);
445                 for (z = 0; z < col->k; z++)
446                         memcpy(darrayGetPr(col->image) + z*col->w*col->h*col->comp + c*col->w*col->h, darrayGetPr(col->values[0]) + z*col->w*col->h, col->w*col->h*sizeof(double));
447         }
448 #endif
449
450         rc = 0;
451
452 error:
453         return rc;
454 }
455
456 void colorize_free(struct colorize *col)
457 {
458         int d;
459
460         /* free memory */
461         for (d = 0; d < col->max_depth; d++) {
462                 if (col->nb_list)
463                         free(col->nb_list[d]);
464                 if (col->marks)
465                         darrayDestroy(col->marks[d]);
466                 if (col->values)
467                         darrayDestroy(col->values[d]);
468                 if (col->flows_xy)
469                         darrayDestroy(col->flows_xy[d]);
470                 if (col->flows)
471                         darrayDestroy(col->flows[d]);
472                 if (col->flows_i)
473                         darrayDestroy(col->flows_i[d]);
474         }
475         free(col->nb_list);
476         free(col->marks);
477         free(col->values);
478         free(col->flows_xy);
479         free(col->flows);
480         free(col->flows_i);
481         darrayDestroy(col->init);
482 }
483
484 int colorize(const darray_t *image, const darray_t *image_mark, const darray_t *flow, const darray_t *flow_i, int inner_iter, int outer_iter, int scalexyz)
485 {
486         struct colorize col;
487         int rc = -1;
488
489         memset(&col, 0, sizeof(col));
490         col.image = image;
491         col.image_mark = image_mark;
492         col.flow = flow;
493         col.flow_i = flow_i;
494         col.inner_iter = inner_iter;
495         col.outer_iter = outer_iter;
496         col.scalexyz = scalexyz;
497
498         rc = colorize_prepare(&col);
499         if (rc < 0)
500                 goto error;
501
502         rc = colorize_solve(&col, 0);
503         if (rc < 0)
504                 goto error;
505
506         rc = 0;
507
508 error:
509         colorize_free(&col);
510
511         return rc;
512 }