Correct PGM height for vertically sampled components.
[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                                         for (unsigned y = 0; y < DCTSIZE; ++y) {
198                                                 unsigned real_x = (mcu_x * image->hsample[cn] + local_xb) * DCTSIZE;
199                                                 unsigned real_y = (mcu_y * image->vsample[cn] + local_yb) * DCTSIZE + y;
200                                         
201                                                 memcpy(image->pixel_data[cn] + real_y * stride + real_x,
202                                                        pixdata + y * DCTSIZE,
203                                                        DCTSIZE);
204                                         }
205                                 }
206                         }
207                 }
208                 
209                 if (++mcu_x == image->num_blocks_horizontal) {
210                         ++mcu_y;
211                         mcu_x = 0;
212
213                         // Some debug code.
214                         const int c = 1;
215                         if (mcu_y == image->num_blocks_vertical) {
216                                 unsigned stride = image->num_blocks_horizontal * image->hsample[c] * DCTSIZE;
217                                 unsigned height = image->num_blocks_vertical * image->vsample[c] * DCTSIZE;
218                                 printf("P5\n%u %u\n255\n", stride, height);
219                                 fwrite(image->pixel_data[c], stride * height, 1, stdout);
220                         }
221                 }
222         }
223         if (len != 0) {
224                 fprintf(stderr, "Error: %u unused bytes at end of SOS segment\n", len);
225         }
226 }
227                         
228 void skip_segment(struct byte_source* source)
229 {
230         uint8_t buf[4096];
231         for ( ;; ) {
232                 ssize_t ret = byte_source_input_func(source, buf, 4096);
233                 if (ret == -1) {
234                         fprintf(stderr, "Input error!\n");
235                         exit(1);
236                 }
237                 if (ret == 0) {
238                         return;
239                 }
240         }
241 }
242         
243 int main(void)
244 {
245         struct jpeg_image jpeg;
246         memset(&jpeg, 0, sizeof(jpeg));
247         init_choices();
248
249         struct byte_source source;
250         init_byte_source(&source, stdio_read, stdin);
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                         exit(0);
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 }