PersonService.java 7.32 KB
package com.viontech.match.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.viontech.keliu.model.FaceFeature;
import com.viontech.keliu.model.Person;
import com.viontech.keliu.model.Pool;
import com.viontech.match.entity.vo.RequestVo;
import com.viontech.match.entity.vo.ResponseVo;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScriptScoreQueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * .
 *
 * @author 谢明辉
 * @date 2020/11/27
 */

@Service
@Slf4j
public class PersonService {
    @Resource
    private RestHighLevelClient client;
    @Resource
    private PoolService poolService;
    @Resource
    private ObjectMapper objectMapper;

    /**
     * 人员比对
     */
    public ResponseVo matchPerson(RequestVo requestVo) throws Exception {
        String rid = requestVo.getRid();
        String poolId = requestVo.getPersonPoolId();
        List<String> unionPersonPoolId = requestVo.getUnionPersonPoolId();
        List<String> poolIds = new ArrayList<>();
        List<Person> result = new ArrayList<>();
        List<Pool> poolStatus = new ArrayList<>();

        if (unionPersonPoolId != null && unionPersonPoolId.size() > 0) {
            poolIds.addAll(unionPersonPoolId);
        }
        if (poolId != null) {
            poolIds.add(poolId);
        }
        log.info("人员匹配操作开始,PoolIds:[{}]", poolIds.toString());
        try {
            for (String id : poolIds) {
                Pool pool = new Pool();
                pool.setPersonPoolId(id);
                if (poolService.existPool(id)) {
                    pool.setStatus(0);
                    List<Person> people = matchPerson(id, requestVo.getPerson());
                    result.addAll(people);
                } else {
                    pool.setStatus(1);
                }
                poolStatus.add(pool);
            }
            ResponseVo success = ResponseVo.success(rid, "success");
            success.setMatchPersons(result);
            success.setMatch(1);
            success.setPersonPoolStatus(poolStatus);
            log.info("人员匹配操作完成,PoolIds:[{}},结果:[{}]", poolIds.toString(), objectMapper.writeValueAsString(success));
            return success;
        } catch (ElasticsearchStatusException e) {
            log.error("matchPerson", e);
            ResponseVo error = ResponseVo.error(rid, e.getDetailedMessage());
            error.setMatch(0);
            return error;
        }
    }

    /**
     * 修改人员特征
     */
    public ResponseVo updatePerson(RequestVo requestVo) {

        return null;
    }

    /**
     * 添加人员
     */
    public BulkResponse addPerson(String poolId, List<Person> personPool) throws IOException {
        BulkRequest bulkRequest = new BulkRequest(poolId);
        for (Person person : personPool) {
            Integer age = person.getAge();
            String gender = person.getGender();
            String personId = person.getPersonId();
            List<FaceFeature> faceFeatures = person.getFaceFeatures();
            if (faceFeatures != null && faceFeatures.size() > 0) {
                for (FaceFeature faceFeature : faceFeatures) {
                    Double[] feature = faceFeature.getFeature();
                    String fid = faceFeature.getFid();

                    IndexRequest indexRequest = new IndexRequest(poolId)
                            .source(XContentType.JSON, "personId", personId, "data", feature, "fid", fid, "age", age, "gender", gender);
                    bulkRequest.add(indexRequest);
                }
            }
        }
        if (bulkRequest.requests().size() == 0) {
            return null;
        }
        return client.bulk(bulkRequest, RequestOptions.DEFAULT);
    }

    /**
     * 删除人员
     */
    public BulkByScrollResponse deletePerson(String poolId, String personId) throws IOException {
        DeleteByQueryRequest deleteByQueryRequest = new DeleteByQueryRequest(poolId)
                .setQuery(new TermQueryBuilder("personId", personId))
                .setRefresh(true);
        return client.deleteByQuery(deleteByQueryRequest, RequestOptions.DEFAULT);
    }


    public List<Person> matchPerson(String poolId, Person person) throws Exception {
        List<FaceFeature> faceFeatures = person.getFaceFeatures();
        List<Person> persons = new ArrayList<>();
        if (faceFeatures != null && faceFeatures.size() > 0) {
            for (FaceFeature faceFeature : faceFeatures) {
                Double[] feature = faceFeature.getFeature();
                if (feature.length < 512) {
                    continue;
                }

                Script script = new Script(
                        ScriptType.INLINE, Script.DEFAULT_SCRIPT_LANG,
                        "(cosineSimilarity(params.data, 'data') + 1) / 2 * 100", Collections.singletonMap("data", feature));
                ScriptScoreQueryBuilder queryBuilder = QueryBuilders.scriptScoreQuery(QueryBuilders.matchAllQuery(), script);

                SearchSourceBuilder builder = new SearchSourceBuilder()
                        .size(5)
                        .query(queryBuilder)
                        .fetchSource(new String[]{"personId", "age", "gender", "fid"}, null);

                SearchRequest searchRequest = new SearchRequest(poolId)
                        .source(builder);

                SearchResponse search = client.search(searchRequest, RequestOptions.DEFAULT);

                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.setScore((int) item.getScore());
                    p.setPersonPoolId(item.getIndex());
                    persons.add(p);
                }
            }
        } else {
            log.info("no face feature");
        }
        if (persons.size() > 5) {
            persons = persons.stream().sorted(Comparator.comparingInt(Person::getScore)).limit(5).collect(Collectors.toList());
        }
        return persons;
    }


}