-! Copyright (C) 2009 Marc Fauconneau.\r
-! See http://factorcode.org/license.txt for BSD license.\r
-USING: accessors arrays assocs byte-arrays\r
-byte-vectors combinators constructors fry grouping hashtables\r
-compression.huffman images io.binary kernel locals\r
-math math.bitwise math.order math.ranges multiline sequences\r
-sorting ;\r
-IN: compression.inflate\r
-\r
-QUALIFIED-WITH: bitstreams bs\r
-\r
-<PRIVATE\r
-\r
-: enum>seq ( assoc -- seq )\r
- dup keys [ ] [ max ] map-reduce 1 + f <array>\r
- [ '[ swap _ set-nth ] assoc-each ] keep ;\r
-\r
-ERROR: zlib-unimplemented ;\r
-ERROR: bad-zlib-data ;\r
-ERROR: bad-zlib-header ;\r
- \r
-:: check-zlib-header ( data -- )\r
- 16 data bs:peek 2 >le be> 31 mod ! checksum\r
- 0 assert= \r
- 4 data bs:read 8 assert= ! compression method: deflate\r
- 4 data bs:read ! log2(max length)-8, 32K max\r
- 7 <= [ bad-zlib-header ] unless \r
- 5 data bs:seek ! drop check bits \r
- 1 data bs:read 0 assert= ! dictionnary - not allowed in png\r
- 2 data bs:seek ! compression level; ignore\r
- ;\r
-\r
-:: default-table ( -- table )\r
- 0 <hashtable> :> table\r
- 0 143 [a,b] 280 287 [a,b] append 8 table set-at\r
- 144 255 [a,b] >array 9 table set-at\r
- 256 279 [a,b] >array 7 table set-at \r
- table enum>seq 1 tail ;\r
-\r
-CONSTANT: clen-shuffle { 16 17 18 0 8 7 9 6 10 5 11 4 12 3 13 2 14 1 15 }\r
-\r
-: get-table ( values size -- table ) \r
- 16 f <array> clone <enum> \r
- [ '[ _ push-at ] 2each ] keep seq>> 1 tail [ natural-sort ] map ;\r
-\r
-:: decode-huffman-tables ( bitstream -- tables )\r
- 5 bitstream bs:read 257 +\r
- 5 bitstream bs:read 1 +\r
- 4 bitstream bs:read 4 +\r
- clen-shuffle swap head\r
- dup [ drop 3 bitstream bs:read ] map\r
- get-table\r
- bitstream swap <huffman-decoder> \r
- [ 2dup + ] dip swap :> k!\r
- '[\r
- _ read1-huff2\r
- {\r
- { [ dup 16 = ] [ 2 bitstream bs:read 3 + 2array ] }\r
- { [ dup 17 = ] [ 3 bitstream bs:read 3 + 2array ] }\r
- { [ dup 18 = ] [ 7 bitstream bs:read 11 + 2array ] }\r
- [ ]\r
- } cond\r
- dup array? [ dup second ] [ 1 ] if\r
- k swap - dup k! 0 >\r
- ] \r
- [ ] produce swap suffix\r
- { } [ dup array? [ dup first 16 = ] [ f ] if [ [ unclip-last ] [ second 1+ swap <repetition> append ] bi* ] [ suffix ] if ] reduce\r
- [ dup array? [ second 0 <repetition> ] [ 1array ] if ] map concat\r
- nip swap cut 2array [ [ length>> [0,b) ] [ ] bi get-table ] map ;\r
- \r
-CONSTANT: length-table\r
- {\r
- 3 4 5 6 7 8 9 10\r
- 11 13 15 17\r
- 19 23 27 31\r
- 35 43 51 59\r
- 67 83 99 115\r
- 131 163 195 227 258\r
- }\r
-\r
-CONSTANT: dist-table\r
- {\r
- 1 2 3 4 \r
- 5 7 9 13 \r
- 17 25 33 49\r
- 65 97 129 193\r
- 257 385 513 769\r
- 1025 1537 2049 3073\r
- 4097 6145 8193 12289\r
- 16385 24577\r
- }\r
-\r
-: nth* ( n seq -- elt )\r
- [ length 1- swap - ] [ nth ] bi ;\r
-\r
-:: inflate-lz77 ( seq -- bytes )\r
- 1000 <byte-vector> :> bytes\r
- seq\r
- [\r
- dup array?\r
- [ first2 '[ _ 1- bytes nth* bytes push ] times ]\r
- [ bytes push ] if\r
- ] each \r
- bytes ;\r
-\r
-:: inflate-dynamic ( bitstream -- bytes )\r
- bitstream decode-huffman-tables\r
- bitstream '[ _ swap <huffman-decoder> ] map :> tables\r
- [\r
- tables first read1-huff2\r
- dup 256 >\r
- [\r
- dup 285 = \r
- [ ]\r
- [ \r
- dup 264 > \r
- [ \r
- dup 261 - 4 /i dup 5 > \r
- [ bad-zlib-data ] when \r
- bitstream bs:read 2array \r
- ]\r
- when \r
- ] if\r
- ! 5 bitstream read-bits ! distance\r
- tables second read1-huff2\r
- dup 3 > \r
- [ \r
- dup 2 - 2 /i dup 13 >\r
- [ bad-zlib-data ] when\r
- bitstream bs:read 2array\r
- ] \r
- when\r
- 2array\r
- ]\r
- when\r
- dup 256 = not\r
- ]\r
- [ ] produce nip\r
- [\r
- dup array? [\r
- first2\r
- [ \r
- dup array? [ first2 ] [ 0 ] if\r
- [ 257 - length-table nth ] [ + ] bi*\r
- ] \r
- [\r
- dup array? [ first2 ] [ 0 ] if\r
- [ dist-table nth ] [ + ] bi*\r
- ] bi*\r
- 2array\r
- ] when\r
- ] map ;\r
- \r
-: inflate-raw ( bitstream -- bytes ) zlib-unimplemented ;\r
-: inflate-static ( bitstream -- bytes ) zlib-unimplemented ;\r
-\r
-:: inflate-loop ( bitstream -- bytes )\r
- [ 1 bitstream bs:read 0 = ]\r
- [\r
- bitstream\r
- 2 bitstream bs:read\r
- { \r
- { 0 [ inflate-raw ] }\r
- { 1 [ inflate-static ] }\r
- { 2 [ inflate-dynamic ] }\r
- { 3 [ bad-zlib-data f ] }\r
- }\r
- case\r
- ]\r
- [ produce ] keep call suffix concat ;\r
- \r
- ! [ produce ] keep dip swap suffix\r
-\r
-:: paeth ( a b c -- p ) \r
- a b + c - { a b c } [ [ - abs ] keep 2array ] with map \r
- sort-keys first second ;\r
- \r
-:: png-unfilter-line ( prev curr filter -- curr' )\r
- prev :> c\r
- prev 3 tail-slice :> b\r
- curr :> a\r
- curr 3 tail-slice :> x\r
- x length [0,b)\r
- filter\r
- {\r
- { 0 [ drop ] }\r
- { 1 [ [| n | n x nth n a nth + 256 wrap n x set-nth ] each ] }\r
- { 2 [ [| n | n x nth n b nth + 256 wrap n x set-nth ] each ] }\r
- { 3 [ [| n | n x nth n a nth n b nth + 2/ + 256 wrap n x set-nth ] each ] }\r
- { 4 [ [| n | n x nth n a nth n b nth n c nth paeth + 256 wrap n x set-nth ] each ] }\r
- \r
- } case \r
- curr 3 tail ;\r
-\r
-PRIVATE>\r
-\r
-! for debug -- shows residual values\r
-: reverse-png-filter' ( lines -- byte-array )\r
- [ first ] [ 1 tail ] [ map ] bi-curry@ bi nip\r
- concat [ 128 + ] B{ } map-as ;\r
- \r
-: reverse-png-filter ( lines -- byte-array )\r
- dup first [ 0 ] replicate prefix\r
- [ { 0 0 } prepend ] map\r
- 2 clump [\r
- first2 dup [ third ] [ 0 2 rot set-nth ] bi png-unfilter-line\r
- ] map B{ } concat-as ;\r
-\r
-: zlib-inflate ( bytes -- bytes )\r
- bs:<lsb0-bit-reader>\r
- [ check-zlib-header ] [ inflate-loop ] bi\r
- inflate-lz77 ;\r
+! Copyright (C) 2009 Marc Fauconneau.
+! See http://factorcode.org/license.txt for BSD license.
+USING: accessors arrays assocs byte-arrays
+byte-vectors combinators constructors fry grouping hashtables
+compression.huffman images io.binary kernel locals
+math math.bitwise math.order math.ranges multiline sequences
+sorting ;
+IN: compression.inflate
+
+QUALIFIED-WITH: bitstreams bs
+
+<PRIVATE
+
+: enum>seq ( assoc -- seq )
+ dup keys [ ] [ max ] map-reduce 1 + f <array>
+ [ '[ swap _ set-nth ] assoc-each ] keep ;
+
+ERROR: zlib-unimplemented ;
+ERROR: bad-zlib-data ;
+ERROR: bad-zlib-header ;
+
+:: check-zlib-header ( data -- )
+ 16 data bs:peek 2 >le be> 31 mod ! checksum
+ 0 assert=
+ 4 data bs:read 8 assert= ! compression method: deflate
+ 4 data bs:read ! log2(max length)-8, 32K max
+ 7 <= [ bad-zlib-header ] unless
+ 5 data bs:seek ! drop check bits
+ 1 data bs:read 0 assert= ! dictionnary - not allowed in png
+ 2 data bs:seek ! compression level; ignore
+ ;
+
+:: default-table ( -- table )
+ 0 <hashtable> :> table
+ 0 143 [a,b] 280 287 [a,b] append 8 table set-at
+ 144 255 [a,b] >array 9 table set-at
+ 256 279 [a,b] >array 7 table set-at
+ table enum>seq 1 tail ;
+
+CONSTANT: clen-shuffle { 16 17 18 0 8 7 9 6 10 5 11 4 12 3 13 2 14 1 15 }
+
+: get-table ( values size -- table )
+ 16 f <array> clone <enum>
+ [ '[ _ push-at ] 2each ] keep seq>> 1 tail [ natural-sort ] map ;
+
+:: decode-huffman-tables ( bitstream -- tables )
+ 5 bitstream bs:read 257 +
+ 5 bitstream bs:read 1 +
+ 4 bitstream bs:read 4 +
+ clen-shuffle swap head
+ dup [ drop 3 bitstream bs:read ] map
+ get-table
+ bitstream swap <huffman-decoder>
+ [ 2dup + ] dip swap :> k!
+ '[
+ _ read1-huff2
+ {
+ { [ dup 16 = ] [ 2 bitstream bs:read 3 + 2array ] }
+ { [ dup 17 = ] [ 3 bitstream bs:read 3 + 2array ] }
+ { [ dup 18 = ] [ 7 bitstream bs:read 11 + 2array ] }
+ [ ]
+ } cond
+ dup array? [ dup second ] [ 1 ] if
+ k swap - dup k! 0 >
+ ]
+ [ ] produce swap suffix
+ { } [ dup array? [ dup first 16 = ] [ f ] if [ [ unclip-last ] [ second 1+ swap <repetition> append ] bi* ] [ suffix ] if ] reduce
+ [ dup array? [ second 0 <repetition> ] [ 1array ] if ] map concat
+ nip swap cut 2array [ [ length>> [0,b) ] [ ] bi get-table ] map ;
+
+CONSTANT: length-table
+ {
+ 3 4 5 6 7 8 9 10
+ 11 13 15 17
+ 19 23 27 31
+ 35 43 51 59
+ 67 83 99 115
+ 131 163 195 227 258
+ }
+
+CONSTANT: dist-table
+ {
+ 1 2 3 4
+ 5 7 9 13
+ 17 25 33 49
+ 65 97 129 193
+ 257 385 513 769
+ 1025 1537 2049 3073
+ 4097 6145 8193 12289
+ 16385 24577
+ }
+
+: nth* ( n seq -- elt )
+ [ length 1- swap - ] [ nth ] bi ;
+
+:: inflate-lz77 ( seq -- bytes )
+ 1000 <byte-vector> :> bytes
+ seq
+ [
+ dup array?
+ [ first2 '[ _ 1- bytes nth* bytes push ] times ]
+ [ bytes push ] if
+ ] each
+ bytes ;
+
+:: inflate-dynamic ( bitstream -- bytes )
+ bitstream decode-huffman-tables
+ bitstream '[ _ swap <huffman-decoder> ] map :> tables
+ [
+ tables first read1-huff2
+ dup 256 >
+ [
+ dup 285 =
+ [ ]
+ [
+ dup 264 >
+ [
+ dup 261 - 4 /i dup 5 >
+ [ bad-zlib-data ] when
+ bitstream bs:read 2array
+ ]
+ when
+ ] if
+ ! 5 bitstream read-bits ! distance
+ tables second read1-huff2
+ dup 3 >
+ [
+ dup 2 - 2 /i dup 13 >
+ [ bad-zlib-data ] when
+ bitstream bs:read 2array
+ ]
+ when
+ 2array
+ ]
+ when
+ dup 256 = not
+ ]
+ [ ] produce nip
+ [
+ dup array? [
+ first2
+ [
+ dup array? [ first2 ] [ 0 ] if
+ [ 257 - length-table nth ] [ + ] bi*
+ ]
+ [
+ dup array? [ first2 ] [ 0 ] if
+ [ dist-table nth ] [ + ] bi*
+ ] bi*
+ 2array
+ ] when
+ ] map ;
+
+:: inflate-raw ( bitstream -- bytes )
+ 8 bitstream bs:align
+ 16 bitstream bs:read :> len
+ 16 bitstream bs:read :> nlen
+ len nlen + 16 >signed -1 assert= ! len + ~len = -1
+ bitstream byte-pos>>
+ bitstream byte-pos>> len +
+ bitstream bytes>> <slice>
+ len 8 * bitstream bs:seek ;
+
+: inflate-static ( bitstream -- bytes ) zlib-unimplemented ;
+
+:: inflate-loop ( bitstream -- bytes )
+ [ 1 bitstream bs:read 0 = ]
+ [
+ bitstream
+ 2 bitstream bs:read
+ {
+ { 0 [ inflate-raw ] }
+ { 1 [ inflate-static ] }
+ { 2 [ inflate-dynamic ] }
+ { 3 [ bad-zlib-data f ] }
+ }
+ case
+ ]
+ [ produce ] keep call suffix concat ;
+
+ ! [ produce ] keep dip swap suffix
+
+:: paeth ( a b c -- p )
+ a b + c - { a b c } [ [ - abs ] keep 2array ] with map
+ sort-keys first second ;
+
+:: png-unfilter-line ( prev curr filter -- curr' )
+ prev :> c
+ prev 3 tail-slice :> b
+ curr :> a
+ curr 3 tail-slice :> x
+ x length [0,b)
+ filter
+ {
+ { 0 [ drop ] }
+ { 1 [ [| n | n x nth n a nth + 256 wrap n x set-nth ] each ] }
+ { 2 [ [| n | n x nth n b nth + 256 wrap n x set-nth ] each ] }
+ { 3 [ [| n | n x nth n a nth n b nth + 2/ + 256 wrap n x set-nth ] each ] }
+ { 4 [ [| n | n x nth n a nth n b nth n c nth paeth + 256 wrap n x set-nth ] each ] }
+
+ } case
+ curr 3 tail ;
+
+PRIVATE>
+
+: reverse-png-filter' ( lines -- byte-array )
+ [ first ] [ 1 tail ] [ map ] bi-curry@ bi nip
+ concat [ 128 + ] B{ } map-as ;
+
+: reverse-png-filter ( lines -- byte-array )
+ dup first [ 0 ] replicate prefix
+ [ { 0 0 } prepend ] map
+ 2 clump [
+ first2 dup [ third ] [ 0 2 rot set-nth ] bi png-unfilter-line
+ ] map B{ } concat-as ;
+
+: zlib-inflate ( bytes -- bytes )
+ bs:<lsb0-bit-reader>
+ [ check-zlib-header ] [ inflate-loop ] bi
+ inflate-lz77 ;
-! Copyright (C) 2009 Marc Fauconneau.\r
-! See http://factorcode.org/license.txt for BSD license.\r
-USING: accessors arrays byte-arrays combinators\r
-constructors grouping compression.huffman images\r
-images.processing io io.binary io.encodings.binary io.files\r
-io.streams.byte-array kernel locals math math.bitwise\r
-math.constants math.functions math.matrices math.order\r
-math.ranges math.vectors memoize multiline namespaces\r
-sequences sequences.deep images.loader ;\r
-QUALIFIED-WITH: bitstreams bs\r
-IN: images.jpeg\r
-\r
-SINGLETON: jpeg-image\r
-{ "jpg" "jpeg" } [ jpeg-image register-image-class ] each\r
-\r
-TUPLE: loading-jpeg < image\r
- { headers }\r
- { bitstream }\r
- { color-info initial: { f f f f } }\r
- { quant-tables initial: { f f } }\r
- { huff-tables initial: { f f f f } }\r
- { components } ;\r
-\r
-<PRIVATE\r
-\r
-CONSTRUCTOR: loading-jpeg ( headers bitstream -- image ) ;\r
-\r
-SINGLETONS: SOF DHT DAC RST SOI EOI SOS DQT DNL DRI DHP EXP\r
-APP JPG COM TEM RES ;\r
-\r
-! ISO/IEC 10918-1 Table B.1\r
-:: >marker ( byte -- marker )\r
- byte\r
- {\r
- { [ dup HEX: CC = ] [ { DAC } ] }\r
- { [ dup HEX: C4 = ] [ { DHT } ] }\r
- { [ dup HEX: C9 = ] [ { JPG } ] }\r
- { [ dup -4 shift HEX: C = ] [ SOF byte 4 bits 2array ] }\r
-\r
- { [ dup HEX: D8 = ] [ { SOI } ] }\r
- { [ dup HEX: D9 = ] [ { EOI } ] }\r
- { [ dup HEX: DA = ] [ { SOS } ] }\r
- { [ dup HEX: DB = ] [ { DQT } ] }\r
- { [ dup HEX: DC = ] [ { DNL } ] }\r
- { [ dup HEX: DD = ] [ { DRI } ] }\r
- { [ dup HEX: DE = ] [ { DHP } ] }\r
- { [ dup HEX: DF = ] [ { EXP } ] }\r
- { [ dup -4 shift HEX: D = ] [ RST byte 4 bits 2array ] }\r
-\r
- { [ dup -4 shift HEX: E = ] [ APP byte 4 bits 2array ] }\r
- { [ dup HEX: FE = ] [ { COM } ] }\r
- { [ dup -4 shift HEX: F = ] [ JPG byte 4 bits 2array ] }\r
-\r
- { [ dup HEX: 01 = ] [ { TEM } ] }\r
- [ { RES } ]\r
- }\r
- cond nip ;\r
-\r
-TUPLE: jpeg-chunk length type data ;\r
-\r
-CONSTRUCTOR: jpeg-chunk ( type length data -- jpeg-chunk ) ;\r
-\r
-TUPLE: jpeg-color-info\r
- h v quant-table dc-huff-table ac-huff-table { diff initial: 0 } id ;\r
-\r
-CONSTRUCTOR: jpeg-color-info ( h v quant-table -- jpeg-color-info ) ;\r
-\r
-: jpeg> ( -- jpeg-image ) loading-jpeg get ;\r
-\r
-: apply-diff ( dc color -- dc' )\r
- [ diff>> + dup ] [ (>>diff) ] bi ;\r
-\r
-: fetch-tables ( component -- )\r
- [ [ jpeg> quant-tables>> nth ] change-quant-table drop ]\r
- [ [ jpeg> huff-tables>> nth ] change-dc-huff-table drop ]\r
- [ [ 2 + jpeg> huff-tables>> nth ] change-ac-huff-table drop ] tri ;\r
-\r
-: read4/4 ( -- a b ) read1 16 /mod ;\r
-\r
-\r
-! headers\r
-\r
-: decode-frame ( header -- )\r
- data>>\r
- binary\r
- [\r
- read1 8 assert=\r
- 2 read be>\r
- 2 read be>\r
- swap 2array jpeg> (>>dim)\r
- read1\r
- [\r
- read1 read4/4 read1 <jpeg-color-info>\r
- swap [ >>id ] keep jpeg> color-info>> set-nth\r
- ] times\r
- ] with-byte-reader ;\r
-\r
-: decode-quant-table ( chunk -- )\r
- dup data>>\r
- binary\r
- [\r
- length>>\r
- 2 - 65 /\r
- [\r
- read4/4 [ 0 assert= ] dip\r
- 64 read\r
- swap jpeg> quant-tables>> set-nth\r
- ] times\r
- ] with-byte-reader ;\r
-\r
-: decode-huff-table ( chunk -- )\r
- data>>\r
- binary\r
- [\r
- 1 ! %fixme: Should handle multiple tables at once\r
- [\r
- read4/4 swap 2 * +\r
- 16 read\r
- dup [ ] [ + ] map-reduce read\r
- binary [ [ read [ B{ } ] unless* ] { } map-as ] with-byte-reader\r
- swap jpeg> huff-tables>> set-nth\r
- ] times\r
- ] with-byte-reader ;\r
-\r
-: decode-scan ( chunk -- )\r
- data>>\r
- binary\r
- [\r
- read1 [0,b)\r
- [ drop\r
- read1 jpeg> color-info>> nth clone\r
- read1 16 /mod [ >>dc-huff-table ] [ >>ac-huff-table ] bi*\r
- ] map jpeg> (>>components)\r
- read1 0 assert=\r
- read1 63 assert=\r
- read1 16 /mod [ 0 assert= ] bi@\r
- ] with-byte-reader ;\r
-\r
-: singleton-first ( seq -- elt )\r
- [ length 1 assert= ] [ first ] bi ;\r
-\r
-: baseline-parse ( -- )\r
- jpeg> headers>>\r
- {\r
- [ [ type>> { SOF 0 } = ] filter singleton-first decode-frame ]\r
- [ [ type>> { DQT } = ] filter [ decode-quant-table ] each ]\r
- [ [ type>> { DHT } = ] filter [ decode-huff-table ] each ]\r
- [ [ type>> { SOS } = ] filter singleton-first decode-scan ]\r
- } cleave ;\r
-\r
-: parse-marker ( -- marker )\r
- read1 HEX: FF assert=\r
- read1 >marker ;\r
-\r
-: parse-headers ( -- chunks )\r
- [ parse-marker dup { SOS } = not ]\r
- [\r
- 2 read be>\r
- dup 2 - read <jpeg-chunk>\r
- ] [ produce ] keep dip swap suffix ;\r
-\r
-MEMO: zig-zag ( -- zz )\r
- {\r
- { 0 1 5 6 14 15 27 28 }\r
- { 2 4 7 13 16 26 29 42 }\r
- { 3 8 12 17 25 30 41 43 }\r
- { 9 11 18 24 31 40 44 53 }\r
- { 10 19 23 32 39 45 52 54 }\r
- { 20 22 33 38 46 51 55 60 }\r
- { 21 34 37 47 50 56 59 61 }\r
- { 35 36 48 49 57 58 62 63 }\r
- } flatten ;\r
-\r
-MEMO: yuv>bgr-matrix ( -- m )\r
- {\r
- { 1 2.03211 0 }\r
- { 1 -0.39465 -0.58060 }\r
- { 1 0 1.13983 }\r
- } ;\r
-\r
-: wave ( x u -- n ) swap 2 * 1 + * pi * 16 / cos ;\r
-\r
-:: dct-vect ( u v -- basis )\r
- { 8 8 } coord-matrix [ { u v } [ wave ] 2map product ] map^2\r
- 1 u v [ 0 = [ 2 sqrt / ] when ] bi@ 4 / m*n ;\r
-\r
-MEMO: dct-matrix ( -- m ) 64 [0,b) [ 8 /mod dct-vect flatten ] map ;\r
-\r
-: mb-dim ( component -- dim ) [ h>> ] [ v>> ] bi 2array ;\r
-\r
-: all-macroblocks ( quot: ( mb -- ) -- )\r
- [\r
- jpeg>\r
- [ dim>> 8 v/n ]\r
- [ color-info>> sift { 0 0 } [ mb-dim vmax ] reduce v/ ] bi\r
- [ ceiling ] map\r
- coord-matrix flip concat\r
- ]\r
- [ each ] bi* ; inline\r
-\r
-: reverse-zigzag ( b -- b' ) zig-zag swap [ nth ] curry map ;\r
-\r
-: idct-factor ( b -- b' ) dct-matrix v.m ;\r
-\r
-USE: math.blas.vectors\r
-USE: math.blas.matrices\r
-\r
-MEMO: dct-matrix-blas ( -- m ) dct-matrix >float-blas-matrix ;\r
-: V.M ( x A -- x.A ) Mtranspose swap M.V ;\r
-: idct-blas ( b -- b' ) >float-blas-vector dct-matrix-blas V.M ;\r
-\r
-: idct ( b -- b' ) idct-blas ;\r
-\r
-:: draw-block ( block x,y color jpeg-image -- )\r
- block dup length>> sqrt >fixnum group flip\r
- dup matrix-dim coord-matrix flip\r
- [\r
- [ first2 spin nth nth ]\r
- [ x,y v+ color id>> 1- jpeg-image draw-color ] bi\r
- ] with each^2 ;\r
-\r
-: sign-extend ( bits v -- v' )\r
- swap [ ] [ 1- 2^ < ] 2bi\r
- [ -1 swap shift 1+ + ] [ drop ] if ;\r
-\r
-: read1-jpeg-dc ( decoder -- dc )\r
- [ read1-huff dup ] [ bs>> bs:read ] bi sign-extend ;\r
-\r
-: read1-jpeg-ac ( decoder -- run/ac )\r
- [ read1-huff 16 /mod dup ] [ bs>> bs:read ] bi sign-extend 2array ;\r
-\r
-:: decode-block ( pos color -- )\r
- color dc-huff-table>> read1-jpeg-dc color apply-diff\r
- 64 0 <array> :> coefs\r
- 0 coefs set-nth\r
- 0 :> k!\r
- [\r
- color ac-huff-table>> read1-jpeg-ac\r
- [ first 1+ k + k! ] [ second k coefs set-nth ] [ ] tri\r
- { 0 0 } = not\r
- k 63 < and\r
- ] loop\r
- coefs color quant-table>> v*\r
- reverse-zigzag idct\r
- ! %fixme: color hack\r
- ! this eat 50% cpu time\r
- color h>> 2 =\r
- [ 8 group 2 matrix-zoom concat ] unless\r
- pos { 8 8 } v* color jpeg> draw-block ;\r
-\r
-: decode-macroblock ( mb -- )\r
- jpeg> components>>\r
- [\r
- [ mb-dim coord-matrix flip concat [ [ { 2 2 } v* ] [ v+ ] bi* ] with map ]\r
- [ [ decode-block ] curry each ] bi\r
- ] with each ;\r
-\r
-: cleanup-bitstream ( bytes -- bytes' )\r
- binary [\r
- [\r
- { HEX: FF } read-until\r
- read1 tuck HEX: 00 = and\r
- ]\r
- [ drop ] produce\r
- swap >marker { EOI } assert=\r
- swap suffix\r
- { HEX: FF } join\r
- ] with-byte-reader ;\r
-\r
-: setup-bitmap ( image -- )\r
- dup dim>> 16 v/n [ ceiling ] map 16 v*n >>dim\r
- BGR >>component-order\r
- f >>upside-down?\r
- dup dim>> first2 * 3 * 0 <array> >>bitmap\r
- drop ;\r
-\r
-: baseline-decompress ( -- )\r
- jpeg> bitstream>> cleanup-bitstream { 255 255 255 255 } append\r
- >byte-array bs:<msb0-bit-reader> jpeg> (>>bitstream)\r
- jpeg> [ bitstream>> ] [ [ [ <huffman-decoder> ] with map ] change-huff-tables drop ] bi\r
- jpeg> components>> [ fetch-tables ] each\r
- jpeg> setup-bitmap\r
- [ decode-macroblock ] all-macroblocks ;\r
-\r
-! this eats ~25% cpu time\r
-: color-transform ( yuv -- rgb )\r
- { 128 0 0 } v+ yuv>bgr-matrix swap m.v\r
- [ 0 max 255 min >fixnum ] map ;\r
-\r
-PRIVATE>\r
-\r
-: load-jpeg ( path -- image )\r
- binary [\r
- parse-marker { SOI } assert=\r
- parse-headers\r
- contents <loading-jpeg>\r
- ] with-file-reader\r
- dup loading-jpeg [\r
- baseline-parse\r
- baseline-decompress\r
- jpeg> bitmap>> 3 <groups> [ color-transform ] change-each\r
- jpeg> [ >byte-array ] change-bitmap drop\r
- ] with-variable ;\r
-\r
-M: jpeg-image load-image* ( path jpeg-image -- bitmap )\r
- drop load-jpeg ;\r
+! Copyright (C) 2009 Marc Fauconneau.
+! See http://factorcode.org/license.txt for BSD license.
+USING: accessors arrays byte-arrays combinators
+constructors grouping compression.huffman images
+images.processing io io.binary io.encodings.binary io.files
+io.streams.byte-array kernel locals math math.bitwise
+math.constants math.functions math.matrices math.order
+math.ranges math.vectors memoize multiline namespaces
+sequences sequences.deep ;
+IN: images.jpeg
+
+QUALIFIED-WITH: bitstreams bs
+
+TUPLE: jpeg-image < image
+ { headers }
+ { bitstream }
+ { color-info initial: { f f f f } }
+ { quant-tables initial: { f f } }
+ { huff-tables initial: { f f f f } }
+ { components } ;
+
+<PRIVATE
+
+CONSTRUCTOR: jpeg-image ( headers bitstream -- image ) ;
+
+SINGLETONS: SOF DHT DAC RST SOI EOI SOS DQT DNL DRI DHP EXP
+APP JPG COM TEM RES ;
+
+! ISO/IEC 10918-1 Table B.1
+:: >marker ( byte -- marker )
+ byte
+ {
+ { [ dup HEX: CC = ] [ { DAC } ] }
+ { [ dup HEX: C4 = ] [ { DHT } ] }
+ { [ dup HEX: C9 = ] [ { JPG } ] }
+ { [ dup -4 shift HEX: C = ] [ SOF byte 4 bits 2array ] }
+
+ { [ dup HEX: D8 = ] [ { SOI } ] }
+ { [ dup HEX: D9 = ] [ { EOI } ] }
+ { [ dup HEX: DA = ] [ { SOS } ] }
+ { [ dup HEX: DB = ] [ { DQT } ] }
+ { [ dup HEX: DC = ] [ { DNL } ] }
+ { [ dup HEX: DD = ] [ { DRI } ] }
+ { [ dup HEX: DE = ] [ { DHP } ] }
+ { [ dup HEX: DF = ] [ { EXP } ] }
+ { [ dup -4 shift HEX: D = ] [ RST byte 4 bits 2array ] }
+
+ { [ dup -4 shift HEX: E = ] [ APP byte 4 bits 2array ] }
+ { [ dup HEX: FE = ] [ { COM } ] }
+ { [ dup -4 shift HEX: F = ] [ JPG byte 4 bits 2array ] }
+
+ { [ dup HEX: 01 = ] [ { TEM } ] }
+ [ { RES } ]
+ }
+ cond nip ;
+
+TUPLE: jpeg-chunk length type data ;
+
+CONSTRUCTOR: jpeg-chunk ( type length data -- jpeg-chunk ) ;
+
+TUPLE: jpeg-color-info
+ h v quant-table dc-huff-table ac-huff-table { diff initial: 0 } id ;
+
+CONSTRUCTOR: jpeg-color-info ( h v quant-table -- jpeg-color-info ) ;
+
+: jpeg> ( -- jpeg-image ) jpeg-image get ;
+
+: apply-diff ( dc color -- dc' )
+ [ diff>> + dup ] [ (>>diff) ] bi ;
+
+: fetch-tables ( component -- )
+ [ [ jpeg> quant-tables>> nth ] change-quant-table drop ]
+ [ [ jpeg> huff-tables>> nth ] change-dc-huff-table drop ]
+ [ [ 2 + jpeg> huff-tables>> nth ] change-ac-huff-table drop ] tri ;
+
+: read4/4 ( -- a b ) read1 16 /mod ;
+
+! headers
+
+: decode-frame ( header -- )
+ data>>
+ binary
+ [
+ read1 8 assert=
+ 2 read be>
+ 2 read be>
+ swap 2array jpeg> (>>dim)
+ read1
+ [
+ read1 read4/4 read1 <jpeg-color-info>
+ swap [ >>id ] keep jpeg> color-info>> set-nth
+ ] times
+ ] with-byte-reader ;
+
+: decode-quant-table ( chunk -- )
+ dup data>>
+ binary
+ [
+ length>>
+ 2 - 65 /
+ [
+ read4/4 [ 0 assert= ] dip
+ 64 read
+ swap jpeg> quant-tables>> set-nth
+ ] times
+ ] with-byte-reader ;
+
+: decode-huff-table ( chunk -- )
+ data>>
+ binary
+ [
+ 1 ! %fixme: Should handle multiple tables at once
+ [
+ read4/4 swap 2 * +
+ 16 read
+ dup [ ] [ + ] map-reduce read
+ binary [ [ read [ B{ } ] unless* ] { } map-as ] with-byte-reader
+ swap jpeg> huff-tables>> set-nth
+ ] times
+ ] with-byte-reader ;
+
+: decode-scan ( chunk -- )
+ data>>
+ binary
+ [
+ read1 [0,b)
+ [ drop
+ read1 jpeg> color-info>> nth clone
+ read1 16 /mod [ >>dc-huff-table ] [ >>ac-huff-table ] bi*
+ ] map jpeg> (>>components)
+ read1 0 assert=
+ read1 63 assert=
+ read1 16 /mod [ 0 assert= ] bi@
+ ] with-byte-reader ;
+
+: singleton-first ( seq -- elt )
+ [ length 1 assert= ] [ first ] bi ;
+
+: baseline-parse ( -- )
+ jpeg> headers>>
+ {
+ [ [ type>> { SOF 0 } = ] filter singleton-first decode-frame ]
+ [ [ type>> { DQT } = ] filter [ decode-quant-table ] each ]
+ [ [ type>> { DHT } = ] filter [ decode-huff-table ] each ]
+ [ [ type>> { SOS } = ] filter singleton-first decode-scan ]
+ } cleave ;
+
+: parse-marker ( -- marker )
+ read1 HEX: FF assert=
+ read1 >marker ;
+
+: parse-headers ( -- chunks )
+ [ parse-marker dup { SOS } = not ]
+ [
+ 2 read be>
+ dup 2 - read <jpeg-chunk>
+ ] [ produce ] keep dip swap suffix ;
+
+MEMO: zig-zag ( -- zz )
+ {
+ { 0 1 5 6 14 15 27 28 }
+ { 2 4 7 13 16 26 29 42 }
+ { 3 8 12 17 25 30 41 43 }
+ { 9 11 18 24 31 40 44 53 }
+ { 10 19 23 32 39 45 52 54 }
+ { 20 22 33 38 46 51 55 60 }
+ { 21 34 37 47 50 56 59 61 }
+ { 35 36 48 49 57 58 62 63 }
+ } flatten ;
+
+MEMO: yuv>bgr-matrix ( -- m )
+ {
+ { 1 2.03211 0 }
+ { 1 -0.39465 -0.58060 }
+ { 1 0 1.13983 }
+ } ;
+
+: wave ( x u -- n ) swap 2 * 1 + * pi * 16 / cos ;
+
+:: dct-vect ( u v -- basis )
+ { 8 8 } coord-matrix [ { u v } [ wave ] 2map product ] map^2
+ 1 u v [ 0 = [ 2 sqrt / ] when ] bi@ 4 / m*n ;
+
+MEMO: dct-matrix ( -- m ) 64 [0,b) [ 8 /mod dct-vect flatten ] map ;
+
+: mb-dim ( component -- dim ) [ h>> ] [ v>> ] bi 2array ;
+
+! : blocks ( component -- seq )
+! mb-dim ! coord-matrix flip concat [ [ { 2 2 } v* ] [ v+ ] bi* ] with map ;
+
+: all-macroblocks ( quot: ( mb -- ) -- )
+ [
+ jpeg>
+ [ dim>> 8 v/n ]
+ [ color-info>> sift { 0 0 } [ mb-dim vmax ] reduce v/ ] bi
+ [ ceiling ] map
+ coord-matrix flip concat
+ ]
+ [ each ] bi* ; inline
+
+: reverse-zigzag ( b -- b' ) zig-zag swap [ nth ] curry map ;
+
+: idct-factor ( b -- b' ) dct-matrix v.m ;
+
+USE: math.blas.vectors
+USE: math.blas.matrices
+
+MEMO: dct-matrix-blas ( -- m ) dct-matrix >float-blas-matrix ;
+: V.M ( x A -- x.A ) Mtranspose swap M.V ;
+: idct-blas ( b -- b' ) >float-blas-vector dct-matrix-blas V.M ;
+
+: idct ( b -- b' ) idct-blas ;
+
+:: draw-block ( block x,y color-id jpeg-image -- )
+ block dup length>> sqrt >fixnum group flip
+ dup matrix-dim coord-matrix flip
+ [
+ [ first2 spin nth nth ]
+ [ x,y v+ color-id jpeg-image draw-color ] bi
+ ] with each^2 ;
+
+: sign-extend ( bits v -- v' )
+ swap [ ] [ 1- 2^ < ] 2bi
+ [ -1 swap shift 1+ + ] [ drop ] if ;
+
+: read1-jpeg-dc ( decoder -- dc )
+ [ read1-huff dup ] [ bs>> bs:read ] bi sign-extend ;
+
+: read1-jpeg-ac ( decoder -- run/ac )
+ [ read1-huff 16 /mod dup ] [ bs>> bs:read ] bi sign-extend 2array ;
+
+:: decode-block ( color -- pixels )
+ color dc-huff-table>> read1-jpeg-dc color apply-diff
+ 64 0 <array> :> coefs
+ 0 coefs set-nth
+ 0 :> k!
+ [
+ color ac-huff-table>> read1-jpeg-ac
+ [ first 1+ k + k! ] [ second k coefs set-nth ] [ ] tri
+ { 0 0 } = not
+ k 63 < and
+ ] loop
+ coefs color quant-table>> v*
+ reverse-zigzag idct ;
+
+:: draw-macroblock-yuv420 ( mb blocks -- )
+ mb { 16 16 } v* :> pos
+ 0 blocks nth pos { 0 0 } v+ 0 jpeg> draw-block
+ 1 blocks nth pos { 8 0 } v+ 0 jpeg> draw-block
+ 2 blocks nth pos { 0 8 } v+ 0 jpeg> draw-block
+ 3 blocks nth pos { 8 8 } v+ 0 jpeg> draw-block
+ 4 blocks nth 8 group 2 matrix-zoom concat pos 1 jpeg> draw-block
+ 5 blocks nth 8 group 2 matrix-zoom concat pos 2 jpeg> draw-block ;
+
+:: draw-macroblock-yuv444 ( mb blocks -- )
+ mb { 8 8 } v* :> pos
+ 3 iota [ [ blocks nth pos ] [ jpeg> draw-block ] bi ] each ;
+
+:: draw-macroblock-y ( mb blocks -- )
+ mb { 8 8 } v* :> pos
+ 0 blocks nth pos 0 jpeg> draw-block
+ 64 0 <array> pos 1 jpeg> draw-block
+ 64 0 <array> pos 2 jpeg> draw-block ;
+
+ ! %fixme: color hack
+ ! color h>> 2 =
+ ! [ 8 group 2 matrix-zoom concat ] unless
+ ! pos { 8 8 } v* color jpeg> draw-block ;
+
+: decode-macroblock ( -- blocks )
+ jpeg> components>>
+ [
+ [ mb-dim first2 * iota ]
+ [ [ decode-block ] curry replicate ] bi
+ ] map concat ;
+
+: cleanup-bitstream ( bytes -- bytes' )
+ binary [
+ [
+ { HEX: FF } read-until
+ read1 tuck HEX: 00 = and
+ ]
+ [ drop ] produce
+ swap >marker { EOI } assert=
+ swap suffix
+ { HEX: FF } join
+ ] with-byte-reader ;
+
+: setup-bitmap ( image -- )
+ dup dim>> 16 v/n [ ceiling ] map 16 v*n >>dim
+ BGR >>component-order
+ f >>upside-down?
+ dup dim>> first2 * 3 * 0 <array> >>bitmap
+ drop ;
+
+ERROR: unsupported-colorspace ;
+SINGLETONS: YUV420 YUV444 Y MAGIC! ;
+
+:: detect-colorspace ( jpeg-image -- csp )
+ jpeg-image color-info>> sift :> colors
+ MAGIC!
+ colors length 1 = [ drop Y ] when
+ colors length 3 =
+ [
+ colors [ mb-dim { 1 1 } = ] all?
+ [ drop YUV444 ] when
+
+ colors unclip
+ [ [ mb-dim { 1 1 } = ] all? ]
+ [ mb-dim { 2 2 } = ] bi* and
+ [ drop YUV420 ] when
+ ] when ;
+
+! this eats ~50% cpu time
+: draw-macroblocks ( mbs -- )
+ jpeg> detect-colorspace
+ {
+ { YUV420 [ [ first2 draw-macroblock-yuv420 ] each ] }
+ { YUV444 [ [ first2 draw-macroblock-yuv444 ] each ] }
+ { Y [ [ first2 draw-macroblock-y ] each ] }
+ [ unsupported-colorspace ]
+ } case ;
+
+! this eats ~25% cpu time
+: color-transform ( yuv -- rgb )
+ { 128 0 0 } v+ yuv>bgr-matrix swap m.v
+ [ 0 max 255 min >fixnum ] map ;
+
+: baseline-decompress ( -- )
+ jpeg> bitstream>> cleanup-bitstream { 255 255 255 255 } append
+ >byte-array bs:<msb0-bit-reader> jpeg> (>>bitstream)
+ jpeg>
+ [ bitstream>> ]
+ [ [ [ <huffman-decoder> ] with map ] change-huff-tables drop ] bi
+ jpeg> components>> [ fetch-tables ] each
+ [ decode-macroblock 2array ] accumulator
+ [ all-macroblocks ] dip
+ jpeg> setup-bitmap draw-macroblocks
+ jpeg> bitmap>> 3 <groups> [ color-transform ] change-each
+ jpeg> [ >byte-array ] change-bitmap drop ;
+
+ERROR: not-a-jpeg-image ;
+
+PRIVATE>
+
+: load-jpeg ( path -- image )
+ binary [
+ parse-marker { SOI } = [ not-a-jpeg-image ] unless
+ parse-headers
+ contents <jpeg-image>
+ ] with-file-reader
+ dup jpeg-image [
+ baseline-parse
+ baseline-decompress
+ ] with-variable ;
+
+M: jpeg-image load-image* ( path jpeg-image -- bitmap )
+ drop load-jpeg ;
+
projects / factor.git / commitdiff