git.sesse.net Git - fjl/blob - driver.c

   1 #include <stdio.h>
   2 #include <string.h>
   3 #include <stdlib.h>
   4
   5 #include "bytesource.h"
   6 #include "choice.h"
   7 #include "dehuff.h"
   8 #include "idct.h"
   9 #include "input.h"
  10 #include "zigzag.h"
  11
  12 struct jpeg_image {
  13         unsigned precision;
  14         unsigned width, height;
  15         unsigned num_components;
  16         unsigned hsample[256], vsample[256], qtable[256];
  17         unsigned max_hsample, max_vsample;
  18         unsigned num_blocks_horizontal, num_blocks_vertical;
  19         uint32_t qvalues[256][DCTSIZE2];
  20         void* idct_data[256];
  21         uint8_t* pixel_data[256];
  22 };
  23
  24 ssize_t stdio_read(void* userdata, uint8_t* buf, size_t count)
  25 {
  26         return fread(buf, 1, count, (FILE*)userdata);
  27 }
  28
  29 void read_dqt(struct byte_source* source, struct jpeg_image* image)
  30 {
  31         unsigned len = read_uint16(byte_source_input_func, source);
  32         assert(len >= 67);
  33         uint8_t precision_table = read_uint8(byte_source_input_func, source);
  34         int precision = precision_table >> 4;  // 0 = 8 bits, otherwise 16 bits.
  35         int table = precision_table & 0x0f;
  36
  37         if (image->idct_data[table] != NULL) {
  38                 idct_choice_free(image->idct_data[table]);
  39         }
  40
  41         if (precision != 0) {
  42                 assert(len == 131);
  43                 fprintf(stderr, "Quantization table %u: 16 bits/entry\n", table);
  44         } else {
  45                 assert(len == 67);
  46                 fprintf(stderr, "Quantization table %u: 8 bits/entry\n", table);
  47         }
  48
  49         for (unsigned i = 0; i < 64; ++i) {
  50                 if (precision != 0) {
  51                         image->qvalues[table][unzigzag[i]] =
  52                                 read_uint16(byte_source_input_func, source);
  53                 } else {
  54                         image->qvalues[table][unzigzag[i]] =
  55                                 read_uint8(byte_source_input_func, source);
  56                 }
  57         }
  58
  59         image->idct_data[table] = idct_choice_alloc(image->qvalues[table]);
  60 }
  61
  62 void read_sof(struct byte_source* source, struct jpeg_image* image)
  63 {
  64         unsigned len = read_uint16(byte_source_input_func, source);
  65         assert(len >= 8);
  66         image->precision = read_uint8(byte_source_input_func, source);
  67         assert(image->precision == 8);
  68         image->height = read_uint16(byte_source_input_func, source);
  69         image->width = read_uint16(byte_source_input_func, source);
  70         image->num_components = read_uint8(byte_source_input_func, source);
  71         len -= 8;
  72
  73         fprintf(stderr, "%u-bit %ux%u JPEG with %u components\n",
  74                 image->precision, image->width, image->height, image->num_components);
  75
  76         for (unsigned i = 0; i < image->num_components; ++i) {
  77                 assert(len >= 3);
  78                 unsigned c = read_uint8(byte_source_input_func, source);
  79                 unsigned sampling_factors = read_uint8(byte_source_input_func, source);
  80                 image->hsample[c] = sampling_factors >> 4;
  81                 image->vsample[c] = sampling_factors & 0x0f;
  82                 image->qtable[c] = read_uint8(byte_source_input_func, source);
  83                 len -= 3;
  84
  85                 if (image->hsample[c] > image->max_hsample) {
  86                         image->max_hsample = image->hsample[c];
  87                 }
  88                 if (image->vsample[c] > image->max_vsample) {
  89                         image->max_vsample = image->vsample[c];
  90                 }
  91
  92                 fprintf(stderr, "Component %u: sampling factors %u x %x, quantization table %u\n",
  93                         c, image->hsample[c], image->vsample[c], image->qtable[c]);
  94         }
  95
  96         image->num_blocks_horizontal = (image->width + image->max_hsample * DCTSIZE - 1) / (image->max_hsample * DCTSIZE);
  97         image->num_blocks_vertical = (image->height + image->max_vsample * DCTSIZE - 1) / (image->max_vsample * DCTSIZE);
  98
  99         for (unsigned c = 0; c < 256; ++c) {
 100                 if (image->hsample[c] == 0) {
 101                         continue;
 102                 }
 103
 104                 unsigned width = image->num_blocks_horizontal * image->hsample[c] * DCTSIZE;
 105                 unsigned height = image->num_blocks_vertical * image->vsample[c] * DCTSIZE;
 106                 image->pixel_data[c] = (uint8_t*)malloc(width * height);
 107                 assert(image->pixel_data[c] != NULL);
 108
 109                 fprintf(stderr, "Component %u: allocating %d x %d\n", c, width, height);
 110         }
 111 }
 112
 113 void read_scan(struct byte_source* source, struct jpeg_image* image, huffman_tables_t* tables)
 114 {
 115         unsigned len = read_uint16(byte_source_input_func, source);
 116         assert(len >= 2);
 117         len -= 2;
 118
 119         assert(len >= 1);
 120         unsigned num_components = read_uint8(byte_source_input_func, source);
 121         --len;
 122
 123         unsigned component_num[256];
 124         unsigned dc_huffman_table[256], ac_huffman_table[256];
 125         unsigned ss, se, ah_al;
 126         int last_dc[256];
 127
 128         for (unsigned i = 0; i < num_components; ++i) {
 129                 unsigned char td_ta;
 130                 assert(len >= 2);
 131                 component_num[i] = read_uint8(byte_source_input_func, source);
 132                 td_ta = read_uint8(byte_source_input_func, source);
 133                 len -= 2;
 134                 dc_huffman_table[i] = td_ta >> 4;
 135                 ac_huffman_table[i] = td_ta & 0x0f;
 136                 last_dc[i] = 0;
 137         }
 138
 139         assert(len >= 3);
 140         ss = read_uint8(byte_source_input_func, source);
 141         se = read_uint8(byte_source_input_func, source);
 142         ah_al = read_uint8(byte_source_input_func, source);
 143         len -= 3;
 144
 145         if (len != 0) {
 146                 fprintf(stderr, "Error: %u unused bytes at end of SOS segment\n", len);
 147         }
 148
 149         struct bit_source bits;
 150         init_bit_source(&bits, byte_source_input_func, 8, source);
 151
 152         unsigned mcu_x = 0, mcu_y = 0;
 153
 154         while (!bits.source_eof) {
 155                 for (unsigned c = 0; c < num_components; ++c) {
 156                         unsigned cn = component_num[c];
 157                         unsigned stride = image->num_blocks_horizontal * image->hsample[cn] * DCTSIZE;
 158                         assert(image->idct_data[image->qtable[cn]] != NULL);
 159
 160                         for (unsigned local_yb = 0; local_yb < image->vsample[cn]; ++local_yb) {
 161                                 for (unsigned local_xb = 0; local_xb < image->hsample[cn]; ++local_xb) {
 162                                         const struct huffman_table* dc_table = &((*tables)[DC_CLASS][dc_huffman_table[c]]);
 163                                         const struct huffman_table* ac_table = &((*tables)[AC_CLASS][ac_huffman_table[c]]);
 164
 165                                         // decode DC component
 166                                         unsigned dc_category = read_huffman_symbol(dc_table, &bits);
 167                                         possibly_refill(&bits, dc_category);
 168                                         last_dc[c] += extend(read_bits(&bits, dc_category), dc_category);
 169
 170                                         int16_t coeff[DCTSIZE2] = { 0 };
 171                                         coeff[0] = last_dc[c];
 172
 173                                         // decode AC components
 174                                         for (unsigned i = 1; i < DCTSIZE2; ++i) {
 175                                                 unsigned rs = read_huffman_symbol(ac_table, &bits);
 176                                                 unsigned r = rs >> 4;
 177                                                 unsigned s = rs & 0xf;
 178
 179                                                 if (rs == 0x00) {
 180                                                         /* end of block */
 181                                                         break;
 182                                                 }
 183                                                 if (rs == 0xf0) {
 184                                                         /* 16 zero coefficients */
 185                                                         i += 15;
 186                                                         continue;
 187                                                 }
 188                                                 i += r;
 189
 190                                                 possibly_refill(&bits, s);
 191                                                 coeff[unzigzag[i]] = extend(read_bits(&bits, s), s);
 192                                         }
 193
 194                                         uint8_t pixdata[DCTSIZE2];
 195                                         idct_choice(coeff, image->idct_data[image->qtable[cn]], pixdata);
 196
 197                                         unsigned real_x = (mcu_x * image->hsample[cn] + local_xb) * DCTSIZE;
 198                                         unsigned real_y = (mcu_y * image->vsample[cn] + local_yb) * DCTSIZE;
 199                                         uint8_t* dest_pixdata = image->pixel_data[cn] + real_y * stride + real_x;
 200                                         for (unsigned y = 0; y < DCTSIZE; ++y, dest_pixdata += stride) {
 201                                                 memcpy(dest_pixdata, pixdata + y * DCTSIZE, DCTSIZE);
 202                                         }
 203                                 }
 204                         }
 205                 }
 206
 207                 if (++mcu_x == image->num_blocks_horizontal) {
 208                         ++mcu_y;
 209                         mcu_x = 0;
 210
 211                         // Some debug code.
 212                         const int c = 1;
 213                         if (mcu_y == image->num_blocks_vertical) {
 214                                 unsigned stride = image->num_blocks_horizontal * image->hsample[c] * DCTSIZE;
 215                                 unsigned height = image->num_blocks_vertical * image->vsample[c] * DCTSIZE;
 216                                 printf("P5\n%u %u\n255\n", stride, height);
 217                                 fwrite(image->pixel_data[c], stride * height, 1, stdout);
 218                         }
 219                 }
 220         }
 221         if (len != 0) {
 222                 fprintf(stderr, "Error: %u unused bytes at end of SOS segment\n", len);
 223         }
 224 }
 225
 226 void skip_segment(struct byte_source* source)
 227 {
 228         uint8_t buf[4096];
 229         for ( ;; ) {
 230                 ssize_t ret = byte_source_input_func(source, buf, 4096);
 231                 if (ret == -1) {
 232                         fprintf(stderr, "Input error!\n");
 233                         exit(1);
 234                 }
 235                 if (ret == 0) {
 236                         return;
 237                 }
 238         }
 239 }
 240
 241 int main(void)
 242 {
 243         struct jpeg_image jpeg;
 244         memset(&jpeg, 0, sizeof(jpeg));
 245         init_choices();
 246
 247         struct byte_source source;
 248         init_byte_source(&source, stdio_read, stdin);
 249
 250         huffman_tables_t tables;
 251
 252         for ( ;; ) {
 253                 uint8_t m2 = byte_source_read_marker(&source);
 254                 assert(m2 != 0x00);
 255
 256                 fprintf(stderr, "Marker 0x%02x, at position %ld\n", m2, ftell(stdin) - source.bytes_available);
 257
 258                 switch (m2) {
 259                 case 0xe0:
 260                 case 0xe1:
 261                 case 0xe2:
 262                 case 0xe3:
 263                 case 0xe4:
 264                 case 0xe5:
 265                 case 0xe6:
 266                 case 0xe7:
 267                 case 0xe8:
 268                 case 0xe9:
 269                 case 0xea:
 270                 case 0xeb:
 271                 case 0xec:
 272                 case 0xed:
 273                 case 0xee:
 274                 case 0xef:
 275                         /* APP0 through APPF */
 276                 case 0xfc:
 277                         /* some EXIF stuff */
 278                 case 0xfe:
 279                         /* comment */
 280                 case 0xff:
 281                         /* ignore */
 282                         skip_segment(&source);
 283                         break;
 284                 case 0xdb:
 285                         /* DQT */
 286                         read_dqt(&source, &jpeg);
 287                         break;
 288                 case 0xc0:
 289                         /* SOF0 (baseline DCT, Huffman encoded) */
 290                         read_sof(&source, &jpeg);
 291                         break;
 292                 case 0xd8:
 293                         /* SOI */
 294                         break;
 295                 case 0xd9:
 296                         /* EOI */
 297                         exit(0);
 298                 case 0xc4:
 299                         /* DHT (define Huffman tables) */
 300                         read_huffman_tables(&tables, byte_source_input_func, &source);
 301                         break;
 302                 case 0xda:
 303                         /* SOS (start of scan) */
 304                         read_scan(&source, &jpeg, &tables);
 305                         break;
 306                 default:
 307                         fprintf(stderr, "Error: Unknown marker 0x%02x\n", m2);
 308                         exit(1);
 309                 }
 310         }
 311 }