First commit

internal/nms/nms.go

@@ -0,0 +1,45 @@
+package nms
+
+import (
+	"math"
+	"sort"
+)
+
+func Apply(boxes [][4]float64, scores []float64, iouThr float64) []int {
+	n := len(scores)
+	order := make([]int, n)
+	for i := range order {
+		order[i] = i
+	}
+	sort.Slice(order, func(i, j int) bool {
+		return scores[order[i]] > scores[order[j]]
+	})
+
+	keep := make([]int, 0)
+	for len(order) > 0 {
+		idx := order[0]
+		keep = append(keep, idx)
+		if len(order) == 1 {
+			break
+		}
+		rest := order[1:]
+		newOrder := make([]int, 0, len(rest))
+		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
+}