-
Notifications
You must be signed in to change notification settings - Fork 0
/
face_recognition.go
127 lines (105 loc) · 3.24 KB
/
face_recognition.go
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
package facerecognition
import (
"github.com/Kagami/go-face"
"github.com/anish-krishnan/Tidepod/tidepod-server/entity"
"gorm.io/gorm"
)
// ClassifyFacesByBoxEngine trains on already labelled faces, and
// classifies all other photos
func ClassifyFacesByBoxEngine(db *gorm.DB, boxes []*entity.Box) map[int]string {
// Mapping photoIDs to respective boxes that are in train or test
var trainSet []int
var testSet []int
var faceMap map[int]string = make(map[int]string)
for j, box := range boxes {
if box.Face.ID != 0 {
trainSet = append(trainSet, j)
faceMap[box.Face.ID] = box.Face.Name
} else {
testSet = append(testSet, j)
}
}
rec, err := face.NewRecognizer("./workflow")
if err != nil {
panic(err)
}
defer rec.Close()
var samples []face.Descriptor
var labels []int32
for _, boxIndex := range trainSet {
face, err := rec.RecognizeSingleFile("./photo_storage/boxes/" + boxes[boxIndex].FilePath)
if face == nil {
continue
} else if err != nil {
panic(err)
}
samples = append(samples, face.Descriptor)
labels = append(labels, int32(boxes[boxIndex].Face.ID))
}
rec.SetSamples(samples, labels)
var result map[int]string = make(map[int]string)
for _, boxIndex := range testSet {
face, err := rec.RecognizeSingleFile("./photo_storage/boxes/" + boxes[boxIndex].FilePath)
if face == nil {
continue
} else if err != nil {
panic(err)
}
label := rec.ClassifyThreshold(face.Descriptor, 0.2)
if label > 0 {
result[boxes[boxIndex].ID] = faceMap[label]
}
}
return result
}
// ClassifyFacesByPhotoEngine trains on already labelled faces, and classifies all other photos
func ClassifyFacesByPhotoEngine(db *gorm.DB, photos []*entity.Photo) {
// Mapping photoIDs to respective boxes that are in train or test
var trainSet map[int][]int = make(map[int][]int)
var testSet map[int][]int = make(map[int][]int)
var photoMap map[int]*entity.Photo = make(map[int]*entity.Photo)
var faceMap map[int]string = make(map[int]string)
for _, photo := range photos {
for j, box := range photo.Boxes {
photoMap[photo.ID] = photo
if box.Face.ID != 0 {
trainSet[photo.ID] = append(trainSet[photo.ID], j)
} else {
testSet[photo.ID] = append(testSet[photo.ID], j)
}
}
}
rec, err := face.NewRecognizer(".")
if err != nil {
panic(err)
}
defer rec.Close()
var samples []face.Descriptor
var labels []int32
for photoID, boxes := range trainSet {
faces, err := rec.RecognizeFile("./photo_storage/saved/" + photoMap[photoID].FilePath)
if err != nil {
panic(err)
}
for _, boxIndex := range boxes {
samples = append(samples, faces[boxIndex].Descriptor)
labels = append(labels, int32(photoMap[photoID].Boxes[boxIndex].Face.ID))
faceMap[photoMap[photoID].Boxes[boxIndex].Face.ID] = photoMap[photoID].Boxes[boxIndex].Face.Name
}
}
rec.SetSamples(samples, labels)
var result map[int]string = make(map[int]string)
for photoID, boxes := range testSet {
for _, boxIndex := range boxes {
face, err := rec.RecognizeSingleFile("./photo_storage/boxes/" + photoMap[photoID].Boxes[boxIndex].FilePath)
if err != nil {
panic(err)
}
if face == nil {
continue
}
label := rec.ClassifyThreshold(face.Descriptor, 0.2)
result[photoMap[photoID].Boxes[boxIndex].ID] = faceMap[label]
}
}
}