AlgUtil.java 3.14 KB
package com.viontech.match.test;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.viontech.keliu.model.Person;
import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.client.RestTemplate;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.HashMap;

/**
 * .
 *
 * @author 谢明辉
 * @date 2020/11/19
 */
@Component
public class AlgUtil {
    @Resource
    private RestTemplateBuilder restTemplateBuilder;

    public Object matchPerson(String poolId, Person person) {
        RestTemplate restTemplate = restTemplateBuilder.build();

        MatchPersonRequestBody matchPersonRequestBody = new MatchPersonRequestBody(Arrays.asList(person.getFaceFeatures().get(0).getFeature()));

        ResponseEntity<String> responseEntity = restTemplate.postForEntity("http://127.0.0.1:9200/" + poolId.toLowerCase() + "/_search", matchPersonRequestBody, String.class);
        return responseEntity.getBody();
    }

    public Object addPerson(String poolId, Person person) throws JsonProcessingException {
        RestTemplate restTemplate = restTemplateBuilder.build();
        HttpHeaders headers = new HttpHeaders();
        headers.add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON.toString());

        HashMap<String, Object> req = new HashMap<>();
        req.put("personId", person.getPersonId());
        req.put("data", person.getFaceFeatures().get(0).getFeature());

        HttpEntity<HashMap> httpEntity = new HttpEntity<>(req, headers);

        ResponseEntity<String> response = restTemplate.exchange("http://127.0.0.1:9200/" + poolId.toLowerCase() + "/_doc/" + person.getPersonId(), HttpMethod.PUT, httpEntity, String.class);
        return response.getBody();
    }

    public Object createPersonPool(String poolId) {
        RestTemplate restTemplate = restTemplateBuilder.build();
        HttpHeaders headers = new HttpHeaders();
        headers.add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON.toString());
        HttpEntity<String> stringHttpEntity = new HttpEntity<>("{\"mappings\":{\"properties\":{\"data\":{\"type\":\"dense_vector\",\"dims\":512},\"personId\":{\"type\":\"keyword\"}}}}", headers);
        ResponseEntity<HashMap> exchange = restTemplate.exchange("http://127.0.0.1:9200/" + poolId.toLowerCase(), HttpMethod.PUT, stringHttpEntity, HashMap.class);
        HashMap body = exchange.getBody();
        return body;
    }

    public Object deletePersonPool(String poolId) {
        try {
            RestTemplate restTemplate = restTemplateBuilder.build();
            ResponseEntity<HashMap> exchange = restTemplate.exchange("http://127.0.0.1:9200/" + poolId.toLowerCase(), HttpMethod.DELETE, null, HashMap.class);
            return exchange.getBody();
        } catch (HttpClientErrorException e) {
            int rawStatusCode = e.getRawStatusCode();
            if (rawStatusCode != HttpStatus.NOT_FOUND.value()) {
                return "failed";
            } else {
                return "success";
            }
        }
    }

}