d7f14c766ff3178a0fc76b8965ad61eef9e3a9e7
[colorize.git] / src / colorize.c
1 #include <stdio.h>
2 #include <string.h>
3 #include <stdlib.h>
4 #include <unistd.h>
5 #include <getopt.h>
6 #include <math.h>
7 #include <time.h>
8 #include "../lib/darray.h"
9 #include "../lib/colorize.h"
10 #include "ppm.h"
11 #include "yuv.h"
12 #include "mark.h"
13 #include "dir_seperator.h"
14 #ifdef WITH_OPENCV
15 #include "opticalflow.h"
16 #endif
17
18 #define min(x,y) ((x < y) ? x : y)
19
20 static void print_help(const char *app);
21 static void print_test_help();
22
23 //#defbug SEQUENCE
24 //#define SINGLE_K_TEST 2
25
26 /*
27  * load sequence
28  */
29
30 struct sequence {
31         char filename[256];
32 };
33 static struct sequence *sequence = NULL;
34
35 /* load sequence list and return number of images in sequence */
36 static int load_sequence(int *start, int *next, char *first_filename)
37 {
38         const char *name = "sequence";
39         int count;
40         int i, j, keyframe_before_start, eof;
41         FILE *fp;
42         char buffer[256], *p;
43
44         /* free previously used sequence */
45         if (sequence)
46                 free(sequence);
47         sequence = NULL;
48
49 again:
50         /* get number of frames until next keyframe or end of squence */
51         fp = fopen(name, "r");
52         if (!fp) {
53                 printf("failed to load sequence '%s'\n", name);
54                 return -1;
55         }
56         count = 0;
57         j = 0;
58         eof = 1;
59         keyframe_before_start = 0;
60         while(fgets(buffer, sizeof(buffer), fp)) {
61                 p = strchr(buffer + 1, '\"');
62                 if (!p)
63                         break;
64                 *p = '\0';
65                 if (j == 0)
66                         strcpy(first_filename, buffer + 1);
67                 p++;
68                 if ((*start) && j == (*start) - 1) {
69                         if (strstr(p, "keyframe"))
70                                 keyframe_before_start = 1;
71                 }
72                 if (j < (*start)) {
73                         j++;
74                         continue;
75                 }
76 #ifdef DEBUG_SEQUENCE
77                 printf("counting %d\n", j);
78 #endif
79                 j++;
80                 count++;
81                 /* skip start frame, since it is always a keyframe */
82                 if (count == 1)
83                         continue;
84                 if (strstr(p, "keyframe")) {
85                         eof = 0;
86                         break;
87                 }
88         }
89         fclose(fp);
90
91         /* end of file case */
92         if (count == 0) {
93 #ifdef DEBUG_SEQUENCE
94                 puts("end of file");
95 #endif
96                 *next = -1;
97                 return 0;
98         } else
99         if (count == 1) {
100 #ifdef DEBUG_SEQUENCE
101                 puts("count=1");
102 #endif
103                 if (eof) {
104 #ifdef DEBUG_SEQUENCE
105                         puts("end of file at last frame");
106 #endif
107                         *next = -1;
108                         return 0;
109                 }
110                 /* we are done, because we already had the last sequence */
111                 if (!keyframe_before_start) {
112 #ifdef DEBUG_SEQUENCE
113                         puts("no keyframe before start");
114 #endif
115                         (*next) = -1;
116                         return 0;
117                 }
118 #ifdef DEBUG_SEQUENCE
119                 puts("keyframe before start");
120 #endif
121                 /* we just have a single last frame */
122                 count = 1;
123                 (*next) = (*start) + 1;
124         } else
125         /* if we have two keyframes side by side, we ignore it */
126         if (count == 2 && !eof) {
127 #ifdef DEBUG_SEQUENCE
128                 puts("count=2");
129 #endif
130                 /* only if the frame before start is not a keyframe
131                  * and we are not the first sequence, because that would mean
132                  * that the first frame is a single sequence */
133                 if (!keyframe_before_start && (*start) > 0) {
134 #ifdef DEBUG_SEQUENCE
135                         puts("no keyframe before start");
136 #endif
137                         (*start)++;
138                         goto again;
139                 }
140 #ifdef DEBUG_SEQUENCE
141                 puts("keyframe before start or first frame");
142 #endif
143                 /* now the start frame has one keyframe before and one after */
144                 count = 1;
145                 (*next) = (*start) + 1;
146         } else {
147                 (*next) = (*start) + count - 1;
148         }
149
150         /* alloc memory for sequence and fill with data */
151         sequence = (struct sequence *)malloc(sizeof(struct sequence) * count);
152         if (!sequence) {
153                 printf("no memory for sequence\n");
154                 return -1;
155         }
156         memset((struct sequence *)sequence, 0, sizeof(struct sequence) * count);
157         fp = fopen(name, "r");
158         if (!fp) {
159                 printf("failed to load sequence '%s'\n", name);
160                 return -1;
161         }
162         i = 0;
163         j = 0;
164         while(fgets(buffer, sizeof(buffer), fp)) {
165                 p = strchr(buffer + 1, '\"');
166                 if (!p)
167                         break;
168                 if (j < (*start)) {
169                         j++;
170                         continue;
171                 }
172                 *p = '\0';
173                 strcpy(sequence[i].filename, buffer + 1);
174                 i++;
175                 if (i == count)
176                         break;
177         }
178         fclose(fp);
179
180         return count;
181 }
182
183 /*
184  * scaling
185  */
186
187 /* scale down image in img_buffer by calculating average */
188 static void scale_img(unsigned char *img_buffer, int width, int height, int scale)
189 {
190         int w, h, i, j, x, y;
191         int r, g, b;
192
193         if (scale == 1)
194                 return;
195
196         w = width / scale;
197         h = height / scale;
198
199         for (i = 0; i < h; i++) {
200                 for (j = 0; j < w; j++) {
201                         r = g = b = 0;
202                         for (y = 0; y < scale; y++) {
203                                 for (x = 0; x < scale; x++) {
204                                         r += img_buffer[((i*scale+y) * width + j*scale+x) * 3 + 0];
205                                         g += img_buffer[((i*scale+y) * width + j*scale+x) * 3 + 1];
206                                         b += img_buffer[((i*scale+y) * width + j*scale+x) * 3 + 2];
207                                 }
208                         }
209                         img_buffer[(i * w + j)*3 + 0] = r / scale / scale;
210                         img_buffer[(i * w + j)*3 + 1] = g / scale / scale;
211                         img_buffer[(i * w + j)*3 + 2] = b / scale / scale;
212                 }
213         }
214 }
215
216 /* scale down mark map in mark_buffer */
217 static void scale_mark(unsigned char *mark_buffer, int width, int height, int scale)
218 {
219         int w, h, i, j, x, y;
220         unsigned char c, temp;
221
222         if (scale == 1)
223                 return;
224
225         w = width / scale;
226         h = height / scale;
227
228         for (i = 0; i < h; i++) {
229                 for (j = 0; j < w; j++) {
230                         c = 0;
231                         /* always use one index other than 0, if there is any in an areaa to be shrinked */
232                         for (y = 0; y < scale; y++) {
233                                 for (x = 0; x < scale; x++) {
234                                         temp = mark_buffer[(i*scale+y) * width + j*scale+x];
235                                         if (temp)
236                                                 c = temp;
237                                 }
238                         }
239                         mark_buffer[i * w + j] = c;
240                 }
241         }
242 }
243
244 /*
245  * options
246  */
247
248 static int in_itr_num = 5, out_itr_num = 1, optical_flow = 1, bright_contrast = 1;
249 int scale = 1, scalexyz = 999;
250
251 static enum test {
252         NO_TEST = 0,
253         FLOW_NEXT,
254         FLOW_PREV,
255         MARKED,
256         MASK,
257         MASK_COLOR,
258         BC_ONLY,
259         BC_IMAGE,
260 } test = NO_TEST;
261
262 static int parse_test(const char *arg)
263 {
264         if (!strcmp(arg, "help")) {
265                 print_test_help();
266                 exit(0);
267         }
268 #ifdef WITH_OPENCV
269         if (!strcmp(arg, "flow-next"))
270                 return FLOW_NEXT;
271         if (!strcmp(arg, "flow-prev"))
272                 return FLOW_PREV;
273 #endif
274         if (!strcmp(arg, "marked"))
275                 return MARKED;
276         if (!strcmp(arg, "mask"))
277                 return MASK;
278         if (!strcmp(arg, "mask+color"))
279                 return MASK_COLOR;
280         if (!strcmp(arg, "bc-only"))
281                 return BC_ONLY;
282         if (!strcmp(arg, "bc-image"))
283                 return BC_IMAGE;
284
285         return NO_TEST;
286 }
287
288 static int handle_options(int argc, char **argv)
289 {
290         int skip_args = 0;
291
292         while (1) {
293                 int option_index = 0, c;
294                 static struct option long_options[] = {
295                         {"help", 0, 0, 'h'},
296                         {"in-itr-num", 1, 0, 'i'},
297                         {"out-itr-num", 1, 0, 'o'},
298                         {"zscale", 1, 0, 'z'},
299                         {"brightness-contrast", 1, 0, 'b'},
300                         {"optical-flow", 1, 0, 'f'},
301                         {"scale", 1, 0, 's'},
302                         {"test", 1, 0, 't'},
303                         {0, 0, 0, 0},
304                 };
305
306                 c = getopt_long(argc, argv, "hi:o:z:b:f:s:t:", long_options, &option_index);
307
308                 if (c == -1)
309                         break;
310
311                 switch (c) {
312                 case 'h':
313                         print_help(argv[0]);
314                         exit(0);
315                 case 'i':
316                         in_itr_num = atoi(optarg);
317                         skip_args += 2;
318                         break;
319                 case 'o':
320                         out_itr_num = atoi(optarg);
321                         skip_args += 2;
322                         break;
323                 case 'z':
324                         scalexyz = atoi(optarg);
325                         skip_args += 2;
326                         break;
327                 case 'b':
328                         bright_contrast = atoi(optarg);
329                         skip_args += 2;
330                         break;
331                 case 'f':
332                         optical_flow = atoi(optarg);
333                         skip_args += 2;
334                         break;
335                 case 's':
336                         scale = atoi(optarg);
337                         skip_args += 2;
338                         break;
339                 case 't':
340                         test = parse_test(optarg);
341                         if (!test) {
342                                 fprintf(stderr, "Invalid test '%s', use '--test help' to get a list of tests\n", optarg);
343                                 exit(-1);
344                         }
345                         skip_args += 2;
346                         break;
347                 default:
348                         break;
349                 }
350         }
351
352         return skip_args;
353 }
354
355 /*
356  * usage
357  */
358
359 static void print_help(const char *app)
360 {
361         printf("Colorize version %s\n\n",
362 #include "../version.h"
363         );
364         printf("Usage: %s [options] <grey ppm image> <marked ppm image> <result ppm image> [<frames> <start>]\n", app);
365         printf("       Colorize grey image using maked image and save to result image.\n");
366         printf("       If frames and start frame is given, image names must include printf integer formatting (e.g. %%04d).\n");
367         printf("Usage: %s [options] <grey ppm image> marked <result ppm image>\n", app);
368         printf("       Colorize grey image using marked mask + palette and save to result image.\n");
369         printf("Usage: %s [options] sequence [list | <start with frame> [<stop with frame>]]\n", app);
370         printf("       Colorize movie sequence (generated by colorize gui) as found in the current directory.\n");
371         printf("       Use list to view sequence segments between keyframes.\n");
372         printf("\nOptions:\n");
373         printf(" -h --help                           This help\n");
374         printf(" -i --in-itr-num <num>               Alter inner iterations (weightening count) of colorization algorithm (default=%d)\n", in_itr_num);
375         printf(" -o --out-itr-num <num>              Alter outer iterations (complete turns) of colorization algorithm (default=%d)\n", out_itr_num);
376         printf(" -z --zscale <levels>                How many grids (staring with the finest) should be scaled in z direction to generate the next coarse grid ");
377         if (scalexyz < 999)
378                 printf("(default=%d)\n", scalexyz);
379         else
380                 printf("(default=infinite)\n");
381         printf(" -b --brightness-contrast [0 | 1]    Apply brightnes and contrast, if defined in palette by GUI (default=%d)\n", bright_contrast);
382 #ifdef WITH_OPENCV
383         printf(" -f --optical-flow [0 | 1]           Apply optical flow, if defined by GUI (default=%d)\n", optical_flow);
384 #endif
385         printf(" -s --scale [1..n]                   Scale down by the given factor for quick and dirty previews (default=%d)\n", scale);
386         printf(" -t --test <test>                    Generate test images. Use 'help' for list of tests\n");
387 }
388
389 static void print_test_help()
390 {
391         printf(" -t --test <test>        Generate test images...\n");
392 #ifdef WITH_OPENCV
393         printf("           flow-next     Optical flow plane to next image\n");
394         printf("           flow-prev     Optical flow plane to previous image\n");
395 #endif
396         printf("           marked        Only apply makred colors to grey image\n");
397         printf("           mask          Show mask of marked areas\n");
398         printf("           mask+color    Show mask of marked areas + color\n");
399         printf("           bc-only       Only apply brightness+contrast, leave colors of grey image as is\n");
400         printf("           bc-image      Show brightness+contrast change on grey image as uv components\n");
401 }
402
403 /*
404  * main function
405  */
406
407 int main(int argc, char *argv[])
408 {
409         darray_t *gI = NULL, *cI = NULL, *markIm = NULL, *ntscIm = NULL;
410         darray_t *flow = NULL, *flow_i = NULL;
411         double *ptr, *ptr2, *ptr3;
412         double diff, sum;
413         int dims[4];
414         int w = 0, h = 0, load_w, load_h, k = 1, index = 0, z;
415         unsigned char *img_buffer, *img_buffer_all = NULL, *mark_buffer = NULL;
416         int rc, x, y, i, ii;
417         time_t start, end;
418         const char* filename;
419         char first_filename[256];
420         int seq_offset = 0, seq_next = 0;
421         int features, change_bc;
422         int skip_args;
423
424         skip_args = handle_options(argc, argv);
425         argc -= skip_args;
426         argv += skip_args;
427
428         if (argc <= 1) {
429                 print_help(argv[0]);
430                 return 0;
431         }
432
433 next_sequence:
434         if (argc > 1 && !strcmp(argv[1], "sequence")) {
435                 k = load_sequence(&seq_offset, &seq_next, first_filename);
436                 if (k == 0)
437                         return 0;
438                 printf("Got %d frames from sequence (frames %d..%d)\n", k, seq_offset, seq_offset + k - 1);
439                 if (argc > 2 && (!strcmp(argv[2], "list") || atoi(argv[2]) > seq_offset)) {
440                         seq_offset = seq_next;
441                         goto next_sequence;
442                 }
443         } else if (argc <= 3) {
444                 print_help(argv[0]);
445                 return 0;
446         } else if (argc > 5) {
447                 k = atoi(argv[4]);
448                 index = atoi(argv[5]);
449         }
450         if (k <= 0)
451                 return 0;
452
453 #ifdef SINGLE_K_TEST
454         if (k > SINGLE_K_TEST)
455                 k = SINGLE_K_TEST;
456 #endif
457
458         time(&start);
459
460         change_bc = 0;
461         for (z = 0; z < k; z++) {
462                 if (sequence) {
463                         filename = sequence[z].filename;
464                         /* first_filename is set by load_sequence */
465                 } else {
466                         filename = argv[1];
467                         strcpy(first_filename, argv[1]);
468                 }
469 #ifdef WITH_OPENCV
470                 // load flow settings
471                 if (sequence || !strcmp(argv[2], "marked")) {
472                         flow_default();
473                         if (optical_flow)
474                                 load_flow(first_filename);
475                         if (flow_enable == 0 && (test == FLOW_NEXT || test == FLOW_PREV)) {
476                                 fprintf(stderr, "Cannot test optical flow, because it is not enabled by GUI.\n");
477                                 exit (0);
478                         }
479                 }
480 #endif
481                 // load original image and convert their RGB components to double RGB array
482                 rc = load_img(-1, &img_buffer, &load_w, &load_h, filename, index + z);
483                 if (rc) {
484                         fprintf(stderr, "Failed to load grey image\n");
485                         return 0;
486                 }
487                 scale_img(img_buffer, load_w, load_h, scale);
488                 if (z == 0) {
489                         w = load_w / scale;
490                         h = load_h / scale;
491                 }
492                 if (load_w/scale != w || load_h/scale != h) {
493                         fprintf(stderr, "Error: All input images must have equal dimenstions.\n");
494                         return 0;
495                 }
496                 // now we know the dimensions, so we can create input arrays
497                 if (!gI) {
498                         dims[0] = w; dims[1] = h; dims[2] = 3; dims[3] = k;
499                         gI = darrayCreate(4, dims);
500                 }
501                 if (!gI) {
502                         printf("failed to create grey image array\n");
503                         exit (0);
504                 }
505                 if (!cI) {
506                         dims[0] = w; dims[1] = h; dims[2] = 3; dims[3] = k;
507                         cI = darrayCreate(4, dims);
508                 }
509                 if (!cI) {
510                         printf("failed to create marked image array\n");
511                         exit (0);
512                 }
513                 ptr = darrayGetPr(gI) + w*h*3*z;
514                 img2array(img_buffer, w, h, ptr, w, h);
515 #ifdef WITH_OPENCV
516                 if (k > 1 && flow_enable) {
517                         if (!img_buffer_all)
518                                 img_buffer_all = malloc(w*h*3*k);
519                         if (!img_buffer_all) {
520                                 printf("failed to create grey image array\n");
521                                 exit (0);
522                         }
523                         memcpy(img_buffer_all + w*h*3*z, img_buffer, w*h*3);
524                 }
525 #endif
526                 free(img_buffer);
527
528                 if (sequence || !strcmp(argv[2], "marked")) {
529                         char name[256];
530                         unsigned char c;
531                         // load marked mask and convert their RGB components to double YUV array
532                         memcpy(darrayGetPr(cI) + w*h*3*z, darrayGetPr(gI) + w*h*3*z, w*h*3 * sizeof(double));
533                         /* add extra memory for unscaled data to prevent buffer overflow */
534                         if (!mark_buffer)
535                                 mark_buffer = (unsigned char *)malloc(w*h*k + load_w*load_h);
536                         if (!mark_buffer) {
537                                 printf("no memory for mark buffer\n");
538                                 exit (0);
539                         }
540                         if (load_palette(first_filename)) {
541                                 printf("failed to load palette for file: '%s'\n", filename);
542                                 exit (0);
543                         }
544                         sprintf(name, filename, index + z);
545                         /* always load full unscaled image, then scale down */
546                         if (load_marked(mark_buffer + w*h*z, load_w, load_h, name) == 0) {
547                                 scale_mark(mark_buffer + w*h*z, load_w, load_h, scale);
548                                 ptr = darrayGetPr(cI) + w*h*3*z;
549                                 for (y = 0; y < h; y++) {
550                                         for (x = 0; x < w; x++) {
551                                                 /* do not apply mask on index 0 */
552                                                 c = mark_buffer[y*w+x + w*h*z];
553                                                 if (c == 0)
554                                                         continue;
555                                                 /* check for any brightness/contrast change */
556                                                 if (bright_contrast && (mark_palette[c-1].bright != 0 || mark_palette[c-1].contrast != 1))
557                                                         change_bc = 1;
558                                                 /* do not apply white pixles, this meas: keep original color */
559                                                 if (mark_palette[c-1].r == 255 && mark_palette[c-1].g == 255 && mark_palette[c-1].b == 255)
560                                                         continue;
561                                                 ptr[y*w+x] = mark_palette[c-1].r / 255.0F;
562                                                 ptr[y*w+x + w*h] = mark_palette[c-1].g / 255.0F;
563                                                 ptr[y*w+x + w*h*2] = mark_palette[c-1].b / 255.0F;
564                                         }
565                                 }
566                         } else
567                                 memset(mark_buffer + w*h*z, 0, w*h);
568                 } else {
569                         // load marked image and convert their RGB components to double YUV array
570                         rc = load_img(-1, &img_buffer, &load_w, &load_h, argv[2], index + z);
571                         if (!rc) {
572                                 if (load_w/scale != w || load_h/scale != h) {
573                                         fprintf(stderr, "Error: All input images must have equal dimenstions.\n");
574                                         return 0;
575                                 }
576                                 ptr = darrayGetPr(cI) + w*h*3*z;
577                                 img2array(img_buffer, w, h, ptr, w, h);
578                                 free(img_buffer);
579                         } else {
580                                 fprintf(stderr, "Failed to load marked image, omitting...\n");
581                                 memcpy(darrayGetPr(cI) + w*h*3*z, darrayGetPr(gI) + w*h*3*z, w*h*3 * sizeof(double));
582                         }
583                 }
584         }
585
586         /* create color mask and ntsc arrays for the colorization process */
587         dims[0] = w; dims[1] = h; dims[2] = k;
588         markIm = darrayCreate(3, dims);
589         if (!markIm) {
590                 printf("failed to create mark array\n");
591                 exit (0);
592         }
593         features = (change_bc) ? 4 : 2;
594         dims[0] = w; dims[1] = h; dims[2] = features+1; dims[3] = k;
595         ntscIm = darrayCreate(4, dims);
596         if (!ntscIm) {
597                 printf("failed to create ntsc array\n");
598                 exit (0);
599         }
600
601         for (z = 0; z < k; z++) {
602                 if (sequence || !strcmp(argv[2], "marked")) {
603                         unsigned char c;
604                         ptr = darrayGetPr(markIm) + w*h*z;
605                         // use marked mask to fill markIm
606                         for (y = 0; y < h; y++) {
607                                 for (x = 0; x < w; x++) {
608                                         if (x < w && y < h) {
609                                                 /* do not apply mask on index 0 */
610                                                 c = mark_buffer[y*w+x + w*h*z];
611                                         } else
612                                                 c = 0;
613                                         if (c)
614                                                 ptr[y*w+x] = 1.0F;
615                                         else
616                                                 ptr[y*w+x] = 0.0F;
617                                 }
618                         }
619                 } else {
620                         // fill color image with marked pixles
621                         // - calculate the difference between two images (original image - color image)
622                         // - convert into absolute (positive values)
623                         // - sum all components to get grey image
624                         // - apply threshold (pixle is 1F, if the absolute difference is > 0.01F)
625                         // original code: markIm=(sum(abs(gI-cI),3)>0.01);
626                         ptr = darrayGetPr(gI) + w*h*3*z;
627                         ptr2 = darrayGetPr(cI) + w*h*3*z;
628                         ptr3 = darrayGetPr(markIm) + w*h*z;
629                         for (i = 0, ii = w * h; i < ii; i++) {
630                                 diff = 0;
631                                 sum = ptr[i] - ptr2[i];
632                                 if (sum < 0)
633                                         diff -= sum;
634                                 else
635                                         diff += sum;
636                                 sum = ptr[i + ii] - ptr2[i + ii];
637                                 if (sum < 0)
638                                         diff -= sum;
639                                 else
640                                         diff += sum;
641                                 sum = ptr[i + ii + ii] - ptr2[i + ii + ii];
642                                 if (sum < 0)
643                                         diff -= sum;
644                                 else
645                                         diff += sum;
646                                 if (diff > 0.01)
647                                         ptr3[i] = 1.0F;
648                                 else
649                                         ptr3[i] = 0.0F;
650                         }
651                 }
652
653                 // convert original image into YUV
654                 ptr = darrayGetPr(gI) + w*h*3*z;
655                 rgb2yuv(ptr, ptr, w, h);
656
657                 // convert maked image into YUV
658                 ptr = darrayGetPr(cI) + w*h*3*z;
659                 rgb2yuv(ptr, ptr, w, h);
660
661                 if (test != BC_ONLY) {
662                         if (sequence || !strcmp(argv[2], "marked")) {
663                                 unsigned char c;
664                                 // generate NTSC image: use luminance from original image and chrominance from original or marked image
665                                 ptr = darrayGetPr(gI) + w*h*3*z;
666                                 ptr2 = darrayGetPr(cI) + w*h*3*z;
667                                 ptr3 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
668                                 /* use original y component */
669                                 memcpy(ptr3, ptr, w * h * sizeof(double));
670                                 /* apply new uv components */
671                                 for (y = 0; y < h; y++) {
672                                         for (x = 0; x < w; x++) {
673                                                 c = mark_buffer[y*w+x + w*h*z];
674                                                 if (c == 0) {
675                                                         ptr3[w * h + w * y + x] = ptr[w * h + w * y + x];
676                                                         ptr3[w * h * 2 + w * y + x] = ptr[w * h * 2 + w * y + x];
677                                                 } else {
678                                                         ptr3[w * h + w * y + x] = ptr2[w * h + w * y + x];
679                                                         ptr3[w * h * 2 + w * y + x] = ptr2[w * h * 2 + w * y + x];
680                                                 }
681                                         }
682                                 }
683                         } else {
684                                 // generate NTSC image: use luminance from original image and chrominance from maked image
685                                 ptr = darrayGetPr(gI) + w*h*3*z;
686                                 ptr2 = darrayGetPr(cI) + w*h*3*z;
687                                 ptr3 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
688                                 memcpy(ptr3, ptr, w * h * sizeof(double));
689                                 memcpy(ptr3 + w * h, ptr2 + w * h, w * h * sizeof(double));
690                                 memcpy(ptr3 + w * h * 2, ptr2 + w * h * 2, w * h * sizeof(double));
691                         }
692                 } else {
693                         /* use grey image as result if BC_ONLY test is selected */
694                         ptr = darrayGetPr(gI) + w*h*3*z;
695                         ptr2 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
696                         memcpy(ptr2, ptr, w * h * 3 * sizeof(double));
697                 }
698
699                 if (change_bc) {
700                         unsigned char c;
701                         ptr2 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
702                         /* use original y component */
703                         memcpy(ptr2, ptr, w * h * sizeof(double));
704                         /* apply brightness and contrast from makred pixles to uv components of grey image */
705                         for (y = 0; y < h; y++) {
706                                 for (x = 0; x < w; x++) {
707                                         /* use unchanged brightness and contrast on index 0 */
708                                         c = mark_buffer[y*w+x + w*h*z];
709                                         if (c == 0) {
710                                                 ptr2[y*w+x + w*h*3] = 0;
711                                                 ptr2[y*w+x + w*h*4] = 0.1;
712                                         } else {
713                                                 ptr2[y*w+x + w*h*3] = mark_palette[c-1].bright / 10.0;
714                                                 ptr2[y*w+x + w*h*4] = mark_palette[c-1].contrast / 10.0;
715                                         }
716                                 }
717                         }
718                 }
719         }
720
721 #ifdef WITH_OPENCV
722         if (k > 1 && flow_enable) {
723                 /* create flow vectors */
724                 dims[0] = w; dims[1] = h; dims[2] = k - 1; dims[3] = 2;
725                 flow = darrayCreate(4, dims);
726                 if (!flow) {
727                         printf("failed to create array\n");
728                         exit (0);
729                 }
730                 flow_i = darrayCreate(4, dims);
731                 if (!flow_i) {
732                         printf("failed to create array\n");
733                         exit (0);
734                 }
735                 printf("Calculating optical flow for %d frames: window=%d\n", k, flow_window/scale);
736         }
737         for (z = 0; z < k-1; z++) {
738                 if (flow)
739                         create_flow_maps(NULL, img_buffer_all + w*h*3*(z+1), img_buffer_all + w*h*3*z, w, h, flow_window/scale, 0, NULL, NULL, darrayGetPr(flow) + w*h*z, darrayGetPr(flow) + w*h*z + w*h*(k-1), NULL);
740                 if (flow_i)
741                         create_flow_maps(img_buffer_all + w*h*3*z, NULL, img_buffer_all + w*h*3*(z+1), w, h, flow_window/scale, 0, darrayGetPr(flow_i) + w*h*z, darrayGetPr(flow_i) + w*h*z + w*h*(k-1), NULL, NULL, NULL);
742         }
743 #endif
744         free(img_buffer_all);
745         img_buffer_all = NULL;
746
747         darrayDestroy(gI);
748         gI = NULL;
749         darrayDestroy(cI);
750         cI = NULL;
751
752         if (test != FLOW_NEXT && test != FLOW_PREV && test != MARKED && test != MASK && test != MASK_COLOR && test != BC_ONLY && test != BC_IMAGE) {
753                 printf("Colorizing %d frames, please wait...\n", k);
754                 rc = colorize(ntscIm, markIm, flow, flow_i, in_itr_num, out_itr_num, scalexyz);
755                 if (rc < 0) {
756                         if (k > 1)
757                                 printf("No memory! Use smaller frames or less frames between key frames or add more memory.");
758                         else
759                                 printf("No memory! Use smaller image or add more memory.");
760                         exit(-1);
761                 }
762         }
763
764         /* if we have a change, we apply brightness+contrast from ntscIm */
765         if (change_bc) {
766                 for (z = 0; z < k; z++) {
767                         ptr = darrayGetPr(ntscIm) + w*h*(features+1)*z;
768                         ptr2 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
769                         for (y = 0; y < h; y++) {
770                                 for (x = 0; x < w; x++) {
771                                         /* apply contrast */
772                                         ptr2[w * y + x] = (ptr2[w * y + x] - 0.5) * ptr[w * h * 4 + w * y + x] * 10.0 + 0.5;
773                                         /* apply brightness */
774                                         ptr2[w * y + x] += ptr[w * h * 3 + w * y + x] * 10.0;
775                                         if (ptr2[w * y + x] < 0)
776                                                 ptr2[w * y + x] = 0;
777                                         if (ptr2[w * y + x] > 1)
778                                                 ptr2[w * y + x] = 1;
779 #if 0
780 #warning TEST: show brightness and contrast change as uv vectors on a grey array */
781 ptr2[w * y + x] = 0.5;
782 ptr2[w * h + w * y + x] = ptr[w * h + w * y + x] * 10;
783 ptr2[w * h * 2 + w * y + x] = ptr[w * h * 2 + w * y + x] * 10 - 1;
784 #endif
785                                 }
786                         }
787                 }
788         }
789
790 #ifdef WITH_OPENCV
791         if (test == FLOW_NEXT || test == FLOW_PREV) {
792                 /* apply flow planes to result image as u and y vector */
793                 for (z = 0; z < k; z++) {
794                         ptr = darrayGetPr(ntscIm) + w*h*(features+1)*z;
795                         if (test == FLOW_NEXT) {
796                                 ptr2 = darrayGetPr(flow) + w*h*z;
797                                 ptr3 = darrayGetPr(flow) + w*h*z*(k-1);
798                         } else {
799                                 ptr2 = darrayGetPr(flow_i) + w*h*z;
800                                 ptr3 = darrayGetPr(flow_i) + w*h*z*(k-1);
801                         }
802                         for (y = 0; y < h; y++) {
803                                 for (x = 0; x < w; x++) {
804                                         ptr[w * y + x] = 0.5;
805                                         if (z < k-1) {
806                                                 ptr[w * y + x + w*h] = ptr2[w * y + x] / 50;
807                                                 ptr[w * y + x + w*h*2] = ptr3[w * y + x] / 50;
808                                         } else {
809                                                 ptr[w * y + x + w*h] = 0;
810                                                 ptr[w * y + x + w*h*2] = 0;
811                                         }
812                                 }
813                         }
814                 }
815         }
816 #endif
817
818         if (test == MASK) {
819                 /* apply maked mask as image */
820                 for (z = 0; z < k; z++) {
821                         ptr = darrayGetPr(ntscIm) + w*h*(features+1)*z;
822                         ptr2 = darrayGetPr(markIm) + w*h*z;
823                         for (y = 0; y < h; y++) {
824                                 for (x = 0; x < w; x++) {
825                                         ptr[w * y + x] = ptr2[w * y + x];
826                                         ptr[w * y + x + w*h] = ptr2[w * y + x];
827                                         ptr[w * y + x + w*h*2] = ptr2[w * y + x];
828                                 }
829                         }
830                 }
831         }
832
833         if (test == MASK_COLOR) {
834                 /* apply maked mask on grey image */
835                 for (z = 0; z < k; z++) {
836                         ptr = darrayGetPr(ntscIm) + w*h*(features+1)*z;
837                         ptr2 = darrayGetPr(markIm) + w*h*z;
838                         for (y = 0; y < h; y++) {
839                                 for (x = 0; x < w; x++) {
840                                         /* darken unmarked areas, make maked areas with uniformed brightness */
841                                         if (ptr2[w * y + x] < 0.5F)
842                                                 ptr[w * y + x] = ptr[w * y + x] / 4;
843                                         else
844                                                 ptr[w * y + x] = 0.5F;
845                                 }
846                         }
847                 }
848         }
849
850         if (test == BC_IMAGE) {
851                 /* apply bc image as result image */
852                 for (z = 0; z < k; z++) {
853                         ptr = darrayGetPr(ntscIm) + w*h*(features+1)*z;
854                         ptr2 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
855                         /* use uniformed brightness as y component */
856                         for (y = 0; y < h; y++) {
857                                 for (x = 0; x < w; x++) {
858                                         ptr[w * y + x] = 0.5F;
859                                 }
860                         }
861                         if (change_bc) {
862                                 ptr2 = darrayGetPr(ntscIm) + w*h*(features+1)*z;
863                                 memcpy(ptr+w*h, ptr2+w*h*3, w * h * sizeof(double));
864                                 memcpy(ptr+w*h*2, ptr2+w*h*4, w * h * sizeof(double));
865                         } else {
866                                 memset(ptr+w*h, 0, w * h * sizeof(double));
867                                 memset(ptr+w*h*2, 0, w * h * sizeof(double));
868                         }
869                 }
870         }
871
872         // save result YUV array to image with RGB components
873         img_buffer = (unsigned char *)malloc(w * h * 3);
874         if (!img_buffer) {
875                 fprintf(stderr, "Failed to allocate image buffer\n");
876                 return 0;
877         }
878         for (z = 0; z < k; z++) {
879                 ptr = darrayGetPr(ntscIm) + w*h*(features+1)*z;
880                 yuv2rgb(ptr, ptr, w, h);
881                 array2img(ptr, w, h, img_buffer, w, h);
882                 if (sequence) {
883                         static char name[256], *p, *q;
884                         p = sequence[z].filename;
885                         while((q = strchr(p, DIR_SEPERATOR)))
886                                 p = q + 1;
887                         strcpy(name, sequence[z].filename);
888                         name[p - sequence[z].filename] = '\0';
889                         strcat(name, "colorized_");
890                         strcat(name, p);
891                         filename = name;
892                 } else
893                         filename = argv[3];
894                 save_img(img_buffer, w, h, filename, index + z);
895         }
896         free(img_buffer);
897
898         time(&end);
899         printf("Elapsed time: %d minutes, %d seconds\n", (int)(end-start)/60, (int)(end-start)%60);
900
901         // destroy
902         darrayDestroy(ntscIm);
903         ntscIm = NULL;
904         darrayDestroy(flow);
905         flow = NULL;
906         darrayDestroy(flow_i);
907         flow_i = NULL;
908         darrayDestroy(markIm);
909         markIm = NULL;
910
911         free(mark_buffer);
912         mark_buffer = NULL;
913
914         darrayDone();
915
916 #ifdef SINGLE_K_TEST
917         exit(0);
918 #endif
919         if (sequence) {
920                 /* if end frame is not given or if not reached */
921                 if (argc <= 3 || atoi(argv[3]) > seq_offset + k) {
922                         seq_offset = seq_offset + k - 1;
923                         goto next_sequence;
924                 }
925         }
926
927         return 0;
928 }
929