111 lines
2.6 KiB
Go
111 lines
2.6 KiB
Go
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
|
|
}
|