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