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