MainController.java
6.92 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
package com.viontech.match.controller;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.viontech.keliu.storage.Storage;
import com.viontech.keliu.util.DateUtil;
import com.viontech.match.PoolEnum;
import com.viontech.match.entity.*;
import com.viontech.match.entity.vo.RequestData;
import com.viontech.match.repository.FaceRecognitionRepository;
import com.viontech.match.service.AlgService;
import com.viontech.match.service.MatchService;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
/**
* .
*
* @author 谢明辉
* @date 2020/11/23
*/
@RestController
@RequestMapping("/alg")
public class MainController {
private final ThreadPoolExecutor poolExecutor = new ThreadPoolExecutor(20, 30, 1, TimeUnit.MINUTES, new LinkedBlockingQueue<>(), new ThreadPoolExecutor.CallerRunsPolicy());
private final AtomicInteger integer = new AtomicInteger();
@Resource
ObjectMapper objectMapper;
@Resource
private MatchService matchService;
@Resource
private FaceRecognitionRepository faceRecognitionRepository;
@Resource
private AlgService algService;
@Resource
private Storage featureStorage;
@PostMapping("/match")
public Object match(RequestData requestData) throws Exception {
MultipartFile file = requestData.getFile();
Double[] data = matchService.extractFeature(file);
Person person = new Person().setFeature(data);
long begin = System.currentTimeMillis();
List<MatchResult> match = algService.matchPerson(person, requestData.getPoolId());
long timeConsumingMatching = System.currentTimeMillis() - begin;
List<String> collect = match.stream().map(MatchResult::getPersonId).collect(Collectors.toList());
List<FaceRecognition> faceRecognitions = faceRecognitionRepository.queryFaceRecognitionByUnids(collect);
if (faceRecognitions == null) {
return null;
}
Map<String, FaceRecognition> map = faceRecognitions.stream().collect(Collectors.toMap(FaceRecognition::getUnid, x -> x, (x, y) -> y));
for (MatchResult item : match) {
FaceRecognition faceRecognition = map.get(item.getPersonId());
if (faceRecognition != null) {
faceRecognition.setMatchScore(item.getScore());
}
}
Collection<FaceRecognition> values = map.values();
List<FaceRecognition> collect1 = values.stream().sorted((o1, o2) -> {
if (o1.getMatchScore() > o2.getMatchScore()) {
return -1;
} else if (o1.getMatchScore().equals(o2.getMatchScore())) {
return 0;
} else {
return 1;
}
}).collect(Collectors.toList());
HashMap<String, Object> resultMap = new HashMap<>(2);
resultMap.put("time", timeConsumingMatching);
resultMap.put("result", collect1);
return resultMap;
}
@GetMapping("/poolIds")
public Object getPoolIds() throws IOException {
return algService.queryPoolList();
}
@GetMapping("/create/{num}")
public Object createPool(@PathVariable("num") Integer num) throws Exception {
PoolEnum pool;
switch (num) {
case 2000:
pool = PoolEnum.POOL_2K;
break;
case 5000:
pool = PoolEnum.POOL_5K;
break;
case 20000:
pool = PoolEnum.POOL_2W;
break;
case 100000:
pool = PoolEnum.POOL_10W;
break;
case 500000:
pool = PoolEnum.POOL_50W;
break;
case 1000000:
pool = PoolEnum.POOL_100W;
break;
case 2000000:
pool = PoolEnum.POOL_200W;
break;
default:
throw new RuntimeException("请输入[2000,5000,20000,100000,500000,1000000,2000000]");
}
algService.deletePool(pool.name);
algService.createPool(pool.name);
Date end = DateUtil.addDays(new Date(), -1);
Date start = DateUtil.addDays(end, -30);
integer.set(0);
List<Date> days = DateUtil.getDaysBetweenDates(start, end);
for (Date day : days) {
List<FaceRecognition> faceRecognitions = faceRecognitionRepository.queryFaceRecognitions(day, 2L);
for (FaceRecognition faceRecognition : faceRecognitions) {
poolExecutor.execute(() -> {
try {
if (integer.get() >= pool.num) {
return;
}
Double[] feature = getFeature(faceRecognition);
if (feature == null) {
return;
}
Person person = new Person();
person.setPersonId(faceRecognition.getPersonUnid());
person.setFeature(feature);
person.setId(faceRecognition.getUnid());
Object o = algService.addPerson(person, pool.name);
System.out.println(objectMapper.writeValueAsString(o));
integer.addAndGet(1);
} catch (Exception e) {
e.printStackTrace();
}
});
if (integer.get() > pool.num) {
return "success";
}
while (poolExecutor.getQueue().size() > 10000) {
System.out.println(poolExecutor.getQueue().size());
TimeUnit.SECONDS.sleep(4);
}
if (integer.get() > pool.num) {
return "success";
}
}
}
return "success";
}
private Double[] getFeature(FaceRecognition face) throws JsonProcessingException {
Short faceFeatureType = face.getFaceFeatureType();
String facePic = face.getFacePic();
if (faceFeatureType == 1) {
facePic = facePic.replace("face-0.jpg", "face-F.jpg");
}
String json = (String) featureStorage.getItem(face.getChannelSerialnum() + "/" + facePic);
Feature feature = objectMapper.readValue(json, Feature.class);
List<Data> datas = feature.getDatas();
if (datas != null && datas.size() > 0) {
Data data = datas.get(0);
Double[] data1 = data.getData();
if (data1.length == 512) {
return data1;
}
}
return null;
}
}