Add SqImage type, which also implements image.Image and image/draw.Image

This is intended as an alternative to the WithSq and ToSqIntegralImg
functionality, though the MeanStdDevWindow function will need to be
changed to use it, and it hasn't been heavily tested yet.

2 files changed, 70 insertions, 9 deletions

diff --git a/integralimg.go b/integralimg.go
index a0cb565..4ae1791 100644
--- a/integralimg.go
+++ b/integralimg.go
@@ -17,15 +17,24 @@ import (
 // I is the Integral Image
 type I [][]uint64
 
-func (i I) ColorModel() color.Model { return color.GrayModel }
+// SqImage is a Square integral Image.
+// A squared integral image is an integral image for which the square of
+// each pixel is saved; this is useful for efficiently calculating
+// Standard Deviation.
+type SqImage [][]uint64
+
+func (i I) ColorModel() color.Model { return color.Gray16Model }
 
 func (i I) Bounds() image.Rectangle {
 	return image.Rectangle {image.Point{0, 0}, image.Point{len(i[0]), len(i)}}
 }
 
-func (i I) At(x, y int) color.Color {
+// at64 is used to return the raw uint64 for a given pixel. Accessing
+// this separately to a (potentially lossy) conversion to a Gray16 is
+// necessary for SqImage to function accurately.
+func (i I) at64(x, y int) uint64 {
 	if !(image.Point{x, y}.In(i.Bounds())) {
-		return color.Gray{}
+		return 0
 	}
 
 	var prevx, prevy, prevxy uint64
@@ -40,11 +49,15 @@ func (i I) At(x, y int) color.Color {
 		prevxy = i[y-1][x-1]
 	}
 	orig := i[y][x] + prevxy - prevx - prevy
+	return orig
+}
 
-	return color.Gray{uint8(orig)}
+func (i I) At(x, y int) color.Color {
+	c := i.at64(x, y)
+	return color.Gray16{uint16(c)}
 }
 
-func (i I) Set(x, y int, c color.Color) {
+func (i I) set64(x, y int, c uint64) {
 	var prevx, prevy, prevxy uint64
 	prevx, prevy, prevxy = 0, 0, 0
 	if x > 0 {
@@ -56,12 +69,16 @@ func (i I) Set(x, y int, c color.Color) {
 	if x > 0 && y > 0 {
 		prevxy = i[y-1][x-1]
 	}
-	gray := color.GrayModel.Convert(c).(color.Gray).Y
-	final := uint64(gray) + prevx + prevy - prevxy
+	final := c + prevx + prevy - prevxy
 	i[y][x] = final
 }
 
-// NewImage returns a new Integral Image with the given bounds.
+func (i I) Set(x, y int, c color.Color) {
+	gray := color.Gray16Model.Convert(c).(color.Gray16).Y
+	i.set64(x, y, uint64(gray))
+}
+
+// NewImage returns a new integral Image with the given bounds.
 func NewImage(r image.Rectangle) *I {
 	w, h := r.Dx(), r.Dy()
 	var rows I
@@ -72,6 +89,30 @@ func NewImage(r image.Rectangle) *I {
 	return &rows
 }
 
+func (i SqImage) ColorModel() color.Model { return I(i).ColorModel() }
+
+func (i SqImage) Bounds() image.Rectangle {
+	return I(i).Bounds()
+}
+
+func (i SqImage) At(x, y int) color.Color {
+	c := I(i).at64(x, y)
+	rt := math.Sqrt(float64(c))
+	return color.Gray16{uint16(rt)}
+}
+
+func (i SqImage) Set(x, y int, c color.Color) {
+	gray := uint64(color.Gray16Model.Convert(c).(color.Gray16).Y)
+	I(i).set64(x, y, gray * gray)
+}
+
+// NewSqImage returns a new squared integral Image with the given bounds.
+func NewSqImage(r image.Rectangle) *SqImage {
+	i := NewImage(r)
+	s := SqImage(*i)
+	return &s
+}
+
 // Sq contains an Integral Image and its Square
 type WithSq struct {
 	Img I
diff --git a/integralimg_test.go b/integralimg_test.go
index df97caa..8b882db 100644
--- a/integralimg_test.go
+++ b/integralimg_test.go
@@ -28,7 +28,27 @@ func TestFromPNG(t *testing.T) {
 	draw.Draw(integral, b, img, b.Min, draw.Src)
 
 	if !imgsequal(img, integral) {
-		t.Errorf("Read png image differs to integral\n")
+		t.Errorf("Read png image differs to integral image\n")
+	}
+}
+
+func TestSqFromPNG(t *testing.T) {
+	f, err := os.Open("testdata/in.png")
+	if err != nil {
+		t.Fatalf("Could not open file %s: %v\n", "testdata/in.png", err)
+	}
+	defer f.Close()
+	img, _, err := image.Decode(f)
+	if err != nil {
+		t.Fatalf("Could not decode image: %v\n", err)
+	}
+	b := img.Bounds()
+
+	integral := NewSqImage(image.Rect(0, 0, b.Dx(), b.Dy()))
+	draw.Draw(integral, b, img, b.Min, draw.Src)
+
+	if !imgsequal(img, integral) {
+		t.Errorf("Read png image differs to square integral image\n")
 	}
 }