Skip to content

谢明辉 / VVAS-Match

Go to a project
Commit fda29510 by xmh
Options

添加是否聚合匹配的选项

1 parent ce426275
Inline Side-by-side
Showing 2 changed files with 67 additions and 58 deletions
src/main/java/com/viontech/match/entity/vo/RequestVo.java
...@@ -43,6 +43,7 @@ public class RequestVo { ...@@ -43,6 +43,7 @@ public class RequestVo {
private String personPoolId; private String personPoolId;
private List<String> unionPersonPoolId; private List<String> unionPersonPoolId;
private Integer size; private Integer size;
private Boolean agg = false;
private FaceFeature newFaceFeature; private FaceFeature newFaceFeature;
......
src/main/java/com/viontech/match/service/PersonService.java
...@@ -29,6 +29,8 @@ import org.elasticsearch.index.reindex.DeleteByQueryRequest; ...@@ -29,6 +29,8 @@ import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.rest.RestStatus; import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.script.Script; import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType; import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.AggregationBuilders; import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations; import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.BucketOrder; import org.elasticsearch.search.aggregations.BucketOrder;
...@@ -41,6 +43,7 @@ import org.springframework.stereotype.Service; ...@@ -41,6 +43,7 @@ import org.springframework.stereotype.Service;
import javax.annotation.Resource; import javax.annotation.Resource;
import java.io.IOException; import java.io.IOException;
import java.text.ParseException;
import java.util.*; import java.util.*;
import java.util.stream.Collectors; import java.util.stream.Collectors;
import java.util.stream.Stream; import java.util.stream.Stream;
...@@ -76,6 +79,7 @@ public class PersonService { ...@@ -76,6 +79,7 @@ public class PersonService {
String poolId = requestVo.getPersonPoolId(); String poolId = requestVo.getPersonPoolId();
List<String> unionPersonPoolId = requestVo.getUnionPersonPoolId(); List<String> unionPersonPoolId = requestVo.getUnionPersonPoolId();
Integer size = requestVo.getSize(); Integer size = requestVo.getSize();
Boolean agg = requestVo.getAgg();
List<String> poolIds = new ArrayList<>(); List<String> poolIds = new ArrayList<>();
List<Person> matchFaces = new ArrayList<>(); List<Person> matchFaces = new ArrayList<>();
List<Person> matchBodies = new ArrayList<>(); List<Person> matchBodies = new ArrayList<>();
...@@ -94,9 +98,9 @@ public class PersonService { ...@@ -94,9 +98,9 @@ public class PersonService {
pool.setPersonPoolId(id); pool.setPersonPoolId(id);
if (poolService.existPool(id)) { if (poolService.existPool(id)) {
pool.setStatus(0); pool.setStatus(0);
List<Person> face = matchPerson(id, requestVo.getPerson(), 0, size); List<Person> face = matchPerson(id, requestVo.getPerson(), 0, size, agg);
matchFaces.addAll(face); matchFaces.addAll(face);
List<Person> body = matchPerson(id, requestVo.getPerson(), 1, size); List<Person> body = matchPerson(id, requestVo.getPerson(), 1, size, agg);
matchBodies.addAll(body); matchBodies.addAll(body);
} else { } else {
pool.setStatus(1); pool.setStatus(1);
...@@ -250,14 +254,14 @@ public class PersonService { ...@@ -250,14 +254,14 @@ public class PersonService {
* @return 匹配结果 * @return 匹配结果
* @throws Exception -- * @throws Exception --
*/ */
public List<Person> matchPerson(String poolId, Person person, int type, Integer size) throws Exception { public List<Person> matchPerson(String poolId, Person person, int type, Integer size, Boolean agg) throws Exception {
List<Person> matchResult = new ArrayList<>(); List<Person> matchResult = new ArrayList<>();
int matchResultSize; int matchResultSize;
if (type == 0) { if (type == 0) {
matchFace(poolId, matchResult, person, size); matchFace(poolId, matchResult, person, size, agg);
matchResultSize = Constant.FACE_MATCH_RESULT_SIZE; matchResultSize = Constant.FACE_MATCH_RESULT_SIZE;
} else { } else {
matchBody(poolId, matchResult, person, size); matchBody(poolId, matchResult, person, size, agg);
matchResultSize = Constant.BODY_MATCH_RESULT_SIZE; matchResultSize = Constant.BODY_MATCH_RESULT_SIZE;
} }
if (size != null) { if (size != null) {
...@@ -272,7 +276,7 @@ public class PersonService { ...@@ -272,7 +276,7 @@ public class PersonService {
return matchResult; return matchResult;
} }
private void matchFace(String poolId, List<Person> matchResult, Person person, Integer size) throws Exception { private void matchFace(String poolId, List<Person> matchResult, Person person, Integer size, Boolean agg) throws Exception {
List<FaceFeature> faceFeatures = person.getFaceFeatures(); List<FaceFeature> faceFeatures = person.getFaceFeatures();
if (faceFeatures != null && faceFeatures.size() > 0) { if (faceFeatures != null && faceFeatures.size() > 0) {
for (FaceFeature faceFeature : faceFeatures) { for (FaceFeature faceFeature : faceFeatures) {
...@@ -282,15 +286,15 @@ public class PersonService { ...@@ -282,15 +286,15 @@ public class PersonService {
continue; continue;
} }
SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.FACE_MATCH_RESULT_SIZE : size, feature, person, 0); SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.FACE_MATCH_RESULT_SIZE : size, feature, person, 0, agg);
matchResult.addAll(match0(searchRequest)); matchResult.addAll(match0(searchRequest, agg));
} }
} else { } else {
log.info("no face feature"); log.info("no face feature");
} }
} }
private void matchBody(String poolId, List<Person> matchResult, Person person, Integer size) throws Exception { private void matchBody(String poolId, List<Person> matchResult, Person person, Integer size, Boolean agg) throws Exception {
List<BodyFeature> bodyFeatures = person.getBodyFeatures(); List<BodyFeature> bodyFeatures = person.getBodyFeatures();
if (bodyFeatures != null && bodyFeatures.size() > 0) { if (bodyFeatures != null && bodyFeatures.size() > 0) {
for (BodyFeature faceFeature : bodyFeatures) { for (BodyFeature faceFeature : bodyFeatures) {
...@@ -306,14 +310,14 @@ public class PersonService { ...@@ -306,14 +310,14 @@ public class PersonService {
feature = Arrays.copyOfRange(feature, 3, Constant.BODY_FEATURE_DIMS_2048 + 3); feature = Arrays.copyOfRange(feature, 3, Constant.BODY_FEATURE_DIMS_2048 + 3);
} }
SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.BODY_MATCH_RESULT_SIZE : size, feature, person, 1); SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.BODY_MATCH_RESULT_SIZE : size, feature, person, 1, agg);
matchResult.addAll(match0(searchRequest)); matchResult.addAll(match0(searchRequest, agg));
} }
} }
} }
private SearchRequest getSearchRequest(String poolId, Integer matchResultSize, Double[] feature, Person person, int type) { private SearchRequest getSearchRequest(String poolId, Integer matchResultSize, Double[] feature, Person person, int type, Boolean agg) {
Script script; Script script;
BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
if (type == 0) { if (type == 0) {
...@@ -353,59 +357,63 @@ public class PersonService { ...@@ -353,59 +357,63 @@ public class PersonService {
ScriptScoreQueryBuilder queryBuilder = QueryBuilders.scriptScoreQuery(boolQuery, script); ScriptScoreQueryBuilder queryBuilder = QueryBuilders.scriptScoreQuery(boolQuery, script);
MaxAggregationBuilder maxScoreAgg = AggregationBuilders.max("max_score").script(new Script("_score")); SearchSourceBuilder builder = new SearchSourceBuilder().query(queryBuilder);
TermsAggregationBuilder personIdAgg = AggregationBuilders.terms("by_personId").field("personId"); if (agg) {
personIdAgg.subAggregation(maxScoreAgg); MaxAggregationBuilder maxScoreAgg = AggregationBuilders.max("max_score").script(new Script("_score"));
personIdAgg.order(BucketOrder.aggregation("max_score", false)); TermsAggregationBuilder personIdAgg = AggregationBuilders.terms("by_personId").field("personId");
personIdAgg.subAggregation(maxScoreAgg);
SearchSourceBuilder builder = new SearchSourceBuilder() personIdAgg.order(BucketOrder.aggregation("max_score", false));
.size(matchResultSize) personIdAgg.size(matchResultSize);
.query(queryBuilder) builder.aggregation(personIdAgg);
// .fetchSource(FETCH_SOURCE, null) } else {
.aggregation(personIdAgg); builder.fetchSource(FETCH_SOURCE, null)
.size(matchResultSize);
}
return new SearchRequest(poolId).source(builder); return new SearchRequest(poolId).source(builder);
} }
private List<Person> match0(SearchRequest searchRequest) throws Exception { private List<Person> match0(SearchRequest searchRequest, Boolean agg) throws Exception {
List<Person> persons = new ArrayList<>(); List<Person> persons = new ArrayList<>();
SearchResponse search = client.search(searchRequest, RequestOptions.DEFAULT); SearchResponse search = client.search(searchRequest, RequestOptions.DEFAULT);
Aggregations aggregations = search.getAggregations(); if (agg) {
Terms byPersonId = aggregations.get("by_personId"); Aggregations aggregations = search.getAggregations();
List<? extends Terms.Bucket> buckets = byPersonId.getBuckets(); Terms byPersonId = aggregations.get("by_personId");
for (Terms.Bucket bucket : buckets) { List<? extends Terms.Bucket> buckets = byPersonId.getBuckets();
String personId = (String) bucket.getKey(); for (Terms.Bucket bucket : buckets) {
Max maxScore = bucket.getAggregations().get("max_score"); String personId = (String) bucket.getKey();
double value = maxScore.getValue(); Max maxScore = bucket.getAggregations().get("max_score");
Person person = new Person().setPersonId(personId).setScore((float) value); double value = maxScore.getValue();
persons.add(person); Person person = new Person().setPersonId(personId).setScore((float) value);
persons.add(person);
}
} else {
SearchHits hits = search.getHits();
SearchHit[] hits1 = hits.getHits();
for (SearchHit item : hits1) {
Map<String, Object> source = item.getSourceAsMap();
Person p = new Person();
p.setPersonId((String) source.get("personId"));
p.setAge((Integer) source.get("age"));
p.setGender((String) source.get("gender"));
p.setChannelSerialNum((String) source.get("channelSerialNum"));
p.setBodyType((Integer) source.get("body_type"));
p.setCounttime(Optional.ofNullable((String) source.get("counttime"))
.map(x -> {
try {
return Constant.DATE_FORMAT.get().parse(x);
} catch (ParseException ignore) {
}
return null;
})
.orElse(null));
p.setScore(item.getScore());
p.setPersonPoolId(item.getIndex());
persons.add(p);
}
} }
// SearchHits hits = search.getHits();
// SearchHit[] hits1 = hits.getHits();
// for (SearchHit item : hits1) {
// Map<String, Object> source = item.getSourceAsMap();
// Person p = new Person();
// p.setPersonId((String) source.get("personId"));
// p.setAge((Integer) source.get("age"));
// p.setGender((String) source.get("gender"));
// p.setChannelSerialNum((String) source.get("channelSerialNum"));
// p.setBodyType((Integer) source.get("body_type"));
//
// p.setCounttime(Optional.ofNullable((String) source.get("counttime"))
// .map(x -> {
// try {
// return Constant.DATE_FORMAT.get().parse(x);
// } catch (ParseException ignore) {
// }
// return null;
// })
// .orElse(null));
//
// p.setScore(item.getScore());
// p.setPersonPoolId(item.getIndex());
// persons.add(p);
// }
return persons; return persons;
} }
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!