Support GRAY & CMYK Jpegs, fix rounding bug

jpeg.c

@@ -31,7 +31,9 @@ void read_jpeg(FILE *in, struct jpeg *jpeg) {
         jpeg->h = d.image_height;
         jpeg->w = d.image_width;
 
-        if(d.num_components != 3) { die("only 3 component jpegs are supported"); }
+        jpeg->c = d.num_components;
+        if(d.num_components < 1 || d.num_components > 4)
+          { die("only jpegs with 1 to 4 components are supported (gray, rgb/yuv or cmyk)"); }
 
         for(int c = 0; c < d.num_components; c++) {
                 unsigned i = d.comp_info[c].quant_tbl_no;
@@ -46,6 +48,8 @@ void read_jpeg(FILE *in, struct jpeg *jpeg) {
                 memcpy(&(jpeg->coefs[c].quant_table), t->quantval, sizeof(uint16_t) * 64);
         }
 
+#define UPDIV(x,y) ((x + ((y) - 1)) / (y))
+
         jvirt_barray_ptr *coefs = jpeg_read_coefficients(&d);
         for(int c = 0; c < d.num_components; c++) {
                 jpeg_component_info *i = &d.comp_info[c];
@@ -56,10 +60,11 @@ void read_jpeg(FILE *in, struct jpeg *jpeg) {
                 coef->h = h;
                 coef->w_samp = d.max_h_samp_factor / i->h_samp_factor;
                 coef->h_samp = d.max_v_samp_factor / i->v_samp_factor;
-                if(coef->h / 8 != (jpeg->h / coef->h_samp + 7) / 8) {
+                if(coef->h / 8 != UPDIV(UPDIV(jpeg->h, coef->h_samp), 8)) {
                         die("jpeg invalid coef h size");
                 }
-                if(coef->w / 8 != (jpeg->w / coef->w_samp + 7) / 8) {
+                if(coef->w / 8 != UPDIV(UPDIV(jpeg->w, coef->w_samp), 8)) {
+                        printf("coeg->w = %d, jpeg->w = %d, coef->w_samp = %d\n", coef->w, jpeg->w, coef->w_samp);
                         die("jpeg invalid coef w size");
                 }
                 if(SIZE_MAX / coef->h / coef->w / coef->h_samp / coef->w_samp < 6) {

jpeg.h

@@ -7,9 +7,10 @@
 #include "jpeg2png.h"
 
 struct jpeg {
+        unsigned c;
         unsigned h;
         unsigned w;
-        struct coef coefs[3];
+        struct coef coefs[4];
 };
 
 void read_jpeg(FILE *in, struct jpeg *jpeg);

jpeg2png.c

@@ -124,11 +124,11 @@ void decode_file(const char* infile, const char *outfile, unsigned iterations[3]
         struct jpeg jpeg;
         read_jpeg(in, &jpeg);
         fclose(in);
-        for(unsigned c = 0; c < 3; c++) {
+        for(unsigned c = 0; c < jpeg.c; c++) {
                 struct coef *coef = &jpeg.coefs[c];
                 decode_coefficients(coef);
         }
-        for(unsigned i = 0; i < 3; i++) {
+        for(unsigned i = 0; i < jpeg.c; i++) {
                 struct coef *coef = &jpeg.coefs[i];
                 float *temp = alloc_simd(sizeof(float) * coef->h * coef->w);
 
@@ -139,13 +139,13 @@ void decode_file(const char* infile, const char *outfile, unsigned iterations[3]
         }
 
         // smooth
-        if(all_together) {
+        if(all_together && jpeg.c == 3) {
                 plog->channel = 3;
                 compute(3, jpeg.coefs, plog, pb, weights[0], pweights, iterations[0]);
         } else {
                 struct logger log = *plog;
                 OPENMP(parallel for schedule(dynamic) firstprivate(log))
-                for(unsigned i = 0; i < 3; i++) {
+                for(unsigned i = 0; i < jpeg.c; i++) {
                         log.channel = i;
                         struct coef *coef = &jpeg.coefs[i];
                         compute(1, coef, &log, pb, weights[i], &pweights[i], iterations[i]);
@@ -161,11 +161,14 @@ void decode_file(const char* infile, const char *outfile, unsigned iterations[3]
         // write png
         FILE *out = fopen(outfile, "wb");
         if(!out) { die_perror("could not open output file `%s`", outfile); }
-        write_png(out, jpeg.w, jpeg.h, png_bits, &jpeg.coefs[0], &jpeg.coefs[1], &jpeg.coefs[2]);
+        if (jpeg.c == 3)
+          write_png(out, jpeg.w, jpeg.h, png_bits, &jpeg.coefs[0], &jpeg.coefs[1], &jpeg.coefs[2]);
+        else
+          write_png(out, jpeg.w, jpeg.h, png_bits, &jpeg.coefs[0], NULL, NULL);
         fclose(out);
 
         // clean up
-        for(unsigned i = 0; i < 3; i++) {
+        for(unsigned i = 0; i < jpeg.c; i++) {
                 free_simd(jpeg.coefs[i].fdata);
                 free(jpeg.coefs[i].data);
         }

png.c

@@ -37,8 +37,8 @@ void write_png(FILE *out, unsigned w, unsigned h, unsigned bits, struct coef *y,
         for(unsigned i = 0; i < h; i++) {
                 for(unsigned j = 0; j < w; j++) {
                         float yi = *p(y->fdata, j, i, y->w, y->h);
-                        float cbi = *p(cb->fdata, j, i, cb->w, cb->h);
-                        float cri = *p(cr->fdata, j, i, cr->w, cr->h);
+                        float cbi = cb ? *p(cb->fdata, j, i, cb->w, cb->h) : 0.0;
+                        float cri = cr ? *p(cr->fdata, j, i, cr->w, cr->h) : 0.0;
 
                         // YCbCr -> RGB
                         float bitfactor = (1 << bits) / 256.;