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