Skip to content

Instantly share code, notes, and snippets.

@anisbhsl
Created March 20, 2020 08:22
Show Gist options
  • Save anisbhsl/7fb2dfac16c51ffc7bc4e0ebbf53437e to your computer and use it in GitHub Desktop.
Save anisbhsl/7fb2dfac16c51ffc7bc4e0ebbf53437e to your computer and use it in GitHub Desktop.
package main
import (
"bufio"
"gocv.io/x/gocv"
"image"
"image/color"
"log"
"os"
"fmt"
)
func main(){
input,_:=gocv.VideoCaptureDevice(0)
defer input.Close()
img:=gocv.NewMat()
defer img.Close()
model:=gocv.ReadNetFromTensorflow("./ssd_mobilenet/saved_model/saved_model.pb")
window:=gocv.NewWindow("Person Tracker")
defer window.Close()
descriptions, err := readDescriptions("./data/coco_labels.txt")
if err != nil {
log.Printf("Error reading descriptions file: %v\n", descriptions)
return
}
model.SetPreferableBackend(3) //openCV backend
model.SetPreferableTarget(0) //0 for CPU and 4 for NVIDIA GPU(needs configuration)
statusColor:=color.RGBA{
R: 0,
G: 255,
B: 0,
A: 0,
}
for{
if ok:=input.Read(&img);!ok{
log.Println("Device closed \n")
return
}
if img.Empty(){
continue
}
//convert image MATRIX to 224 x 224 blob
blob:=gocv.BlobFromImage(img,1.0,image.Pt(224,224),gocv.NewScalar(0, 0, 0, 0), true, false)
//feed blob into classifier
model.SetInput(blob,"input")
//Run a forward pass
prob:=model.Forward("softmax2")
//reshape results into 1*1000 matrix
probMat:=prob.Reshape(1,1)
//determine the most probable classificaton
_,maxVal,_,maxLoc:=gocv.MinMaxLoc(probMat)
//display classification
desc:="Unknown"
if maxLoc.X<1000{
desc=descriptions[maxLoc.X]
}
status := fmt.Sprintf("description: %v, maxVal: %v\n", desc, maxVal)
gocv.PutText(&img, status, image.Pt(10, 20), gocv.FontHersheyPlain, 1.2, statusColor, 2)
blob.Close()
prob.Close()
probMat.Close()
window.IMShow(img)
window.IMShow(img)
if window.WaitKey(1) >= 0 {
break
}
}
}
// readDescriptions reads the descriptions from a file
// and returns a slice of its lines.
func readDescriptions(path string) ([]string, error) {
file, err := os.Open(path)
if err != nil {
return nil, err
}
defer file.Close()
var lines []string
scanner := bufio.NewScanner(file)
for scanner.Scan() {
lines = append(lines, scanner.Text())
}
return lines, scanner.Err()
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment