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