octree-quantize.ss added.
svn: r10992
This commit is contained in:
Danny Yoo
parent
6b42b0dfe6
commit
6b2e5f4014
393
collects/file/octree-quantize.ss
Normal file
393
collects/file/octree-quantize.ss
Normal file
|
|
@ -0,0 +1,393 @@
|
|||
#lang scheme/base
|
||||
(require scheme/class
|
||||
scheme/contract
|
||||
scheme/gui/base)
|
||||
|
||||
|
||||
;; quantize: bytes (integer-in 1 255) -> (values bytes gif-colormap color)
|
||||
;; Adaptively quantizes 24-bit image data to fit within, at most, 256 colors
|
||||
;; (including the transparent color.) The signature is meant to match the
|
||||
;; function provided by file/gif.
|
||||
(define (quantize argb #:palette (palette 255))
|
||||
(let* ([an-octree (make-octree-from-argb argb palette)]
|
||||
[len (bytes-length argb)]
|
||||
[dest (make-bytes (/ len 4))])
|
||||
(for ([i (in-range 0 len 4)])
|
||||
(let ([a (bytes-ref argb i)]
|
||||
[r (bytes-ref argb (+ i 1))]
|
||||
[g (bytes-ref argb (+ i 2))]
|
||||
[b (bytes-ref argb (+ i 3))]
|
||||
[n (arithmetic-shift i -2)])
|
||||
(cond
|
||||
[(alpha-opaque? a)
|
||||
(let ([palette-index (octree-lookup-index an-octree r g b)])
|
||||
(bytes-set! dest n palette-index))]
|
||||
[else
|
||||
(bytes-set! dest n TRANSPARENT-INDEX)])))
|
||||
(values dest (vector->list (octree-palette an-octree)) TRANSPARENT-INDEX)))
|
||||
|
||||
|
||||
|
||||
|
||||
;; make-octree-from-argb: bytes positive-number -> octree
|
||||
;; Constructs an octree ready to quantize the colors from an-argb.
|
||||
(define (make-octree-from-argb an-argb n)
|
||||
(unless (> n 0)
|
||||
(raise (make-exn:fail:octree "n must be positive" (current-continuation-marks))))
|
||||
(let ([an-octree (new-octree)]
|
||||
[len (bytes-length an-argb)])
|
||||
(let make-octree-loop ([i 0])
|
||||
(when (< i len)
|
||||
(let ([a (bytes-ref an-argb i)]
|
||||
[r (bytes-ref an-argb (+ i 1))]
|
||||
[g (bytes-ref an-argb (+ i 2))]
|
||||
[b (bytes-ref an-argb (+ i 3))])
|
||||
(when (alpha-opaque? a)
|
||||
(octree-insert-color! an-octree r g b)
|
||||
(let reduction-loop ()
|
||||
(when (> (octree-leaf-count an-octree) n)
|
||||
(octree-reduce! an-octree)
|
||||
(reduction-loop)))))
|
||||
(make-octree-loop (+ i 4))))
|
||||
(octree-finalize! an-octree)
|
||||
an-octree))
|
||||
|
||||
|
||||
;; alpha-opaque? byte -> boolean
|
||||
;; Returns true if the alpha value is considered opaque.
|
||||
(define (alpha-opaque? a)
|
||||
(>= a 128))
|
||||
|
||||
|
||||
(define TRANSPARENT-INDEX 0)
|
||||
|
||||
|
||||
|
||||
;; quantize-bitmap: bitmap positive-number -> bitmap
|
||||
;; Given a bitmap, returns a new bitmap quantized to, at most, n colors.
|
||||
(define (quantize-bitmap bm n)
|
||||
(let* ([width (send bm get-width)]
|
||||
[height (send bm get-height)]
|
||||
[len (* width height 4)]
|
||||
[source-buffer (make-bytes len)]
|
||||
[_ (send bm get-argb-pixels 0 0 width height source-buffer)]
|
||||
[an-octree (make-octree-from-argb source-buffer n)]
|
||||
[dest-buffer (make-bytes len)])
|
||||
(let quantize-bitmap-loop ([i 0])
|
||||
(when (< i len)
|
||||
(let* ([i+1 (+ i 1)]
|
||||
[i+2 (+ i 2)]
|
||||
[i+3 (+ i 3)]
|
||||
[a (bytes-ref source-buffer i)]
|
||||
[r (bytes-ref source-buffer i+1)]
|
||||
[g (bytes-ref source-buffer i+2)]
|
||||
[b (bytes-ref source-buffer i+3)])
|
||||
(cond
|
||||
[(alpha-opaque? a)
|
||||
(let-values ([(new-r new-g new-b)
|
||||
(octree-lookup an-octree r g b)])
|
||||
(bytes-set! dest-buffer i 255)
|
||||
(bytes-set! dest-buffer i+1 new-r)
|
||||
(bytes-set! dest-buffer i+2 new-g)
|
||||
(bytes-set! dest-buffer i+3 new-b))]
|
||||
[else
|
||||
(bytes-set! dest-buffer i 0)
|
||||
(bytes-set! dest-buffer i+1 0)
|
||||
(bytes-set! dest-buffer i+2 0)
|
||||
(bytes-set! dest-buffer i+3 0)]))
|
||||
(quantize-bitmap-loop (+ i 4))))
|
||||
(let* ([new-bm (make-object bitmap% width height)]
|
||||
[dc (make-object bitmap-dc% new-bm)])
|
||||
(send dc set-argb-pixels 0 0 width height dest-buffer)
|
||||
(send dc set-bitmap #f)
|
||||
new-bm)))
|
||||
|
||||
|
||||
;; The maximum level height of an octree.
|
||||
(define MAX-LEVEL 7)
|
||||
|
||||
|
||||
(define byte/c (integer-in 0 255))
|
||||
(define octet/c (integer-in 0 7))
|
||||
|
||||
|
||||
|
||||
(define-struct (exn:fail:octree exn:fail) ())
|
||||
|
||||
|
||||
;; A color is an (vector byte byte byte)
|
||||
|
||||
|
||||
;; An octree is a:
|
||||
(define-struct octree (root ;; node
|
||||
leaf-count ;; number
|
||||
reduction-heads ;; (vectorof (or/c node #f))
|
||||
palette ;; (vectorof (or/c color #f))
|
||||
)
|
||||
#:mutable)
|
||||
;; reduction-heads is used to accelerate the search for a reduction candidate.
|
||||
|
||||
|
||||
;; A subtree node is a:
|
||||
(define-struct node (leaf? ;; bool
|
||||
npixels ;; number -- number of pixels this subtree node represents
|
||||
redsum ;; number
|
||||
greensum ;; number
|
||||
bluesum ;; number
|
||||
children ;; (vectorof (or/c #f node))
|
||||
next ;; (or/c #f node)
|
||||
palette-index ;; (or/c #f (integer-in 0 255))
|
||||
)
|
||||
#:mutable)
|
||||
;; node-next is used to accelerate the search for a reduction candidate.
|
||||
|
||||
|
||||
;; new-octree: -> octree
|
||||
(define (new-octree)
|
||||
(let* ([root-node (make-node #f ;; not a leaf
|
||||
0 ;; no pixels under us yet
|
||||
0 ;; red sum
|
||||
0 ;; green sum
|
||||
0 ;; blue sum
|
||||
(make-vector 8 #f) ;; no children so far
|
||||
#f ;; next
|
||||
#f ;; palette-index
|
||||
)]
|
||||
[an-octree
|
||||
(make-octree root-node
|
||||
0 ;; no leaves so far
|
||||
(make-vector (add1 MAX-LEVEL) #f) ;; no reductions so far
|
||||
(make-vector 256 #(0 0 0)) ;; the palette
|
||||
)])
|
||||
;; Although we'll almost never reduce to this level, initialize the first
|
||||
;; reducible node to the root, for completeness sake.
|
||||
(vector-set! (octree-reduction-heads an-octree) 0 root-node)
|
||||
an-octree))
|
||||
|
||||
|
||||
|
||||
;; rgb->index: natural-number byte byte byte -> octet
|
||||
;; Given a level and an (r,g,b) triplet, returns an octet that can be used
|
||||
;; as an index into our octree structure.
|
||||
(define (rgb->index level r g b)
|
||||
(bitwise-ior
|
||||
(bitwise-and 4 (arithmetic-shift r (- level 5)))
|
||||
(bitwise-and 2 (arithmetic-shift g (- level 6)))
|
||||
(bitwise-and 1 (arithmetic-shift b (- level 7)))))
|
||||
|
||||
|
||||
;; octree-insert-color!: octree byte byte byte -> void
|
||||
;; Accumulates a new r,g,b triplet into the octree.
|
||||
(define (octree-insert-color! an-octree r g b)
|
||||
(node-insert-color! (octree-root an-octree) an-octree r g b 0))
|
||||
|
||||
|
||||
;; node-insert-color!: node octree byte byte byte natural-number -> void
|
||||
;; Adds a color to the node subtree. While we hit #f, we create new nodes.
|
||||
;; If we hit an existing leaf, we accumulate our color into it.
|
||||
(define (node-insert-color! a-node an-octree r g b level)
|
||||
(let insert-color-loop ([a-node a-node]
|
||||
[level level])
|
||||
(cond [(node-leaf? a-node)
|
||||
;; update the leaf with the new color
|
||||
(set-node-npixels! a-node (add1 (node-npixels a-node)))
|
||||
(set-node-redsum! a-node (+ (node-redsum a-node) r))
|
||||
(set-node-greensum! a-node (+ (node-greensum a-node) g))
|
||||
(set-node-bluesum! a-node (+ (node-bluesum a-node) b))]
|
||||
[else
|
||||
;; create the child node if necessary
|
||||
(let ([index (rgb->index level r g b)])
|
||||
(unless (vector-ref (node-children a-node) index)
|
||||
(let ([new-node (make-node (= level MAX-LEVEL) ; leaf?
|
||||
0 ; npixels
|
||||
0 ; redsum
|
||||
0 ; greensum
|
||||
0 ; bluesum
|
||||
(make-vector 8 #f) ; no children yet
|
||||
#f ; and no next node yet
|
||||
#f ; or palette index
|
||||
)])
|
||||
(vector-set! (node-children a-node) index new-node)
|
||||
(cond
|
||||
[(= level MAX-LEVEL)
|
||||
;; If we added a leaf, mark it in the octree.
|
||||
(set-octree-leaf-count! an-octree
|
||||
(add1 (octree-leaf-count an-octree)))]
|
||||
[else
|
||||
;; Attach the node as a reducible node if it's interior.
|
||||
(set-node-next!
|
||||
new-node (vector-ref (octree-reduction-heads an-octree)
|
||||
(add1 level)))
|
||||
(vector-set! (octree-reduction-heads an-octree)
|
||||
(add1 level)
|
||||
new-node)])))
|
||||
;; and recur on the child node.
|
||||
(insert-color-loop (vector-ref (node-children a-node) index)
|
||||
(add1 level)))])))
|
||||
|
||||
|
||||
;; octree-reduce!: octree -> void
|
||||
;; Reduces one of the subtrees, collapsing the children into a single node.
|
||||
(define (octree-reduce! an-octree)
|
||||
(let ([candidate (pop-reduction-candidate! an-octree)])
|
||||
(node-reduce! candidate an-octree)))
|
||||
|
||||
|
||||
;; node-reduce!: node octree -> void
|
||||
;; Reduces the interior node.
|
||||
(define (node-reduce! a-node an-octree)
|
||||
(for ([child (in-vector (node-children a-node))]
|
||||
#:when child)
|
||||
(set-node-npixels! a-node (+ (node-npixels a-node)
|
||||
(node-npixels child)))
|
||||
(set-node-redsum! a-node (+ (node-redsum a-node)
|
||||
(node-redsum child)))
|
||||
(set-node-greensum! a-node (+ (node-greensum a-node)
|
||||
(node-greensum child)))
|
||||
(set-node-bluesum! a-node (+ (node-bluesum a-node)
|
||||
(node-bluesum child)))
|
||||
(set-octree-leaf-count! an-octree (sub1 (octree-leaf-count an-octree))))
|
||||
(set-node-leaf?! a-node #t)
|
||||
(set-octree-leaf-count! an-octree (add1 (octree-leaf-count an-octree))))
|
||||
|
||||
|
||||
;; find-reduction-candidate!: octree -> node
|
||||
;; Returns a bottom-level interior node for reduction. Also takes the
|
||||
;; candidate out of the conceptual queue of reduction candidates.
|
||||
(define (pop-reduction-candidate! an-octree)
|
||||
(let loop ([i MAX-LEVEL])
|
||||
(cond
|
||||
[(vector-ref (octree-reduction-heads an-octree) i)
|
||||
=>
|
||||
(lambda (candidate-node)
|
||||
(when (> i 0)
|
||||
(vector-set! (octree-reduction-heads an-octree) i
|
||||
(node-next candidate-node)))
|
||||
candidate-node)]
|
||||
[else
|
||||
(loop (sub1 i))])))
|
||||
|
||||
|
||||
;; octree-finalize!: octree -> void
|
||||
;; Finalization does a few things:
|
||||
;; * Walks through the octree and reduces any interior nodes with just one leaf child.
|
||||
;; Optimizes future lookups.
|
||||
;; * Fills in the palette of the octree and the palette indexes of the leaf nodes.
|
||||
;; * Note: palette index 0 is always reserved for the transparent color.
|
||||
(define (octree-finalize! an-octree)
|
||||
;; Collapse one-leaf interior nodes.
|
||||
(let loop ([a-node (octree-root an-octree)])
|
||||
(for ([child (in-vector (node-children a-node))]
|
||||
#:when (and child (not (node-leaf? child))))
|
||||
(loop child)
|
||||
|
||||
(when (interior-node-one-leaf-child? a-node)
|
||||
(node-reduce! a-node an-octree))))
|
||||
|
||||
;; Attach palette entries.
|
||||
(let ([current-palette-index 1])
|
||||
(let loop ([a-node (octree-root an-octree)])
|
||||
(cond [(node-leaf? a-node)
|
||||
(let ([n (node-npixels a-node)])
|
||||
(vector-set! (octree-palette an-octree) current-palette-index
|
||||
(vector (quotient (node-redsum a-node) n)
|
||||
(quotient (node-greensum a-node) n)
|
||||
(quotient (node-bluesum a-node) n)))
|
||||
(set-node-palette-index! a-node current-palette-index)
|
||||
(set! current-palette-index (add1 current-palette-index)))]
|
||||
[else
|
||||
(for ([child (in-vector (node-children a-node))]
|
||||
#:when child)
|
||||
(loop child))]))))
|
||||
|
||||
|
||||
;; interior-node-one-leaf-child?: node -> boolean
|
||||
(define (interior-node-one-leaf-child? a-node)
|
||||
(let ([child-list (filter values (vector->list (node-children a-node)))])
|
||||
(and (= (length child-list) 1)
|
||||
(node-leaf? (car child-list)))))
|
||||
|
||||
|
||||
|
||||
;; octree-lookup: octree byte byte byte -> (values byte byte byte)
|
||||
;; Returns the palettized color.
|
||||
(define (octree-lookup an-octree r g b)
|
||||
(let ([index
|
||||
(node-lookup-index (octree-root an-octree) an-octree r g b 0)])
|
||||
(let ([vec (vector-ref (octree-palette an-octree) index)])
|
||||
(values (vector-ref vec 0)
|
||||
(vector-ref vec 1)
|
||||
(vector-ref vec 2)))))
|
||||
|
||||
|
||||
;; octree-lookup-index: octree byte byte byte -> byte
|
||||
;; Returns the paletized color as an index into the octree-palette vector.
|
||||
(define (octree-lookup-index an-octree r g b)
|
||||
(node-lookup-index (octree-root an-octree) an-octree r g b 0))
|
||||
|
||||
|
||||
;; node-lookup-index: node byte byte byte natural-number -> byte
|
||||
;; Returns the palettized color index.
|
||||
(define (node-lookup-index a-node an-octree r g b level)
|
||||
(let loop ([a-node a-node]
|
||||
[level level])
|
||||
(cond [(node-leaf? a-node)
|
||||
(node-palette-index a-node)]
|
||||
[else
|
||||
(let ([child (vector-ref (node-children a-node) (rgb->index level r g b))])
|
||||
(unless child
|
||||
(raise (make-exn:fail:octree
|
||||
(format
|
||||
"octree-lookup: Color (~a, ~a, ~a) not previously inserted"
|
||||
r g b)
|
||||
(current-continuation-marks))))
|
||||
(loop child (add1 level)))])))
|
||||
|
||||
|
||||
|
||||
(define color? byte?)
|
||||
|
||||
(define (gif-colormap? l)
|
||||
(and (list? l)
|
||||
(member (length l) '(2 4 8 16 32 64 128 256))
|
||||
(andmap (lambda (c)
|
||||
(and (vector? c)
|
||||
(= 3 (vector-length c))
|
||||
(color? (vector-ref c 0))
|
||||
(color? (vector-ref c 1))
|
||||
(color? (vector-ref c 2))))
|
||||
l)))
|
||||
|
||||
(define (argb-bytes? b)
|
||||
(and (bytes? b)
|
||||
(zero? (remainder (bytes-length b) 4))))
|
||||
|
||||
|
||||
(provide/contract [quantize ((argb-bytes?) (#:palette (integer-in 1 255))
|
||||
. ->* . (values bytes? gif-colormap? color?))]
|
||||
[quantize-bitmap ((is-a?/c bitmap%) natural-number/c
|
||||
. -> .
|
||||
(is-a?/c bitmap%))]
|
||||
|
||||
[struct (exn:fail:octree exn:fail) ((message string?)
|
||||
(continuation-marks continuation-mark-set?))]
|
||||
[struct octree
|
||||
((root node?)
|
||||
(leaf-count natural-number/c)
|
||||
(reduction-heads (vectorof (or/c node? false/c)))
|
||||
(palette (vectorof (vector/c byte/c byte/c byte/c))))]
|
||||
[struct node
|
||||
((leaf? boolean?)
|
||||
(npixels natural-number/c)
|
||||
(redsum natural-number/c)
|
||||
(greensum natural-number/c)
|
||||
(bluesum natural-number/c)
|
||||
(children (vectorof (or/c false/c node?)))
|
||||
(next (or/c node? false/c))
|
||||
(palette-index byte/c))]
|
||||
[new-octree (-> octree?)]
|
||||
[octree-insert-color! (octree? byte/c byte/c byte/c . -> . any)]
|
||||
[octree-lookup (octree? byte/c byte/c byte/c . -> . (values byte/c byte/c byte/c))]
|
||||
[octree-reduce! (octree? . -> . any)]
|
||||
[octree-finalize! (octree? . -> . any)]
|
||||
[rgb->index (natural-number/c byte/c byte/c byte/c . -> . octet/c)])
|
||||
Loading…
Reference in New Issue
Block a user