yolov5

README

Go YOLO V5

This repository is a proof of concept and should at this moment be considered as work in progress.

Please do not use this in production as it's not properly tested yet.

Feel free to contribute!

Milestone v1.0.0

• Fix/Add tests
• Re-add filter functionality
• Re-think net interface
• Add proper mock definition for regen

Documentation

Overview

Package yolov5 provides a Go implementation of the YOLO V5 object detection system: https://pjreddie.com/darknet/yolo/.

The yolov5 package leverages gocv(https://github.com/hybridgroup/gocv) for a neural net able to detect object.

In order for the neural net to be able to detect objects, it needs the pre-trained network model consisting of a .cfg file and a .weights file. Using the Makefile provied by the library, these models can simply be downloaded by running 'make models'.

In order to use the package, make sure you've checked the prerequisites in the README: https://github.com/wimspaargaren/yolov5#prerequisites

Index

• Constants
• func DrawDetections(frame *gocv.Mat, detections []ObjectDetection)
• type Config
  • func DefaultConfig() Config
• type Net
  • func NewNet(modelPath, cocoNamePath string) (Net, error)
  • func NewNetWithConfig(modelPath, cocoNamePath string, config Config) (Net, error)
• type ObjectDetection

Examples

• NewNet
• NewNetWithConfig

Constants

const (
	DefaultRows        = 25200
	DefaultStepSize    = 85
	DefaultYoloVersion = 5
	DefaultInputWidth  = 1280
	DefaultInputHeight = 1280

	DefaultConfThreshold float32 = 0.5
	DefaultNMSThreshold  float32 = 0.4
)

Default constants for initialising the yolov5 net.

Functions

func DrawDetections

func DrawDetections(frame *gocv.Mat, detections []ObjectDetection)

DrawDetections draws a given list of object detections on a gocv Matrix.

Types

type Config

type Config struct {
	Rows     int
	StepSize int
	// InputWidth & InputHeight are used to determine the input size of the image for the network
	InputWidth  int
	InputHeight int
	YoloVersion int
	// ConfidenceThreshold can be used to determine the minimum confidence before an object is considered to be "detected"
	ConfidenceThreshold float32
	// Non-maximum suppression threshold used for removing overlapping bounding boxes
	NMSThreshold float32

	// Type on which the network will be executed
	NetTargetType  gocv.NetTargetType
	NetBackendType gocv.NetBackendType

	// NewNet function can be used to inject a custom neural net
	NewNet func(modelPath string) ml.NeuralNet
}

Config can be used to customise the settings of the neural network used for object detection.

func DefaultConfig

func DefaultConfig() Config

DefaultConfig used to create a working yolov5 net out of the box.

type Net

type Net interface {
	Close() error
	GetDetections(gocv.Mat) ([]ObjectDetection, error)
	GetDetectionsWithFilter(gocv.Mat, map[string]bool) ([]ObjectDetection, error)
}

Net the yolov5 net.

func NewNet

func NewNet(modelPath, cocoNamePath string) (Net, error)

NewNet creates new yolo net for given weight path, config and coconames list.

Example

yolov5Model := path.Join(os.Getenv("GOPATH"), "src/github.com/wimspaargaren/yolov5/data/yolov5/yolov5s.onnx")
cocoNamesPath := path.Join(os.Getenv("GOPATH"), "src/github.com/wimspaargaren/data/yolov5/coco.names")

yolonet, err := NewNet(yolov5Model, cocoNamesPath)
if err != nil {
	log.WithError(err).Fatal("unable to create yolo net")
}

// Gracefully close the net when the program is done
defer func() {
	err := yolonet.Close()
	if err != nil {
		log.WithError(err).Error("unable to gracefully close yolo net")
	}
}()

imagePath := path.Join(os.Getenv("GOPATH"), "src/github.com/wimspaargaren/yolov5/data/example_images/bird.jpg")
frame := gocv.IMRead(imagePath, gocv.IMReadColor)

detections, err := yolonet.GetDetections(frame)
if err != nil {
	log.WithError(err).Fatal("unable to retrieve predictions")
}

DrawDetections(&frame, detections)

window := gocv.NewWindow("Result Window")
defer func() {
	err := window.Close()
	if err != nil {
		log.WithError(err).Error("unable to close window")
	}
}()

window.IMShow(frame)
window.ResizeWindow(872, 585)

window.WaitKey(10000000000)

func NewNetWithConfig

func NewNetWithConfig(modelPath, cocoNamePath string, config Config) (Net, error)

NewNetWithConfig creates new yolo net with given config.

Example

yolov5Model := path.Join(os.Getenv("GOPATH"), "src/github.com/wimspaargaren/yolov5/data/yolov5/yolov5s.onnx")
cocoNamesPath := path.Join(os.Getenv("GOPATH"), "src/github.com/wimspaargaren/data/yolov5/coco.names")

conf := DefaultConfig()
// Set the neural net to use CUDA
conf.NetBackendType = gocv.NetBackendCUDA
conf.NetTargetType = gocv.NetTargetCUDA

yolonet, err := NewNetWithConfig(yolov5Model, cocoNamesPath, conf)
if err != nil {
	log.WithError(err).Fatal("unable to create yolo net")
}

// Gracefully close the net when the program is done
defer func() {
	err := yolonet.Close()
	if err != nil {
		log.WithError(err).Error("unable to gracefully close yolo net")
	}
}()

// ...

type ObjectDetection

type ObjectDetection struct {
	ClassID     int
	ClassName   string
	BoundingBox image.Rectangle
	Confidence  float32
}

ObjectDetection represents information of an object detected by the neural net.