Test0.java 6.27 KB
package com.viontech.match;

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.entity.*;
import com.viontech.match.repository.FaceRecognitionRepository;
import com.viontech.match.service.AlgService;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.annotation.Resource;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.text.ParseException;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

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

@SpringBootTest
@RunWith(SpringRunner.class)
public class Test0 {

    private final ThreadPoolExecutor poolExecutor = new ThreadPoolExecutor(20, 30, 1, TimeUnit.MINUTES, new LinkedBlockingQueue<>(), new ThreadPoolExecutor.CallerRunsPolicy());
    private final AtomicInteger integer = new AtomicInteger();
    private final PoolEnum pool = PoolEnum.POOL_50W;
    @Resource
    ObjectMapper objectMapper;
    @Resource
    private AlgService algService;
    @Resource
    private FaceRecognitionRepository faceRecognitionRepository;
    @Resource
    private Storage featureStorage;
    @Resource
    private RestHighLevelClient client;

    @Test
    public void createPoolTest() throws IOException {
        Object test2 = algService.createPool(pool.name);
        System.out.println(test2);

    }

    @Test
    public void deleteIndexTest() throws IOException {
        algService.deletePool(pool.name);
    }

    @Test
    public void queryIndex() throws IOException {
        String[] strings = algService.queryPoolList();
        HashMap<String, Object> result = new HashMap<>();
        for (String item : strings) {
            CountResponse count = client.count(new CountRequest(item), RequestOptions.DEFAULT);
            long count1 = count.getCount();
            result.put(item, count1);
        }
        System.out.println(result.toString());
    }

    @Test
    public void addPersonTest() throws IOException {
        Person person1 = getPerson("3");
        Object test2 = algService.addPerson(person1, pool.name);
        System.out.println(objectMapper.writeValueAsString(test2));
    }

    @Test
    public void matchPersonTest() throws Exception {
        Person person = getPerson("3");
        List<MatchResult> test2 = algService.matchPerson(person, pool.name);
        System.out.println(objectMapper.writeValueAsString(test2));

    }

    @Test
    public void batchAdd() throws ParseException, IOException, InterruptedException {
        algService.deletePool(pool.name);
        algService.createPool(pool.name);

        Date start = DateUtil.parse(DateUtil.FORMAT_SHORT, "2020-11-01");
        Date end = DateUtil.parse(DateUtil.FORMAT_SHORT, "2020-11-24");
        int i = 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;
                }
                while (poolExecutor.getQueue().size() > 1000) {
                    System.out.println(poolExecutor.getQueue().size());
                    TimeUnit.SECONDS.sleep(10);
                }
                if (integer.get() > pool.num) {
                    return;
                }
            }
        }

    }

    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;
    }


    public Person getPerson(String id) throws IOException {
        Path path = Paths.get("C:\\Users\\vion\\Desktop\\" + id + ".feature");
        String featureStr = Files.readAllLines(path).get(0);
        Feature feature = objectMapper.readValue(featureStr, Feature.class);
        Double[] data = feature.getDatas().get(0).getData();

        Person person = new Person();
        person.setPersonId(UUID.randomUUID().toString());
        person.setFeature(data);
        return person;
    }
}