mirror of https://bitbucket.org/ausocean/av.git
129 lines
3.6 KiB
Go
129 lines
3.6 KiB
Go
//go:build !nocv
|
|
// +build !nocv
|
|
|
|
/*
|
|
DESCRIPTION
|
|
Provides the methods for the turbidity probe using GoCV. Turbidity probe
|
|
will collect the most recent frames in a buffer and write the latest sharpness
|
|
and contrast scores to the probe.
|
|
|
|
AUTHORS
|
|
Russell Stanley <russell@ausocean.org>
|
|
|
|
LICENSE
|
|
Copyright (C) 2021-2022 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 (
|
|
"os"
|
|
"time"
|
|
|
|
"gocv.io/x/gocv"
|
|
"gonum.org/v1/gonum/stat"
|
|
|
|
"bitbucket.org/ausocean/av/turbidity"
|
|
"bitbucket.org/ausocean/utils/logger"
|
|
)
|
|
|
|
// Turbidity sensor constants.
|
|
const (
|
|
k1, k2 = 8, 8 // Block size, must be divisible by the size template with no remainder.
|
|
filterSize = 3 // Sobel filter size.
|
|
scale = 1.0 // Amount of scale applied to sobel filter values.
|
|
alpha = 1.0 // Paramater for contrast equation.
|
|
)
|
|
|
|
// Misc constants.
|
|
const (
|
|
maxImages = 10 // Max number of images read when evaluating turbidity.
|
|
)
|
|
|
|
type turbidityProbe struct {
|
|
sharpness, contrast float64
|
|
delay time.Duration
|
|
ticker time.Ticker
|
|
ts *turbidity.TurbiditySensor
|
|
log logger.Logger
|
|
}
|
|
|
|
// NewTurbidityProbe returns a new turbidity probe.
|
|
func NewTurbidityProbe(log logger.Logger, delay time.Duration) (*turbidityProbe, error) {
|
|
tp := new(turbidityProbe)
|
|
tp.log = log
|
|
tp.delay = delay
|
|
tp.ticker = *time.NewTicker(delay)
|
|
|
|
// Create the turbidity sensor.
|
|
standard := gocv.IMRead("../../turbidity/images/template.jpg", gocv.IMReadGrayScale)
|
|
template := gocv.IMRead("../../turbidity/images/template.jpg", gocv.IMReadGrayScale)
|
|
ts, err := turbidity.NewTurbiditySensor(template, standard, k1, k2, filterSize, scale, alpha)
|
|
if err != nil {
|
|
log.Error("failed create turbidity sensor", "error", err.Error())
|
|
}
|
|
tp.ts = ts
|
|
return tp, nil
|
|
}
|
|
|
|
// Write, reads input h264 frames in the form of a byte stream and writes the the sharpness and contrast
|
|
// scores of a video to the the turbidity probe.
|
|
func (tp *turbidityProbe) Write(p []byte) (int, error) {
|
|
select {
|
|
case <-tp.ticker.C:
|
|
var imgs []gocv.Mat
|
|
img := gocv.NewMat()
|
|
|
|
// Write byte array to a temp file.
|
|
file, err := os.CreateTemp("temp", "video*.h264")
|
|
if err != nil {
|
|
tp.log.Error("failed to create temp file", "error", err.Error())
|
|
return 0, err
|
|
}
|
|
defer os.Remove(file.Name())
|
|
n, err := file.Write(p)
|
|
if err != nil {
|
|
tp.log.Error("failed to write to temporary file", "error", err.Error())
|
|
return n, err
|
|
}
|
|
|
|
// Read the file and store each frame.
|
|
vc, err := gocv.VideoCaptureFile(file.Name())
|
|
if err != nil {
|
|
tp.log.Error("failed to open video file", "error", err.Error())
|
|
return len(p), err
|
|
}
|
|
for vc.Read(&img) && len(imgs) < maxImages {
|
|
imgs = append(imgs, img.Clone())
|
|
}
|
|
|
|
// Process video data to get saturation and contrast scores.
|
|
res, err := tp.ts.Evaluate(imgs)
|
|
if err != nil {
|
|
tp.log.Error("evaluate failed", "errror", err.Error())
|
|
return len(p), err
|
|
}
|
|
tp.contrast = stat.Mean(res.Contrast, nil)
|
|
tp.sharpness = stat.Mean(res.Sharpness, nil)
|
|
default:
|
|
}
|
|
return len(p), nil
|
|
}
|
|
|
|
func (tp *turbidityProbe) Close() error {
|
|
return nil
|
|
}
|