More 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 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 read_scan(struct byte_source* source, struct jpeg_image* image, huffman_tables_t* tables)
118 {
119         unsigned len = read_uint16(byte_source_input_func, source);
120         assert(len >= 2);
121         len -= 2;
122
123         assert(len >= 1);
124         unsigned num_components = read_uint8(byte_source_input_func, source);
125         --len;
126
127         unsigned component_num[256];
128         unsigned dc_huffman_table[256], ac_huffman_table[256];
129         unsigned ss, se, ah_al;
130         int last_dc[256];
131
132         for (unsigned i = 0; i < num_components; ++i) {
133                 unsigned char td_ta;
134                 assert(len >= 2);
135                 component_num[i] = read_uint8(byte_source_input_func, source);
136                 td_ta = read_uint8(byte_source_input_func, source);
137                 len -= 2;
138                 dc_huffman_table[i] = td_ta >> 4;
139                 ac_huffman_table[i] = td_ta & 0x0f;
140                 last_dc[i] = 0;
141         }
142
143         assert(len >= 3);
144         ss = read_uint8(byte_source_input_func, source);
145         se = read_uint8(byte_source_input_func, source);
146         ah_al = read_uint8(byte_source_input_func, source);
147         len -= 3;
148
149         if (len != 0) {
150                 fprintf(stderr, "Error: %u unused bytes at end of SOS segment\n", len);
151         }
152
153         struct bit_source bits;
154         init_bit_source(&bits, byte_source_input_func, 8, source);
155                 
156         unsigned mcu_x = 0, mcu_y = 0;
157
158         while (!bits.source_eof) {
159                 for (unsigned c = 0; c < num_components; ++c) {
160                         unsigned cn = component_num[c];
161                         assert(image->idct_data[image->qtable[cn]] != NULL);
162
163                         uint8_t* pixel_write_pointer_y = image->pixel_write_pointer[cn];
164                         for (unsigned local_yb = 0; local_yb < image->vsample[cn]; ++local_yb, pixel_write_pointer_y += image->stride[cn] * DCTSIZE) {
165                                 uint8_t* pixel_write_pointer = pixel_write_pointer_y;
166                                 for (unsigned local_xb = 0; local_xb < image->hsample[cn]; ++local_xb, pixel_write_pointer += DCTSIZE) {
167                                         const struct huffman_table* dc_table = &((*tables)[DC_CLASS][dc_huffman_table[c]]);
168                                         const struct huffman_table* ac_table = &((*tables)[AC_CLASS][ac_huffman_table[c]]);
169
170                                         // decode DC component
171                                         unsigned dc_category = read_huffman_symbol(dc_table, &bits);
172                                         possibly_refill(&bits, dc_category + DEHUF_TABLE_BITS);
173                                         last_dc[c] += extend(read_bits(&bits, dc_category), dc_category);
174                                         
175                                         int16_t coeff[DCTSIZE2] = { 0 };
176                                         coeff[0] = last_dc[c];
177
178                                         // decode AC components
179                                         for (unsigned i = 1; i < DCTSIZE2; ++i) {
180                                                 unsigned rs = read_huffman_symbol_no_refill(ac_table, &bits);
181                                                 unsigned r = rs >> 4;
182                                                 unsigned s = rs & 0xf;
183
184                                                 if (rs == 0x00) {
185                                                         /* end of block */
186                                                         break;
187                                                 }
188                                                 if (rs == 0xf0) {
189                                                         /* 16 zero coefficients */
190                                                         possibly_refill(&bits, DEHUF_TABLE_BITS);
191                                                         i += 15;
192                                                         continue;
193                                                 }
194                                                 i += r;
195
196                                                 possibly_refill(&bits, s + DEHUF_TABLE_BITS);
197                                                 coeff[unzigzag[i]] = extend(read_bits(&bits, s), s);
198                                         }
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
221                         // Some debug code.
222                         const int c = 1;
223                         if (mcu_y == image->num_blocks_vertical) {
224                                 unsigned stride = image->num_blocks_horizontal * image->hsample[c] * DCTSIZE;
225                                 unsigned height = image->num_blocks_vertical * image->vsample[c] * DCTSIZE;
226                                 printf("P5\n%u %u\n255\n", stride, height);
227                                 fwrite(image->pixel_data[c], stride * height, 1, stdout);
228                         }
229                 }
230         }
231         if (len != 0) {
232                 fprintf(stderr, "Error: %u unused bytes at end of SOS segment\n", len);
233         }
234 }
235                         
236 void skip_segment(struct byte_source* source)
237 {
238         uint8_t buf[4096];
239         for ( ;; ) {
240                 ssize_t ret = byte_source_input_func(source, buf, 4096);
241                 if (ret == -1) {
242                         fprintf(stderr, "Input error!\n");
243                         exit(1);
244                 }
245                 if (ret == 0) {
246                         return;
247                 }
248         }
249 }
250         
251 int main(void)
252 {
253         struct jpeg_image jpeg;
254         memset(&jpeg, 0, sizeof(jpeg));
255         init_choices();
256
257         struct byte_source source;
258         init_byte_source(&source, stdio_read, stdin);
259
260         huffman_tables_t tables;
261
262         for ( ;; ) {
263                 uint8_t m2 = byte_source_read_marker(&source);
264                 assert(m2 != 0x00);
265
266                 fprintf(stderr, "Marker 0x%02x, at position %ld\n", m2, ftell(stdin) - source.bytes_available);
267
268                 switch (m2) {
269                 case 0xe0:
270                 case 0xe1:
271                 case 0xe2:
272                 case 0xe3:
273                 case 0xe4:
274                 case 0xe5:
275                 case 0xe6:
276                 case 0xe7:
277                 case 0xe8:
278                 case 0xe9:
279                 case 0xea:
280                 case 0xeb:
281                 case 0xec:
282                 case 0xed:
283                 case 0xee:
284                 case 0xef:
285                         /* APP0 through APPF */
286                 case 0xfc:
287                         /* some EXIF stuff */
288                 case 0xfe:
289                         /* comment */
290                 case 0xff:
291                         /* ignore */
292                         skip_segment(&source);
293                         break;
294                 case 0xdb:
295                         /* DQT */
296                         read_dqt(&source, &jpeg);
297                         break;
298                 case 0xc0:
299                         /* SOF0 (baseline DCT, Huffman encoded) */
300                         read_sof(&source, &jpeg);
301                         break;
302                 case 0xd8:
303                         /* SOI */
304                         break;
305                 case 0xd9:
306                         /* EOI */
307                         exit(0);
308                 case 0xc4:
309                         /* DHT (define Huffman tables) */
310                         read_huffman_tables(&tables, byte_source_input_func, &source);
311                         break;
312                 case 0xda:
313                         /* SOS (start of scan) */
314                         read_scan(&source, &jpeg, &tables);
315                         break;
316                 default:
317                         fprintf(stderr, "Error: Unknown marker 0x%02x\n", m2);
318                         exit(1);
319                 }
320         }
321 }