First commit

pkg/yolo/yolo.go

@@ -0,0 +1,236 @@
+package yolo
+
+import (
+	"image"
+	"math"
+
+	"github.com/anomalyco/bytetrack-counter-go/pkg/rknn"
+	"gocv.io/x/gocv"
+)
+
+type Detection struct {
+	BBox      [4]float64
+	Score     float64
+	Class     int
+	Keypoints [][2]float64
+}
+
+type Detector struct {
+	rknn         *rknn.Context
+	imgSz        int
+	conf         float64
+	iouThr       float64
+	numClasses   int
+	scoreSigmoid bool
+}
+
+func NewDetector(modelPath string, imgSz int, conf float64, numClasses int, scoreSigmoid bool) (*Detector, error) {
+	ctx, err := rknn.LoadModel(modelPath)
+	if err != nil {
+		return nil, err
+	}
+	return &Detector{
+		rknn:         ctx,
+		imgSz:        imgSz,
+		conf:         conf,
+		iouThr:       0.45,
+		numClasses:   numClasses,
+		scoreSigmoid: scoreSigmoid,
+	}, nil
+}
+
+func (d *Detector) Detect(frame gocv.Mat) ([]Detection, error) {
+	h0, w0 := frame.Rows(), frame.Cols()
+
+	scale := math.Min(float64(d.imgSz)/float64(h0), float64(d.imgSz)/float64(w0))
+	nh, nw := int(float64(h0)*scale), int(float64(w0)*scale)
+
+	resized := gocv.NewMat()
+	defer resized.Close()
+	gocv.Resize(frame, &resized, image.Point{X: nw, Y: nh}, 0, 0, gocv.InterpolationLinear)
+
+	letterbox := gocv.NewMatWithSize(d.imgSz, d.imgSz, gocv.MatTypeCV8UC3)
+	defer letterbox.Close()
+	letterbox.SetTo(gocv.NewScalar(114, 114, 114, 0))
+
+	dy := (d.imgSz - nh) / 2
+	dx := (d.imgSz - nw) / 2
+	roi := letterbox.Region(image.Rect(dx, dy, dx+nw, dy+nh))
+	resized.CopyTo(&roi)
+
+	rgb := gocv.NewMat()
+	defer rgb.Close()
+	gocv.CvtColor(letterbox, &rgb, gocv.ColorBGRToRGB)
+
+	data := rgb.ToBytes()
+
+	outputs, err := d.rknn.InferenceRGB(data, d.imgSz, d.imgSz)
+	if err != nil {
+		return nil, err
+	}
+	if len(outputs) == 0 {
+		return nil, nil
+	}
+
+	out := outputs[0]
+	return d.decodeOutput(out, h0, w0, scale, float64(dy), float64(dx))
+}
+
+func (d *Detector) decodeOutput(out []float32, h0, w0 int, scale, padY, padX float64) ([]Detection, error) {
+	if len(out) == 0 {
+		return nil, nil
+	}
+
+	stride := d.numClasses + 4
+	if len(out)%stride != 0 {
+		return nil, nil
+	}
+
+	numDet := len(out) / stride
+	if numDet == 0 {
+		return nil, nil
+	}
+
+	var detections []Detection
+
+	for i := 0; i < numDet; i++ {
+		off := i * stride
+
+		cx := float64(out[off+0])
+		cy := float64(out[off+1])
+		w := float64(out[off+2])
+		h := float64(out[off+3])
+
+		x1 := cx - w/2
+		y1 := cy - h/2
+		x2 := cx + w/2
+		y2 := cy + h/2
+
+		var maxScore float64
+		var bestClass int
+		for c := 0; c < d.numClasses; c++ {
+			score := float64(out[off+4+c])
+			if d.scoreSigmoid {
+				score = sigmoid(clamp(score, -10, 10))
+			}
+			if score > maxScore {
+				maxScore = score
+				bestClass = c
+			}
+		}
+
+		if maxScore <= d.conf {
+			continue
+		}
+
+		x1 = (x1 - padX) / scale
+		y1 = (y1 - padY) / scale
+		x2 = (x2 - padX) / scale
+		y2 = (y2 - padY) / scale
+
+		x1 = clamp(x1, 0, float64(w0))
+		y1 = clamp(y1, 0, float64(h0))
+		x2 = clamp(x2, 0, float64(w0))
+		y2 = clamp(y2, 0, float64(h0))
+
+		detections = append(detections, Detection{
+			BBox:  [4]float64{x1, y1, x2, y2},
+			Score: maxScore,
+			Class: bestClass,
+		})
+	}
+
+	return nmsByClass(detections, d.iouThr, d.numClasses), nil
+}
+
+func nmsByClass(dets []Detection, iouThr float64, numClasses int) []Detection {
+	if len(dets) == 0 {
+		return nil
+	}
+
+	byClass := make([][]int, numClasses)
+	for i, d := range dets {
+		byClass[d.Class] = append(byClass[d.Class], i)
+	}
+
+	var result []Detection
+	for cls := 0; cls < numClasses; cls++ {
+		idxs := byClass[cls]
+		if len(idxs) == 0 {
+			continue
+		}
+
+		scores := make([]float64, len(idxs))
+		boxes := make([][4]float64, len(idxs))
+		for i, idx := range idxs {
+			scores[i] = dets[idx].Score
+			boxes[i] = dets[idx].BBox
+		}
+
+		keep := nmsIndices(boxes, scores, iouThr)
+		for _, k := range keep {
+			result = append(result, dets[idxs[k]])
+		}
+	}
+	return result
+}
+
+func nmsIndices(boxes [][4]float64, scores []float64, iouThr float64) []int {
+	if len(scores) == 0 {
+		return nil
+	}
+
+	order := make([]int, len(scores))
+	for i := range order {
+		order[i] = i
+	}
+
+	for i := 0; i < len(scores); i++ {
+		for j := i + 1; j < len(scores); j++ {
+			if scores[order[i]] < scores[order[j]] {
+				order[i], order[j] = order[j], order[i]
+			}
+		}
+	}
+
+	var keep []int
+	for len(order) > 0 {
+		idx := order[0]
+		keep = append(keep, idx)
+		if len(order) == 1 {
+			break
+		}
+
+		rest := order[1:]
+		var newOrder []int
+		for _, r := range rest {
+			xx1 := math.Max(boxes[idx][0], boxes[r][0])
+			yy1 := math.Max(boxes[idx][1], boxes[r][1])
+			xx2 := math.Min(boxes[idx][2], boxes[r][2])
+			yy2 := math.Min(boxes[idx][3], boxes[r][3])
+			w := math.Max(0, xx2-xx1)
+			h := math.Max(0, yy2-yy1)
+			inter := w * h
+			areaI := (boxes[idx][2] - boxes[idx][0]) * (boxes[idx][3] - boxes[idx][1])
+			areaR := (boxes[r][2] - boxes[r][0]) * (boxes[r][3] - boxes[r][1])
+			iou := inter / (areaI + areaR - inter + 1e-16)
+			if iou < iouThr {
+				newOrder = append(newOrder, r)
+			}
+		}
+		order = newOrder
+	}
+	return keep
+}
+
+func sigmoid(x float64) float64 {
+	return 1.0 / (1.0 + math.Exp(-x))
+}
+
+func clamp(x, lo, hi float64) float64 {
+	return math.Max(lo, math.Min(hi, x))
+}
+
+func (d *Detector) Release() {
+	d.rknn.Release()
+}