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