image-pool-service/src/main/java/com/mokasz/image_pool/restful/controller/ImageController.java

package com.mokasz.image_pool.restful.controller;

import com.mokasz.image_pool.algorithm.FeatureExtractor;
import com.mokasz.image_pool.restful.entity.ImageQueryResponse;
import com.mokasz.image_pool.restful.entity.ImageSaveResponse;
import org.bytedeco.javacpp.Loader;
import org.bytedeco.opencv.opencv_java;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.img_hash.PHash;
import org.opencv.imgcodecs.Imgcodecs;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.sql.ResultSet;
import java.util.List;

@RestController
@RequestMapping("/service/image")
public class ImageController {
    // 静态初始化 OpenCV
    static {
        Loader.load(opencv_java.class);
    }

    private final Logger logger = LoggerFactory.getLogger(this.getClass());
    private final PHash pHash = PHash.create();
    @Autowired
    private JdbcTemplate jdbcTemplate;

    private void calculate(byte[] fileBytes, float[] vector) {
        Mat input = Imgcodecs.imdecode(new MatOfByte(fileBytes), Imgcodecs.IMREAD_UNCHANGED);
        Mat output = new Mat();
        pHash.compute(input, output);
        transformMatToArrayOfFloat(output, vector);
    }

    private void transformMatToArrayOfFloat(Mat mat, float[] floats) {
        if (mat.width() != 8 || mat.height() != 1) {
            throw new RuntimeException("Matrix dimension is not 1x8");
        }
        for (int i = 0; i < 8; i++) {
            floats[i] = (float) mat.get(0, i)[0];
        }
    }

    private void transformBinaryArray(float[] octets, float[] binaries) {
        for (int i = 0; i < 8; i++) {
            for (int j = 0; j < 8; j++) {
                int v = (int) octets[i];
                binaries[i * 8 + j] = (v & (1 << j)) >> j;
            }
        }
    }

    private String getReadableHash(float[] floats) {
        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < 8; i++) {
            stringBuilder.append(String.format("%02X", (int) floats[i]));
        }
        return stringBuilder.toString();
    }

    @PostMapping("/store2")
    public ImageSaveResponse store2(@RequestParam("file") MultipartFile file) throws IOException {
        FeatureExtractor.compute(file.getBytes());
        return new ImageSaveResponse("OK", 1L);
    }

    @PostMapping("/store")
    public ImageSaveResponse store(@RequestParam("file") MultipartFile file, @RequestParam("pool") int pool) throws IOException {
        logger.info("File size: {} Name: {}", file.getSize(), file.getOriginalFilename());
        float[] floats = new float[8];
        calculate(file.getBytes(), floats);
        float[] binaries = new float[64];
        transformBinaryArray(floats, binaries);
        logger.info("{}", binaries);
        String hash = getReadableHash(floats);
        int updated = jdbcTemplate.update("INSERT INTO image_pool.image_hash VALUES (?, NOW(), ?) ON CONFLICT (hash) DO NOTHING", hash, binaries);
        logger.info("{}", updated);
        jdbcTemplate.update("INSERT INTO image_pool.image_to_pool VALUES (?, ?, NOW()) ON CONFLICT (hash, pool_id) DO NOTHING", hash, pool);
        return new ImageSaveResponse(hash, 1L);
    }

    @PostMapping("/match")
    public List<ImageQueryResponse> match(@RequestParam("file") MultipartFile file,
                                          @RequestParam("pool") int pool,
                                          @RequestParam(value = "max", defaultValue = "1") float max,
                                          @RequestParam(value = "min", defaultValue = "1") float min,
                                          @RequestParam(value = "limit", defaultValue = "10") int limit) throws IOException {
        float[] floats = new float[8];
        calculate(file.getBytes(), floats);
        float[] binaries = new float[64];
        transformBinaryArray(floats, binaries);
        logger.info("{}", floats);
        List<ImageQueryResponse> list = jdbcTemplate.query("SELECT * FROM (SELECT hash, (64 - l2_distance(?::FLOAT2[], vector)) / 64 AS similarity FROM image_pool.image_hash WHERE hash IN (SELECT hash FROM image_pool.image_to_pool WHERE pool_id = ?) ORDER BY vector <-> ?::FLOAT2[] ASC) t WHERE similarity >=? and similarity<=? ORDER BY similarity DESC LIMIT ?",
                (ResultSet rs, int rowNum) -> new ImageQueryResponse(rs.getFloat("similarity"), rs.getString("hash")), binaries, pool, binaries, min, max, limit);
        return list;
    }
}