image/draw: add Drawer, FloydSteinberg and the op.Draw method.

R=r, andybons
CC=andybons, golang-dev
https://golang.org/cl/10977043

src/pkg/image/color/color.go

@@ -253,13 +253,6 @@ func gray16Model(c Color) Color {
 // Palette is a palette of colors.
 type Palette []Color
 
-func diff(a, b uint32) uint32 {
-	if a > b {
-		return a - b
-	}
-	return b - a
-}
-
 // Convert returns the palette color closest to c in Euclidean R,G,B space.
 func (p Palette) Convert(c Color) Color {
 	if len(p) == 0 {
@@ -271,19 +264,20 @@ func (p Palette) Convert(c Color) Color {
 // Index returns the index of the palette color closest to c in Euclidean
 // R,G,B space.
 func (p Palette) Index(c Color) int {
+	// A batch version of this computation is in image/draw/draw.go.
+
 	cr, cg, cb, _ := c.RGBA()
-	// Shift by 1 bit to avoid potential uint32 overflow in sum-squared-difference.
-	cr >>= 1
-	cg >>= 1
-	cb >>= 1
 	ret, bestSSD := 0, uint32(1<<32-1)
 	for i, v := range p {
 		vr, vg, vb, _ := v.RGBA()
-		vr >>= 1
-		vg >>= 1
-		vb >>= 1
-		dr, dg, db := diff(cr, vr), diff(cg, vg), diff(cb, vb)
-		ssd := (dr * dr) + (dg * dg) + (db * db)
+		// We shift by 1 bit to avoid potential uint32 overflow in
+		// sum-squared-difference.
+		delta := (int32(cr) - int32(vr)) >> 1
+		ssd := uint32(delta * delta)
+		delta = (int32(cg) - int32(vg)) >> 1
+		ssd += uint32(delta * delta)
+		delta = (int32(cb) - int32(vb)) >> 1
+		ssd += uint32(delta * delta)
 		if ssd < bestSSD {
 			if ssd == 0 {
 				return i

src/pkg/image/draw/draw.go

@@ -16,6 +16,12 @@ import (
 // m is the maximum color value returned by image.Color.RGBA.
 const m = 1<<16 - 1
 
+// A draw.Image is an image.Image with a Set method to change a single pixel.
+type Image interface {
+	image.Image
+	Set(x, y int, c color.Color)
+}
+
 // Op is a Porter-Duff compositing operator.
 type Op int
 
@@ -26,15 +32,31 @@ const (
 	Src
 )
 
-// A draw.Image is an image.Image with a Set method to change a single pixel.
-type Image interface {
-	image.Image
-	Set(x, y int, c color.Color)
+// Draw implements the Drawer interface by calling the Draw function with this
+// Op.
+func (op Op) Draw(dst Image, r image.Rectangle, src image.Image, sp image.Point) {
+	DrawMask(dst, r, src, sp, nil, image.Point{}, op)
 }
 
-// Draw calls DrawMask with a nil mask.
-func Draw(dst Image, r image.Rectangle, src image.Image, sp image.Point, op Op) {
-	DrawMask(dst, r, src, sp, nil, image.ZP, op)
+// Drawer contains the Draw method.
+type Drawer interface {
+	// Draw aligns r.Min in dst with sp in src and then replaces the
+	// rectangle r in dst with the result of drawing src on dst.
+	Draw(dst Image, r image.Rectangle, src image.Image, sp image.Point)
+}
+
+// FloydSteinberg is a Drawer that is the Src Op with Floyd-Steinberg error
+// diffusion.
+var FloydSteinberg Drawer = floydSteinberg{}
+
+type floydSteinberg struct{}
+
+func (floydSteinberg) Draw(dst Image, r image.Rectangle, src image.Image, sp image.Point) {
+	clip(dst, &r, src, &sp, nil, nil)
+	if r.Empty() {
+		return
+	}
+	drawPaletted(dst, r, src, sp, true)
 }
 
 // clip clips r against each image's bounds (after translating into the
@@ -58,6 +80,17 @@ func clip(dst Image, r *image.Rectangle, src image.Image, sp *image.Point, mask
 	(*mp).Y += dy
 }
 
+func processBackward(dst Image, r image.Rectangle, src image.Image, sp image.Point) bool {
+	return image.Image(dst) == src &&
+		r.Overlaps(r.Add(sp.Sub(r.Min))) &&
+		(sp.Y < r.Min.Y || (sp.Y == r.Min.Y && sp.X < r.Min.X))
+}
+
+// Draw calls DrawMask with a nil mask.
+func Draw(dst Image, r image.Rectangle, src image.Image, sp image.Point, op Op) {
+	DrawMask(dst, r, src, sp, nil, image.Point{}, op)
+}
+
 // DrawMask aligns r.Min in dst with sp in src and mp in mask and then replaces the rectangle r
 // in dst with the result of a Porter-Duff composition. A nil mask is treated as opaque.
 func DrawMask(dst Image, r image.Rectangle, src image.Image, sp image.Point, mask image.Image, mp image.Point, op Op) {
@@ -67,7 +100,8 @@ func DrawMask(dst Image, r image.Rectangle, src image.Image, sp image.Point, mas
 	}
 
 	// Fast paths for special cases. If none of them apply, then we fall back to a general but slow implementation.
-	if dst0, ok := dst.(*image.RGBA); ok {
+	switch dst0 := dst.(type) {
+	case *image.RGBA:
 		if op == Over {
 			if mask == nil {
 				switch src0 := src.(type) {
@@ -113,19 +147,20 @@ func DrawMask(dst Image, r image.Rectangle, src image.Image, sp image.Point, mas
 		}
 		drawRGBA(dst0, r, src, sp, mask, mp, op)
 		return
+	case *image.Paletted:
+		if op == Src && mask == nil && !processBackward(dst, r, src, sp) {
+			drawPaletted(dst0, r, src, sp, false)
+		}
 	}
 
 	x0, x1, dx := r.Min.X, r.Max.X, 1
 	y0, y1, dy := r.Min.Y, r.Max.Y, 1
-	if image.Image(dst) == src && r.Overlaps(r.Add(sp.Sub(r.Min))) {
-		// Rectangles overlap: process backward?
-		if sp.Y < r.Min.Y || sp.Y == r.Min.Y && sp.X < r.Min.X {
-			x0, x1, dx = x1-1, x0-1, -1
-			y0, y1, dy = y1-1, y0-1, -1
-		}
+	if processBackward(dst, r, src, sp) {
+		x0, x1, dx = x1-1, x0-1, -1
+		y0, y1, dy = y1-1, y0-1, -1
 	}
 
-	var out *color.RGBA64
+	var out color.RGBA64
 	sy := sp.Y + y0 - r.Min.Y
 	my := mp.Y + y0 - r.Min.Y
 	for y := y0; y != y1; y, sy, my = y+dy, sy+dy, my+dy {
@@ -147,9 +182,6 @@ func DrawMask(dst Image, r image.Rectangle, src image.Image, sp image.Point, mas
 				dst.Set(x, y, src.At(sx, sy))
 			default:
 				sr, sg, sb, sa := src.At(sx, sy).RGBA()
-				if out == nil {
-					out = new(color.RGBA64)
-				}
 				if op == Over {
 					dr, dg, db, da := dst.At(x, y).RGBA()
 					a := m - (sa * ma / m)
@@ -163,7 +195,11 @@ func DrawMask(dst Image, r image.Rectangle, src image.Image, sp image.Point, mas
 					out.B = uint16(sb * ma / m)
 					out.A = uint16(sa * ma / m)
 				}
-				dst.Set(x, y, out)
+				// The third argument is &out instead of out (and out is
+				// declared outside of the inner loop) to avoid the implicit
+				// conversion to color.Color here allocating memory in the
+				// inner loop if sizeof(color.RGBA64) > sizeof(uintptr).
+				dst.Set(x, y, &out)
 			}
 		}
 	}
@@ -500,3 +536,131 @@ func drawRGBA(dst *image.RGBA, r image.Rectangle, src image.Image, sp image.Poin
 		i0 += dy * dst.Stride
 	}
 }
+
+// clamp clamps i to the interval [0, 0xffff].
+func clamp(i int32) int32 {
+	if i < 0 {
+		return 0
+	}
+	if i > 0xffff {
+		return 0xffff
+	}
+	return i
+}
+
+func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point, floydSteinberg bool) {
+	// TODO(nigeltao): handle the case where the dst and src overlap.
+	// Does it even make sense to try and do Floyd-Steinberg whilst
+	// walking the image backward (right-to-left bottom-to-top)?
+
+	// If dst is an *image.Paletted, we have a fast path for dst.Set and
+	// dst.At. The dst.Set equivalent is a batch version of the algorithm
+	// used by color.Palette's Index method in image/color/color.go, plus
+	// optional Floyd-Steinberg error diffusion.
+	palette, pix, stride := [][3]int32(nil), []byte(nil), 0
+	if p, ok := dst.(*image.Paletted); ok {
+		palette = make([][3]int32, len(p.Palette))
+		for i, col := range p.Palette {
+			r, g, b, _ := col.RGBA()
+			palette[i][0] = int32(r)
+			palette[i][1] = int32(g)
+			palette[i][2] = int32(b)
+		}
+		pix, stride = p.Pix[p.PixOffset(r.Min.X, r.Min.Y):], p.Stride
+	}
+
+	// quantErrorCurr and quantErrorNext are the Floyd-Steinberg quantization
+	// errors that have been propagated to the pixels in the current and next
+	// rows. The +2 simplifies calculation near the edges.
+	var quantErrorCurr, quantErrorNext [][3]int32
+	if floydSteinberg {
+		quantErrorCurr = make([][3]int32, r.Dx()+2)
+		quantErrorNext = make([][3]int32, r.Dx()+2)
+	}
+
+	// Loop over each source pixel.
+	out := color.RGBA64{A: 0xffff}
+	for y := 0; y != r.Dy(); y++ {
+		for x := 0; x != r.Dx(); x++ {
+			// er, eg and eb are the pixel's R,G,B values plus the
+			// optional Floyd-Steinberg error.
+			sr, sg, sb, _ := src.At(sp.X+x, sp.Y+y).RGBA()
+			er, eg, eb := int32(sr), int32(sg), int32(sb)
+			if floydSteinberg {
+				er = clamp(er + quantErrorCurr[x+1][0]/16)
+				eg = clamp(eg + quantErrorCurr[x+1][1]/16)
+				eb = clamp(eb + quantErrorCurr[x+1][2]/16)
+			}
+
+			if palette != nil {
+				// Find the closest palette color in Euclidean R,G,B space: the
+				// one that minimizes sum-squared-difference. We shift by 1 bit
+				// to avoid potential uint32 overflow in sum-squared-difference.
+				// TODO(nigeltao): consider smarter algorithms.
+				bestIndex, bestSSD := 0, uint32(1<<32-1)
+				for index, p := range palette {
+					delta := (er - p[0]) >> 1
+					ssd := uint32(delta * delta)
+					delta = (eg - p[1]) >> 1
+					ssd += uint32(delta * delta)
+					delta = (eb - p[2]) >> 1
+					ssd += uint32(delta * delta)
+					if ssd < bestSSD {
+						bestIndex, bestSSD = index, ssd
+						if ssd == 0 {
+							break
+						}
+					}
+				}
+				pix[y*stride+x] = byte(bestIndex)
+
+				if !floydSteinberg {
+					continue
+				}
+				er -= int32(palette[bestIndex][0])
+				eg -= int32(palette[bestIndex][1])
+				eb -= int32(palette[bestIndex][2])
+
+			} else {
+				out.R = uint16(er)
+				out.G = uint16(eg)
+				out.B = uint16(eb)
+				// The third argument is &out instead of out (and out is
+				// declared outside of the inner loop) to avoid the implicit
+				// conversion to color.Color here allocating memory in the
+				// inner loop if sizeof(color.RGBA64) > sizeof(uintptr).
+				dst.Set(r.Min.X+x, r.Min.Y+y, &out)
+
+				if !floydSteinberg {
+					continue
+				}
+				sr, sg, sb, _ = dst.At(r.Min.X+x, r.Min.Y+y).RGBA()
+				er -= int32(sr)
+				eg -= int32(sg)
+				eb -= int32(sb)
+			}
+
+			// Propagate the Floyd-Steinberg quantization error.
+			quantErrorNext[x+0][0] += er * 3
+			quantErrorNext[x+0][1] += eg * 3
+			quantErrorNext[x+0][2] += eb * 3
+			quantErrorNext[x+1][0] += er * 5
+			quantErrorNext[x+1][1] += eg * 5
+			quantErrorNext[x+1][2] += eb * 5
+			quantErrorNext[x+2][0] += er * 1
+			quantErrorNext[x+2][1] += eg * 1
+			quantErrorNext[x+2][2] += eb * 1
+			quantErrorCurr[x+2][0] += er * 7
+			quantErrorCurr[x+2][1] += eg * 7
+			quantErrorCurr[x+2][2] += eb * 7
+		}
+
+		// Recycle the quantization error buffers.
+		if floydSteinberg {
+			quantErrorCurr, quantErrorNext = quantErrorNext, quantErrorCurr
+			for i := range quantErrorNext {
+				quantErrorNext[i] = [3]int32{}
+			}
+		}
+	}
+}

src/pkg/image/draw/draw_test.go

@@ -7,6 +7,8 @@ package draw
 import (
 	"image"
 	"image/color"
+	"image/png"
+	"os"
 	"testing"
 )
 
@@ -352,3 +354,76 @@ func TestFill(t *testing.T) {
 		check("whole")
 	}
 }
+
+// TestFloydSteinbergCheckerboard tests that the result of Floyd-Steinberg
+// error diffusion of a uniform 50% gray source image with a black-and-white
+// palette is a checkerboard pattern.
+func TestFloydSteinbergCheckerboard(t *testing.T) {
+	b := image.Rect(0, 0, 640, 480)
+	// We can't represent 50% exactly, but 0x7fff / 0xffff is close enough.
+	src := &image.Uniform{color.Gray16{0x7fff}}
+	dst := image.NewPaletted(b, color.Palette{color.Black, color.White})
+	FloydSteinberg.Draw(dst, b, src, image.Point{})
+	nErr := 0
+	for y := b.Min.Y; y < b.Max.Y; y++ {
+		for x := b.Min.X; x < b.Max.X; x++ {
+			got := dst.Pix[dst.PixOffset(x, y)]
+			want := uint8(x+y) % 2
+			if got != want {
+				t.Errorf("at (%d, %d): got %d, want %d", x, y, got, want)
+				if nErr++; nErr == 10 {
+					t.Fatal("there may be more errors")
+				}
+			}
+		}
+	}
+}
+
+// embeddedPaletted is an Image that behaves like an *image.Paletted but whose
+// type is not *image.Paletted.
+type embeddedPaletted struct {
+	*image.Paletted
+}
+
+// TestPaletted tests that the drawPaletted function behaves the same
+// regardless of whether dst is an *image.Paletted.
+func TestPaletted(t *testing.T) {
+	f, err := os.Open("../testdata/video-001.png")
+	if err != nil {
+		t.Fatal("open: %v", err)
+	}
+	defer f.Close()
+	src, err := png.Decode(f)
+	if err != nil {
+		t.Fatal("decode: %v", err)
+	}
+	b := src.Bounds()
+
+	cgaPalette := color.Palette{
+		color.RGBA{0x00, 0x00, 0x00, 0xff},
+		color.RGBA{0x55, 0xff, 0xff, 0xff},
+		color.RGBA{0xff, 0x55, 0xff, 0xff},
+		color.RGBA{0xff, 0xff, 0xff, 0xff},
+	}
+	drawers := map[string]Drawer{
+		"src":             Src,
+		"floyd-steinberg": FloydSteinberg,
+	}
+
+loop:
+	for dName, d := range drawers {
+		dst0 := image.NewPaletted(b, cgaPalette)
+		dst1 := image.NewPaletted(b, cgaPalette)
+		d.Draw(dst0, b, src, image.Point{})
+		d.Draw(embeddedPaletted{dst1}, b, src, image.Point{})
+		for y := b.Min.Y; y < b.Max.Y; y++ {
+			for x := b.Min.X; x < b.Max.X; x++ {
+				if !eq(dst0.At(x, y), dst1.At(x, y)) {
+					t.Errorf("%s: at (%d, %d), %v versus %v",
+						dName, x, y, dst0.At(x, y), dst1.At(x, y))
+					continue loop
+				}
+			}
+		}
+	}
+}