goyolov5

README

YOLOV5 for golang

Introduction

Basic example

package main

import (
	"github.com/danhilltech/goyolov5"
)

func main() {
	yolov5, err := goyolov5.NewYoloV5("yolov5n.torchscript.gpu.batch1.pt", DeviceGPU, 640, false)
	if err != nil {
		panic(err)
	}

	f, err := os.Open("path/to/my/image.png")
	if err != nil {
		panic(err)
	}
	defer f.Close()

	input, _, err := image.Decode(f)
	if err != nil {
		t.Fatal(err)
	}
	tensor := goyolov5.NewTensorFromImage(input)

	outTensor := goyolov5.NewTensorFromImage(tensor)

	predictions, err := yolov5.Infer(tensor, 0.5, 0.4, outTensor)
	if err != nil {
		t.Fatal(err)
	}
}

CUDA

Weights

This library uses traced torchscript versions of YOLOv5. Instructions on exporting can be found in the YOLOv5 repository. Alternatively, there's a simple dockerfile plus script to generate CPU and GPU versions of the v6 release models:

make weights

Developing

Inside the .devcontainer directory is an example devcontainer.json for use with VSCode. Duplicate the example and edit accordingly. Typically, this would mean adding or removing CUDA support depending on your hardware. E.g. add --gpus all to runArgs and -tags=cuda to gopls, testFlags, toolsEnvVars etc.

Tests

Basic test coverage run

go test

or

go test -tags=cuda

CUDA

Make sure you've installed nvidia-docker2 and the nvidia-container-toolkit.

Documentation

Index

• Constants
• Variables
• func FlattenDim(shape []int64) int
• func GetAndResetLastErr() *C.char
• func Iou(b1, b2 Bbox) (retVal float64)
• func TorchErr() error
• type Bbox
• type ByConfBbox
  • func (bb ByConfBbox) Len() int
  • func (bb ByConfBbox) Less(i, j int) bool
  • func (bb ByConfBbox) Swap(i, j int)
• type Cmodule
• type Ctensor
• type DeviceType
  • func DeviceCudaIfAvailable() DeviceType
• type Prediction
• type RGB
  • func (c RGB) RGBA() (r, g, b, a uint32)
• type Tensor
  • func NewTensor(r image.Rectangle) *Tensor
  • func NewTensorFromImage(i image.Image) *Tensor
  • func (p *Tensor) At(x, y int) color.Color
  • func (p *Tensor) Bounds() image.Rectangle
  • func (p *Tensor) ColorModel() color.Model
  • func (img *Tensor) DrawRect(rect image.Rectangle, col color.Color)
  • func (ts *Tensor) Drop() error
  • func (img *Tensor) HLine(x1, y, x2 int, col color.Color)
  • func (ts *Tensor) Numel() uint
  • func (p *Tensor) PixOffset(x, y int) int
  • func (p *Tensor) RGBAAt(x, y int) color.RGBA
  • func (img *Tensor) Resize(targetSize int) (*Tensor, float64, error)
  • func (p *Tensor) Set(x, y int, c color.Color)
  • func (ts *Tensor) Size() ([]int64, error)
  • func (p *Tensor) ToSquareShape() (*Tensor, int, int, error)
  • func (img *Tensor) VLine(x, y1, y2 int, col color.Color)
• type YoloV5
  • func NewYoloV5(path string, device DeviceType, size int, half bool) (*YoloV5, error)
  • func (yolov5 *YoloV5) Infer(inputTensorRaw *Tensor, confidenceThreshold float32, nmsThreshold float64, ...) ([][]Prediction, error)

Constants

const (
	ClassesOffset = 5
	NPreds        = 25200
	PredSize      = 85
	NClasses      = PredSize - ClassesOffset
)

const (
	COCO_PERSON = 0
)

Variables

var ColorModel = color.ModelFunc(rgbModel)

ColorModel is RGB color model instance

Functions

func FlattenDim

func FlattenDim(shape []int64) int

FlattenDim counts number of elements with given shape

func GetAndResetLastErr

func GetAndResetLastErr() *C.char

func Iou

func Iou(b1, b2 Bbox) (retVal float64)

Intersection over union of two bounding boxes.

func TorchErr

func TorchErr() error

TorchErr checks and retrieves last error message from C `thread_local` if existing and frees up C memory the C pointer points to.

