1 ! Copyright (C) 2009 Marc Fauconneau.
2 ! See http://factorcode.org/license.txt for BSD license.
3 USING: accessors arrays byte-arrays combinators
4 grouping compression.huffman images fry
5 images.processing io io.binary io.encodings.binary io.files
6 io.streams.byte-array kernel locals math math.bitwise
7 math.constants math.functions math.matrices math.order
8 math.ranges math.vectors memoize multiline namespaces
9 sequences sequences.deep images.loader io.streams.limited ;
12 QUALIFIED-WITH: bitstreams bs
16 TUPLE: loading-jpeg < image
19 { color-info initial: { f f f f } }
20 { quant-tables initial: { f f } }
21 { huff-tables initial: { f f f f } }
24 "jpg" jpeg-image register-image-class
25 "jpeg" jpeg-image register-image-class
29 : <loading-jpeg> ( headers bitstream -- image )
30 loading-jpeg new swap >>bitstream swap >>headers ;
32 SINGLETONS: SOF DHT DAC RST SOI EOI SOS DQT DNL DRI DHP EXP
35 ! ISO/IEC 10918-1 Table B.1
36 :: >marker ( byte -- marker )
39 { [ dup HEX: CC = ] [ { DAC } ] }
40 { [ dup HEX: C4 = ] [ { DHT } ] }
41 { [ dup HEX: C9 = ] [ { JPG } ] }
42 { [ dup -4 shift HEX: C = ] [ SOF byte 4 bits 2array ] }
44 { [ dup HEX: D8 = ] [ { SOI } ] }
45 { [ dup HEX: D9 = ] [ { EOI } ] }
46 { [ dup HEX: DA = ] [ { SOS } ] }
47 { [ dup HEX: DB = ] [ { DQT } ] }
48 { [ dup HEX: DC = ] [ { DNL } ] }
49 { [ dup HEX: DD = ] [ { DRI } ] }
50 { [ dup HEX: DE = ] [ { DHP } ] }
51 { [ dup HEX: DF = ] [ { EXP } ] }
52 { [ dup -4 shift HEX: D = ] [ RST byte 4 bits 2array ] }
54 { [ dup -4 shift HEX: E = ] [ APP byte 4 bits 2array ] }
55 { [ dup HEX: FE = ] [ { COM } ] }
56 { [ dup -4 shift HEX: F = ] [ JPG byte 4 bits 2array ] }
58 { [ dup HEX: 01 = ] [ { TEM } ] }
63 TUPLE: jpeg-chunk length type data ;
65 : <jpeg-chunk> ( type length data -- jpeg-chunk )
71 TUPLE: jpeg-color-info
72 h v quant-table dc-huff-table ac-huff-table { diff initial: 0 } id ;
74 : <jpeg-color-info> ( h v quant-table -- jpeg-color-info )
80 : jpeg> ( -- jpeg-image ) jpeg-image get ;
82 : apply-diff ( dc color -- dc' )
83 [ diff>> + dup ] [ (>>diff) ] bi ;
85 : fetch-tables ( component -- )
86 [ [ jpeg> quant-tables>> nth ] change-quant-table drop ]
87 [ [ jpeg> huff-tables>> nth ] change-dc-huff-table drop ]
88 [ [ 2 + jpeg> huff-tables>> nth ] change-ac-huff-table drop ] tri ;
90 : read4/4 ( -- a b ) read1 16 /mod ;
94 : decode-frame ( header -- )
101 swap 2array jpeg> (>>dim)
104 read1 read4/4 read1 <jpeg-color-info>
105 swap [ >>id ] keep jpeg> color-info>> set-nth
109 : decode-quant-table ( chunk -- )
116 read4/4 [ 0 assert= ] dip
118 swap jpeg> quant-tables>> set-nth
122 : decode-huff-table ( chunk -- )
123 data>> [ binary <byte-reader> ] [ length ] bi
126 [ input-stream get [ count>> ] [ limit>> ] bi < ]
130 dup [ ] [ + ] map-reduce read
131 binary [ [ read [ B{ } ] unless* ] { } map-as ] with-byte-reader
132 swap jpeg> huff-tables>> set-nth
134 ] with-input-stream* ;
136 : decode-scan ( chunk -- )
142 read1 jpeg> color-info>> nth clone
143 read1 16 /mod [ >>dc-huff-table ] [ >>ac-huff-table ] bi*
144 ] map jpeg> (>>components)
147 read1 16 /mod [ 0 assert= ] bi@
150 : singleton-first ( seq -- elt )
151 [ length 1 assert= ] [ first ] bi ;
153 ERROR: not-a-baseline-jpeg-image ;
155 : baseline-parse ( -- )
156 jpeg> headers>> [ type>> { SOF 0 } = ] any? [ not-a-baseline-jpeg-image ] unless
159 [ [ type>> { SOF 0 } = ] filter singleton-first decode-frame ]
160 [ [ type>> { DQT } = ] filter [ decode-quant-table ] each ]
161 [ [ type>> { DHT } = ] filter [ decode-huff-table ] each ]
162 [ [ type>> { SOS } = ] filter singleton-first decode-scan ]
165 : parse-marker ( -- marker )
166 read1 HEX: FF assert=
169 : parse-headers ( -- chunks )
170 [ parse-marker dup { SOS } = not ]
173 dup 2 - read <jpeg-chunk>
174 ] [ produce ] keep dip swap suffix ;
176 MEMO: zig-zag ( -- zz )
178 { 0 1 5 6 14 15 27 28 }
179 { 2 4 7 13 16 26 29 42 }
180 { 3 8 12 17 25 30 41 43 }
181 { 9 11 18 24 31 40 44 53 }
182 { 10 19 23 32 39 45 52 54 }
183 { 20 22 33 38 46 51 55 60 }
184 { 21 34 37 47 50 56 59 61 }
185 { 35 36 48 49 57 58 62 63 }
188 MEMO: yuv>bgr-matrix ( -- m )
191 { 1 -0.39465 -0.58060 }
195 : wave ( x u -- n ) swap 2 * 1 + * pi * 16 / cos ;
197 :: dct-vect ( u v -- basis )
198 { 8 8 } coord-matrix [ { u v } [ wave ] 2map product ] map^2
199 1 u v [ 0 = [ 2 sqrt / ] when ] bi@ 4 / m*n ;
201 MEMO: dct-matrix ( -- m ) 64 [0,b) [ 8 /mod dct-vect flatten ] map ;
203 : mb-dim ( component -- dim ) [ h>> ] [ v>> ] bi 2array ;
205 ! : blocks ( component -- seq )
206 ! mb-dim ! coord-matrix flip concat [ [ { 2 2 } v* ] [ v+ ] bi* ] with map ;
208 : all-macroblocks ( quot: ( mb -- ) -- )
212 [ color-info>> sift { 0 0 } [ mb-dim vmax ] reduce v/ ] bi
214 coord-matrix flip concat
216 [ each ] bi* ; inline
218 : reverse-zigzag ( b -- b' ) zig-zag swap [ nth ] curry map ;
220 : idct-factor ( b -- b' ) dct-matrix v.m ;
222 USE: math.blas.vectors
223 USE: math.blas.matrices
225 MEMO: dct-matrix-blas ( -- m ) dct-matrix >float-blas-matrix ;
226 : V.M ( x A -- x.A ) Mtranspose swap M.V ;
227 : idct-blas ( b -- b' ) >float-blas-vector dct-matrix-blas V.M ;
229 : idct ( b -- b' ) idct-factor ;
231 :: draw-block ( block x,y color-id jpeg-image -- )
232 block dup length>> sqrt >fixnum group flip
233 dup matrix-dim coord-matrix flip
235 [ '[ _ [ second ] [ first ] bi ] dip nth nth ]
236 [ x,y v+ color-id jpeg-image draw-color ] bi
239 : sign-extend ( bits v -- v' )
240 swap [ ] [ 1 - 2^ < ] 2bi
241 [ -1 swap shift 1 + + ] [ drop ] if ;
243 : read1-jpeg-dc ( decoder -- dc )
244 [ read1-huff dup ] [ bs>> bs:read ] bi sign-extend ;
246 : read1-jpeg-ac ( decoder -- run/ac )
247 [ read1-huff 16 /mod dup ] [ bs>> bs:read ] bi sign-extend 2array ;
249 :: decode-block ( color -- pixels )
250 color dc-huff-table>> read1-jpeg-dc color apply-diff
251 64 0 <array> :> coefs
255 color ac-huff-table>> read1-jpeg-ac
256 [ first 1 + k + k! ] [ second k coefs set-nth ] [ ] tri
260 coefs color quant-table>> v*
261 reverse-zigzag idct ;
263 :: draw-macroblock-yuv420 ( mb blocks -- )
264 mb { 16 16 } v* :> pos
265 0 blocks nth pos { 0 0 } v+ 0 jpeg> draw-block
266 1 blocks nth pos { 8 0 } v+ 0 jpeg> draw-block
267 2 blocks nth pos { 0 8 } v+ 0 jpeg> draw-block
268 3 blocks nth pos { 8 8 } v+ 0 jpeg> draw-block
269 4 blocks nth 8 group 2 matrix-zoom concat pos 1 jpeg> draw-block
270 5 blocks nth 8 group 2 matrix-zoom concat pos 2 jpeg> draw-block ;
272 :: draw-macroblock-yuv444 ( mb blocks -- )
274 3 iota [ [ blocks nth pos ] [ jpeg> draw-block ] bi ] each ;
276 :: draw-macroblock-y ( mb blocks -- )
278 0 blocks nth pos 0 jpeg> draw-block
279 64 0 <array> pos 1 jpeg> draw-block
280 64 0 <array> pos 2 jpeg> draw-block ;
284 ! [ 8 group 2 matrix-zoom concat ] unless
285 ! pos { 8 8 } v* color jpeg> draw-block ;
287 : decode-macroblock ( -- blocks )
290 [ mb-dim first2 * iota ]
291 [ [ decode-block ] curry replicate ] bi
294 : cleanup-bitstream ( bytes -- bytes' )
297 { HEX: FF } read-until
298 read1 [ HEX: 00 = and ] keep swap
301 swap >marker { EOI } assert=
306 : setup-bitmap ( image -- )
307 dup dim>> 16 v/n [ ceiling ] map 16 v*n >>dim
308 BGR >>component-order
309 ubyte-components >>component-type
311 dup dim>> first2 * 3 * 0 <array> >>bitmap
314 ERROR: unsupported-colorspace ;
315 SINGLETONS: YUV420 YUV444 Y MAGIC! ;
317 :: detect-colorspace ( jpeg-image -- csp )
318 jpeg-image color-info>> sift :> colors
320 colors length 1 = [ drop Y ] when
323 colors [ mb-dim { 1 1 } = ] all?
327 [ [ mb-dim { 1 1 } = ] all? ]
328 [ mb-dim { 2 2 } = ] bi* and
332 ! this eats ~50% cpu time
333 : draw-macroblocks ( mbs -- )
334 jpeg> detect-colorspace
336 { YUV420 [ [ first2 draw-macroblock-yuv420 ] each ] }
337 { YUV444 [ [ first2 draw-macroblock-yuv444 ] each ] }
338 { Y [ [ first2 draw-macroblock-y ] each ] }
339 [ unsupported-colorspace ]
342 ! this eats ~25% cpu time
343 : color-transform ( yuv -- rgb )
344 { 128 0 0 } v+ yuv>bgr-matrix swap m.v
345 [ 0 max 255 min >fixnum ] map ;
347 : baseline-decompress ( -- )
348 jpeg> bitstream>> cleanup-bitstream { 255 255 255 255 } append
349 >byte-array bs:<msb0-bit-reader> jpeg> (>>bitstream)
352 [ [ [ <huffman-decoder> ] with map ] change-huff-tables drop ] bi
353 jpeg> components>> [ fetch-tables ] each
354 [ decode-macroblock 2array ] accumulator
355 [ all-macroblocks ] dip
356 jpeg> setup-bitmap draw-macroblocks
357 jpeg> bitmap>> 3 <groups> [ color-transform ] map! drop
358 jpeg> [ >byte-array ] change-bitmap drop ;
360 ERROR: not-a-jpeg-image ;
362 : loading-jpeg>image ( loading-jpeg -- image )
368 : load-jpeg ( stream -- loading-jpeg )
370 parse-marker { SOI } = [ not-a-jpeg-image ] unless
372 unlimited-input contents <loading-jpeg>
373 ] with-input-stream ;
377 M: jpeg-image stream>image ( stream jpeg-image -- bitmap )
378 drop load-jpeg loading-jpeg>image ;