Reduce memory usage in image processing pipeline (#14)

* Implement a cached kernel scaler

Optimized for scaling a lot of images in a few fixed destination and
source sizes.

* test: Add benchmark for entire conversion pipeline

* Update processing pipeline to use new scaler

* test: Avoid extra copying in Grayscale benchmarks

* Remove unnecessary mutex pointer from cache scaler

* Decouple grayscale func

* Implement an image pool

* Utilize image pool in processing pipeline

* Prevent new allocations if decoded image is already gray

cmd/mangaconv/main.go

@@ -56,18 +56,20 @@ If provided directory does not exist, mangaconv will attempt to create it. (defa
 		}
 	}()
 
+	converter := mangaconv.New(mangaconv.Params{
+		Cutoff: *cutoff,
+		Gamma:  *gamma,
+		Height: *height,
+		Width:  *width,
+	})
+
 	var wg sync.WaitGroup
 	for i := 0; i < 2; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
 			for t := range targets {
-				if err := mangaconv.Convert(t.in, t.out, mangaconv.Params{
-					Cutoff: *cutoff,
-					Gamma:  *gamma,
-					Height: *height,
-					Width:  *width,
-				}); err != nil {
+				if err := converter.Convert(t.in, t.out); err != nil {
 					fmt.Println("Failed to convert", filepath.Base(t.in), err)
 					return
 				}

imgutil/grayscale.go

@@ -8,49 +8,50 @@ import (
 // Grayscale returns an image converted to grayscale.
 // It always returns a copy.
 func Grayscale(img image.Image) *image.Gray {
-	switch i := img.(type) {
+	b := img.Bounds()
+	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
+	grayscale(dst, img)
+	return dst
+}
+
+func grayscale(dst *image.Gray, src image.Image) {
+	switch i := src.(type) {
 	case *image.Gray:
-		return clone(i)
+		clone(dst, i)
 	case *image.RGBA:
-		return rgbaToGray(i)
+		rgbaToGray(dst, i)
 	case *image.RGBA64:
-		return rgba64ToGray(i)
+		rgba64ToGray(dst, i)
 	case *image.NRGBA:
-		return nrgbaToGray(i)
+		nrgbaToGray(dst, i)
 	case *image.NRGBA64:
-		return nrgba64ToGray(i)
+		nrgba64ToGray(dst, i)
 	case *image.YCbCr:
-		return ycbcrToGray(i)
+		ycbcrToGray(dst, i)
 	default:
-		return drawGray(img)
+		drawGray(dst, src)
 	}
 }
 
 // drawGray uses draw.Draw as a slow conversion fallback for unsupported image types.
-func drawGray(img image.Image) *image.Gray {
-	b := img.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
-	draw.Draw(dst, dst.Bounds(), img, b.Min, draw.Src)
-	return dst
+func drawGray(dst *image.Gray, src image.Image) {
+	draw.Draw(dst, dst.Bounds(), src, src.Bounds().Min, draw.Src)
 }
 
 // clone returns a copy of a grayscale image. It additionally corrects negative bounds and removes
 // dangling bytes from the underlaying pixel slice, if any are present.
-func clone(src *image.Gray) *image.Gray {
-	b := src.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
+func clone(dst, src *image.Gray) {
 	if dst.Stride == src.Stride {
 		// no need to correct stride, simply copy pixels.
 		copy(dst.Pix, src.Pix)
-		return dst
+		return
 	}
 	// need to correct stride.
-	for i := 0; i < b.Dy(); i++ {
+	for i := 0; i < src.Rect.Dy(); i++ {
 		dstH := i * dst.Stride
 		srcH := i * src.Stride
 		copy(dst.Pix[dstH:dstH+dst.Stride], src.Pix[srcH:srcH+dst.Stride])
 	}
-	return dst
 }
 
 // rgbToGray returns a grayscale value from alpha premultiplied red, green and blue values.
@@ -69,10 +70,8 @@ func rgbToGray(r, g, b uint32) uint8 {
 }
 
 // rgbaToGray converts an RGBA image to grayscale.
-func rgbaToGray(src *image.RGBA) *image.Gray {
-	b := src.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
-	concurrentIterate(b.Dy(), func(y int) {
+func rgbaToGray(dst *image.Gray, src *image.RGBA) {
+	concurrentIterate(src.Rect.Dy(), func(y int) {
 		for x := 0; x < dst.Stride; x++ {
 			i := y*src.Stride + x*4
 			s := src.Pix[i : i+3 : i+3]
@@ -87,14 +86,11 @@ func rgbaToGray(src *image.RGBA) *image.Gray {
 			dst.Pix[y*dst.Stride+x] = rgbToGray(r, g, b)
 		}
 	})
-	return dst
 }
 
 // rgba64ToGray lossily converts a 64 bit RGBA image to grayscale.
-func rgba64ToGray(src *image.RGBA64) *image.Gray {
-	b := src.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
-	concurrentIterate(b.Dy(), func(y int) {
+func rgba64ToGray(dst *image.Gray, src *image.RGBA64) {
+	concurrentIterate(src.Rect.Dy(), func(y int) {
 		for x := 0; x < dst.Stride; x++ {
 			i := y*src.Stride + x*8
 			s := src.Pix[i : i+6 : i+6]
@@ -106,14 +102,11 @@ func rgba64ToGray(src *image.RGBA64) *image.Gray {
 			dst.Pix[y*dst.Stride+x] = rgbToGray(r, g, b)
 		}
 	})
-	return dst
 }
 
 // nrgbaToGray converts an NRGBA image to grayscale.
-func nrgbaToGray(src *image.NRGBA) *image.Gray {
-	b := src.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
-	concurrentIterate(b.Dy(), func(y int) {
+func nrgbaToGray(dst *image.Gray, src *image.NRGBA) {
+	concurrentIterate(src.Rect.Dy(), func(y int) {
 		for x := 0; x < dst.Stride; x++ {
 			i := y*src.Stride + x*4
 			s := src.Pix[i : i+4 : i+4]
@@ -133,14 +126,11 @@ func nrgbaToGray(src *image.NRGBA) *image.Gray {
 			dst.Pix[y*dst.Stride+x] = rgbToGray(r, g, b)
 		}
 	})
-	return dst
 }
 
 // nrgba64ToGray lossily converts an NRGBA64 image to grayscale.
-func nrgba64ToGray(src *image.NRGBA64) *image.Gray {
-	b := src.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
-	concurrentIterate(b.Dy(), func(y int) {
+func nrgba64ToGray(dst *image.Gray, src *image.NRGBA64) {
+	concurrentIterate(src.Rect.Dy(), func(y int) {
 		for x := 0; x < dst.Stride; x++ {
 			i := y*src.Stride + x*8
 			s := src.Pix[i : i+8 : i+8]
@@ -156,17 +146,13 @@ func nrgba64ToGray(src *image.NRGBA64) *image.Gray {
 			dst.Pix[y*dst.Stride+x] = rgbToGray(r, g, b)
 		}
 	})
-	return dst
 }
 
 // ycbcrToGray converts a YCbCr image to grayscale.
-func ycbcrToGray(src *image.YCbCr) *image.Gray {
-	b := src.Bounds()
-	dst := image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
-	for i := 0; i < b.Dy(); i++ {
+func ycbcrToGray(dst *image.Gray, src *image.YCbCr) {
+	for i := 0; i < src.Rect.Dy(); i++ {
 		dstH := i * dst.Stride
 		srcH := i * src.YStride
 		copy(dst.Pix[dstH:dstH+dst.Stride], src.Y[srcH:srcH+dst.Stride])
 	}
-	return dst
 }

imgutil/grayscale_test.go

@@ -190,11 +190,7 @@ func BenchmarkGrayscale(b *testing.B) {
 		}
 		b.Run(bb.name, func(b *testing.B) {
 			for i := 0; i < b.N; i++ {
-				b.StopTimer()
-				img := cloneImg(bb.img)
-				b.StartTimer()
-
-				imgutil.Grayscale(img)
+				imgutil.Grayscale(bb.img)
 			}
 		})
 	}

imgutil/imgutil.go

@@ -131,11 +131,3 @@ func FitRect(rect image.Rectangle, x, y int) image.Rectangle {
 	}
 	return image.Rect(0, 0, int(math.Round(scale*width)), int(math.Round(scale*height)))
 }
-
-// Fit returns an image scaled to fit the specified bounding box without changing the aspect ratio.
-// It returns a copy of the image.
-func Fit(img *image.Gray, x, y int) *image.Gray {
-	dst := image.NewGray(FitRect(img.Rect, x, y))
-	CatmullRom.Scale(dst, img)
-	return dst
-}

imgutil/imgutil_test.go

@@ -190,50 +190,3 @@ func BenchmarkAutoContrast(b *testing.B) {
 		imgutil.AutoContrast(img, 1)
 	}
 }
-
-func TestFit(t *testing.T) {
-	tests := []struct {
-		name  string
-		image *image.Gray
-		x     int
-		y     int
-		want  *image.Gray
-	}{
-		{
-			name:  "fit height",
-			image: mustBeGray(mustReadImg("testdata/wikipe-tan-Gray.png")),
-			x:     1000,
-			y:     100,
-			want:  mustBeGray(mustReadImg("testdata/wikipe-tan-100h.png")),
-		},
-		{
-			name:  "fit width",
-			image: mustBeGray(mustReadImg("testdata/wikipe-tan-Gray.png")),
-			x:     100,
-			y:     1000,
-			want:  mustBeGray(mustReadImg("testdata/wikipe-tan-100w.png")),
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got := imgutil.Fit(tt.image, tt.x, tt.y)
-			if got != tt.want {
-				if !isWithinDeltaDiff(tt.want.Pix, got.Pix, 1) {
-					t.Errorf("Fit() difference above acceptable delta")
-				}
-			}
-		})
-	}
-}
-
-func BenchmarkFit(b *testing.B) {
-	src := mustBeGray(mustReadImg("testdata/wikipe-tan-Gray.png"))
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		b.StopTimer()
-		img := cloneGray(src)
-		b.StartTimer()
-
-		imgutil.Fit(img, 150, 150)
-	}
-}

imgutil/pool.go

@@ -0,0 +1,60 @@
+package imgutil
+
+import (
+	"image"
+	"sync"
+)
+
+// NewImagePool creates an ImagePool.
+func NewImagePool() *ImagePool {
+	return &ImagePool{
+		cache: make(map[int]*sync.Pool),
+	}
+}
+
+// ImagePool maintains a sync.Pool of pixel arrays for each image resolution gotten from it.
+type ImagePool struct {
+	cache map[int]*sync.Pool
+	mu    sync.Mutex
+}
+
+// GetFromImage converts an image into a grayscale image with pixel slice taken from the pool.
+func (p *ImagePool) GetFromImage(img image.Image) *image.Gray {
+	if i, ok := img.(*image.Gray); ok {
+		return i
+	}
+	dst := p.Get(img.Bounds().Dx(), img.Bounds().Dy())
+	grayscale(dst, img)
+	return dst
+}
+
+func (p *ImagePool) getPool(pixLen int) *sync.Pool {
+	p.mu.Lock()
+	pool, ok := p.cache[pixLen]
+	if !ok {
+		pool = &sync.Pool{
+			New: func() interface{} {
+				tmp := make([]uint8, pixLen)
+				return &tmp
+			},
+		}
+		p.cache[pixLen] = pool
+	}
+	p.mu.Unlock()
+	return pool
+}
+
+// Get gets a grayscale image of specified width and height with pixel slice taken from the pool.
+func (p *ImagePool) Get(width, height int) *image.Gray {
+	tmp := p.getPool(width * height).Get().(*[]uint8)
+	return &image.Gray{
+		Pix:    *tmp,
+		Stride: width,
+		Rect:   image.Rect(0, 0, width, height),
+	}
+}
+
+// Put puts an images pixel slice back into the pool.
+func (p *ImagePool) Put(img *image.Gray) {
+	p.getPool(len(img.Pix)).Put(&img.Pix)
+}

imgutil/scale.go

@@ -9,6 +9,42 @@ import (
 	"sync"
 )
 
+type cacheScaler struct {
+	kernel *Kernel
+	cache  map[cacheKey]Scaler
+	mu     sync.Mutex
+}
+
+type cacheKey struct {
+	dw, dh, sw, sh int
+}
+
+// Scale implements the Scaler interface.
+func (z *cacheScaler) Scale(dst, src *image.Gray) {
+	key := cacheKey{dst.Rect.Dx(), dst.Rect.Dy(), src.Rect.Dx(), src.Rect.Dy()}
+
+	z.mu.Lock()
+	scaler, ok := z.cache[key]
+	if !ok {
+		scaler = z.kernel.NewScaler(key.dw, key.dh, key.sw, key.sh)
+		z.cache[key] = scaler
+	}
+	z.mu.Unlock()
+
+	scaler.Scale(dst, src)
+}
+
+// NewCacheScaler creates, caches and reuses kernel scalers optimized for each unique combination
+// of destination and source width and height. It is mostly useful when scaling a large quantity of
+// images in a few fixed sizes.
+func NewCacheScaler(k *Kernel) Scaler {
+	return &cacheScaler{
+		kernel: k,
+		cache:  make(map[cacheKey]Scaler),
+		mu:     sync.Mutex{},
+	}
+}
+
 // Scaler scales the source image to the destination image.
 type Scaler interface {
 	Scale(dst, src *image.Gray)