c5ae87e3b003cc2b060ccfe3fa1571ff817a2246
[fjl] / 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 }