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 {
private String personPoolId;
private List<String> unionPersonPoolId;
private Integer size;
private Boolean agg = false;
private FaceFeature newFaceFeature;
......
src/main/java/com/viontech/match/service/PersonService.java
......@@ -29,6 +29,8 @@ import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.script.Script;
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.Aggregations;
import org.elasticsearch.search.aggregations.BucketOrder;
......@@ -41,6 +43,7 @@ import org.springframework.stereotype.Service;
import javax.annotation.Resource;
import java.io.IOException;
import java.text.ParseException;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;
......@@ -76,6 +79,7 @@ public class PersonService {
String poolId = requestVo.getPersonPoolId();
List<String> unionPersonPoolId = requestVo.getUnionPersonPoolId();
Integer size = requestVo.getSize();
Boolean agg = requestVo.getAgg();
List<String> poolIds = new ArrayList<>();
List<Person> matchFaces = new ArrayList<>();
List<Person> matchBodies = new ArrayList<>();
......@@ -94,9 +98,9 @@ public class PersonService {
pool.setPersonPoolId(id);
if (poolService.existPool(id)) {
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);
List<Person> body = matchPerson(id, requestVo.getPerson(), 1, size);
List<Person> body = matchPerson(id, requestVo.getPerson(), 1, size, agg);
matchBodies.addAll(body);
} else {
pool.setStatus(1);
......@@ -250,14 +254,14 @@ public class PersonService {
* @return 匹配结果
* @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<>();
int matchResultSize;
if (type == 0) {
matchFace(poolId, matchResult, person, size);
matchFace(poolId, matchResult, person, size, agg);
matchResultSize = Constant.FACE_MATCH_RESULT_SIZE;
} else {
matchBody(poolId, matchResult, person, size);
matchBody(poolId, matchResult, person, size, agg);
matchResultSize = Constant.BODY_MATCH_RESULT_SIZE;
}
if (size != null) {
......@@ -272,7 +276,7 @@ public class PersonService {
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();
if (faceFeatures != null && faceFeatures.size() > 0) {
for (FaceFeature faceFeature : faceFeatures) {
......@@ -282,15 +286,15 @@ public class PersonService {
continue;
}
SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.FACE_MATCH_RESULT_SIZE : size, feature, person, 0);
matchResult.addAll(match0(searchRequest));
SearchRequest searchRequest = getSearchRequest(poolId, size == null ? Constant.FACE_MATCH_RESULT_SIZE : size, feature, person, 0, agg);
matchResult.addAll(match0(searchRequest, agg));
}
} else {
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();
if (bodyFeatures != null && bodyFeatures.size() > 0) {
for (BodyFeature faceFeature : bodyFeatures) {
......@@ -306,14 +310,14 @@ public class PersonService {
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;
BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
if (type == 0) {
......@@ -353,59 +357,63 @@ public class PersonService {
ScriptScoreQueryBuilder queryBuilder = QueryBuilders.scriptScoreQuery(boolQuery, script);
MaxAggregationBuilder maxScoreAgg = AggregationBuilders.max("max_score").script(new Script("_score"));
TermsAggregationBuilder personIdAgg = AggregationBuilders.terms("by_personId").field("personId");
personIdAgg.subAggregation(maxScoreAgg);
personIdAgg.order(BucketOrder.aggregation("max_score", false));
SearchSourceBuilder builder = new SearchSourceBuilder()
.size(matchResultSize)
.query(queryBuilder)
// .fetchSource(FETCH_SOURCE, null)
.aggregation(personIdAgg);
SearchSourceBuilder builder = new SearchSourceBuilder().query(queryBuilder);
if (agg) {
MaxAggregationBuilder maxScoreAgg = AggregationBuilders.max("max_score").script(new Script("_score"));
TermsAggregationBuilder personIdAgg = AggregationBuilders.terms("by_personId").field("personId");
personIdAgg.subAggregation(maxScoreAgg);
personIdAgg.order(BucketOrder.aggregation("max_score", false));
personIdAgg.size(matchResultSize);
builder.aggregation(personIdAgg);
} else {
builder.fetchSource(FETCH_SOURCE, null)
.size(matchResultSize);
}
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<>();
SearchResponse search = client.search(searchRequest, RequestOptions.DEFAULT);
Aggregations aggregations = search.getAggregations();
Terms byPersonId = aggregations.get("by_personId");
List<? extends Terms.Bucket> buckets = byPersonId.getBuckets();
for (Terms.Bucket bucket : buckets) {
String personId = (String) bucket.getKey();
Max maxScore = bucket.getAggregations().get("max_score");
double value = maxScore.getValue();
Person person = new Person().setPersonId(personId).setScore((float) value);
persons.add(person);
if (agg) {
Aggregations aggregations = search.getAggregations();
Terms byPersonId = aggregations.get("by_personId");
List<? extends Terms.Bucket> buckets = byPersonId.getBuckets();
for (Terms.Bucket bucket : buckets) {
String personId = (String) bucket.getKey();
Max maxScore = bucket.getAggregations().get("max_score");
double value = maxScore.getValue();
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;
}
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!