implement avg in C

internal/pkg/stitch/auto.go

@@ -244,7 +244,7 @@ func (r *AutoStitcher) Frame(frameColor image.Image, ts time.Time) *Train {
 	}
 
 	// Check for minimal contrast and brightness.
-	avg, avgDev := avg.RGBA(frameRGBA)
+	avg, avgDev := avg.RGBAC(frameRGBA)
 	if sum3(avg)/3 < minContrastAvg || sum3(avgDev)/3 < minContrastAvgDev {
 		log.Trace().Interface("avgDev", avgDev).Interface("avg", avg).Msg("contrast too low, discarding")
 		return nil

pkg/avg/c.c

@@ -0,0 +1,55 @@
+#include "c.h"
+
+static const int four = 4;
+
+static int64_t iabs(int64_t a) {
+  if (a < 0) {
+    return -a;
+  }
+  return a;
+}
+
+void RGBAC(const int m, const int n, const int s,
+           /* pixels */
+           const uint8_t* const pix,
+           /* return parameters */
+           retData* ret) {
+  uint64_t sum[3] = {0};
+
+  for (int y = 0; y < n; y++) {
+    const int ys = y * s;
+    for (int x = 0; x < m; x++) {
+      const int ix = ys + x * four;
+      sum[0] += (uint64_t)(pix[ix + 0]);
+      sum[1] += (uint64_t)(pix[ix + 1]);
+      sum[2] += (uint64_t)(pix[ix + 2]);
+    }
+  }
+
+  const uint64_t cnt = (uint64_t)m * (uint64_t)n;
+  int64_t avgPx[3] = {
+      sum[0] / cnt,
+      sum[1] / cnt,
+      sum[2] / cnt,
+  };
+
+  sum[0] = sum[1] = sum[2] = 0;
+  for (int y = 0; y < n; y++) {
+    const int ys = y * s;
+    for (int x = 0; x < m; x++) {
+      const int ix = ys + x * four;
+
+      sum[0] += iabs((int64_t)pix[ix + 0] - avgPx[0]);
+      sum[1] += iabs((int64_t)pix[ix + 1] - avgPx[1]);
+      sum[2] += iabs((int64_t)pix[ix + 2] - avgPx[2]);
+    }
+  }
+
+  ret->avg[0] = (float64)avgPx[0] / 255.;
+  ret->avg[1] = (float64)avgPx[1] / 255.;
+  ret->avg[2] = (float64)avgPx[2] / 255.;
+
+  ret->avgDev[0] = (float64)sum[0] / (float64)cnt / 255.;
+  ret->avgDev[1] = (float64)sum[1] / (float64)cnt / 255.;
+  ret->avgDev[2] = (float64)sum[2] / (float64)cnt / 255.;
+}

pkg/avg/c.go

@@ -0,0 +1,40 @@
+package avg
+
+// See pkg/pmatch/c.go for docs on CC flags.
+
+// #cgo CFLAGS: -Wall -Wextra -pedantic -std=c99
+// #cgo CFLAGS: -O2
+//
+// #cgo amd64 CFLAGS: -march=native
+//
+// #cgo arm64 CFLAGS: -mcpu=cortex-a72 -mtune=cortex-a72
+//
+// #include "c.h"
+import "C"
+import "image"
+
+// RGBAC computes the pixel average, and pixel mean deviation from average,
+// on an RGBA image, per channel.
+// Note that the alpha channel is ignored.
+// Scaled to [0, 1].
+// Implemented in Cgo.
+func RGBAC(img *image.RGBA) ([3]float64, [3]float64) {
+
+	m, n := img.Bounds().Dx(), img.Bounds().Dy()
+	s := img.Stride
+
+	ret := C.retData{}
+	C.RGBAC(C.int(m), C.int(n), C.int(s),
+		(*C.uint8_t)(&img.Pix[0]), (*C.retData)(&ret))
+
+	return [3]float64{
+			float64(ret.avg[0]),
+			float64(ret.avg[1]),
+			float64(ret.avg[2]),
+		},
+		[3]float64{
+			float64(ret.avgDev[0]),
+			float64(ret.avgDev[1]),
+			float64(ret.avgDev[2]),
+		}
+}

pkg/avg/c.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <stdint.h>
+
+#define float64 double
+
+typedef struct retData {
+  float64 avg[3];
+  float64 avgDev[3];
+} retData;
+
+void RGBAC(const int m, const int n, const int s,
+           /* pixels */
+           const uint8_t* const pix,
+           /* return parameters */
+           retData* ret);

pkg/avg/c_test.go

@@ -0,0 +1,56 @@
+package avg
+
+import (
+	"testing"
+
+	"github.com/jo-m/trainbot/pkg/imutil"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_RGBAC(t *testing.T) {
+	high, err := imutil.Load("testdata/high.jpg")
+	require.NoError(t, err)
+	highRGB := imutil.ToRGBA(high)
+	avg, avgDev := RGBAC(highRGB)
+	assert.InDelta(t, 0.45, avg[0], 0.01)
+	assert.InDelta(t, 0.38, avg[1], 0.01)
+	assert.InDelta(t, 0.44, avg[2], 0.01)
+	assert.InDelta(t, 0.25, avgDev[0], 0.01)
+	assert.InDelta(t, 0.22, avgDev[1], 0.01)
+	assert.InDelta(t, 0.18, avgDev[2], 0.01)
+
+	mid, err := imutil.Load("testdata/mid.jpg")
+	require.NoError(t, err)
+	midRGB := imutil.ToRGBA(mid)
+	avg, avgDev = RGBAC(midRGB)
+	assert.InDelta(t, 0.019, avg[0], 0.001)
+	assert.InDelta(t, 0.015, avg[1], 0.001)
+	assert.InDelta(t, 0.007, avg[2], 0.001)
+	assert.InDelta(t, 0.023, avgDev[0], 0.001)
+	assert.InDelta(t, 0.016, avgDev[1], 0.001)
+	assert.InDelta(t, 0.009, avgDev[2], 0.001)
+
+	low, err := imutil.Load("testdata/low.jpg")
+	require.NoError(t, err)
+	lowRGB := imutil.ToRGBA(low)
+	avg, avgDev = RGBAC(lowRGB)
+	assert.InDelta(t, 0., avg[0], 0.0001)
+	assert.InDelta(t, 0., avg[1], 0.0001)
+	assert.InDelta(t, 0., avg[2], 0.0001)
+	assert.InDelta(t, 0.0038, avgDev[0], 0.0001)
+	assert.InDelta(t, 0.0029, avgDev[1], 0.0001)
+	assert.InDelta(t, 0.0022, avgDev[2], 0.0001)
+}
+
+func Benchmark_RGBAC(b *testing.B) {
+	high, err := imutil.Load("testdata/high.jpg")
+	if err != nil {
+		b.Error(err)
+	}
+	highRGB := imutil.ToRGBA(high)
+
+	for i := 0; i < b.N; i++ {
+		RGBAC(highRGB)
+	}
+}

pkg/pmatch/c.c

@@ -1,6 +1,6 @@
 #include "c.h"
 
-const int four = 4;
+static const int four = 4;
 
 void SearchRGBAC(const int m, const int n, const int du, const int dv,
                  const int is, const int ps,