av/turbidity/main.go

// +build !nocv

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
	}

	for i := range slice {
		if slice[i] > max {
			max = slice[i]
		}
		if slice[i] < min {
			min = slice[i]
		}
	}

	for i := range slice {
		out[i] = (slice[i] - min) / (max - min)
	}
	return out
}

// Return the average of a slice.
func average(slice []float64) float64 {

	out := 0.0

	for i := range slice {
		out += slice[i]
	}
	return out / float64(len(slice))
}

func plotResults(x, saturation, contrast []float64) error {

	err := plotToFile(
		"Results",
		"Almond Milk (ml)",
		"Score",
		func(p *plot.Plot) error {
			return plotutil.AddLinePoints(p,
				"Contrast", plotterXY(x, contrast),
				"Saturation", plotterXY(x, saturation),
			)
		},
	)
	if err != nil {
		return fmt.Errorf("Could not plot results: %w", err)
	}
	return nil
}

// plotToFile creates a plot with a specified name and x&y titles using the
// provided draw function, and then saves to a PNG file with filename of name.
func plotToFile(name, xTitle, yTitle string, draw func(*plot.Plot) error) error {
	p := plot.New()
	p.Title.Text = name
	p.X.Label.Text = xTitle
	p.Y.Label.Text = yTitle
	err := draw(p)
	if err != nil {
		return fmt.Errorf("could not draw plot contents: %w", err)
	}
	if err := p.Save(15*vg.Centimeter, 15*vg.Centimeter, "plots/"+name+".png"); err != nil {
		return fmt.Errorf("could not save plot: %w", err)
	}
	return nil
}

// plotterXY provides a plotter.XYs type value based on the given x and y data.
func plotterXY(x, y []float64) plotter.XYs {
	xy := make(plotter.XYs, len(x))
	for i := range x {
		xy[i].X = x[i]
		xy[i].Y = y[i]
	}
	return xy
}