improved NMSBoxes code

hybridgroup · May 14, 2023 · fd3de44 · fd3de44
1 parent cca4917
commit fd3de44
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/dnn.go b/dnn.go
@@ -522,7 +522,7 @@ func (l *Layer) OutputNameToIndex(name string) int {
 //
 // For futher details, please see:
 // https://docs.opencv.org/4.4.0/d6/d0f/group__dnn.html#ga9d118d70a1659af729d01b10233213ee
-func NMSBoxes(bboxes []image.Rectangle, scores []float32, scoreThreshold float32, nmsThreshold float32, indices []int) {
+func NMSBoxes(bboxes []image.Rectangle, scores []float32, scoreThreshold float32, nmsThreshold float32) (indices []int) {
  bboxesRectArr := []C.struct_Rect{}
  for _, v := range bboxes {
  bbox := C.struct_Rect{
@@ -560,6 +560,8 @@ func NMSBoxes(bboxes []image.Rectangle, scores []float32, scoreThreshold float32
 
  ptr := *(*[]C.int)(unsafe.Pointer(h))
 
+
+ indices = make([]int,indicesVector.length)
  for i := 0; i < int(indicesVector.length); i++ {
  indices[i] = int(ptr[i])
  }
@@ -570,7 +572,7 @@ func NMSBoxes(bboxes []image.Rectangle, scores []float32, scoreThreshold float32
 //
 // For futher details, please see:
 // https://docs.opencv.org/4.4.0/d6/d0f/group__dnn.html#ga9d118d70a1659af729d01b10233213ee
-func NMSBoxesWithParams(bboxes []image.Rectangle, scores []float32, scoreThreshold float32, nmsThreshold float32, indices []int, eta float32, topK int) {
+func NMSBoxesWithParams(bboxes []image.Rectangle, scores []float32, scoreThreshold float32, nmsThreshold float32, eta float32, topK int) (indices []int) {
  bboxesRectArr := []C.struct_Rect{}
  for _, v := range bboxes {
  bbox := C.struct_Rect{
@@ -608,6 +610,7 @@ func NMSBoxesWithParams(bboxes []image.Rectangle, scores []float32, scoreThresho
 
  ptr := *(*[]C.int)(unsafe.Pointer(h))
 
+ indices = make([]int,indicesVector.length)
  for i := 0; i < int(indicesVector.length); i++ {
  indices[i] = int(ptr[i])
  }

diff --git a/dnn_test.go b/dnn_test.go
@@ -591,11 +591,10 @@ func TestNMSBoxes(t *testing.T) {
  image.Rect(156, 51, 640, 480),
  }
  scores := []float32{0.82094115, 0.7998236, 0.9809663, 0.99717456, 0.89628726}
- indices := make([]int, 10)
  scoreThreshold := float32(0.5)
  nmsThreshold := float32(0.4)
 
- NMSBoxes(bboxes, scores, scoreThreshold, nmsThreshold, indices)
+ indices := NMSBoxes(bboxes, scores, scoreThreshold, nmsThreshold)
 
  if indices[0] != 3 {
  t.Errorf("Invalid NMSBoxes test indices: %v", indices)
@@ -619,11 +618,10 @@ func TestNMSBoxesWithParams(t *testing.T) {
  image.Rect(156, 51, 640, 480),
  }
  scores := []float32{0.82094115, 0.7998236, 0.9809663, 0.99717456, 0.89628726}
- indices := make([]int, 10)
  scoreThreshold := float32(0.5)
  nmsThreshold := float32(0.4)
 
- NMSBoxesWithParams(bboxes, scores, scoreThreshold, nmsThreshold, indices, float32(1.0), 0)
+ indices := NMSBoxesWithParams(bboxes, scores, scoreThreshold, nmsThreshold, float32(1.0), 0)
 
  if indices[0] != 3 {
  t.Errorf("Invalid NMSBoxesWithParams test indices: %v", indices)