code cleanup, improved corner detection in transform function, fixed some comments

turbidity/plot.go

@@ -1,12 +1,9 @@
-//go:build !nocv
-// +build !nocv
-
 /*
 DESCRIPTION
-Turbidity is a program to measure water clarity using computer vison
+  Plotting functions for the turbidity sensor results.
 
 AUTHORS
-Russell Stanley <russell@ausocean.org>
+  Russell Stanley <russell@ausocean.org>
 
 LICENSE
   Copyright (C) 2020 the Australian Ocean Lab (AusOcean)
@@ -29,79 +26,25 @@ package main
 
 import (
 	"fmt"
-	"log"
 	"math"
 
 	"gonum.org/v1/plot"
 	"gonum.org/v1/plot/plotter"
 	"gonum.org/v1/plot/plotutil"
 	"gonum.org/v1/plot/vg"
-
-	"gocv.io/x/gocv"
 )
 
-const (
-	nImages  = 13
-	nSamples = 10
-)
-
-func main() {
-	// Load template and standard image.
-	template := gocv.IMRead("template.jpg", gocv.IMReadGrayScale)
-	standard := gocv.IMRead("default.jpg", gocv.IMReadGrayScale)
-
-	imgs := make([][]gocv.Mat, nImages)
-
-	// Load test images.
-	for i := range imgs {
-		imgs[i] = make([]gocv.Mat, nSamples)
-		for j := range imgs[i] {
-			imgs[i][j] = gocv.IMRead(fmt.Sprintf("images/t-%v/000%v.jpg", i, j), gocv.IMReadGrayScale)
-		}
-	}
-
-	// Create turbidity sensor.
-	ts := TurbiditySensor{template: template, standard: standard, k1: 8, k2: 8, sobelFilterSize: 3, scale: 1.0, alpha: 1.0}
-
-	var finalRes Results
-	finalRes.new(nImages)
-
-	// Score each image by calculating the average score from camera burst.
-	for i := range imgs {
-		// Evaluate camera burst.
-		sample_result, err := ts.Evaluate(imgs[i])
-		if err != nil {
-			log.Fatalf("Evaluation Failed: %v", err)
-		}
-
-		// Add the average result from camera burst.
-		finalRes.update(average(sample_result.saturation), average(sample_result.contrast), float64(i)*2.5, i)
-	}
-
-	// Plot the final results.
-	err := plotResults(finalRes.turbidity, normalize(finalRes.saturation), normalize(finalRes.contrast))
-	if err != nil {
-		log.Fatalf("Plotting Failed: %v", err)
-	}
-
-	log.Printf("Saturation: %v", finalRes.saturation)
-	log.Printf("Contrast: %v", finalRes.contrast)
-}
-
-// Plotting Functions.
-
 // Normalize values in a slice between 0 and 1.
 func normalize(slice []float64) []float64 {
-
 	max := -math.MaxFloat64
 	min := math.MaxFloat64
-
 	out := make([]float64, len(slice))
 
 	if len(slice) <= 1 {
 		return slice
 	}
 
+	// Find the max and min values of the slice.
 	for i := range slice {
 		if slice[i] > max {
 			max = slice[i]
@@ -119,9 +62,9 @@ func normalize(slice []float64) []float64 {
 
 // Return the average of a slice.
 func average(slice []float64) float64 {
+	var out float64
 
-	out := 0.0
+	// Sum all elements in the slice.
-
 	for i := range slice {
 		out += slice[i]
 	}
@@ -129,7 +72,6 @@ func average(slice []float64) float64 {
 }
 
 func plotResults(x, saturation, contrast []float64) error {
-
 	err := plotToFile(
 		"Results",
 		"Almond Milk (ml)",

turbidity/results.go

@@ -3,10 +3,10 @@
 
 /*
 DESCRIPTION
-Results struct used to store results from the turbidity sensor
+  Results struct used to store results from the turbidity sensor.
 
 AUTHORS
-Russell Stanley <russell@ausocean.org>
+  Russell Stanley <russell@ausocean.org>
 
 LICENSE
   Copyright (C) 2020 the Australian Ocean Lab (AusOcean)
@@ -27,22 +27,33 @@ LICENSE
 
 package main
 
-// struct to hold the results of the turbidity sensor.
+import "fmt"
+
+// Results holds the results of the turbidity sensor.
 type Results struct {
-	turbidity  []float64
+	Turbidity  []float64
-	saturation []float64
+	Saturation []float64
-	contrast   []float64
+	Contrast   []float64
 }
 
-func (r *Results) new(n int) {
+// NewResults constructs the results object.
-	r.turbidity = make([]float64, n)
+func NewResults(n int) (*Results, error) {
-	r.saturation = make([]float64, n)
+
-	r.contrast = make([]float64, n)
+	if n <= 0 {
+		return nil, fmt.Errorf("invalid result size: %v.", n)
+	}
+
+	r := new(Results)
+	r.Turbidity = make([]float64, n)
+	r.Saturation = make([]float64, n)
+	r.Contrast = make([]float64, n)
+
+	return r, nil
 }
 
-// Update results to add new values at specified index.
+// Update adds new values to slice at specified index.
-func (r *Results) update(saturation, contrast, turbidity float64, index int) {
+func (r *Results) Update(newSat, newCont, newTurb float64, index int) {
-	r.saturation[index] = saturation
+	r.Saturation[index] = newSat
-	r.contrast[index] = contrast
+	r.Contrast[index] = newCont
-	r.turbidity[index] = turbidity
+	r.Turbidity[index] = newTurb
 }

turbidity/turbidity.go

@@ -3,11 +3,11 @@
 
 /*
 DESCRIPTION
-Holds the turbidity sensor struct. Can evaluate 4x4 chessboard markers
+  Holds the turbidity sensor struct. Can evaluate 4x4 chessboard markers
-in an image to measure the sharpness and contrast.
+  in an image to measure the sharpness and contrast.
 
 AUTHORS
-Russell Stanley <russell@ausocean.org>
+  Russell Stanley <russell@ausocean.org>
 
 LICENSE
   Copyright (C) 2020 the Australian Ocean Lab (AusOcean)
@@ -37,89 +37,116 @@ import (
 	"gocv.io/x/gocv"
 )
 
-// Turbidity Sensor.
+// TurbiditySensor is a software based turbidity sensor that uses CV to determine saturation and constrast level
+// of a chessboard-like target submerged in water that can be correlated to turbidity/visibility values.
 type TurbiditySensor struct {
-	template, standard      gocv.Mat
+	template, templateCorners gocv.Mat
+	standard, standardCorners gocv.Mat
 	k1, k2, sobelFilterSize   int
-	alpha, scale            float64
+	scale, alpha              float64
 }
 
-// Given a slice of test images, return the sharpness and contrast scores.
+// NewTurbiditySensor constructor for a turbidity sensor.
-func (ts TurbiditySensor) Evaluate(imgs []gocv.Mat) (Results, error) {
+func NewTurbiditySensor(template, standard gocv.Mat, k1, k2, sobelFilterSize int, scale, alpha float64) (*TurbiditySensor, error) {
+	ts := new(TurbiditySensor)
+	templateCorners := gocv.NewMat()
+	standardCorners := gocv.NewMat()
 
-	var result Results
+	// Validate template image is not empty and has valid corners.
-	result.new(len(imgs))
+	if template.Empty() {
-
+		return nil, errors.New("template image is empty.")
-	for i := range imgs {
-
-		// Transform image.
-		marker, err := ts.Transform(imgs[i])
-		if err != nil {
-			return result, fmt.Errorf("Image %v: %w", i, err)
 	}
+	if !gocv.FindChessboardCorners(template, image.Pt(3, 3), &templateCorners, gocv.CalibCBNormalizeImage) {
+		return nil, errors.New("could not find corners in template image")
+	}
+	ts.template = template
+	ts.templateCorners = templateCorners
 
-		// Apply sobel filter.
+	// Validate standard image is not empty and has valid corners.
-		edge := ts.Sobel(marker)
+	if standard.Empty() {
+		return nil, errors.New("standard image is empty.")
+	}
+	if !gocv.FindChessboardCorners(standard, image.Pt(3, 3), &standardCorners, gocv.CalibCBNormalizeImage) {
+		return nil, errors.New("could not find corners in standard image")
+	}
+	ts.standard = standard
+	ts.standardCorners = standardCorners
 
-		// Evaluate image.
+	ts.k1, ts.k2, ts.sobelFilterSize = k1, k2, sobelFilterSize
-		scores, err := ts.EvaluateImage(marker, edge)
+	ts.alpha, ts.scale = alpha, scale
+
+	return ts, nil
+}
+
+// Evaluate, given a slice of images, return the sharpness and contrast scores.
+func (ts TurbiditySensor) Evaluate(imgs []gocv.Mat) (*Results, error) {
+	result, err := NewResults(len(imgs))
 	if err != nil {
 		return result, err
 	}
 
-		result.update(scores[0], scores[1], float64(i*10), i)
+	for i := range imgs {
+		// Transform image.
+		marker, err := ts.transform(imgs[i])
+		if err != nil {
+			return result, fmt.Errorf("image %v: %w", i, err)
+		}
+
+		// Apply sobel filter.
+		edge := ts.sobel(marker)
+
+		// Evaluate image.
+		sharpScore, contScore, err := ts.EvaluateImage(marker, edge)
+		if err != nil {
+			return result, err
+		}
+
+		result.Update(sharpScore, contScore, float64(i), i)
 	}
 
 	return result, nil
 }
 
-// Evaluate image sharpness and contrast using blocks of size k1 by k2. Return a slice of the respective scores.
+// EvaluateImage will evaluate image sharpness and contrast using blocks of size k1 by k2. Return the respective scores.
-func (ts TurbiditySensor) EvaluateImage(img, edge gocv.Mat) ([]float64, error) {
+func (ts TurbiditySensor) EvaluateImage(img, edge gocv.Mat) (float64, float64, error) {
-
+	var sharpness float64
-	result := make([]float64, 2) // [0.0, 0.0]
+	var contrast float64
 
 	if img.Rows()%ts.k1 != 0 || img.Cols()%ts.k2 != 0 {
-		return nil, fmt.Errorf("Dimensions not compatible (%v, %v)", ts.k1, ts.k2)
+		return math.NaN(), math.NaN(), fmt.Errorf("dimensions not compatible (%v, %v)", ts.k1, ts.k2)
 	}
 
-	lStep := int(img.Rows() / ts.k1)
+	lStep := img.Rows() / ts.k1
-	kStep := int(img.Cols() / ts.k2)
+	kStep := img.Cols() / ts.k2
-
-	for l := 0; l < img.Rows(); l = l + lStep {
-		for k := 0; k < img.Cols(); k = k + kStep {
 
+	for l := 0; l < img.Rows(); l += lStep {
+		for k := 0; k < img.Cols(); k += kStep {
 			// Enhancement Measure Estimation (EME), provides a measure of the sharpness.
-			err := ts.EvaluateBlock(edge, l, k, l+lStep, k+kStep, result, "EME", ts.alpha)
+			sharpValue := ts.evaluateBlockEME(edge, l, k, l+lStep, k+kStep)
-			if err != nil {
+			sharpness += sharpValue
-				return nil, err
-			}
 
 			// AMEE, provides a measure of the contrast.
-			err = ts.EvaluateBlock(img, l, k, l+lStep, k+kStep, result, "AMEE", ts.alpha)
+			contValue := ts.evaluateBlockAMEE(img, l, k, l+lStep, k+kStep)
-			if err != nil {
+			contrast += contValue
-				return nil, err
-			}
 		}
 	}
 
-	// EME.
+	// Scale EME based on block size.
-	result[0] = 2.0 / (float64(ts.k1) * float64(ts.k2)) * result[0]
+	sharpness = 2.0 / (float64(ts.k1 * ts.k2)) * sharpness
 
-	// AMEE.
+	// Scale and flip AMEE based on block size.
-	result[1] = -1.0 / (float64(ts.k1) * float64(ts.k2)) * result[1]
+	contrast = -1.0 / (float64(ts.k1 * ts.k2)) * contrast
 
-	return result, nil
+	return sharpness, contrast, nil
 }
 
-// Evaluate a block within an image and add to to the result slice.
+// minMax returns the max and min pixel values of an image block.
-func (ts TurbiditySensor) EvaluateBlock(img gocv.Mat, xStart, yStart, xEnd, yEnd int, result []float64, operation string, alpha float64) error {
+func (ts TurbiditySensor) minMax(img gocv.Mat, xStart, yStart, xEnd, yEnd int) (float64, float64) {
-
 	max := -math.MaxFloat64
 	min := math.MaxFloat64
 
 	for i := xStart; i < xEnd; i++ {
 		for j := yStart; j < yEnd; j++ {
-
 			value := float64(img.GetUCharAt(i, j))
 
 			// Check max/min conditions, zero values are ignored.
@@ -131,52 +158,56 @@ func (ts TurbiditySensor) EvaluateBlock(img gocv.Mat, xStart, yStart, xEnd, yEnd
 			}
 		}
 	}
+	return max, min
+}
+
+// evaluateBlockEME will evaluate an image block and return the value to be added to the sharpness result.
+func (ts TurbiditySensor) evaluateBlockEME(img gocv.Mat, xStart, yStart, xEnd, yEnd int) float64 {
+	max, min := ts.minMax(img, xStart, yStart, xEnd, yEnd)
 
 	// Blocks which have no information are ignored.
 	if max != -math.MaxFloat64 && min != math.MaxFloat64 && max != min {
-		if operation == "EME" {
+		return math.Log(max / min)
-			result[0] += math.Log(max / min)
-		} else if operation == "AMEE" {
-			contrast := (max + min) / (max - min)
-			result[1] += math.Pow(alpha*(contrast), alpha) * math.Log(contrast)
-		} else {
-			return fmt.Errorf("Invalid operation: %v", operation)
 	}
-	}
+	return 0.0
-	return nil
 }
 
-// Search image for matching template. Returns the transformed image  which best match the template.
+// evaluateBlockAMEE will evaluate an image block and return the value to be added to the contrast result.
-func (ts TurbiditySensor) Transform(img gocv.Mat) (gocv.Mat, error) {
+func (ts TurbiditySensor) evaluateBlockAMEE(img gocv.Mat, xStart, yStart, xEnd, yEnd int) float64 {
+	max, min := ts.minMax(img, xStart, yStart, xEnd, yEnd)
+
+	// Blocks which have no information are ignored.
+	if max != -math.MaxFloat64 && min != math.MaxFloat64 && max != min {
+		contrast := (max + min) / (max - min)
+		return math.Pow(ts.alpha*(contrast), ts.alpha) * math.Log(contrast)
+	}
+	return 0.0
+}
+
+// transform will search img for matching template. Returns the transformed image which best match the template.
+func (ts TurbiditySensor) transform(img gocv.Mat) (gocv.Mat, error) {
 	out := gocv.NewMat()
 	mask := gocv.NewMat()
-	corners_img := gocv.NewMat()
+	imgCorners := gocv.NewMat()
-	corners_template := gocv.NewMat()
 
-	// Find corners in image.
+	// Check image is valid.
-	if !gocv.FindChessboardCorners(ts.standard, image.Pt(3, 3), &corners_img, gocv.CalibCBNormalizeImage) {
+	if img.Empty() {
-		// Apply default if transformation fails.
+		return out, errors.New("image is empty, cannot transform")
-		fmt.Println("Corner detection failed applying standard transformation")
-		if !gocv.FindChessboardCorners(ts.standard, image.Pt(3, 3), &corners_img, gocv.CalibCBNormalizeImage) {
-			return out, errors.New("Could not find corners in default image")
 	}
-	}
+	// Check image for corners, if non can be found corners will be set to default value.
-
+	if !gocv.FindChessboardCorners(img, image.Pt(3, 3), &imgCorners, gocv.CalibCBFastCheck) {
-	// Find corners in template.
+		imgCorners = ts.standardCorners
-	if !gocv.FindChessboardCorners(ts.template, image.Pt(3, 3), &corners_template, gocv.CalibCBNormalizeImage) {
-		return out, errors.New("Could not find corners in template")
 	}
 
 	// Find and apply transformation.
-	H := gocv.FindHomography(corners_img, &corners_template, gocv.HomograpyMethodRANSAC, 3.0, &mask, 2000, 0.995)
+	H := gocv.FindHomography(imgCorners, &ts.templateCorners, gocv.HomograpyMethodRANSAC, 3.0, &mask, 2000, 0.995)
 	gocv.WarpPerspective(img, &out, H, image.Pt(ts.template.Rows(), ts.template.Cols()))
 
 	return out, nil
 }
 
-// Apply sobel filter to an image with a given scale and return the result.
+// sobel will apply sobel filter to an image and return the result.
-func (ts TurbiditySensor) Sobel(img gocv.Mat) gocv.Mat {
+func (ts TurbiditySensor) sobel(img gocv.Mat) gocv.Mat {
-
 	dx := gocv.NewMat()
 	dy := gocv.NewMat()
 	sobel := gocv.NewMat()

turbidity/turbidity_test.go

@@ -0,0 +1,88 @@
+/*
+DESCRIPTION
+  Testing functions for the turbidity sensor using images from
+  previous experiment.
+
+AUTHORS
+  Russell Stanley <russell@ausocean.org>
+
+LICENSE
+  Copyright (C) 2020 the Australian Ocean Lab (AusOcean)
+
+  It is free software: you can redistribute it and/or modify them
+  under the terms of the GNU General Public License as published by the
+  Free Software Foundation, either version 3 of the License, or (at your
+  option) any later version.
+
+  It is distributed in the hope that it will be useful, but WITHOUT
+  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
+  for more details.
+
+  You should have received a copy of the GNU General Public License
+  in gpl.txt.  If not, see http://www.gnu.org/licenses.
+*/
+
+package main
+
+import (
+	"fmt"
+	"testing"
+
+	"gocv.io/x/gocv"
+)
+
+const (
+	nImages   = 13 // Number of images to test. (Max 13)
+	nSamples  = 10 // Number of samples for each image. (Max 10)
+	increment = 2.5
+)
+
+func TestImages(t *testing.T) {
+	// Load template and standard image.
+	template := gocv.IMRead("images/template.jpg", gocv.IMReadGrayScale)
+	standard := gocv.IMRead("images/default.jpg", gocv.IMReadGrayScale)
+
+	imgs := make([][]gocv.Mat, nImages)
+
+	// Load test images.
+	for i := range imgs {
+		imgs[i] = make([]gocv.Mat, nSamples)
+		for j := range imgs[i] {
+			imgs[i][j] = gocv.IMRead(fmt.Sprintf("images/t-%v/000%v.jpg", i, j), gocv.IMReadGrayScale)
+		}
+	}
+
+	// Create turbidity sensor.
+	ts, err := NewTurbiditySensor(template, standard, 8, 8, 3, 1.0, 1.0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create results.
+	results, err := NewResults(nImages)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Score each image by calculating the average score from camera burst.
+	for i := range imgs {
+		// Evaluate camera burst.
+		sample_result, err := ts.Evaluate(imgs[i])
+		if err != nil {
+			t.Fatalf("Evaluation Failed: %v", err)
+		}
+
+		// Add the average result from camera burst.
+		results.Update(average(sample_result.Saturation), average(sample_result.Contrast), float64(i)*increment, i)
+	}
+
+	// Plot the final results.
+	err = plotResults(results.Turbidity, normalize(results.Saturation), normalize(results.Contrast))
+	if err != nil {
+		t.Fatalf("Plotting Failed: %v", err)
+	}
+
+	t.Logf("Saturation: %v", results.Saturation)
+	t.Logf("Contrast: %v", results.Contrast)
+}