MainController.java 6.92 KB
package com.viontech.match.controller;

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.PoolEnum;
import com.viontech.match.entity.*;
import com.viontech.match.entity.vo.RequestData;
import com.viontech.match.repository.FaceRecognitionRepository;
import com.viontech.match.service.AlgService;
import com.viontech.match.service.MatchService;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * .
 *
 * @author 谢明辉
 * @date 2020/11/23
 */
@RestController
@RequestMapping("/alg")
public class MainController {

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

    @PostMapping("/match")
    public Object match(RequestData requestData) throws Exception {
        MultipartFile file = requestData.getFile();
        Double[] data = matchService.extractFeature(file);
        Person person = new Person().setFeature(data);
        long begin = System.currentTimeMillis();
        List<MatchResult> match = algService.matchPerson(person, requestData.getPoolId());
        long timeConsumingMatching = System.currentTimeMillis() - begin;
        List<String> collect = match.stream().map(MatchResult::getPersonId).collect(Collectors.toList());
        List<FaceRecognition> faceRecognitions = faceRecognitionRepository.queryFaceRecognitionByUnids(collect);
        if (faceRecognitions == null) {
            return null;
        }
        Map<String, FaceRecognition> map = faceRecognitions.stream().collect(Collectors.toMap(FaceRecognition::getUnid, x -> x, (x, y) -> y));
        for (MatchResult item : match) {
            FaceRecognition faceRecognition = map.get(item.getPersonId());
            if (faceRecognition != null) {
                faceRecognition.setMatchScore(item.getScore());
            }
        }
        Collection<FaceRecognition> values = map.values();
        List<FaceRecognition> collect1 = values.stream().sorted((o1, o2) -> {
            if (o1.getMatchScore() > o2.getMatchScore()) {
                return -1;
            } else if (o1.getMatchScore().equals(o2.getMatchScore())) {
                return 0;
            } else {
                return 1;
            }
        }).collect(Collectors.toList());

        HashMap<String, Object> resultMap = new HashMap<>(2);
        resultMap.put("time", timeConsumingMatching);
        resultMap.put("result", collect1);
        return resultMap;
    }

    @GetMapping("/poolIds")
    public Object getPoolIds() throws IOException {
        return algService.queryPoolList();
    }

    @GetMapping("/create/{num}")
    public Object createPool(@PathVariable("num") Integer num) throws Exception {
        PoolEnum pool;
        switch (num) {
            case 2000:
                pool = PoolEnum.POOL_2K;
                break;
            case 5000:
                pool = PoolEnum.POOL_5K;
                break;
            case 20000:
                pool = PoolEnum.POOL_2W;
                break;
            case 100000:
                pool = PoolEnum.POOL_10W;
                break;
            case 500000:
                pool = PoolEnum.POOL_50W;
                break;
            case 1000000:
                pool = PoolEnum.POOL_100W;
                break;
            case 2000000:
                pool = PoolEnum.POOL_200W;
                break;
            default:
                throw new RuntimeException("请输入[2000,5000,20000,100000,500000,1000000,2000000]");
        }

        algService.deletePool(pool.name);
        algService.createPool(pool.name);

        Date end = DateUtil.addDays(new Date(), -1);
        Date start = DateUtil.addDays(end, -30);
        integer.set(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 "success";
                }
                while (poolExecutor.getQueue().size() > 10000) {
                    System.out.println(poolExecutor.getQueue().size());
                    TimeUnit.SECONDS.sleep(4);
                }
                if (integer.get() > pool.num) {
                    return "success";
                }
            }
        }
        return "success";
    }

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