46 lines
1.0 KiB
Go
46 lines
1.0 KiB
Go
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
|
|
}
|