Do some manual strength reduction.
[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 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 }