├── .gitignore ├── obama.jpg ├── reg_embeddings.npy ├── go.mod ├── LICENSE ├── model_converter.py ├── iresnet.py ├── utils.go ├── README.md ├── go.sum └── main.go /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | tmp* 3 | -------------------------------------------------------------------------------- /obama.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/modanesh/GoFaceRec/HEAD/obama.jpg -------------------------------------------------------------------------------- /reg_embeddings.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/modanesh/GoFaceRec/HEAD/reg_embeddings.npy -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module test 2 | 3 | go 1.17 4 | 5 | require ( 6 | github.com/galeone/tensorflow/tensorflow/go v0.0.0-20221023090153-6b7fa0680c3e 7 | github.com/galeone/tfgo v0.0.0-20221023090852-d89a5c7e31e1 8 | github.com/sbinet/npyio v0.7.0 9 | gocv.io/x/gocv v0.32.1 10 | gonum.org/v1/gonum v0.13.0 11 | ) 12 | 13 | require ( 14 | github.com/campoy/embedmd v1.0.0 // indirect 15 | github.com/pmezard/go-difflib v1.0.0 // indirect 16 | google.golang.org/protobuf v1.28.1 // indirect 17 | ) 18 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Mohamad H. Danesh 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /model_converter.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os.path 3 | import sys 4 | from collections import OrderedDict 5 | from collections import namedtuple 6 | 7 | import numpy as np 8 | import onnx 9 | import onnxruntime 10 | import torch 11 | import torch.nn as nn 12 | from google.protobuf.json_format import MessageToDict 13 | from onnx_tf.backend import prepare 14 | 15 | from preprocessing.magface import iresnet 16 | 17 | 18 | class NetworkBuilder_inf(nn.Module): 19 | def __init__(self, args): 20 | super(NetworkBuilder_inf, self).__init__() 21 | self.features = iresnet.iresnet100(pretrained=False, 22 | num_classes=args.embedding_size) 23 | 24 | def forward(self, input): 25 | # add Fp, a pose feature 26 | x = self.features(input) 27 | return x 28 | 29 | 30 | def arg_parser(): 31 | parser = argparse.ArgumentParser() 32 | parser.add_argument('--model_path', required=True) 33 | parser.add_argument('--img_size', required=True) 34 | args = parser.parse_args() 35 | return args 36 | 37 | 38 | def to_numpy(tensor): 39 | return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() 40 | 41 | 42 | def load_dict_inf(args, model): 43 | if os.path.isfile(args.resume): 44 | print('=> loading pth from {} ...'.format(args.resume)) 45 | if args.cpu_mode: 46 | checkpoint = torch.load(args.resume, map_location=torch.device("cpu")) 47 | else: 48 | checkpoint = torch.load(args.resume) 49 | _state_dict = clean_dict_inf(model, checkpoint['state_dict']) 50 | model_dict = model.state_dict() 51 | model_dict.update(_state_dict) 52 | model.load_state_dict(model_dict) 53 | # delete to release more space 54 | del checkpoint 55 | del _state_dict 56 | else: 57 | sys.exit("=> No checkpoint found at '{}'".format(args.resume)) 58 | return model 59 | 60 | 61 | def clean_dict_inf(model, state_dict): 62 | _state_dict = OrderedDict() 63 | for k, v in state_dict.items(): 64 | # # assert k[0:1] == 'features.module.' 65 | new_k = 'features.' + '.'.join(k.split('.')[2:]) 66 | if new_k in model.state_dict().keys() and \ 67 | v.size() == model.state_dict()[new_k].size(): 68 | _state_dict[new_k] = v 69 | # assert k[0:1] == 'module.features.' 70 | new_kk = '.'.join(k.split('.')[1:]) 71 | if new_kk in model.state_dict().keys() and \ 72 | v.size() == model.state_dict()[new_kk].size(): 73 | _state_dict[new_kk] = v 74 | num_model = len(model.state_dict().keys()) 75 | num_ckpt = len(_state_dict.keys()) 76 | if num_model != num_ckpt: 77 | sys.exit("=> Not all weights loaded, model params: {}, loaded params: {}".format( 78 | num_model, num_ckpt)) 79 | return _state_dict 80 | 81 | 82 | if __name__ == '__main__': 83 | args = arg_parser() 84 | torch_model_path = args.model_path + '.pth' 85 | onnx_model_path = args.model_path + '.onnx' 86 | 87 | dummy = torch.randn(64, 3, args.img_size, args.img_size) 88 | 89 | Torch_args = namedtuple('Args', ['arch', 'resume', 'embedding_size', 'cpu_mode']) 90 | torch_args = Torch_args('iresnet100', torch_model_path, 512, True) 91 | 92 | torch_model = NetworkBuilder_inf(torch_args) 93 | torch_model = load_dict_inf(torch_args, torch_model) 94 | 95 | torch_model.eval() 96 | torch_out = torch_model(dummy) 97 | if not os.path.exists(onnx_model_path): 98 | torch.onnx.export(torch_model, dummy, onnx_model_path) 99 | print("Torch model converted to ONNX and saved to file!") 100 | else: 101 | print("Torch model was already converted to ONNX!") 102 | 103 | # ONNX 104 | onnx_model = onnx.load(onnx_model_path) 105 | onnx.checker.check_model(onnx_model) 106 | print("ONNX model checked and it is fine!") 107 | 108 | print("ONNX graph configurations:") 109 | for _input in onnx_model.graph.input: 110 | print(MessageToDict(_input)) 111 | ort_session = onnxruntime.InferenceSession(onnx_model_path) 112 | 113 | # compute ONNX Runtime output prediction 114 | ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy)} 115 | ort_outs = ort_session.run(None, ort_inputs) 116 | 117 | # compare ONNX Runtime and PyTorch results 118 | np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05) 119 | print("Exported ONNX model has been tested with ONNXRuntime, and the result looks good!") 120 | 121 | # TensorFlow 122 | print("Converting the ONNX model to TF:") 123 | tf_rep = prepare(onnx_model) 124 | tf_rep.export_graph(args.model_path + "_pb") 125 | print("TF model saved!") 126 | -------------------------------------------------------------------------------- /iresnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | # from torchvision.models.utils import load_state_dict_from_url 4 | 5 | 6 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): 7 | """3x3 convolution with padding""" 8 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 9 | padding=dilation, groups=groups, bias=False, dilation=dilation) 10 | 11 | 12 | def conv1x1(in_planes, out_planes, stride=1): 13 | """1x1 convolution""" 14 | return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) 15 | 16 | 17 | class IBasicBlock(nn.Module): 18 | expansion = 1 19 | 20 | def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, 21 | base_width=64, dilation=1): 22 | super(IBasicBlock, self).__init__() 23 | if groups != 1 or base_width != 64: 24 | raise ValueError( 25 | 'BasicBlock only supports groups=1 and base_width=64') 26 | if dilation > 1: 27 | raise NotImplementedError( 28 | "Dilation > 1 not supported in BasicBlock") 29 | # Both self.conv1 and self.downsample layers downsample the input when stride != 1 30 | self.bn1 = nn.BatchNorm2d(inplanes, eps=2e-05, momentum=0.9) 31 | self.conv1 = conv3x3(inplanes, planes) 32 | self.bn2 = nn.BatchNorm2d(planes, eps=2e-05, momentum=0.9) 33 | self.prelu = nn.PReLU(planes) 34 | self.conv2 = conv3x3(planes, planes, stride) 35 | self.bn3 = nn.BatchNorm2d(planes, eps=2e-05, momentum=0.9) 36 | self.downsample = downsample 37 | self.stride = stride 38 | 39 | def forward(self, x): 40 | identity = x 41 | 42 | out = self.bn1(x) 43 | out = self.conv1(out) 44 | out = self.bn2(out) 45 | out = self.prelu(out) 46 | out = self.conv2(out) 47 | out = self.bn3(out) 48 | 49 | if self.downsample is not None: 50 | identity = self.downsample(x) 51 | 52 | out += identity 53 | 54 | return out 55 | 56 | 57 | class IResNet(nn.Module): 58 | fc_scale = 7 * 7 59 | 60 | def __init__(self, block, layers, num_classes=512, zero_init_residual=False, 61 | groups=1, width_per_group=64, replace_stride_with_dilation=None): 62 | super(IResNet, self).__init__() 63 | 64 | self.inplanes = 64 65 | self.dilation = 1 66 | if replace_stride_with_dilation is None: 67 | # each element in the tuple indicates if we should replace 68 | # the 2x2 stride with a dilated convolution instead 69 | replace_stride_with_dilation = [False, False, False] 70 | if len(replace_stride_with_dilation) != 3: 71 | raise ValueError("replace_stride_with_dilation should be None " 72 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) 73 | self.groups = groups 74 | self.base_width = width_per_group 75 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, 76 | bias=False) 77 | self.bn1 = nn.BatchNorm2d(self.inplanes, eps=2e-05, momentum=0.9) 78 | self.prelu = nn.PReLU(self.inplanes) 79 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2) 80 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, 81 | dilate=replace_stride_with_dilation[0]) 82 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, 83 | dilate=replace_stride_with_dilation[1]) 84 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, 85 | dilate=replace_stride_with_dilation[2]) 86 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 87 | 88 | self.bn2 = nn.BatchNorm2d( 89 | 512 * block.expansion, eps=2e-05, momentum=0.9) 90 | self.dropout = nn.Dropout2d(p=0.4, inplace=True) 91 | self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_classes) 92 | self.features = nn.BatchNorm1d(num_classes, eps=2e-05, momentum=0.9) 93 | 94 | for m in self.modules(): 95 | if isinstance(m, nn.Conv2d): 96 | nn.init.kaiming_normal_( 97 | m.weight, mode='fan_out', nonlinearity='relu') 98 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): 99 | nn.init.constant_(m.weight, 1) 100 | nn.init.constant_(m.bias, 0) 101 | 102 | if zero_init_residual: 103 | for m in self.modules(): 104 | if isinstance(m, IBasicBlock): 105 | nn.init.constant_(m.bn2.weight, 0) 106 | 107 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 108 | downsample = None 109 | previous_dilation = self.dilation 110 | if dilate: 111 | self.dilation *= stride 112 | stride = 1 113 | if stride != 1 or self.inplanes != planes * block.expansion: 114 | downsample = nn.Sequential( 115 | conv1x1(self.inplanes, planes * block.expansion, stride), 116 | nn.BatchNorm2d(planes * block.expansion, 117 | eps=2e-05, momentum=0.9), 118 | ) 119 | 120 | layers = [] 121 | layers.append(block(self.inplanes, planes, stride, downsample, self.groups, 122 | self.base_width, previous_dilation)) 123 | self.inplanes = planes * block.expansion 124 | for _ in range(1, blocks): 125 | layers.append(block(self.inplanes, planes, groups=self.groups, 126 | base_width=self.base_width, dilation=self.dilation)) 127 | 128 | return nn.Sequential(*layers) 129 | 130 | def forward(self, x): 131 | x = self.conv1(x) 132 | x = self.bn1(x) 133 | x = self.prelu(x) 134 | 135 | x = self.layer1(x) 136 | x = self.layer2(x) 137 | x = self.layer3(x) 138 | x = self.layer4(x) 139 | 140 | x = self.bn2(x) 141 | x = self.dropout(x) 142 | x = x.view(x.size(0), -1) 143 | x = self.fc(x) 144 | x = self.features(x) 145 | 146 | return x 147 | 148 | 149 | def _iresnet(arch, block, layers, pretrained, progress, **kwargs): 150 | model = IResNet(block, layers, **kwargs) 151 | # if pretrained: 152 | # state_dict = load_state_dict_from_url(model_urls[arch], 153 | # progress=progress) 154 | # model.load_state_dict(state_dict) 155 | return model 156 | 157 | 158 | def iresnet100(pretrained=False, progress=True, **kwargs): 159 | return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained, progress, 160 | **kwargs) 161 | -------------------------------------------------------------------------------- /utils.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | tf "github.com/galeone/tensorflow/tensorflow/go" 7 | "gocv.io/x/gocv" 8 | "image" 9 | "image/color" 10 | "io/ioutil" 11 | "math" 12 | "os" 13 | "reflect" 14 | "strings" 15 | ) 16 | 17 | // FloatImage represents a custom image type that satisfies the image.Image interface 18 | type FloatImage struct { 19 | data [][]float32 20 | } 21 | 22 | func (f FloatImage) ColorModel() color.Model { 23 | return color.Gray16Model 24 | } 25 | 26 | func (f FloatImage) Bounds() image.Rectangle { 27 | height := len(f.data) 28 | width := len(f.data[0]) 29 | return image.Rect(0, 0, width, height) 30 | } 31 | 32 | func (f FloatImage) At(x, y int) color.Color { 33 | return color.Gray16{Y: uint16(f.data[y][x])} 34 | } 35 | 36 | // ConvertFloatImage converts [][]float32 to FloatImage 37 | func ConvertFloatImage(data [][]float32) *FloatImage { 38 | return &FloatImage{data: data} 39 | } 40 | 41 | func matToFloat32Slice(mat gocv.Mat) [][]float32 { 42 | rows, cols := mat.Rows(), mat.Cols() 43 | slice := make([][]float32, rows) 44 | 45 | for i := 0; i < rows; i++ { 46 | rowSlice := make([]float32, cols) 47 | for j := 0; j < cols; j++ { 48 | val := mat.GetFloatAt(i, j) 49 | rowSlice[j] = float32(val) 50 | } 51 | slice[i] = rowSlice 52 | } 53 | return slice 54 | } 55 | 56 | func findExtremeValue(slice interface{}, operation string) float32 { 57 | switch reflect.TypeOf(slice).Kind() { 58 | case reflect.Slice: 59 | val := reflect.ValueOf(slice) 60 | if val.Len() == 0 { 61 | return 0 // Return a default value when the slice is empty 62 | } 63 | extremeVal := initializeExtremeValue(operation) // Initialize extreme value based on operation 64 | for i := 0; i < val.Len(); i++ { 65 | elem := val.Index(i).Interface() 66 | extreme := findExtremeValue(elem, operation) 67 | if isBetter(extreme, extremeVal, operation) { 68 | extremeVal = extreme 69 | } 70 | } 71 | return extremeVal 72 | default: 73 | // Handle non-slice types (e.g., single element) 74 | switch slice := slice.(type) { 75 | case float32: 76 | return slice 77 | default: 78 | // Return a default value when the element is not a float32 79 | return 0 80 | } 81 | } 82 | } 83 | 84 | func storeSliceToFile(slice interface{}, filename string) error { 85 | // Get the value and kind of the input slice 86 | value := reflect.ValueOf(slice) 87 | kind := value.Kind() 88 | 89 | // Ensure the input is a slice 90 | if kind != reflect.Slice { 91 | return fmt.Errorf("input is not a slice") 92 | } 93 | 94 | // Flatten the slice 95 | flattened, err := flattenSlice(slice) 96 | if err != nil { 97 | return err 98 | } 99 | 100 | // Convert the flattened slice to a string representation 101 | dataStr := strings.Trim(strings.Join(strings.Fields(fmt.Sprint(flattened)), " "), "[]") 102 | 103 | // Write the string data to a text file 104 | if _, err := os.Stat(filename); err == nil { 105 | fmt.Println("File already exists!") 106 | } else if errors.Is(err, os.ErrNotExist) { 107 | _ = ioutil.WriteFile(filename, []byte(dataStr), 0644) 108 | } else { 109 | fmt.Println("Schrodinger: file may or may not exist. See err for details.", err) 110 | } 111 | return nil 112 | } 113 | 114 | func tensorsToFloat64Slices(tensors []*tf.Tensor) ([][]float64, error) { 115 | result := make([][]float64, len(tensors)) 116 | 117 | for i, t := range tensors { 118 | // Get the data from the *tf.Tensor as a 1D []float32 slice 119 | data, ok := t.Value().([]float32) 120 | if !ok { 121 | return nil, fmt.Errorf("expected tensor to be of type []float32, but got %T", t.Value()) 122 | } 123 | 124 | // Convert the []float32 to []float64 125 | float64Data := make([]float64, len(data)) 126 | for j, v := range data { 127 | float64Data[j] = float64(v) 128 | } 129 | 130 | // Append the []float64 to the result 131 | result[i] = float64Data 132 | } 133 | 134 | return result, nil 135 | } 136 | 137 | func float32SliceToPoint2fSlice(float32Slice []float32) []gocv.Point2f { 138 | if len(float32Slice)%2 != 0 { 139 | panic("float32Slice length must be even.") 140 | } 141 | 142 | point2fSlice := make([]gocv.Point2f, len(float32Slice)/2) 143 | 144 | for i := 0; i < len(float32Slice); i += 2 { 145 | point2fSlice[i/2] = gocv.Point2f{X: float32Slice[i], Y: float32Slice[i+1]} 146 | } 147 | 148 | return point2fSlice 149 | } 150 | 151 | // Reshape a 1D slice into a 2D slice of the given dimensions 152 | func reshape2D(data []float64, rows, cols int) [][]float64 { 153 | result := make([][]float64, rows) 154 | for i := 0; i < rows; i++ { 155 | result[i] = data[i*cols : (i+1)*cols] 156 | } 157 | return result 158 | } 159 | 160 | func svdGolubReinsch(A, U, S, Vt [][]float64) { 161 | m := len(A) 162 | n := len(A[0]) 163 | 164 | // Use A^T * A to compute eigenvectors and eigenvalues. 165 | AtA := make([][]float64, n) 166 | for i := range AtA { 167 | AtA[i] = make([]float64, n) 168 | for j := range AtA[i] { 169 | sum := 0.0 170 | for k := 0; k < m; k++ { 171 | sum += A[k][i] * A[k][j] 172 | } 173 | AtA[i][j] = sum 174 | } 175 | } 176 | 177 | // Compute eigenvalues and eigenvectors of A^T * A. 178 | eigenValues, eigenVectors := eigen(AtA) 179 | 180 | // Compute singular values and singular vectors. 181 | for i := 0; i < n; i++ { 182 | S[i][i] = math.Sqrt(eigenValues[i][0]) 183 | for j := 0; j < n; j++ { 184 | Vt[i][j] = eigenVectors[i][j] 185 | } 186 | } 187 | 188 | // Compute U from A * V. 189 | for i := 0; i < m; i++ { 190 | for j := 0; j < m; j++ { 191 | sum := 0.0 192 | for k := 0; k < n; k++ { 193 | sum += A[i][k] * Vt[k][j] 194 | } 195 | U[i][j] = sum / S[j][j] 196 | } 197 | } 198 | } 199 | 200 | func scaleVector(v []float64, s float64) []float64 { 201 | result := make([]float64, len(v)) 202 | for i := range v { 203 | result[i] = v[i] * s 204 | } 205 | return result 206 | } 207 | 208 | func float32ToMat(data [][][]float32) gocv.Mat { 209 | height := len(data) 210 | width := len(data[0]) 211 | channels := len(data[0][0]) 212 | 213 | // Create a new Mat from the flat data. 214 | sizes := []int{height, width, channels} 215 | mat := gocv.NewMatWithSizes(sizes, gocv.MatTypeCV32FC3) 216 | 217 | for i := 0; i < height; i++ { 218 | for j := 0; j < width; j++ { 219 | for k := 0; k < channels; k++ { 220 | mat.SetFloatAt3(i, j, k, data[i][j][k]) 221 | } 222 | } 223 | } 224 | return mat 225 | } 226 | 227 | func ConvertToFloats(img image.Image) [][]float32 { 228 | bounds := img.Bounds() 229 | height := bounds.Dy() 230 | width := bounds.Dx() 231 | 232 | data := make([][]float32, height) 233 | for y := 0; y < height; y++ { 234 | data[y] = make([]float32, width) 235 | for x := 0; x < width; x++ { 236 | grayColor := color.GrayModel.Convert(img.At(x, y)).(color.Gray) 237 | data[y][x] = float32(grayColor.Y) 238 | } 239 | } 240 | 241 | return data 242 | } 243 | 244 | func reshape(slice []float32) [][]float32 { 245 | var reshaped [][]float32 246 | for i := 0; i < len(slice); i += 9 { 247 | end := i + 9 248 | if end > len(slice) { 249 | end = len(slice) 250 | } 251 | reshaped = append(reshaped, slice[i:end]) 252 | } 253 | return reshaped 254 | } 255 | 256 | func getShape(slice interface{}) []int { 257 | var shape []int 258 | val := reflect.ValueOf(slice) 259 | for val.Kind() == reflect.Slice { 260 | shape = append(shape, val.Len()) 261 | val = val.Index(0) 262 | } 263 | return shape 264 | } 265 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## GoFaceRec 2 | 3 | This repository uses [tfgo](https://github.com/galeone/tfgo) to perform face recognition on an image from file. After 4 | a lot of efforts, I came to the conclusion that if one wants to load a deep learning model in PyTorch or Jax in Go, they better think twice before committing a good amount of effort into it. Instead, they are better to first convert their models to TensorFlow and then work with tfgo. 5 | 6 | In this repo, the input image is first processed, and then its embeddings are compared against the ones already computed from our dataset. In order to compute and save embeddings from an arbitrary dataset, one can use the [QMagFace's repo](https://github.com/pterhoer/QMagFace). Once the embeddings are ready, this repo uses Go in order to do face recognition. If the distance between embeddings falls bellow a specific threshold, then the face is considered as unknown. Otherwise, the proper label will be printed. 7 | 8 | ### Requirements 9 | 10 | This project is tested using `Go 1.17` on Ubuntu 20.04. Except for `tfgo`, latest version of other packages have been used and installed. 11 | 12 | For `gocv`, the version of `OpenCV` installed is `4.7`. And for `tfgo`, I installed [this version](https://github.com/galeone/tfgo) instead of the official one. 13 | 14 | 15 | ### Installing 16 | Just run the following command in you project in order to install this package: 17 | ```shell 18 | go get github.com/modanesh/GoFaceRec@v0.1.1 19 | ``` 20 | 21 | ### Converting models 22 | There are many ways to convert a non-TF model to a TF one. For that purpose, I used ONNX as an intermediary to convert 23 | the [QMagFace's model](https://github.com/pterhoer/QMagFace) from PyTorch to TF. 24 | 25 | Use the [model_converter.py](model_converter.py) script to convert the PyTorch model to ONNX first, and then the ONNX 26 | model to TF. 27 | 28 | Some of the code in the [model_converter.py](model_converter.py) is taken from the official [QMagFace's implementation](https://github.com/pterhoer/QMagFace). 29 | 30 | For this project, you may download the `MTCNN` and `MagFace` tensorflow models from the following URL: 31 | 32 | | Model | URL | 33 | |---------|----------------| 34 | | MTCNN | [Google Drive](https://drive.google.com/drive/folders/1_H_2qAdUfyKT8yyP9cxn0m20mMZe_yjg?usp=sharing)| 35 | | MagFace | [Google Drive](https://drive.google.com/drive/folders/136zEa3-6ve31HFse2sVgTi-eWkiLtfpk?usp=sharing)| 36 | 37 | ### Extracting layers 38 | 39 | In order to run the model using tfgo, you should know the input and output layers' names. In order to extract such 40 | information, the `saved_model_cli` command could be useful. A model exported with `tf.saved_model.save()` automatically 41 | comes with the "serve" tag because the SavedModel file format is designed for serving. This tag contains the various 42 | functions exported. Among these, there is always present the "serving_default" `signature_def`. This signature def 43 | works exactly like the TF 1.x graph. Get the input tensor and the output tensor, and use them as placeholder to feed 44 | and output to get, respectively. 45 | 46 | To get info inside a SavedModel the best tool is `saved_model_cli` that comes with the TensorFlow Python package, for 47 | example: 48 | ``` 49 | saved_model_cli show --all --dir output/keras 50 | gives, among the others, this info: 51 | 52 | signature_def['serving_default']: 53 | The given SavedModel SignatureDef contains the following input(s): 54 | inputs['inputs_input'] tensor_info: 55 | dtype: DT_FLOAT 56 | shape: (-1, 28, 28, 1) 57 | name: serving_default_inputs_input:0 58 | The given SavedModel SignatureDef contains the following output(s): 59 | outputs['logits'] tensor_info: 60 | dtype: DT_FLOAT 61 | shape: (-1, 10) 62 | name: StatefulPartitionedCall:0 63 | Method name is: tensorflow/serving/predict 64 | ``` 65 | 66 | Knowing the input and output layers' names, `serving_default_inputs_input:0` and `StatefulPartitionedCall:0`, is 67 | essential to run the model in tfgo. 68 | 69 | 70 | ### Running the model 71 | 72 | This project uses MTCNN for face detection and QMagFace for face recognition. For MTCNN, three stages (PNet, RNet, ONet) have been used in a close fashion similar to [FaceNet](https://github.com/davidsandberg/facenet). Each stage is done in its corresponding function: 73 | - First stage (PNet): [`totalBoxes := firstStage(scales, img, pnetModel)`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L1658) 74 | - Second stage (RNet): [`squaredBoxes := secondStage(totalBoxes, width, height, img, rnetModel)`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L1667) 75 | - Third stage (ONet): [`thirdPickedBoxes, pickedPoints := thirdStage(squaredBoxes, width, height, img, onetModel)`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L1676) 76 | 77 | You may download the models from the available [Google Drive URLs](#converting-models). 78 | 79 | After the face detection stage, there is a face alignment. The function to perform face alignment is [`pImgs := alignFace(thirdPickedBoxes, pickedPoints, img)`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L1685) which imitates the steps from [here](https://github.com/pterhoer/QMagFace/blob/main/preprocessing/insightface/src/face_preprocess.py#L195). 80 | 81 | Finally, once the face is detected and aligned, the recognition phase can start. It happens at this line: [`recognizeFace(pImgs, qmfModel, regEmbeddings, bSize, regFiles)`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L1694). 82 | 83 | Use the bellow command to run the code: 84 | ```shell 85 | go run main.go IMAGE.jpg path/to/REGISTERED_IMAGES path/to/EMBEDDINGS.npy path/to/MTCNN_MODELS_DIR path/to/MAGFACE_MODEL_DIR 86 | ``` 87 | where: 88 | - `IMAGE.jpg`: path to the given image 89 | - `path/to/REGISTERED_IMAGES`: directory containing register images 90 | - `path/to/EMBEDDINGS.npy`: the embeddings extracted from the register images using the [Python's QMagFace implementation](https://github.com/pterhoer/QMagFace) 91 | - `path/to/MTCNN_MODELS_DIR`: directory containing tensorflow models for MTCNN 92 | - `path/to/MAGFACE_MODEL_DIR`: directory containing tensorflow model for MagFace 93 | 94 | ### Challenges 95 | 96 | The main challenge thus far was the conversion between [`gocv.Mat`](https://github.com/hybridgroup/gocv), [`tfgo.Tensor`](https://github.com/galeone/tfgo), [`gonum`](https://github.com/gonum/gonum/), and Go's native slice. The conversion is required as some matrix transformations are only available in `gocv` and some in `tfgo`. Also, the input to the `tfgo` model should be of type `tfgo.Tensor`, so inevitably one needs to convert the image read by `gocv` to `tfgo`. Also, some matrix operations are not available in any of these packages, so I had to implement them myself from scratch. To do so, I had to use Go's native slice. So inevitable conversions between these types are frequent throughout the code. 97 | 98 | For example, the function [`adjustInput()`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L502) besides doing some scaling, it also converts a `gocv.Mat` to Go's `[][][][]float32`. In addition, at this line: [`inputBufTensor, _ := tf.NewTensor(inputBuf)`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L402) a `[][][][]float32` slice is converted to a `tfgo.Tensor`. 99 | 100 | In contrast, these type conversions are done pretty easy and fast in Python. 101 | 102 | ### ToDo 103 | - [X] Check why recognition model takes so long for a forward pass. In Python, it takes about 0.5 milliseconds while in Go it takes about 5500 milliseconds. For the first run, in Go, the session instantiation takes a long time. For next runs, Go runs pretty fast. [Take a look at this issue](https://github.com/galeone/tfgo/issues/4). The [`fakeRun()`](https://github.com/modanesh/GoFaceRec/blob/main/main.go?plain=1#L1582) function is for that purpose. 104 | - [X] Upload the models 105 | - [ ] Create a Go package 106 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= 2 | gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8= 3 | git.sr.ht/~sbinet/gg v0.3.1/go.mod h1:KGYtlADtqsqANL9ueOFkWymvzUvLMQllU5Ixo+8v3pc= 4 | github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= 5 | github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= 6 | github.com/ajstarks/deck v0.0.0-20200831202436-30c9fc6549a9/go.mod h1:JynElWSGnm/4RlzPXRlREEwqTHAN3T56Bv2ITsFT3gY= 7 | github.com/ajstarks/deck/generate v0.0.0-20210309230005-c3f852c02e19/go.mod h1:T13YZdzov6OU0A1+RfKZiZN9ca6VeKdBdyDV+BY97Tk= 8 | github.com/ajstarks/svgo v0.0.0-20180226025133-644b8db467af/go.mod h1:K08gAheRH3/J6wwsYMMT4xOr94bZjxIelGM0+d/wbFw= 9 | github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b/go.mod h1:1KcenG0jGWcpt8ov532z81sp/kMMUG485J2InIOyADM= 10 | github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8= 11 | github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8= 12 | github.com/campoy/embedmd v1.0.0 h1:V4kI2qTJJLf4J29RzI/MAt2c3Bl4dQSYPuflzwFH2hY= 13 | github.com/campoy/embedmd v1.0.0/go.mod h1:oxyr9RCiSXg0M3VJ3ks0UGfp98BpSSGr0kpiX3MzVl8= 14 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 15 | github.com/fogleman/gg v1.2.1-0.20190220221249-0403632d5b90/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k= 16 | github.com/fogleman/gg v1.3.0/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k= 17 | github.com/galeone/tensorflow/tensorflow/go v0.0.0-20221023090153-6b7fa0680c3e h1:9+2AEFZymTi25FIIcDwuzcOPH04z9+fV6XeLiGORPDI= 18 | github.com/galeone/tensorflow/tensorflow/go v0.0.0-20221023090153-6b7fa0680c3e/go.mod h1:TelZuq26kz2jysARBwOrTv16629hyUsHmIoj54QqyFo= 19 | github.com/galeone/tfgo v0.0.0-20221023090852-d89a5c7e31e1 h1:fsU+Je1A2kNu7e1LUAyOUrzXk50/2yP1tFGqKyNkOsI= 20 | github.com/galeone/tfgo v0.0.0-20221023090852-d89a5c7e31e1/go.mod h1:3YgYBeIX42t83uP27Bd4bSMxTnQhSbxl0pYSkCDB1tc= 21 | github.com/go-fonts/dejavu v0.1.0/go.mod h1:4Wt4I4OU2Nq9asgDCteaAaWZOV24E+0/Pwo0gppep4g= 22 | github.com/go-fonts/latin-modern v0.2.0/go.mod h1:rQVLdDMK+mK1xscDwsqM5J8U2jrRa3T0ecnM9pNujks= 23 | github.com/go-fonts/latin-modern v0.3.0/go.mod h1:ysEQXnuT/sCDOAONxC7ImeEDVINbltClhasMAqEtRK0= 24 | github.com/go-fonts/liberation v0.1.1/go.mod h1:K6qoJYypsmfVjWg8KOVDQhLc8UDgIK2HYqyqAO9z7GY= 25 | github.com/go-fonts/liberation v0.2.0/go.mod h1:K6qoJYypsmfVjWg8KOVDQhLc8UDgIK2HYqyqAO9z7GY= 26 | github.com/go-fonts/liberation v0.3.0/go.mod h1:jdJ+cqF+F4SUL2V+qxBth8fvBpBDS7yloUL5Fi8GTGY= 27 | github.com/go-fonts/stix v0.1.0/go.mod h1:w/c1f0ldAUlJmLBvlbkvVXLAD+tAMqobIIQpmnUIzUY= 28 | github.com/go-gl/glfw v0.0.0-20190409004039-e6da0acd62b1/go.mod h1:vR7hzQXu2zJy9AVAgeJqvqgH9Q5CA+iKCZ2gyEVpxRU= 29 | github.com/go-latex/latex v0.0.0-20210118124228-b3d85cf34e07/go.mod h1:CO1AlKB2CSIqUrmQPqA0gdRIlnLEY0gK5JGjh37zN5U= 30 | github.com/go-latex/latex v0.0.0-20210823091927-c0d11ff05a81/go.mod h1:SX0U8uGpxhq9o2S/CELCSUxEWWAuoCUcVCQWv7G2OCk= 31 | github.com/go-latex/latex v0.0.0-20230307184459-12ec69307ad9/go.mod h1:gWuR/CrFDDeVRFQwHPvsv9soJVB/iqymhuZQuJ3a9OM= 32 | github.com/go-pdf/fpdf v0.5.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhOh5M= 33 | github.com/go-pdf/fpdf v0.6.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhOh5M= 34 | github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k= 35 | github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= 36 | github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 37 | github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg= 38 | github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 39 | github.com/hybridgroup/mjpeg v0.0.0-20140228234708-4680f319790e/go.mod h1:eagM805MRKrioHYuU7iKLUyFPVKqVV6um5DAvCkUtXs= 40 | github.com/jung-kurt/gofpdf v1.0.0/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes= 41 | github.com/jung-kurt/gofpdf v1.0.3-0.20190309125859-24315acbbda5/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes= 42 | github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= 43 | github.com/pascaldekloe/goe v0.1.0/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc= 44 | github.com/phpdave11/gofpdf v1.4.2/go.mod h1:zpO6xFn9yxo3YLyMvW8HcKWVdbNqgIfOOp2dXMnm1mY= 45 | github.com/phpdave11/gofpdi v1.0.12/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk/7bXwjDoI= 46 | github.com/phpdave11/gofpdi v1.0.13/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk/7bXwjDoI= 47 | github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 48 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 49 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 50 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 51 | github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfFZQK844Gfx8o5WFuvpxWRwnSoipWe/p622j1v06w= 52 | github.com/ruudk/golang-pdf417 v0.0.0-20201230142125-a7e3863a1245/go.mod h1:pQAZKsJ8yyVxGRWYNEm9oFB8ieLgKFnamEyDmSA0BRk= 53 | github.com/sbinet/npyio v0.7.0 h1:KH8n5VrI1O2FeNAHwa0WmC1f9nGNtXNzQHBkyoU8tuE= 54 | github.com/sbinet/npyio v0.7.0/go.mod h1:4jmxspVr/RFRPc6zSGR/8FP6nb9m7EpypUXrU/cf/nU= 55 | github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= 56 | github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= 57 | github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= 58 | gocv.io/x/gocv v0.32.1 h1:BC9hHs5+47nVgySUFVKntc6RsF3SULFzqk6OV9xz+C0= 59 | gocv.io/x/gocv v0.32.1/go.mod h1:oc6FvfYqfBp99p+yOEzs9tbYF9gOrAQSeL/dyIPefJU= 60 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 61 | golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= 62 | golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= 63 | golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= 64 | golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= 65 | golang.org/x/crypto v0.1.0/go.mod h1:RecgLatLF4+eUMCP1PoPZQb+cVrJcOPbHkTkbkB9sbw= 66 | golang.org/x/exp v0.0.0-20180321215751-8460e604b9de/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= 67 | golang.org/x/exp v0.0.0-20180807140117-3d87b88a115f/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= 68 | golang.org/x/exp v0.0.0-20190125153040-c74c464bbbf2/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= 69 | golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= 70 | golang.org/x/exp v0.0.0-20191002040644-a1355ae1e2c3/go.mod h1:NOZ3BPKG0ec/BKJQgnvsSFpcKLM5xXVWnvZS97DWHgE= 71 | golang.org/x/exp v0.0.0-20230321023759-10a507213a29 h1:ooxPy7fPvB4kwsA2h+iBNHkAbp/4JxTSwCmvdjEYmug= 72 | golang.org/x/exp v0.0.0-20230321023759-10a507213a29/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= 73 | golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs= 74 | golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js= 75 | golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 76 | golang.org/x/image v0.0.0-20190910094157-69e4b8554b2a/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 77 | golang.org/x/image v0.0.0-20200119044424-58c23975cae1/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 78 | golang.org/x/image v0.0.0-20200430140353-33d19683fad8/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 79 | golang.org/x/image v0.0.0-20200618115811-c13761719519/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 80 | golang.org/x/image v0.0.0-20201208152932-35266b937fa6/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 81 | golang.org/x/image v0.0.0-20210216034530-4410531fe030/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= 82 | golang.org/x/image v0.0.0-20210607152325-775e3b0c77b9/go.mod h1:023OzeP/+EPmXeapQh35lcL3II3LrY8Ic+EFFKVhULM= 83 | golang.org/x/image v0.0.0-20210628002857-a66eb6448b8d/go.mod h1:023OzeP/+EPmXeapQh35lcL3II3LrY8Ic+EFFKVhULM= 84 | golang.org/x/image v0.0.0-20211028202545-6944b10bf410/go.mod h1:023OzeP/+EPmXeapQh35lcL3II3LrY8Ic+EFFKVhULM= 85 | golang.org/x/image v0.5.0/go.mod h1:FVC7BI/5Ym8R25iw5OLsgshdUBbT1h5jZTpA+mvAdZ4= 86 | golang.org/x/image v0.6.0/go.mod h1:MXLdDR43H7cDJq5GEGXEVeeNhPgi+YYEQ2pC1byI1x0= 87 | golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028/go.mod h1:E/iHnbuqvinMTCcRqshq8CkpyQDoeVncDDYHnLhea+o= 88 | golang.org/x/mod v0.1.0/go.mod h1:0QHyrYULN0/3qlju5TqG8bIK38QM8yzMo5ekMj3DlcY= 89 | golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= 90 | golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= 91 | golang.org/x/mod v0.6.0/go.mod h1:4mET923SAdbXp2ki8ey+zGs1SLqsuM2Y0uvdZR/fUNI= 92 | golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= 93 | golang.org/x/mod v0.9.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= 94 | golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= 95 | golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= 96 | golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= 97 | golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= 98 | golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= 99 | golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco= 100 | golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= 101 | golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc= 102 | golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 103 | golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 104 | golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 105 | golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 106 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 107 | golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 108 | golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 109 | golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 110 | golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 111 | golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 112 | golang.org/x/sys v0.0.0-20210304124612-50617c2ba197/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 113 | golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 114 | golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 115 | golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 116 | golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 117 | golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 118 | golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 119 | golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= 120 | golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= 121 | golang.org/x/term v0.1.0/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= 122 | golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= 123 | golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U= 124 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 125 | golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= 126 | golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= 127 | golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= 128 | golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= 129 | golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= 130 | golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= 131 | golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= 132 | golang.org/x/tools v0.0.0-20180525024113-a5b4c53f6e8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 133 | golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 134 | golang.org/x/tools v0.0.0-20190206041539-40960b6deb8e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 135 | golang.org/x/tools v0.0.0-20190927191325-030b2cf1153e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= 136 | golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= 137 | golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= 138 | golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= 139 | golang.org/x/tools v0.2.0/go.mod h1:y4OqIKeOV/fWJetJ8bXPU1sEVniLMIyDAZWeHdV+NTA= 140 | golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= 141 | golang.org/x/tools v0.7.0/go.mod h1:4pg6aUX35JBAogB10C9AtvVL+qowtN4pT3CGSQex14s= 142 | golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 143 | golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 144 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 145 | golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 146 | gonum.org/v1/gonum v0.0.0-20180816165407-929014505bf4/go.mod h1:Y+Yx5eoAFn32cQvJDxZx5Dpnq+c3wtXuadVZAcxbbBo= 147 | gonum.org/v1/gonum v0.8.2/go.mod h1:oe/vMfY3deqTw+1EZJhuvEW2iwGF1bW9wwu7XCu0+v0= 148 | gonum.org/v1/gonum v0.9.3/go.mod h1:TZumC3NeyVQskjXqmyWt4S3bINhy7B4eYwW69EbyX+0= 149 | gonum.org/v1/gonum v0.13.0 h1:a0T3bh+7fhRyqeNbiC3qVHYmkiQgit3wnNan/2c0HMM= 150 | gonum.org/v1/gonum v0.13.0/go.mod h1:/WPYRckkfWrhWefxyYTfrTtQR0KH4iyHNuzxqXAKyAU= 151 | gonum.org/v1/netlib v0.0.0-20190313105609-8cb42192e0e0/go.mod h1:wa6Ws7BG/ESfp6dHfk7C6KdzKA7wR7u/rKwOGE66zvw= 152 | gonum.org/v1/plot v0.0.0-20190515093506-e2840ee46a6b/go.mod h1:Wt8AAjI+ypCyYX3nZBvf6cAIx93T+c/OS2HFAYskSZc= 153 | gonum.org/v1/plot v0.9.0/go.mod h1:3Pcqqmp6RHvJI72kgb8fThyUnav364FOsdDo2aGW5lY= 154 | gonum.org/v1/plot v0.10.1/go.mod h1:VZW5OlhkL1mysU9vaqNHnsy86inf6Ot+jB3r+BczCEo= 155 | google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= 156 | google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= 157 | google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= 158 | honnef.co/go/tools v0.1.3/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las= 159 | rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | tf "github.com/galeone/tensorflow/tensorflow/go" 6 | tg "github.com/galeone/tfgo" 7 | "github.com/sbinet/npyio" 8 | "gocv.io/x/gocv" 9 | "gonum.org/v1/gonum/blas/blas64" 10 | "gonum.org/v1/gonum/lapack" 11 | "gonum.org/v1/gonum/lapack/lapack64" 12 | "image" 13 | "math" 14 | "os" 15 | "path/filepath" 16 | "reflect" 17 | "sort" 18 | "strings" 19 | "time" 20 | ) 21 | 22 | func sliceIndex(number int) []int { 23 | indexes := make([]int, number) 24 | for i := 0; i < number; i++ { 25 | indexes[i] = i 26 | } 27 | return indexes 28 | } 29 | 30 | func removeIndex(s []int, idx []int) []int { 31 | sort.Ints(idx) 32 | offset := 0 33 | for _, i := range idx { 34 | i -= offset 35 | s = append(s[:i], s[i+1:]...) 36 | offset++ 37 | } 38 | return s 39 | } 40 | 41 | func nms(boxes [][]float32, overlapThreshold float64, mode string) []int { 42 | if len(boxes) == 0 { 43 | return []int{} 44 | } 45 | 46 | // initialize the list of picked indexes 47 | pick := []int{} 48 | 49 | // grab the coordinates of the bounding boxes 50 | var x1, y1, x2, y2, score []float32 51 | for _, box := range boxes { 52 | x1 = append(x1, box[0]) 53 | y1 = append(y1, box[1]) 54 | x2 = append(x2, box[2]) 55 | y2 = append(y2, box[3]) 56 | score = append(score, box[4]) 57 | } 58 | 59 | area := make([]float32, len(boxes)) 60 | for i := range boxes { 61 | area[i] = (x2[i] - x1[i] + 1) * (y2[i] - y1[i] + 1) 62 | } 63 | 64 | idxs := make([]int, len(score)) 65 | for i := 0; i < len(score); i++ { 66 | idxs[i] = i 67 | } 68 | 69 | sort.Slice(idxs, func(i, j int) bool { return score[idxs[i]] < score[idxs[j]] }) 70 | 71 | // keep looping while some indexes still remain in the indexes list 72 | for len(idxs) > 0 { 73 | // grab the last index in the indexes list and add the index value to the list of picked indexes 74 | last := len(idxs) - 1 75 | i := idxs[last] 76 | pick = append(pick, i) 77 | 78 | xx1 := make([]float64, last) 79 | yy1 := make([]float64, last) 80 | xx2 := make([]float64, last) 81 | yy2 := make([]float64, last) 82 | for j := 0; j < last; j++ { 83 | xx1[j] = math.Max(float64(x1[i]), float64(x1[idxs[j]])) 84 | yy1[j] = math.Max(float64(y1[i]), float64(y1[idxs[j]])) 85 | xx2[j] = math.Min(float64(x2[i]), float64(x2[idxs[j]])) 86 | yy2[j] = math.Min(float64(y2[i]), float64(y2[idxs[j]])) 87 | } 88 | 89 | // compute the width and height of the bounding box 90 | w := make([]float64, last) 91 | h := make([]float64, last) 92 | for j := 0; j < last; j++ { 93 | w[j] = math.Max(0, xx2[j]-xx1[j]+1) 94 | h[j] = math.Max(0, yy2[j]-yy1[j]+1) 95 | } 96 | inter := make([]float64, last) 97 | overlap := make([]float64, last) 98 | for j := 0; j < last; j++ { 99 | inter[j] = w[j] * h[j] 100 | if mode == "Min" { 101 | overlap[j] = inter[j] / math.Min(float64(area[i]), float64(area[idxs[j]])) 102 | } else { 103 | overlap[j] = inter[j] / (float64(area[i]) + float64(area[idxs[j]]) - inter[j]) 104 | } 105 | } 106 | 107 | // delete all indexes from the index list that have 108 | toDelete := []int{} 109 | for j := 0; j < last; j++ { 110 | if overlap[j] > overlapThreshold { 111 | toDelete = append(toDelete, j) 112 | } 113 | } 114 | toDelete = append(toDelete, last) 115 | 116 | deleteIdxs := make([]int, len(toDelete)) 117 | for j := 0; j < len(toDelete); j++ { 118 | deleteIdxs[j] = idxs[toDelete[j]] 119 | } 120 | idxs = removeIndex(idxs, toDelete) 121 | } 122 | return pick 123 | } 124 | 125 | func extractData(npArray [][][]float32, index int) [][]float32 { 126 | // Assuming the dimensions of your slice are as follows: 127 | var dim1, dim2 = len(npArray), len(npArray[0]) 128 | 129 | // Creating a 2D slice to store the extracted values 130 | var extractedData = make([][]float32, dim1) 131 | for i := 0; i < dim1; i++ { 132 | extractedData[i] = make([]float32, dim2) 133 | } 134 | 135 | // Extracting the desired values 136 | for i := 0; i < dim1; i++ { 137 | for j := 0; j < dim2; j++ { 138 | extractedData[i][j] = npArray[i][j][index] 139 | } 140 | } 141 | return extractedData 142 | } 143 | 144 | func fix(val float32) float32 { 145 | if val < 0 { 146 | return float32(math.Ceil(float64(val))) 147 | } else { 148 | return float32(math.Floor(float64(val))) 149 | } 150 | } 151 | 152 | func computeQ1(bb [][]float32, stride, scale float32) [][]float32 { 153 | q1 := make([][]float32, len(bb)) 154 | for i, pair := range bb { 155 | q1[i] = make([]float32, len(pair)) 156 | for j, val := range pair { 157 | q1[i][j] = fix((stride*val + 1) / scale) 158 | } 159 | } 160 | return q1 161 | } 162 | 163 | func computeQ2(bb [][]float32, stride, scale, cellsize float32) [][]float32 { 164 | q2 := make([][]float32, len(bb)) 165 | for i, pair := range bb { 166 | q2[i] = make([]float32, len(pair)) 167 | for j, val := range pair { 168 | q2[i][j] = fix((stride*val + cellsize) / scale) 169 | } 170 | } 171 | return q2 172 | } 173 | 174 | func generateBBox(imap [][]float32, reg [][][]float32, scale float64, t float32) [][]float32 { 175 | stride := 2 176 | cellsize := 12 177 | 178 | imap = transpose2D(imap) 179 | dx1 := transpose2D(extractData(reg, 0)) 180 | dy1 := transpose2D(extractData(reg, 1)) 181 | dx2 := transpose2D(extractData(reg, 2)) 182 | dy2 := transpose2D(extractData(reg, 3)) 183 | 184 | var ys []int 185 | var xs []int 186 | for i := range imap { 187 | for j := range imap[i] { 188 | if imap[i][j] >= t { 189 | ys = append(ys, i) 190 | xs = append(xs, j) 191 | } 192 | } 193 | } 194 | 195 | if len(ys) == 1 { 196 | dx1 = flip2D(dx1) 197 | dy1 = flip2D(dy1) 198 | dx2 = flip2D(dx2) 199 | dy2 = flip2D(dy2) 200 | } 201 | 202 | scores := make([]float32, len(ys)) 203 | for i, y := range ys { 204 | scores[i] = imap[y][xs[i]] 205 | } 206 | regResult := make([][]float32, len(ys)) 207 | for i := range regResult { 208 | regResult[i] = []float32{ 209 | dx1[ys[i]][xs[i]], 210 | dy1[ys[i]][xs[i]], 211 | dx2[ys[i]][xs[i]], 212 | dy2[ys[i]][xs[i]], 213 | } 214 | } 215 | if len(regResult) == 0 { 216 | regResult = make([][]float32, 0) 217 | } 218 | 219 | bb := make([][]float32, len(xs)) 220 | for i := 0; i < len(xs); i++ { 221 | bb[i] = []float32{float32(ys[i]), float32(xs[i])} 222 | } 223 | 224 | q1 := computeQ1(bb, float32(stride), float32(scale)) 225 | q2 := computeQ2(bb, float32(stride), float32(scale), float32(cellsize)) 226 | 227 | boundingbox := make([][]float32, len(ys)) 228 | for i := range boundingbox { 229 | boundingbox[i] = append(append(q1[i], q2[i]...), append([]float32{scores[i]}, regResult[i]...)...) 230 | } 231 | 232 | return boundingbox 233 | } 234 | 235 | func transpose2D(matrix [][]float32) [][]float32 { 236 | rows := len(matrix) 237 | if rows == 0 { 238 | return matrix 239 | } 240 | 241 | cols := len(matrix[0]) 242 | result := make([][]float32, cols) 243 | for i := range result { 244 | result[i] = make([]float32, rows) 245 | } 246 | 247 | for i, row := range matrix { 248 | for j, val := range row { 249 | result[j][i] = val 250 | } 251 | } 252 | 253 | return result 254 | } 255 | 256 | func transpose3D(slice [][][]float32) [][][]float32 { 257 | var ( 258 | x = len(slice) 259 | y = len(slice[0]) 260 | z = len(slice[0][0]) 261 | ) 262 | 263 | newSlice := make([][][]float32, z) 264 | for i := range newSlice { 265 | newSlice[i] = make([][]float32, x) 266 | for j := range newSlice[i] { 267 | newSlice[i][j] = make([]float32, y) 268 | } 269 | } 270 | 271 | for i, s := range slice { 272 | for j, ss := range s { 273 | for k, v := range ss { 274 | newSlice[k][i][j] = v 275 | } 276 | } 277 | } 278 | 279 | return newSlice 280 | } 281 | 282 | func flip2D(matrix [][]float32) [][]float32 { 283 | for i := 0; i < len(matrix)/2; i++ { 284 | matrix[i], matrix[len(matrix)-1-i] = matrix[len(matrix)-1-i], matrix[i] 285 | } 286 | return matrix 287 | } 288 | 289 | func flatten4DTo2D(data [][][][]float32) [][]float32 { 290 | var dim2, dim3 = len(data[0]), len(data[0][0]) 291 | 292 | var extractedData = make([][]float32, dim2) 293 | for i := 0; i < dim2; i++ { 294 | extractedData[i] = make([]float32, dim3) 295 | } 296 | 297 | for i := 0; i < dim2; i++ { 298 | for j := 0; j < dim3; j++ { 299 | extractedData[i][j] = data[0][i][j][1] 300 | } 301 | } 302 | return extractedData 303 | } 304 | 305 | func transpose(x [][][][]float32, order []int) [][][][]float32 { 306 | if len(order) != 4 { 307 | panic("order must have a length of 4") 308 | } 309 | 310 | dims := []int{len(x), len(x[0]), len(x[0][0]), len(x[0][0][0])} 311 | newDims := []int{dims[order[0]], dims[order[1]], dims[order[2]], dims[order[3]]} 312 | 313 | // Create output slice with transposed dimensions 314 | out := make([][][][]float32, newDims[0]) 315 | for i := range out { 316 | out[i] = make([][][]float32, newDims[1]) 317 | for j := range out[i] { 318 | out[i][j] = make([][]float32, newDims[2]) 319 | for k := range out[i][j] { 320 | out[i][j][k] = make([]float32, newDims[3]) 321 | } 322 | } 323 | } 324 | 325 | // Transpose elements 326 | for i := 0; i < dims[0]; i++ { 327 | for j := 0; j < dims[1]; j++ { 328 | for k := 0; k < dims[2]; k++ { 329 | for l := 0; l < dims[3]; l++ { 330 | xIndex := []int{i, j, k, l} 331 | outIndex := []int{xIndex[order[0]], xIndex[order[1]], xIndex[order[2]], xIndex[order[3]]} 332 | out[outIndex[0]][outIndex[1]][outIndex[2]][outIndex[3]] = x[xIndex[0]][xIndex[1]][xIndex[2]][xIndex[3]] 333 | } 334 | } 335 | } 336 | } 337 | 338 | return out 339 | } 340 | 341 | func initializeExtremeValue(operation string) float32 { 342 | switch operation { 343 | case "min": 344 | return float32(math.Inf(1)) // Initialize with positive infinity for min operation 345 | case "max": 346 | return float32(math.Inf(-1)) // Initialize with negative infinity for max operation 347 | default: 348 | return 0 349 | } 350 | } 351 | 352 | func isBetter(candidate float32, currentExtreme float32, operation string) bool { 353 | switch operation { 354 | case "min": 355 | return candidate < currentExtreme 356 | case "max": 357 | return candidate > currentExtreme 358 | default: 359 | return false 360 | } 361 | } 362 | 363 | func flattenSlice(slice interface{}) ([]interface{}, error) { 364 | value := reflect.ValueOf(slice) 365 | kind := value.Kind() 366 | 367 | // Ensure the input is a slice 368 | if kind != reflect.Slice { 369 | return nil, fmt.Errorf("input is not a slice") 370 | } 371 | 372 | // Flatten the slice recursively 373 | var flattened []interface{} 374 | for i := 0; i < value.Len(); i++ { 375 | element := value.Index(i) 376 | elementKind := element.Kind() 377 | 378 | if elementKind == reflect.Slice { 379 | subSlice, err := flattenSlice(element.Interface()) 380 | if err != nil { 381 | return nil, err 382 | } 383 | flattened = append(flattened, subSlice...) 384 | } else { 385 | flattened = append(flattened, element.Interface()) 386 | } 387 | } 388 | 389 | return flattened, nil 390 | } 391 | 392 | func detectFirstStage(img gocv.Mat, net *tg.Model, scale float64, threshold float32) [][]float32 { 393 | height, width := img.Size()[0], img.Size()[1] 394 | ws := int(math.Ceil(float64(width) * scale)) 395 | hs := int(math.Ceil(float64(height) * scale)) 396 | 397 | imData := gocv.NewMat() 398 | defer imData.Close() 399 | gocv.Resize(img, &imData, image.Point{X: ws, Y: hs}, 0, 0, gocv.InterpolationLinear) 400 | 401 | inputBuf := adjustInput(imData) 402 | inputBufTensor, _ := tf.NewTensor(inputBuf) 403 | newShape := []int64{1, int64(ws), int64(hs), 3} 404 | inputBufTensor.Reshape(newShape) 405 | 406 | netOutput := net.Exec([]tf.Output{ 407 | net.Op("PartitionedCall", 0), 408 | net.Op("PartitionedCall", 1), 409 | }, map[tf.Output]*tf.Tensor{ 410 | net.Op("serving_default_input_1", 0): inputBufTensor, 411 | }) 412 | 413 | reg, ok := netOutput[0].Value().([][][][]float32) 414 | if !ok { 415 | fmt.Println("Failed to convert reg to [][][][]float64") 416 | } 417 | heatmap, ok := netOutput[1].Value().([][][][]float32) 418 | if !ok { 419 | fmt.Println("Failed to convert heatmap to [][]float64") 420 | } 421 | order := []int{0, 2, 1, 3} 422 | reg = transpose(reg, order) 423 | heatmap2d := flatten4DTo2D(transpose(heatmap, order)) 424 | 425 | boxes := generateBBox(heatmap2d, reg[0], scale, threshold) 426 | if len(boxes) == 0 { 427 | return nil 428 | } 429 | 430 | // nms 431 | pick := nms(boxes, 0.5, "Union") 432 | 433 | var pickedBoxes [][]float32 434 | for _, index := range pick { 435 | pickedBoxes = append(pickedBoxes, boxes[index]) 436 | } 437 | return pickedBoxes 438 | } 439 | 440 | func convertToSquare(bbox [][]float32) [][]float32 { 441 | squareBbox := make([][]float32, len(bbox)) 442 | for i := 0; i < len(bbox); i++ { 443 | squareBbox[i] = make([]float32, len(bbox[i])) 444 | copy(squareBbox[i], bbox[i]) 445 | h := bbox[i][3] - bbox[i][1] + 1 446 | w := bbox[i][2] - bbox[i][0] + 1 447 | maxSide := float32(math.Max(float64(h), float64(w))) 448 | squareBbox[i][0] = bbox[i][0] + w*0.5 - maxSide*0.5 449 | squareBbox[i][1] = bbox[i][1] + h*0.5 - maxSide*0.5 450 | squareBbox[i][2] = squareBbox[i][0] + maxSide - 1 451 | squareBbox[i][3] = squareBbox[i][1] + maxSide - 1 452 | } 453 | return squareBbox 454 | } 455 | 456 | func pad(bboxes [][]float32, w float32, h float32) ([]float32, []float32, []float32, []float32, []float32, []float32, []float32, []float32, []float32, []float32) { 457 | numBox := len(bboxes) 458 | tmpw := make([]float32, numBox) 459 | tmph := make([]float32, numBox) 460 | 461 | for i := range bboxes { 462 | tmpw[i] = bboxes[i][2] - bboxes[i][0] + 1 463 | tmph[i] = bboxes[i][3] - bboxes[i][1] + 1 464 | } 465 | 466 | dx := make([]float32, numBox) 467 | dy := make([]float32, numBox) 468 | edx := make([]float32, numBox) 469 | edy := make([]float32, numBox) 470 | 471 | x := make([]float32, numBox) 472 | y := make([]float32, numBox) 473 | ex := make([]float32, numBox) 474 | ey := make([]float32, numBox) 475 | 476 | for i := range bboxes { 477 | dx[i], dy[i] = 0, 0 478 | edx[i], edy[i] = tmpw[i]-1, tmph[i]-1 479 | x[i], y[i], ex[i], ey[i] = bboxes[i][0], bboxes[i][1], bboxes[i][2], bboxes[i][3] 480 | 481 | if ex[i] > w-1 { 482 | edx[i] = tmpw[i] + w - 2 - ex[i] 483 | ex[i] = w - 1 484 | } 485 | if ey[i] > h-1 { 486 | edy[i] = tmph[i] + h - 2 - ey[i] 487 | ey[i] = h - 1 488 | } 489 | if x[i] < 0 { 490 | dx[i] = 0 - x[i] 491 | x[i] = 0 492 | } 493 | if y[i] < 0 { 494 | dy[i] = 0 - y[i] 495 | y[i] = 0 496 | } 497 | } 498 | 499 | return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph 500 | } 501 | 502 | func adjustInput(inData gocv.Mat) [][][][]float32 { 503 | width := inData.Cols() 504 | height := inData.Rows() 505 | channels := inData.Channels() 506 | outData := make([][][][]float32, 1) 507 | outData[0] = make([][][]float32, width) 508 | 509 | // Scale the data and expand dimensions 510 | for w := 0; w < width; w++ { 511 | outData[0][w] = make([][]float32, height) 512 | for h := 0; h < height; h++ { 513 | outData[0][w][h] = make([]float32, channels) 514 | val := inData.GetVecbAt(h, w) 515 | for c := 0; c < channels; c++ { 516 | outData[0][w][h][c] = (float32(val[c]) - 127.5) * 0.0078125 517 | } 518 | } 519 | } 520 | return outData 521 | } 522 | 523 | func matToSlice(inData gocv.Mat) [][][]float32 { 524 | width := inData.Rows() 525 | height := inData.Cols() 526 | channels := inData.Channels() 527 | outData := make([][][]float32, width) 528 | 529 | for w := 0; w < width; w++ { 530 | outData[w] = make([][]float32, height) 531 | for h := 0; h < height; h++ { 532 | outData[w][h] = make([]float32, channels) 533 | val := inData.GetVecbAt(w, h) 534 | for c := 0; c < channels; c++ { 535 | outData[w][h][c] = float32(val[c]) / 255 536 | } 537 | } 538 | } 539 | return outData 540 | } 541 | 542 | func CalibrateBox(bbox [][]float32, reg [][]float32) [][]float32 { 543 | n := len(bbox) 544 | w := make([]float32, n) 545 | h := make([]float32, n) 546 | for i := 0; i < n; i++ { 547 | w[i] = bbox[i][2] - bbox[i][0] + 1 548 | h[i] = bbox[i][3] - bbox[i][1] + 1 549 | } 550 | regM := make([][]float32, n) 551 | for i := range regM { 552 | regM[i] = make([]float32, 4) 553 | regM[i][0] = w[i] 554 | regM[i][1] = h[i] 555 | regM[i][2] = w[i] 556 | regM[i][3] = h[i] 557 | } 558 | aug := make([][]float32, n) 559 | for i := range aug { 560 | aug[i] = make([]float32, 4) 561 | aug[i][0] = regM[i][0] * reg[i][0] 562 | aug[i][1] = regM[i][1] * reg[i][1] 563 | aug[i][2] = regM[i][2] * reg[i][2] 564 | aug[i][3] = regM[i][3] * reg[i][3] 565 | } 566 | for i := 0; i < n; i++ { 567 | bbox[i][0] += aug[i][0] 568 | bbox[i][1] += aug[i][1] 569 | bbox[i][2] += aug[i][2] 570 | bbox[i][3] += aug[i][3] 571 | } 572 | return bbox 573 | } 574 | 575 | func preprocessMat(img gocv.Mat, bbox []float32, landmark [][]float32) gocv.Mat { 576 | var M [][]float64 = nil 577 | var det, bb []float32 = nil, nil 578 | 579 | if landmark != nil { 580 | src := [][]float64{ 581 | {30.2946, 51.6963}, 582 | {65.5318, 51.5014}, 583 | {48.0252, 71.7366}, 584 | {33.5493, 92.3655}, 585 | {62.7299, 92.2041}, 586 | } 587 | 588 | for i := 0; i < len(src); i++ { 589 | src[i][0] += 8.0 590 | } 591 | dst := float32ToFloat64(landmark) 592 | M = umeyama(dst, src, true) 593 | M = M[:len(M)-1] 594 | fmt.Printf("Matrix M: %v\n", M) 595 | } 596 | 597 | if M == nil { 598 | if bbox == nil { 599 | det = make([]float32, 4) 600 | det[0] = float32(img.Cols()) * 0.0625 601 | det[1] = float32(img.Rows()) * 0.0625 602 | det[2] = float32(img.Cols()) - det[0] 603 | det[3] = float32(img.Rows()) - det[1] 604 | } else { 605 | det = bbox 606 | } 607 | 608 | bb = make([]float32, 4) 609 | bb[0] = float32(math.Max(float64(det[0])-float64(44.0/2), 0)) 610 | bb[1] = float32(math.Max(float64(det[1])-float64(44.0/2), 0)) 611 | bb[2] = float32(math.Min(float64(det[2])+float64(44.0/2), float64(img.Cols()))) 612 | bb[3] = float32(math.Min(float64(det[3])+float64(44.0/2), float64(img.Rows()))) 613 | 614 | ret := img.Region(image.Rect(int(bb[0]), int(bb[1]), int(bb[2]), int(bb[3]))) 615 | gocv.Resize(ret, &ret, image.Point{X: 112, Y: 112}, 0, 0, gocv.InterpolationLinear) 616 | return ret 617 | } else { 618 | warped := gocv.NewMat() 619 | mMat := float64ToMat(M) 620 | val := img.GetVecbAt(0, 0) 621 | fmt.Printf("img[0,0]: %v\n", val) 622 | val = img.GetVecbAt(24, 10) 623 | fmt.Printf("img[24,10]: %v\n", val) 624 | 625 | gocv.WarpAffine(img, &warped, mMat, image.Point{X: 112, Y: 112}) 626 | val = warped.GetVecbAt(0, 0) 627 | fmt.Printf("warped[0,0]: %v\n", val) 628 | val = warped.GetVecbAt(0, 1) 629 | fmt.Printf("warped[0,1]: %v\n", val) 630 | 631 | return warped 632 | } 633 | } 634 | 635 | func umeyama(src, dst [][]float64, estimateScale bool) [][]float64 { 636 | num := len(src) 637 | dim := len(src[0]) 638 | 639 | // Compute mean of src and dst. 640 | srcMean := make([]float64, dim) 641 | dstMean := make([]float64, dim) 642 | for i := 0; i < num; i++ { 643 | for j := 0; j < dim; j++ { 644 | srcMean[j] += src[i][j] 645 | dstMean[j] += dst[i][j] 646 | } 647 | } 648 | for j := 0; j < dim; j++ { 649 | srcMean[j] /= float64(num) 650 | dstMean[j] /= float64(num) 651 | } 652 | 653 | // Subtract mean from src and dst. 654 | srcDemean := make([][]float64, num) 655 | dstDemean := make([][]float64, num) 656 | for i := range srcDemean { 657 | srcDemean[i] = make([]float64, dim) 658 | dstDemean[i] = make([]float64, dim) 659 | } 660 | for i := 0; i < num; i++ { 661 | for j := 0; j < dim; j++ { 662 | srcDemean[i][j] = src[i][j] - srcMean[j] 663 | dstDemean[i][j] = dst[i][j] - dstMean[j] 664 | } 665 | } 666 | fmt.Printf("src: %v\n", src) 667 | fmt.Printf("srcMean: %v\n", srcMean) 668 | fmt.Printf("srcDemean: %v\n", srcDemean) 669 | 670 | // Eq. (38). 671 | A := make([][]float64, dim) 672 | for i := range A { 673 | A[i] = make([]float64, dim) 674 | } 675 | for i := 0; i < dim; i++ { 676 | for j := 0; j < dim; j++ { 677 | for k := 0; k < num; k++ { 678 | A[i][j] += dstDemean[k][i] * srcDemean[k][j] 679 | } 680 | A[i][j] /= float64(num) 681 | } 682 | } 683 | 684 | // Eq. (39). 685 | d := make([]float64, dim) 686 | for i := range d { 687 | d[i] = 1.0 688 | } 689 | if determinant(A) < 0 { 690 | d[dim-1] = -1 691 | } 692 | 693 | T := make([][]float64, dim+1) 694 | for i := range T { 695 | T[i] = make([]float64, dim+1) 696 | for j := range T[i] { 697 | if i == j { 698 | T[i][j] = 1.0 699 | } 700 | } 701 | } 702 | 703 | V, S, U := svd(A) 704 | fmt.Printf("Matrix A: %v\n", A) 705 | fmt.Printf("Matrix U: %v\n", U) 706 | fmt.Printf("Matrix V: %v\n", V) 707 | fmt.Printf("Vector d: %v\n", d) 708 | fmt.Printf("Matrix S: %v\n", S) 709 | 710 | // Eq. (40) and (43). 711 | rank := matrixRank(A) 712 | if rank == 0 { 713 | for i := range T { 714 | for j := range T[i] { 715 | T[i][j] = math.NaN() 716 | } 717 | } 718 | return T 719 | } else if rank == dim-1 { 720 | if determinant(U)*determinant(V) > 0 { 721 | for i := 0; i < dim; i++ { 722 | for j := 0; j < dim; j++ { 723 | T[i][j] = 0 724 | for k := 0; k < dim; k++ { 725 | T[i][j] += U[i][k] * V[k][j] 726 | } 727 | } 728 | } 729 | } else { 730 | s := d[dim-1] 731 | d[dim-1] = -1 732 | Temp := make([][]float64, dim) 733 | for i := range Temp { 734 | Temp[i] = make([]float64, dim) 735 | } 736 | for i := 0; i < dim; i++ { 737 | for j := 0; j < dim; j++ { 738 | if i == j { 739 | Temp[i][j] = d[i] 740 | } 741 | } 742 | } 743 | for i := 0; i < dim; i++ { 744 | for j := 0; j < dim; j++ { 745 | T[i][j] = 0 746 | for k := 0; k < dim; k++ { 747 | T[i][j] += U[i][k] * Temp[k][j] 748 | } 749 | } 750 | } 751 | for i := 0; i < dim; i++ { 752 | for j := 0; j < dim; j++ { 753 | Temp[i][j] = 0 754 | for k := 0; k < dim; k++ { 755 | Temp[i][j] += T[i][k] * V[k][j] 756 | } 757 | } 758 | } 759 | T = Temp 760 | d[dim-1] = s 761 | } 762 | } else { 763 | Temp := make([][]float64, dim) 764 | for i := range Temp { 765 | Temp[i] = make([]float64, dim) 766 | } 767 | for i := 0; i < dim; i++ { 768 | for j := 0; j < dim; j++ { 769 | if i == j { 770 | Temp[i][j] = d[i] 771 | } 772 | } 773 | } 774 | for i := 0; i < dim; i++ { 775 | for j := 0; j < dim; j++ { 776 | T[i][j] = 0 777 | for k := 0; k < dim; k++ { 778 | T[i][j] += U[i][k] * Temp[k][j] 779 | } 780 | } 781 | } 782 | for i := 0; i < dim; i++ { 783 | for j := 0; j < dim; j++ { 784 | Temp[i][j] = 0 785 | for k := 0; k < dim; k++ { 786 | Temp[i][j] += T[i][k] * V[k][j] 787 | } 788 | } 789 | } 790 | T = Temp 791 | } 792 | fmt.Printf("Matrix T 0: %v\n", T) 793 | 794 | var scale float64 795 | if estimateScale { 796 | // Eq. (41) and (42). 797 | var srcVar float64 798 | for i := 0; i < dim; i++ { 799 | var temp float64 800 | for j := 0; j < num; j++ { 801 | temp += (src[j][i] - srcMean[i]) * (src[j][i] - srcMean[i]) 802 | } 803 | srcVar += temp / float64(num) 804 | } 805 | var sumSD float64 806 | for i := 0; i < dim; i++ { 807 | sumSD += S[0][i] * d[i] 808 | } 809 | scale = 1.0 / srcVar * sumSD 810 | fmt.Printf("srcVar: %v\n", srcVar) 811 | fmt.Printf("sumSD: %v\n", sumSD) 812 | } else { 813 | scale = 1.0 814 | } 815 | fmt.Printf("Scale: %v\n", scale) 816 | 817 | for i := 0; i < dim; i++ { 818 | var temp float64 819 | for j := 0; j < dim; j++ { 820 | temp += T[i][j] * srcMean[j] 821 | } 822 | T[i][dim] = dstMean[i] - scale*temp 823 | } 824 | fmt.Printf("Matrix T 1: %v\n", T) 825 | 826 | for i := 0; i < dim; i++ { 827 | for j := 0; j < dim; j++ { 828 | T[i][j] *= scale 829 | } 830 | } 831 | fmt.Printf("Matrix T 2: %v\n", T) 832 | return T 833 | } 834 | 835 | func matrixRank(A [][]float64) int { 836 | U, _, _ := svd(A) 837 | rank := 0 838 | for i := range U { 839 | if U[i][i] > 1e-10 { 840 | rank++ 841 | } 842 | } 843 | return rank 844 | } 845 | 846 | func determinant(A [][]float64) float64 { 847 | n := len(A) 848 | if n == 2 { 849 | return A[0][0]*A[1][1] - A[0][1]*A[1][0] 850 | } 851 | 852 | det := 0.0 853 | for i := 0; i < n; i++ { 854 | subMatrix := make([][]float64, n-1) 855 | for j := range subMatrix { 856 | subMatrix[j] = make([]float64, n-1) 857 | copy(subMatrix[j], A[j+1]) 858 | subMatrix[j] = append(subMatrix[j][:i], subMatrix[j][i+1:]...) 859 | } 860 | 861 | det += math.Pow(-1, float64(i)) * A[0][i] * determinant(subMatrix) 862 | } 863 | 864 | return det 865 | } 866 | 867 | func svd(A [][]float64) (U, S, Vt [][]float64) { 868 | m := len(A) 869 | n := len(A[0]) 870 | 871 | // Convert A to column-major order. 872 | data := make([]float64, m*n) 873 | for i := 0; i < n; i++ { 874 | for j := 0; j < m; j++ { 875 | data[i*m+j] = A[j][i] 876 | } 877 | } 878 | 879 | U = make([][]float64, m) 880 | for i := range U { 881 | U[i] = make([]float64, m) 882 | } 883 | 884 | S = make([][]float64, m) 885 | for i := range S { 886 | S[i] = make([]float64, n) 887 | } 888 | 889 | Vt = make([][]float64, n) 890 | for i := range Vt { 891 | Vt[i] = make([]float64, n) 892 | } 893 | 894 | // Compute SVD using gonum/lapack. 895 | blasData := blas64.General{ 896 | Rows: m, 897 | Cols: n, 898 | Stride: n, 899 | Data: data, 900 | } 901 | blasU := blas64.General{ 902 | Rows: m, 903 | Cols: m, 904 | Stride: m, 905 | Data: make([]float64, m*m), 906 | } 907 | blasVt := blas64.General{ 908 | Rows: n, 909 | Cols: n, 910 | Stride: n, 911 | Data: make([]float64, n*n), 912 | } 913 | work := make([]float64, 1) 914 | lwork := -1 915 | 916 | // Compute work size. 917 | ok := lapack64.Gesvd(lapack.SVDAll, lapack.SVDAll, blasData, blasU, blasVt, S[0], work, lwork) 918 | 919 | if !ok { 920 | panic("SVD failed") 921 | } 922 | 923 | // Allocate work with the correct size. 924 | lwork = int(work[0]) 925 | work = make([]float64, lwork) 926 | 927 | // Compute SVD. 928 | ok = lapack64.Gesvd(lapack.SVDAll, lapack.SVDAll, blasData, blasU, blasVt, S[0], work, lwork) 929 | 930 | if !ok { 931 | panic("SVD failed") 932 | } 933 | 934 | // Convert U, S, Vt back to row-major order. 935 | for i := 0; i < m; i++ { 936 | for j := 0; j < m; j++ { 937 | U[i][j] = blasU.Data[j*m+i] 938 | } 939 | } 940 | 941 | for i := 0; i < m; i++ { 942 | for j := 0; j < n; j++ { 943 | S[i][j] = S[i][j] 944 | } 945 | } 946 | 947 | for i := 0; i < n; i++ { 948 | for j := 0; j < n; j++ { 949 | Vt[i][j] = blasVt.Data[j*n+i] 950 | } 951 | } 952 | 953 | return U, S, Vt 954 | } 955 | 956 | func eigen(A [][]float64) (values, vectors [][]float64) { 957 | // Perform eigenvalue decomposition using QR algorithm with shifts. 958 | n := len(A) 959 | 960 | values = make([][]float64, n) 961 | vectors = make([][]float64, n) 962 | for i := range values { 963 | values[i] = make([]float64, 1) 964 | vectors[i] = make([]float64, n) 965 | } 966 | 967 | B := make([][]float64, n) 968 | for i := range B { 969 | B[i] = make([]float64, n) 970 | copy(B[i], A[i]) 971 | } 972 | 973 | for iter := 0; iter < 50; iter++ { 974 | // Check if B is diagonal. 975 | sumOffDiagonal := 0.0 976 | for i := 0; i < n; i++ { 977 | for j := 0; j < n; j++ { 978 | if i != j { 979 | sumOffDiagonal += B[i][j] * B[i][j] 980 | } 981 | } 982 | } 983 | 984 | if sumOffDiagonal <= 1e-12 { 985 | break 986 | } 987 | 988 | // QR decomposition of B. 989 | Q, R := qr(B) 990 | 991 | // B = R * Q 992 | B = matmul(R, Q) 993 | } 994 | 995 | for i := 0; i < n; i++ { 996 | values[i][0] = B[i][i] 997 | for j := 0; j < n; j++ { 998 | vectors[i][j] = B[i][j] 999 | } 1000 | } 1001 | 1002 | return values, vectors 1003 | } 1004 | 1005 | func qr(A [][]float64) (Q, R [][]float64) { 1006 | m := len(A) 1007 | n := len(A[0]) 1008 | 1009 | Q = make([][]float64, m) 1010 | for i := range Q { 1011 | Q[i] = make([]float64, n) 1012 | } 1013 | 1014 | R = make([][]float64, n) 1015 | for i := range R { 1016 | R[i] = make([]float64, n) 1017 | } 1018 | 1019 | for j := 0; j < n; j++ { 1020 | v := make([]float64, m) 1021 | for i := 0; i < m; i++ { 1022 | v[i] = A[i][j] 1023 | } 1024 | 1025 | for k := 0; k < j; k++ { 1026 | R[k][j] = dotproduct(Q[k], v) 1027 | for i := 0; i < m; i++ { 1028 | v[i] -= R[k][j] * Q[i][k] 1029 | } 1030 | } 1031 | 1032 | R[j][j] = norm(v) 1033 | for i := 0; i < m; i++ { 1034 | Q[i][j] = v[i] / R[j][j] 1035 | } 1036 | } 1037 | 1038 | return Q, R 1039 | } 1040 | 1041 | func dotproduct(u, v []float64) float64 { 1042 | sum := 0.0 1043 | for i := range u { 1044 | sum += u[i] * v[i] 1045 | } 1046 | return sum 1047 | } 1048 | 1049 | func norm(v []float64) float64 { 1050 | sum := 0.0 1051 | for _, x := range v { 1052 | sum += x * x 1053 | } 1054 | return math.Sqrt(sum) 1055 | } 1056 | 1057 | func matmul(A, B [][]float64) [][]float64 { 1058 | m := len(A) 1059 | n := len(B[0]) 1060 | p := len(B) 1061 | 1062 | C := make([][]float64, m) 1063 | for i := range C { 1064 | C[i] = make([]float64, n) 1065 | for j := 0; j < n; j++ { 1066 | sum := 0.0 1067 | for k := 0; k < p; k++ { 1068 | sum += A[i][k] * B[k][j] 1069 | } 1070 | C[i][j] = sum 1071 | } 1072 | } 1073 | 1074 | return C 1075 | } 1076 | 1077 | func float32ToFloat64(data32 [][]float32) [][]float64 { 1078 | data64 := make([][]float64, len(data32)) 1079 | for i := range data32 { 1080 | data64[i] = make([]float64, len(data32[i])) 1081 | for j := range data32[i] { 1082 | data64[i][j] = float64(data32[i][j]) 1083 | } 1084 | } 1085 | return data64 1086 | } 1087 | 1088 | func float64ToMat(data [][]float64) gocv.Mat { 1089 | rows := len(data) 1090 | cols := len(data[0]) 1091 | 1092 | sizes := []int{rows, cols} 1093 | mat := gocv.NewMatWithSizes(sizes, gocv.MatTypeCV32F) 1094 | 1095 | for i := 0; i < rows; i++ { 1096 | for j := 0; j < cols; j++ { 1097 | mat.SetFloatAt(i, j, float32(data[i][j])) 1098 | } 1099 | } 1100 | return mat 1101 | } 1102 | 1103 | func generateEmbeddings(imgs [][][]float32) *tf.Tensor { 1104 | transposedImgs := transpose3D(imgs) 1105 | permutedImgs, _ := tf.NewTensor(transposedImgs) 1106 | permutedImgs.Reshape([]int64{1, 3, 112, 112}) 1107 | return permutedImgs 1108 | } 1109 | 1110 | func normalize(vecs [][]float32) ([][]float32, []float32) { 1111 | r := len(vecs) 1112 | c := len(vecs[0]) 1113 | norms := make([]float32, r) 1114 | for i := 0; i < r; i++ { 1115 | norm := 0.0 1116 | for _, v := range vecs[i] { 1117 | norm += float64(v * v) 1118 | } 1119 | norms[i] = float32(math.Sqrt(norm)) 1120 | } 1121 | normed := make([][]float32, r) 1122 | for i := 0; i < r; i++ { 1123 | normed[i] = make([]float32, c) 1124 | for j := 0; j < c; j++ { 1125 | normed[i][j] = vecs[i][j] / norms[i] 1126 | } 1127 | } 1128 | return normed, norms 1129 | } 1130 | 1131 | func cosineSimilarityNoPair(fAnch, fTest [][]float32, isNormed bool) []float32 { 1132 | if !isNormed { 1133 | fAnch, _ = normalize(fAnch) 1134 | fTest, _ = normalize(fTest) 1135 | } 1136 | r1 := len(fAnch) 1137 | r2 := len(fTest) 1138 | 1139 | result := make([]float32, r1*r2) 1140 | counter := 0 1141 | 1142 | for i := 0; i < r1; i++ { 1143 | for j := 0; j < r2; j++ { 1144 | sum := float32(0.0) 1145 | for k := range fAnch[i] { 1146 | sum += fAnch[i][k] * fTest[j][k] 1147 | } 1148 | result[counter] = sum 1149 | counter++ 1150 | } 1151 | } 1152 | 1153 | return result 1154 | } 1155 | 1156 | func computeSQNoPair(fAnchor, fTest [][]float32) ([]float32, []float32) { 1157 | fAnchor, qAnchor := normalize(fAnchor) 1158 | fTest, qTest := normalize(fTest) 1159 | s := cosineSimilarityNoPair(fAnchor, fTest, true) 1160 | 1161 | q := make([]float32, len(qAnchor)*len(qTest)) 1162 | counter := 0 1163 | for _, i := range qAnchor { 1164 | for _, j := range qTest { 1165 | q[counter] = float32(math.Min(float64(i), float64(j))) 1166 | counter++ 1167 | } 1168 | } 1169 | return s, q 1170 | } 1171 | 1172 | func similarityNoPair(fAnch, fTest [][]float32) []float32 { 1173 | s, q := computeSQNoPair(fAnch, fTest) 1174 | alpha := float32(0.077428) 1175 | beta := float32(0.125926) 1176 | omega := make([]float32, len(s)) 1177 | for i, v := range s { 1178 | omega[i] = beta*v - alpha 1179 | if omega[i] >= 0 { 1180 | omega[i] = 0 1181 | } 1182 | } 1183 | 1184 | result := make([]float32, len(s)) 1185 | for i := range s { 1186 | result[i] = omega[i]*q[i] + s[i] 1187 | } 1188 | return result 1189 | } 1190 | 1191 | func loadNpy(filePath string) ([][]float32, error) { 1192 | file, err := os.Open(filePath) 1193 | if err != nil { 1194 | return nil, err 1195 | } 1196 | defer file.Close() 1197 | 1198 | var flattened []float32 1199 | if err := npyio.Read(file, &flattened); err != nil { 1200 | return nil, err 1201 | } 1202 | 1203 | // Assuming the second dimension is 512 1204 | const dim2 = 512 1205 | dim1 := len(flattened) / dim2 1206 | 1207 | matrix := make([][]float32, dim1) 1208 | for i := range matrix { 1209 | matrix[i] = make([]float32, dim2) 1210 | for j := range matrix[i] { 1211 | matrix[i][j] = flattened[i*dim2+j] 1212 | } 1213 | } 1214 | 1215 | return matrix, nil 1216 | } 1217 | 1218 | func getRegFiles(regDataPath string) ([]string, error) { 1219 | var regFiles []string 1220 | 1221 | err := filepath.Walk(regDataPath, func(path string, info os.FileInfo, err error) error { 1222 | if err != nil { 1223 | return err 1224 | } 1225 | 1226 | if !info.IsDir() && (strings.HasSuffix(path, ".png") || strings.HasSuffix(path, ".jpg")) { 1227 | regFiles = append(regFiles, path) 1228 | } 1229 | 1230 | return nil 1231 | }) 1232 | 1233 | if err != nil { 1234 | return nil, err 1235 | } 1236 | 1237 | return regFiles, nil 1238 | } 1239 | 1240 | func refineBoxes(totalBoxes [][]float32) [][]float32 { 1241 | bbW := make([]float32, len(totalBoxes)) 1242 | bbH := make([]float32, len(totalBoxes)) 1243 | 1244 | for i := range totalBoxes { 1245 | bbW[i] = totalBoxes[i][2] - totalBoxes[i][0] + 1 1246 | bbH[i] = totalBoxes[i][3] - totalBoxes[i][1] + 1 1247 | } 1248 | 1249 | refined := make([][]float32, len(totalBoxes)) 1250 | for i := range totalBoxes { 1251 | refined[i] = make([]float32, 5) 1252 | refined[i][0] = totalBoxes[i][0] + totalBoxes[i][5]*bbW[i] 1253 | refined[i][1] = totalBoxes[i][1] + totalBoxes[i][6]*bbH[i] 1254 | refined[i][2] = totalBoxes[i][2] + totalBoxes[i][7]*bbW[i] 1255 | refined[i][3] = totalBoxes[i][3] + totalBoxes[i][8]*bbH[i] 1256 | refined[i][4] = totalBoxes[i][4] 1257 | } 1258 | 1259 | return refined 1260 | } 1261 | 1262 | func firstStage(scales []float64, img gocv.Mat, pnetModel *tg.Model) [][]float32 { 1263 | slicedIndex := sliceIndex(len(scales)) 1264 | var totalBoxes [][]float32 1265 | for _, batch := range slicedIndex { 1266 | localBoxes := detectFirstStage(img, pnetModel, scales[batch], 0.6) 1267 | totalBoxes = append(totalBoxes, localBoxes...) 1268 | } 1269 | 1270 | // merge the detection from first stage 1271 | fiveTotalBoxes := make([][]float32, len(totalBoxes)) 1272 | for i := range fiveTotalBoxes { 1273 | fiveTotalBoxes[i] = totalBoxes[i][:5] // Using slice notation a[i:j] 1274 | } 1275 | mergedBoxes := nms(fiveTotalBoxes, 0.7, "Union") 1276 | 1277 | var pickedBoxes [][]float32 1278 | for _, idx := range mergedBoxes { 1279 | pickedBoxes = append(pickedBoxes, totalBoxes[idx]) 1280 | } 1281 | 1282 | // refine the boxes 1283 | refinedBoxes := refineBoxes(pickedBoxes) 1284 | 1285 | totalBoxes = convertToSquare(refinedBoxes) 1286 | for i := range totalBoxes { 1287 | totalBoxes[i][0] = float32(math.Round(float64(totalBoxes[i][0]))) 1288 | totalBoxes[i][1] = float32(math.Round(float64(totalBoxes[i][1]))) 1289 | totalBoxes[i][2] = float32(math.Round(float64(totalBoxes[i][2]))) 1290 | totalBoxes[i][3] = float32(math.Round(float64(totalBoxes[i][3]))) 1291 | } 1292 | return totalBoxes 1293 | } 1294 | 1295 | func secondStage(totalBoxes [][]float32, width, height int, img gocv.Mat, rnetModel *tg.Model) [][]float32 { 1296 | numBox := len(totalBoxes) 1297 | 1298 | // pad the bbox 1299 | dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph := pad(totalBoxes, float32(width), float32(height)) 1300 | inputBuf := make([][][][]float32, numBox) 1301 | for i := 0; i < numBox; i++ { 1302 | tmp := gocv.NewMatWithSize(int(tmph[i]), int(tmpw[i]), gocv.MatTypeCV8UC3) 1303 | defer tmp.Close() 1304 | 1305 | imgRoi := img.Region(image.Rect(int(x[i]), int(y[i]), int(ex[i])+1, int(ey[i])+1)) 1306 | defer imgRoi.Close() 1307 | 1308 | tmpRoi := tmp.Region(image.Rect(int(dx[i]), int(dy[i]), int(edx[i])+1, int(edy[i])+1)) 1309 | defer tmpRoi.Close() 1310 | 1311 | imgRoi.CopyTo(&tmpRoi) 1312 | 1313 | resized := gocv.NewMat() 1314 | defer resized.Close() 1315 | gocv.Resize(tmp, &resized, image.Pt(24, 24), 0, 0, gocv.InterpolationLinear) 1316 | 1317 | inputBuf[i] = adjustInput(resized)[0] 1318 | } 1319 | 1320 | inputBufTensor, _ := tf.NewTensor(inputBuf) 1321 | output := rnetModel.Exec([]tf.Output{ 1322 | rnetModel.Op("PartitionedCall", 0), 1323 | rnetModel.Op("PartitionedCall", 1), 1324 | }, map[tf.Output]*tf.Tensor{ 1325 | rnetModel.Op("serving_default_input_2", 0): inputBufTensor, 1326 | }) 1327 | rNetOutput0, ok := output[0].Value().([][]float32) 1328 | if !ok { 1329 | fmt.Println("Failed to convert rNetOutput to [][]float64") 1330 | } 1331 | rNetOutput1, ok := output[1].Value().([][]float32) 1332 | if !ok { 1333 | fmt.Println("Failed to convert rNetOutput to [][]float64") 1334 | } 1335 | 1336 | score := make([]float32, len(rNetOutput1)) 1337 | for i, v := range rNetOutput1 { 1338 | score[i] = v[1] 1339 | } 1340 | 1341 | passed := make([]int, 0) 1342 | for i, v := range score { 1343 | if v > 0.7 { 1344 | passed = append(passed, i) 1345 | } 1346 | } 1347 | 1348 | totalBoxesNew := make([][]float32, 0) 1349 | for _, i := range passed { 1350 | totalBoxesNew = append(totalBoxesNew, totalBoxes[i]) 1351 | } 1352 | totalBoxes = totalBoxesNew 1353 | if len(totalBoxes) == 0 { 1354 | fmt.Println("No face detected! Stop!") 1355 | } 1356 | 1357 | for i, idx := range passed { 1358 | totalBoxes[i][4] = rNetOutput1[idx][1] 1359 | } 1360 | 1361 | reg := make([][]float32, len(passed)) 1362 | for i, idx := range passed { 1363 | reg[i] = rNetOutput0[idx] 1364 | } 1365 | 1366 | // nms 1367 | pick := nms(totalBoxes, 0.7, "Union") 1368 | var newPickedBoxes [][]float32 1369 | var pickedReg [][]float32 1370 | for _, i := range pick { 1371 | newPickedBoxes = append(newPickedBoxes, totalBoxes[i]) 1372 | pickedReg = append(pickedReg, reg[i]) 1373 | } 1374 | 1375 | calibratedBoxes := CalibrateBox(newPickedBoxes, pickedReg) 1376 | 1377 | squaredBoxes := convertToSquare(calibratedBoxes) 1378 | for i := range squaredBoxes { 1379 | for j := 0; j < 4; j++ { 1380 | squaredBoxes[i][j] = float32(math.Round(float64(squaredBoxes[i][j]))) 1381 | } 1382 | } 1383 | return squaredBoxes 1384 | } 1385 | 1386 | func thirdStage(squaredBoxes [][]float32, width, height int, img gocv.Mat, onetModel *tg.Model) ([][]float32, [][]float32) { 1387 | numBox := len(squaredBoxes) 1388 | totalBoxes := squaredBoxes 1389 | // pad the bbox 1390 | dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph := pad(totalBoxes, float32(width), float32(height)) 1391 | 1392 | // (3, 48, 48) is the input shape for ONet 1393 | inputBuf := make([][][][]float32, numBox) 1394 | for i := 0; i < numBox; i++ { 1395 | tmp := gocv.NewMatWithSize(int(tmph[i]), int(tmpw[i]), gocv.MatTypeCV8UC3) 1396 | defer tmp.Close() 1397 | 1398 | imgRoi := img.Region(image.Rect(int(x[i]), int(y[i]), int(ex[i])+1, int(ey[i])+1)) 1399 | defer imgRoi.Close() 1400 | 1401 | tmpRoi := tmp.Region(image.Rect(int(dx[i]), int(dy[i]), int(edx[i])+1, int(edy[i])+1)) 1402 | defer tmpRoi.Close() 1403 | 1404 | imgRoi.CopyTo(&tmpRoi) 1405 | 1406 | resized := gocv.NewMat() 1407 | defer resized.Close() 1408 | gocv.Resize(tmp, &resized, image.Pt(48, 48), 0, 0, gocv.InterpolationLinear) 1409 | 1410 | inputBuf[i] = adjustInput(resized)[0] 1411 | } 1412 | 1413 | inputBufTensor, _ := tf.NewTensor(inputBuf) 1414 | output := onetModel.Exec([]tf.Output{ 1415 | onetModel.Op("PartitionedCall", 0), 1416 | onetModel.Op("PartitionedCall", 1), 1417 | onetModel.Op("PartitionedCall", 2), 1418 | }, map[tf.Output]*tf.Tensor{ 1419 | onetModel.Op("serving_default_input_3", 0): inputBufTensor, 1420 | }) 1421 | oNetOutput0, ok := output[1].Value().([][]float32) 1422 | if !ok { 1423 | fmt.Println("Failed to convert rNetOutput to [][]float64") 1424 | } 1425 | oNetOutput1, ok := output[0].Value().([][]float32) 1426 | if !ok { 1427 | fmt.Println("Failed to convert rNetOutput to [][]float64") 1428 | } 1429 | oNetOutput2, ok := output[2].Value().([][]float32) 1430 | if !ok { 1431 | fmt.Println("Failed to convert rNetOutput to [][]float64") 1432 | } 1433 | 1434 | score := make([]float32, len(oNetOutput2)) 1435 | for i, v := range oNetOutput2 { 1436 | score[i] = v[1] 1437 | } 1438 | 1439 | passed := make([]int, 0) 1440 | for i, v := range score { 1441 | if v > 0.8 { 1442 | passed = append(passed, i) 1443 | } 1444 | } 1445 | 1446 | totalBoxesNew := make([][]float32, 0) 1447 | for _, i := range passed { 1448 | totalBoxesNew = append(totalBoxesNew, totalBoxes[i]) 1449 | } 1450 | totalBoxes = totalBoxesNew 1451 | if len(totalBoxes) == 0 { 1452 | fmt.Println("No face detected! Stop!") 1453 | } 1454 | 1455 | for i, idx := range passed { 1456 | totalBoxes[i][4] = oNetOutput2[idx][1] 1457 | } 1458 | 1459 | reg := make([][]float32, len(passed)) 1460 | for i, idx := range passed { 1461 | reg[i] = oNetOutput1[idx] 1462 | } 1463 | 1464 | points := make([][]float32, len(passed)) 1465 | for i, idx := range passed { 1466 | points[i] = oNetOutput0[idx] 1467 | } 1468 | 1469 | bbW := make([]float32, len(totalBoxes)) 1470 | bbH := make([]float32, len(totalBoxes)) 1471 | 1472 | for i := range totalBoxes { 1473 | bbW[i] = totalBoxes[i][2] - totalBoxes[i][0] + 1 1474 | bbH[i] = totalBoxes[i][3] - totalBoxes[i][1] + 1 1475 | } 1476 | 1477 | for i := 0; i < len(points); i++ { 1478 | for j := 0; j < 5; j++ { 1479 | points[i][j] = totalBoxes[i][0] + bbW[i]*points[i][j] 1480 | points[i][j+5] = totalBoxes[i][1] + bbH[i]*points[i][j+5] 1481 | } 1482 | } 1483 | 1484 | // nms 1485 | calibratedBoxes := CalibrateBox(totalBoxes, reg) 1486 | pick := nms(calibratedBoxes, 0.7, "Min") 1487 | 1488 | var thirdPickedBoxes [][]float32 1489 | var pickedPoints [][]float32 1490 | for _, i := range pick { 1491 | thirdPickedBoxes = append(thirdPickedBoxes, calibratedBoxes[i]) 1492 | pickedPoints = append(pickedPoints, points[i]) 1493 | } 1494 | return thirdPickedBoxes, pickedPoints 1495 | } 1496 | 1497 | func alignFace(thirdPickedBoxes, pickedPoints [][]float32, img gocv.Mat) [][][][]float32 { 1498 | if len(thirdPickedBoxes) == 0 || len(pickedPoints) == 0 { 1499 | fmt.Println("return nil") 1500 | } 1501 | var pImgs [][][][]float32 1502 | for i := 0; i < len(pickedPoints); i++ { 1503 | p := make([][]float32, 5) 1504 | for j := 0; j < 5; j++ { 1505 | p[j] = make([]float32, 2) 1506 | p[j][0] = pickedPoints[i][j] 1507 | p[j][1] = pickedPoints[i][j+5] 1508 | } 1509 | b := thirdPickedBoxes[i] 1510 | matFace := preprocessMat(img, b, p) 1511 | 1512 | val := matFace.GetVecbAt(0, 0) 1513 | fmt.Printf("matFace[0,0]: %v\n", val) 1514 | val = matFace.GetVecbAt(12, 23) 1515 | fmt.Printf("matFace[12,23]: %v\n", val) 1516 | 1517 | sliceFace := matToSlice(matFace) 1518 | pImgs = append(pImgs, sliceFace) 1519 | } 1520 | return pImgs 1521 | } 1522 | 1523 | func recognizeFace(pImgs [][][][]float32, qmfModel *tg.Model, regEmbeddings [][]float32, bSize int, regFiles []string) { 1524 | if len(pImgs) == 0 { 1525 | fmt.Println("return nil") 1526 | } 1527 | 1528 | for _, pImg := range pImgs { 1529 | transformedFaces := generateEmbeddings(pImg) 1530 | recognitionStart_1 := time.Now() 1531 | frameEmbeddings := qmfModel.Exec([]tf.Output{ 1532 | qmfModel.Op("PartitionedCall", 0), 1533 | }, map[tf.Output]*tf.Tensor{ 1534 | qmfModel.Op("serving_default_input.1", 0): transformedFaces, 1535 | }) 1536 | elapsed := time.Since(recognitionStart_1) 1537 | elapsedMilliseconds := elapsed.Milliseconds() 1538 | fmt.Printf("----------------> ForwardPass execution time: %d milliseconds\n", elapsedMilliseconds) 1539 | 1540 | frameEmbeddingsFloat32, ok := frameEmbeddings[0].Value().([][]float32) 1541 | if !ok { 1542 | fmt.Println("Failed to convert rNetOutput to [][]float32") 1543 | } 1544 | 1545 | qmfScores := similarityNoPair(frameEmbeddingsFloat32, regEmbeddings) 1546 | nB := int(math.Ceil(float64(len(qmfScores)) / float64(bSize))) 1547 | 1548 | classIDs := make([]string, nB) 1549 | recScores := make([]float32, nB) 1550 | targetTh := float32(-0.4) 1551 | 1552 | for i := 0; i < nB; i++ { 1553 | startIndex := i * bSize 1554 | endIndex := (i + 1) * bSize 1555 | if endIndex > len(qmfScores) { 1556 | endIndex = len(qmfScores) 1557 | } 1558 | qmfSlice := qmfScores[startIndex:endIndex] 1559 | 1560 | maxScore := qmfSlice[0] 1561 | maxIndex := 0 1562 | for j, score := range qmfSlice { 1563 | if score > maxScore { 1564 | maxScore = score 1565 | maxIndex = j 1566 | } 1567 | } 1568 | 1569 | if maxScore > targetTh { 1570 | classIDs[i] = filepath.Base(filepath.Dir(regFiles[maxIndex])) 1571 | } else { 1572 | classIDs[i] = "unknown" 1573 | } 1574 | recScores[i] = maxScore 1575 | } 1576 | 1577 | fmt.Println("----------------- classIDs ------->", classIDs) 1578 | fmt.Println("----------------- recScores ------->", recScores) 1579 | } 1580 | } 1581 | 1582 | func fakeRun(net *tg.Model, outputOpName string, s0, s1, s2 int) { 1583 | fakeData := make([][][][]float32, 1) 1584 | fakeData[0] = make([][][]float32, s0) 1585 | for i := 0; i < s0; i++ { 1586 | fakeData[0][i] = make([][]float32, s1) 1587 | for j := 0; j < s1; j++ { 1588 | fakeData[0][i][j] = make([]float32, s2) 1589 | for k := 0; k < s2; k++ { 1590 | fakeData[0][i][j][k] = 1.0 1591 | } 1592 | } 1593 | } 1594 | fakeInput, _ := tf.NewTensor(fakeData) 1595 | 1596 | _ = net.Exec([]tf.Output{ 1597 | net.Op("PartitionedCall", 0), 1598 | }, map[tf.Output]*tf.Tensor{ 1599 | net.Op(outputOpName, 0): fakeInput, 1600 | }) 1601 | } 1602 | 1603 | func main() { 1604 | 1605 | //************************************************************************************ 1606 | // preprocessing and loading 1607 | //************************************************************************************ 1608 | imagePath := os.Args[1] 1609 | regFilesPath := os.Args[2] 1610 | embeddingsFilePath := os.Args[3] 1611 | mtcnnModelPath := os.Args[4] 1612 | magfaceModelPath := os.Args[5] 1613 | 1614 | img := gocv.IMRead(imagePath, gocv.IMReadColor) 1615 | defer img.Close() 1616 | 1617 | height := img.Size()[0] 1618 | width := img.Size()[1] 1619 | minDetSize := 12 1620 | minSize := 50 1621 | var scales []float64 1622 | m := float64(minDetSize) / float64(minSize) 1623 | minL := math.Min(float64(height), float64(width)) * m 1624 | factorCount := 0 1625 | factor := 0.709 1626 | for minL > float64(minDetSize) { 1627 | scales = append(scales, m*math.Pow(factor, float64(factorCount))) 1628 | minL *= factor 1629 | factorCount++ 1630 | } 1631 | pnetModel := tg.LoadModel(mtcnnModelPath+"/pnet_pb", []string{"serve"}, nil) 1632 | rnetModel := tg.LoadModel(mtcnnModelPath+"/rnet_pb", []string{"serve"}, nil) 1633 | onetModel := tg.LoadModel(mtcnnModelPath+"/onet_pb", []string{"serve"}, nil) 1634 | qmfModel := tg.LoadModel(magfaceModelPath, []string{"serve"}, nil) 1635 | 1636 | regEmbeddings, err := loadNpy(embeddingsFilePath) 1637 | if err != nil { 1638 | fmt.Println("Error:", err) 1639 | return 1640 | } 1641 | 1642 | regFiles, _ := getRegFiles(regFilesPath) 1643 | bSize := len(regFiles) 1644 | 1645 | fakeRun(pnetModel, "serving_default_input_1", 12, 12, 3) 1646 | fakeRun(rnetModel, "serving_default_input_2", 24, 24, 3) 1647 | fakeRun(onetModel, "serving_default_input_3", 48, 48, 3) 1648 | fakeRun(qmfModel, "serving_default_input.1", 3, 112, 112) 1649 | 1650 | //************************************************************************************ 1651 | // detect face 1652 | //************************************************************************************ 1653 | 1654 | //////////////////////////////////////////// 1655 | // first stage 1656 | //////////////////////////////////////////// 1657 | start := time.Now() 1658 | totalBoxes := firstStage(scales, img, pnetModel) 1659 | elapsed := time.Since(start) 1660 | elapsedMilliseconds := elapsed.Milliseconds() 1661 | fmt.Printf("----------------> First detection execution time: %d milliseconds\n", elapsedMilliseconds) 1662 | 1663 | //////////////////////////////////////////// 1664 | // second stage 1665 | //////////////////////////////////////////// 1666 | detection2Start := time.Now() 1667 | squaredBoxes := secondStage(totalBoxes, width, height, img, rnetModel) 1668 | elapsed = time.Since(detection2Start) 1669 | elapsedMilliseconds = elapsed.Milliseconds() 1670 | fmt.Printf("----------------> Second detection execution time: %d milliseconds\n", elapsedMilliseconds) 1671 | 1672 | ////////////////////////////////////////////// 1673 | //// third stage 1674 | ////////////////////////////////////////////// 1675 | detection3Start := time.Now() 1676 | thirdPickedBoxes, pickedPoints := thirdStage(squaredBoxes, width, height, img, onetModel) 1677 | elapsed = time.Since(detection3Start) 1678 | elapsedMilliseconds = elapsed.Milliseconds() 1679 | fmt.Printf("----------------> Third detection execution time: %d milliseconds\n", elapsedMilliseconds) 1680 | 1681 | //************************************************************************************ 1682 | // align face 1683 | //************************************************************************************ 1684 | alignStart := time.Now() 1685 | pImgs := alignFace(thirdPickedBoxes, pickedPoints, img) 1686 | elapsed = time.Since(alignStart) 1687 | elapsedMilliseconds = elapsed.Milliseconds() 1688 | fmt.Printf("----------------> Face alignment execution time: %d milliseconds\n", elapsedMilliseconds) 1689 | 1690 | ////************************************************************************************ 1691 | //// recognize face 1692 | ////************************************************************************************ 1693 | recognitionStart := time.Now() 1694 | recognizeFace(pImgs, qmfModel, regEmbeddings, bSize, regFiles) 1695 | elapsed = time.Since(recognitionStart) 1696 | elapsedMilliseconds = elapsed.Milliseconds() 1697 | fmt.Printf("----------------> Face recognition execution time: %d milliseconds\n", elapsedMilliseconds) 1698 | elapsed = time.Since(start) 1699 | elapsedMilliseconds = elapsed.Milliseconds() 1700 | fmt.Printf("----------------> Total execution time: %d milliseconds\n", elapsedMilliseconds) 1701 | } 1702 | --------------------------------------------------------------------------------