NOTE: Go language atm does not have generic function something similar to `macro` in Rust language, does it? So we have to wrap this function to any Libtorch C function call to check error instead of doing the other way around. See Go2 proposal: https://github.com/golang/go/issues/32620

Types

type Bbox

type Bbox struct {
	// contains filtered or unexported fields
}

type ByConfBbox

type ByConfBbox []Bbox

func (ByConfBbox) Len

func (bb ByConfBbox) Len() int

Implement sort.Interface for []Bbox on Bbox.confidence: =====================================================

func (ByConfBbox) Less

func (bb ByConfBbox) Less(i, j int) bool

func (ByConfBbox) Swap

func (bb ByConfBbox) Swap(i, j int)

type Cmodule

type Cmodule = C.module

type Ctensor

type Ctensor = C.tensor

type DeviceType

type DeviceType = int32

const (
	DeviceCPU DeviceType = -1
	DeviceGPU DeviceType = 0
)

func DeviceCudaIfAvailable

func DeviceCudaIfAvailable() DeviceType

type Prediction

type Prediction struct {
	Rect            image.Rectangle
	Confidence      float64
	ClassIndex      uint
	ClassConfidence float64
}

type RGB

type RGB struct {
	R, G, B uint8
}

RGB color

func (RGB) RGBA

func (c RGB) RGBA() (r, g, b, a uint32)

RGBA implements Color.RGBA

type Tensor

type Tensor struct {
	// Pix holds the image's stream, in R, G, B order.
	Pix []uint8
	// Stride is the Pix stride (in bytes) between vertically adjacent pixels.
	Stride int
	// Rect is the image's bounds.
	Rect image.Rectangle
	// contains filtered or unexported fields
}

Tensor represent image data which has RGB colors. Tensor is compatible with image.RGBA, but does not have alpha channel to reduce using memory.

func NewTensor

func NewTensor(r image.Rectangle) *Tensor

NewTensor allocates and returns RGB image

func NewTensorFromImage

func NewTensorFromImage(i image.Image) *Tensor

func (*Tensor) At

func (p *Tensor) At(x, y int) color.Color

At implements image.Image.At

func (*Tensor) Bounds

func (p *Tensor) Bounds() image.Rectangle

Bounds implements image.Image.At

func (*Tensor) ColorModel

func (p *Tensor) ColorModel() color.Model

ColorModel returns RGB color model.

func (*Tensor) DrawRect

func (img *Tensor) DrawRect(rect image.Rectangle, col color.Color)

Rect draws a rectangle utilizing HLine() and VLine()

func (*Tensor) Drop

func (ts *Tensor) Drop() error

Drop drops (frees) the tensor

func (*Tensor) HLine

func (img *Tensor) HLine(x1, y, x2 int, col color.Color)

HLine draws a horizontal line

func (*Tensor) Numel

func (ts *Tensor) Numel() uint

Numel returns the total number of elements stored in a tensor.

func (*Tensor) PixOffset

func (p *Tensor) PixOffset(x, y int) int

PixOffset returns the index of the first element of Pix that corresponds to the pixel at (x, y).

func (*Tensor) RGBAAt

func (p *Tensor) RGBAAt(x, y int) color.RGBA

RGBAAt returns the color of the pixel at (x, y) as RGBA.

func (*Tensor) Resize

func (img *Tensor) Resize(targetSize int) (*Tensor, float64, error)

func (*Tensor) Set

func (p *Tensor) Set(x, y int, c color.Color)

func (*Tensor) Size

func (ts *Tensor) Size() ([]int64, error)

func (*Tensor) ToSquareShape

func (p *Tensor) ToSquareShape() (*Tensor, int, int, error)

func (*Tensor) VLine

func (img *Tensor) VLine(x, y1, y2 int, col color.Color)

VLine draws a veritcal line

type YoloV5

type YoloV5 struct {
	// contains filtered or unexported fields
}

func NewYoloV5

func NewYoloV5(path string, device DeviceType, size int, half bool) (*YoloV5, error)

func (*YoloV5) Infer

func (yolov5 *YoloV5) Infer(inputTensorRaw *Tensor, confidenceThreshold float32, nmsThreshold float64, annotateTensor *Tensor) ([][]Prediction, error)