basis/images/jpeg/jpeg.factor

   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
   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 ;
  10 IN: images.jpeg
  11
  12 QUALIFIED-WITH: bitstreams bs
  13
  14 TUPLE: jpeg-image < image
  15     { headers }
  16     { bitstream }
  17     { color-info initial: { f f f f } }
  18     { quant-tables initial: { f f } }
  19     { huff-tables initial: { f f f f } }
  20     { components } ;
  21
  22 <PRIVATE
  23
  24 : <jpeg-image> ( headers bitstream -- image )
  25     jpeg-image new swap >>bitstream swap >>headers ;
  26
  27 SINGLETONS: SOF DHT DAC RST SOI EOI SOS DQT DNL DRI DHP EXP
  28 APP JPG COM TEM RES ;
  29
  30 ! ISO/IEC 10918-1 Table B.1
  31 :: >marker ( byte -- marker )
  32     byte
  33     {
  34       { [ dup HEX: CC = ] [ { DAC } ] }
  35       { [ dup HEX: C4 = ] [ { DHT } ] }
  36       { [ dup HEX: C9 = ] [ { JPG } ] }
  37       { [ dup -4 shift HEX: C = ] [ SOF byte 4 bits 2array ] }
  38
  39       { [ dup HEX: D8 = ] [ { SOI } ] }
  40       { [ dup HEX: D9 = ] [ { EOI } ] }
  41       { [ dup HEX: DA = ] [ { SOS } ] }
  42       { [ dup HEX: DB = ] [ { DQT } ] }
  43       { [ dup HEX: DC = ] [ { DNL } ] }
  44       { [ dup HEX: DD = ] [ { DRI } ] }
  45       { [ dup HEX: DE = ] [ { DHP } ] }
  46       { [ dup HEX: DF = ] [ { EXP } ] }
  47       { [ dup -4 shift HEX: D = ] [ RST byte 4 bits 2array ] }
  48
  49       { [ dup -4 shift HEX: E = ] [ APP byte 4 bits 2array ] }
  50       { [ dup HEX: FE = ] [ { COM } ] }
  51       { [ dup -4 shift HEX: F = ] [ JPG byte 4 bits 2array ] }
  52
  53       { [ dup HEX: 01 = ] [ { TEM } ] }
  54       [ { RES } ]
  55     }
  56     cond nip ;
  57
  58 TUPLE: jpeg-chunk length type data ;
  59
  60 : <jpeg-chunk> ( type length data -- jpeg-chunk )
  61     jpeg-chunk new
  62         swap >>data
  63         swap >>length
  64         swap >>type ;
  65
  66 TUPLE: jpeg-color-info
  67     h v quant-table dc-huff-table ac-huff-table { diff initial: 0 } id ;
  68
  69 : <jpeg-color-info> ( h v quant-table -- jpeg-color-info )
  70     jpeg-color-info new
  71         swap >>quant-table
  72         swap >>v
  73         swap >>h ;
  74
  75 : jpeg> ( -- jpeg-image ) jpeg-image get ;
  76
  77 : apply-diff ( dc color -- dc' )
  78     [ diff>> + dup ] [ (>>diff) ] bi ;
  79
  80 : fetch-tables ( component -- )
  81     [ [ jpeg> quant-tables>> nth ] change-quant-table drop ]
  82     [ [ jpeg> huff-tables>> nth ] change-dc-huff-table drop ]
  83     [ [ 2 + jpeg> huff-tables>> nth ] change-ac-huff-table drop ] tri ;
  84
  85 : read4/4 ( -- a b ) read1 16 /mod ;
  86
  87 ! headers
  88
  89 : decode-frame ( header -- )
  90     data>>
  91     binary
  92     [
  93         read1 8 assert=
  94         2 read be>
  95         2 read be>
  96         swap 2array jpeg> (>>dim)
  97         read1
  98         [
  99             read1 read4/4 read1 <jpeg-color-info>
 100             swap [ >>id ] keep jpeg> color-info>> set-nth
 101         ] times
 102     ] with-byte-reader ;
 103
 104 : decode-quant-table ( chunk -- )
 105     dup data>>
 106     binary
 107     [
 108         length>>
 109         2 - 65 /
 110         [
 111             read4/4 [ 0 assert= ] dip
 112             64 read
 113             swap jpeg> quant-tables>> set-nth
 114         ] times
 115     ] with-byte-reader ;
 116
 117 : decode-huff-table ( chunk -- )
 118     data>>
 119     binary
 120     [
 121         1 ! %fixme: Should handle multiple tables at once
 122         [
 123             read4/4 swap 2 * +
 124             16 read
 125             dup [ ] [ + ] map-reduce read
 126             binary [ [ read [ B{ } ] unless* ] { } map-as ] with-byte-reader
 127             swap jpeg> huff-tables>> set-nth
 128         ] times
 129     ] with-byte-reader ;
 130
 131 : decode-scan ( chunk -- )
 132     data>>
 133     binary
 134     [
 135         read1 [0,b)
 136         [   drop
 137             read1 jpeg> color-info>> nth clone
 138             read1 16 /mod [ >>dc-huff-table ] [ >>ac-huff-table ] bi*
 139         ] map jpeg> (>>components)
 140         read1 0 assert=
 141         read1 63 assert=
 142         read1 16 /mod [ 0 assert= ] bi@
 143     ] with-byte-reader ;
 144
 145 : singleton-first ( seq -- elt )
 146     [ length 1 assert= ] [ first ] bi ;
 147
 148 : baseline-parse ( -- )
 149     jpeg> headers>>
 150     {
 151         [ [ type>> { SOF 0 } = ] filter singleton-first decode-frame ]
 152         [ [ type>> { DQT } = ] filter [ decode-quant-table ] each ]
 153         [ [ type>> { DHT } = ] filter [ decode-huff-table ] each ]
 154         [ [ type>> { SOS } = ] filter singleton-first decode-scan ]
 155     } cleave ;
 156
 157 : parse-marker ( -- marker )
 158     read1 HEX: FF assert=
 159     read1 >marker ;
 160
 161 : parse-headers ( -- chunks )
 162     [ parse-marker dup { SOS } = not ]
 163     [
 164         2 read be>
 165         dup 2 - read <jpeg-chunk>
 166     ] [ produce ] keep dip swap suffix ;
 167
 168 MEMO: zig-zag ( -- zz )
 169     {
 170         {  0  1  5  6 14 15 27 28 }
 171         {  2  4  7 13 16 26 29 42 }
 172         {  3  8 12 17 25 30 41 43 }
 173         {  9 11 18 24 31 40 44 53 }
 174         { 10 19 23 32 39 45 52 54 }
 175         { 20 22 33 38 46 51 55 60 }
 176         { 21 34 37 47 50 56 59 61 }
 177         { 35 36 48 49 57 58 62 63 }
 178     } flatten ;
 179
 180 MEMO: yuv>bgr-matrix ( -- m )
 181     {
 182         { 1  2.03211  0       }
 183         { 1 -0.39465 -0.58060 }
 184         { 1  0        1.13983 }
 185     } ;
 186
 187 : wave ( x u -- n ) swap 2 * 1 + * pi * 16 / cos ;
 188
 189 :: dct-vect ( u v -- basis )
 190     { 8 8 } coord-matrix [ { u v } [ wave ] 2map product ] map^2
 191     1 u v [ 0 = [ 2 sqrt / ] when ] bi@ 4 / m*n ;
 192
 193 MEMO: dct-matrix ( -- m ) 64 [0,b) [ 8 /mod dct-vect flatten ] map ;
 194
 195 : mb-dim ( component -- dim )  [ h>> ] [ v>> ] bi 2array ;
 196
 197 ! : blocks ( component -- seq )
 198 !    mb-dim ! coord-matrix flip concat [ [ { 2 2 } v* ] [ v+ ] bi* ] with map ;
 199
 200 : all-macroblocks ( quot: ( mb -- ) -- )
 201     [
 202         jpeg>
 203         [ dim>> 8 v/n ]
 204         [ color-info>> sift { 0 0 } [ mb-dim vmax ] reduce v/ ] bi
 205         [ ceiling ] map
 206         coord-matrix flip concat
 207     ]
 208     [ each ] bi* ; inline
 209
 210 : reverse-zigzag ( b -- b' ) zig-zag swap [ nth ] curry map ;
 211
 212 : idct-factor ( b -- b' ) dct-matrix v.m ;
 213
 214 USE: math.blas.vectors
 215 USE: math.blas.matrices
 216
 217 MEMO: dct-matrix-blas ( -- m ) dct-matrix >float-blas-matrix ;
 218 : V.M ( x A -- x.A ) Mtranspose swap M.V ;
 219 : idct-blas ( b -- b' ) >float-blas-vector dct-matrix-blas V.M ;
 220
 221 : idct ( b -- b' ) idct-blas ;
 222
 223 :: draw-block ( block x,y color-id jpeg-image -- )
 224     block dup length>> sqrt >fixnum group flip
 225     dup matrix-dim coord-matrix flip
 226     [
 227         [ first2 spin nth nth ]
 228         [ x,y v+ color-id jpeg-image draw-color ] bi
 229     ] with each^2 ;
 230
 231 : sign-extend ( bits v -- v' )
 232     swap [ ] [ 1 - 2^ < ] 2bi
 233     [ -1 swap shift 1 + + ] [ drop ] if ;
 234
 235 : read1-jpeg-dc ( decoder -- dc )
 236     [ read1-huff dup ] [ bs>> bs:read ] bi sign-extend ;
 237
 238 : read1-jpeg-ac ( decoder -- run/ac )
 239     [ read1-huff 16 /mod dup ] [ bs>> bs:read ] bi sign-extend 2array ;
 240
 241 :: decode-block ( color -- pixels )
 242     color dc-huff-table>> read1-jpeg-dc color apply-diff
 243     64 0 <array> :> coefs
 244     0 coefs set-nth
 245     0 :> k!
 246     [
 247         color ac-huff-table>> read1-jpeg-ac
 248         [ first 1 + k + k! ] [ second k coefs set-nth ] [ ] tri
 249         { 0 0 } = not
 250         k 63 < and
 251     ] loop
 252     coefs color quant-table>> v*
 253     reverse-zigzag idct ;
 254
 255 :: draw-macroblock-yuv420 ( mb blocks -- )
 256     mb { 16 16 } v* :> pos
 257     0 blocks nth pos { 0 0 } v+ 0 jpeg> draw-block
 258     1 blocks nth pos { 8 0 } v+ 0 jpeg> draw-block
 259     2 blocks nth pos { 0 8 } v+ 0 jpeg> draw-block
 260     3 blocks nth pos { 8 8 } v+ 0 jpeg> draw-block
 261     4 blocks nth 8 group 2 matrix-zoom concat pos 1 jpeg> draw-block
 262     5 blocks nth 8 group 2 matrix-zoom concat pos 2 jpeg> draw-block ;
 263
 264 :: draw-macroblock-yuv444 ( mb blocks -- )
 265     mb { 8 8 } v* :> pos
 266     3 iota [ [ blocks nth pos ] [ jpeg> draw-block ] bi ] each ;
 267
 268 :: draw-macroblock-y ( mb blocks -- )
 269     mb { 8 8 } v* :> pos
 270     0 blocks nth pos 0 jpeg> draw-block
 271     64 0 <array> pos 1 jpeg> draw-block
 272     64 0 <array> pos 2 jpeg> draw-block ;
 273
 274     ! %fixme: color hack
 275  !   color h>> 2 =
 276  !   [ 8 group 2 matrix-zoom concat ] unless
 277  !   pos { 8 8 } v* color jpeg> draw-block ;
 278
 279 : decode-macroblock ( -- blocks )
 280     jpeg> components>>
 281     [
 282         [ mb-dim first2 * iota ]
 283         [ [ decode-block ] curry replicate ] bi
 284     ] map concat ;
 285
 286 : cleanup-bitstream ( bytes -- bytes' )
 287     binary [
 288         [
 289             { HEX: FF } read-until
 290             read1 tuck HEX: 00 = and
 291         ]
 292         [ drop ] produce
 293         swap >marker {  EOI } assert=
 294         swap suffix
 295         { HEX: FF } join
 296     ] with-byte-reader ;
 297
 298 : setup-bitmap ( image -- )
 299     dup dim>> 16 v/n [ ceiling ] map 16 v*n >>dim
 300     BGR >>component-order
 301     ubyte-components >>component-type
 302     f >>upside-down?
 303     dup dim>> first2 * 3 * 0 <array> >>bitmap
 304     drop ;
 305
 306 ERROR: unsupported-colorspace ;
 307 SINGLETONS: YUV420 YUV444 Y MAGIC! ;
 308
 309 :: detect-colorspace ( jpeg-image -- csp )
 310     jpeg-image color-info>> sift :> colors
 311     MAGIC!
 312     colors length 1 = [ drop Y ] when
 313     colors length 3 =
 314     [
 315         colors [ mb-dim { 1 1 } = ] all?
 316         [ drop YUV444 ] when
 317
 318         colors unclip
 319         [ [ mb-dim { 1 1 } = ] all? ]
 320         [ mb-dim { 2 2 } =  ] bi* and
 321         [ drop YUV420 ] when
 322     ] when ;
 323
 324 ! this eats ~50% cpu time
 325 : draw-macroblocks ( mbs -- )
 326     jpeg> detect-colorspace
 327     {
 328         { YUV420 [ [ first2 draw-macroblock-yuv420 ] each ] }
 329         { YUV444 [ [ first2 draw-macroblock-yuv444 ] each ] }
 330         { Y      [ [ first2 draw-macroblock-y ] each ] }
 331         [ unsupported-colorspace ]
 332     } case ;
 333
 334 ! this eats ~25% cpu time
 335 : color-transform ( yuv -- rgb )
 336     { 128 0 0 } v+ yuv>bgr-matrix swap m.v
 337     [ 0 max 255 min >fixnum ] map ;
 338
 339 : baseline-decompress ( -- )
 340     jpeg> bitstream>> cleanup-bitstream { 255 255 255 255 } append
 341     >byte-array bs:<msb0-bit-reader> jpeg> (>>bitstream)
 342     jpeg>
 343     [ bitstream>> ]
 344     [ [ [ <huffman-decoder> ] with map ] change-huff-tables drop ] bi
 345     jpeg> components>> [ fetch-tables ] each
 346     [ decode-macroblock 2array ] accumulator
 347     [ all-macroblocks ] dip
 348     jpeg> setup-bitmap draw-macroblocks
 349     jpeg> bitmap>> 3 <groups> [ color-transform ] change-each
 350     jpeg> [ >byte-array ] change-bitmap drop ;
 351
 352 ERROR: not-a-jpeg-image ;
 353
 354 PRIVATE>
 355
 356 : load-jpeg ( path -- image )
 357     binary [
 358         parse-marker { SOI } = [ not-a-jpeg-image ] unless
 359         parse-headers
 360         contents <jpeg-image>
 361     ] with-file-reader
 362     dup jpeg-image [
 363         baseline-parse
 364         baseline-decompress
 365     ] with-variable ;
 366
 367 M: jpeg-image load-image* ( path jpeg-image -- bitmap )
 368     drop load-jpeg ;
 369