First commit

internal/iou/iou.go

@@ -0,0 +1,110 @@
+package iou
+
+import (
+	"math"
+	"sort"
+)
+
+func NMS(boxes [][4]float64, scores []float64, iouThr float64) []int {
+	n := len(scores)
+	idx := make([]int, n)
+	for i := range idx {
+		idx[i] = i
+	}
+	sort.Slice(idx, func(i, j int) bool {
+		return scores[idx[i]] > scores[idx[j]]
+	})
+
+	keep := make([]int, 0, n)
+	for len(idx) > 0 {
+		cur := idx[0]
+		keep = append(keep, cur)
+		if len(idx) == 1 {
+			break
+		}
+		rest := idx[1:]
+		newOrder := make([]int, 0, len(rest))
+		bCur := boxes[cur]
+		areaCur := (bCur[2] - bCur[0]) * (bCur[3] - bCur[1])
+		for _, ri := range rest {
+			bRi := boxes[ri]
+			xx1 := math.Max(bCur[0], bRi[0])
+			yy1 := math.Max(bCur[1], bRi[1])
+			xx2 := math.Min(bCur[2], bRi[2])
+			yy2 := math.Min(bCur[3], bRi[3])
+			w := math.Max(0, xx2-xx1)
+			h := math.Max(0, yy2-yy1)
+			inter := w * h
+			areaRi := (bRi[2] - bRi[0]) * (bRi[3] - bRi[1])
+			ov := inter / (areaCur + areaRi - inter + 1e-16)
+			if ov < iouThr {
+				newOrder = append(newOrder, ri)
+			}
+		}
+		idx = newOrder
+	}
+	return keep
+}
+
+func PairwiseIoU(boxesA, boxesB [][4]float64) [][]float64 {
+	n, m := len(boxesA), len(boxesB)
+	mat := make([][]float64, n)
+	for i := range mat {
+		mat[i] = make([]float64, m)
+	}
+	if n == 0 || m == 0 {
+		return mat
+	}
+	for i := 0; i < n; i++ {
+		areaA := (boxesA[i][2] - boxesA[i][0]) * (boxesA[i][3] - boxesA[i][1])
+		for j := 0; j < m; j++ {
+			xx1 := math.Max(boxesA[i][0], boxesB[j][0])
+			yy1 := math.Max(boxesA[i][1], boxesB[j][1])
+			xx2 := math.Min(boxesA[i][2], boxesB[j][2])
+			yy2 := math.Min(boxesA[i][3], boxesB[j][3])
+			iw := math.Max(0, xx2-xx1)
+			ih := math.Max(0, yy2-yy1)
+			inter := iw * ih
+			areaB := (boxesB[j][2] - boxesB[j][0]) * (boxesB[j][3] - boxesB[j][1])
+			mat[i][j] = inter / (areaA + areaB - inter + 1e-16)
+		}
+	}
+	return mat
+}
+
+func GreedyMatchRowCol(costMatrix [][]float64, threshold float64) []struct{ Row, Col int } {
+	if len(costMatrix) == 0 || len(costMatrix[0]) == 0 {
+		return nil
+	}
+	n, m := len(costMatrix), len(costMatrix[0])
+
+	type entry struct {
+		cost     float64
+		row, col int
+	}
+	entries := make([]entry, 0, n*m)
+	for i := 0; i < n; i++ {
+		for j := 0; j < m; j++ {
+			entries = append(entries, entry{costMatrix[i][j], i, j})
+		}
+	}
+	sort.Slice(entries, func(a, b int) bool {
+		return entries[a].cost < entries[b].cost
+	})
+
+	rowUsed := make(map[int]bool)
+	colUsed := make(map[int]bool)
+	var pairs []struct{ Row, Col int }
+	for _, e := range entries {
+		if e.cost >= threshold {
+			break
+		}
+		if rowUsed[e.row] || colUsed[e.col] {
+			continue
+		}
+		rowUsed[e.row] = true
+		colUsed[e.col] = true
+		pairs = append(pairs, struct{ Row, Col int }{e.row, e.col})
+	}
+	return pairs
+}