goreco/libprocessor.go at master · nefeg/goreco · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
package goreco

import (
	"fmt"
	"gocv.io/x/gocv"
)

type processor struct {
	Processor
	process func(detectionResult gocv.Mat, threshold float32, all bool) []Box
}

func (p *processor) Process(detectionResult gocv.Mat, threshold float32, all bool) []Box {
	return p.process(detectionResult, threshold, all)
}

func NewProcessorSSD() Processor {
	p := &processor{}
	p.process = postProcessSSD
	return p
}

func NewProcessorYOLO() Processor {
	p := &processor{}
	p.process = postProcessYOLO
	return p
}

//func (p *processor)Process(detectionResult gocv.Mat, threshold float32 ) []Box{
//
//}

// performDetection analyzes the results from the detector network,
// which produces an output blob with a shape 1x1xNx7
// where N is the number of detections, and each detection
// is a vector of float values
// [batchId, classId, confidence, left, top, right, bottom]
func postProcessSSD(results gocv.Mat, threshold float32, all bool) (b []Box) {

	for i := 0; i < results.Total(); i += 7 {

		confidence := results.GetFloatAt(0, i+2)
		classId := int(results.GetFloatAt(0, i+1))

		//fmt.Println(results.Cols(), results.Size())
		//fmt.Println(results.GetFloatAt(0, i+1), results.GetFloatAt(0, i+2), results.GetFloatAt(0, i+3), results.GetFloatAt(0, i+4))
		//fmt.Println(results.GetFloatAt(0, i+5), results.GetFloatAt(0, i+6), results.GetFloatAt(0, i+7), results.GetFloatAt(0, i+8))
		//fmt.Println(results.GetFloatAt(0, i+9), results.GetFloatAt(0, i+10), results.GetFloatAt(0, i+11), results.GetFloatAt(0, i+12))

		//if confidence > threshold {
		if confidence > threshold {
			if !all && classId != 1 {
				continue
			}
			box := Box{
				class: classId,
				rect: struct {
					min [2]float32
					max [2]float32
				}{
					[2]float32{results.GetFloatAt(0, i+3), results.GetFloatAt(0, i+4)},
					[2]float32{results.GetFloatAt(0, i+5), results.GetFloatAt(0, i+6)},
				},
				conf: confidence,
			}
			//fmt.Println(box)
			b = append(b, box)
		}
	}
	return b
}

/**
The outputs object are vectors of lenght 85

4x the bounding box (centerx, centery, width, height)
1x box confidence
80x class confidence
*/
func postProcessYOLO(results gocv.Mat, threshold float32, all bool) (b []Box) {

	fts, _ := results.DataPtrFloat32()

	fmt.Println(len(fts))

	fmt.Println(results.Size(), results.Total())
	fmt.Println(results.GetFloatAt(0, 0), results.GetFloatAt(0, 1), results.GetFloatAt(0, 2), results.GetFloatAt(0, 3))
	fmt.Println(results.GetFloatAt(0, 4), results.GetFloatAt(0, 5))
	threshold = 0.2

	for i := 0; i < results.Total(); i += 85 {
		tensor := fts[i : i+85]

		for i, v := range tensor[5:] {
			if v != 0 {
				fmt.Println(i, v)
			}
		}
		//boxConfidence := tensor[4]
		classConfidence, classId := FindMaxFloat32(tensor[5:])

		if classId == 0 {
			//if classId ==0 ||classConfidence <threshold{
			continue
		}

		box := Box{
			class: classId,
			rect: struct {
				min [2]float32
				max [2]float32
			}{
				[2]float32{tensor[0] - tensor[2]/2, tensor[1] - tensor[3]/2},
				[2]float32{tensor[0] + tensor[2]/2, tensor[1] + tensor[3]/2},
			},
			conf: classConfidence,
		}

		fmt.Println(box)
		b = append(b, box)

		//fmt.Println( boxConfidence, box, classId, classConf )

		//confidence := results.GetFloatAt(0, i+2)
		//classId := int(results.GetFloatAt(0, i+1))
		//
		//fmt.Println(confidence, classId)

		//if confidence > threshold {
		//	//if confidence > threshold && classId ==1 {
		//	box := Box{
		//		class: classId,
		//		rect: struct{
		//			min	[2]float32
		//			max [2]float32
		//		}{
		//			[2]float32{results.GetFloatAt(0, i+3), results.GetFloatAt(0, i+4)},
		//			[2]float32{results.GetFloatAt(0, i+5), results.GetFloatAt(0, i+6)},
		//		},
		//		conf: confidence,
		//	}
		//	fmt.Println(box)
		//	b =  append( b, box )
		//}
	}
	return b
}