Browse Source

feat: 相似度min

zhangyuting 5 năm trước cách đây
mục cha
commit
0c88e34ead

+ 10 - 10
src/main/java/com/mokasz/image_pool/restful/controller/ImageController.java

@@ -21,6 +21,7 @@ import org.springframework.web.multipart.MultipartFile;
 
 import java.io.IOException;
 import java.sql.ResultSet;
+import java.util.ArrayList;
 import java.util.List;
 
 @RestController
@@ -36,22 +37,21 @@ public class ImageController {
     @Autowired
     private JdbcTemplate jdbcTemplate;
 
-    private void calculate (byte[] fileBytes, float[] vector)
-    {
+    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) {
+    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) {
+    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];
@@ -60,7 +60,7 @@ public class ImageController {
         }
     }
 
-    private String getReadableHash (float[] floats) {
+    private String getReadableHash(float[] floats) {
         StringBuilder stringBuilder = new StringBuilder();
         for (int i = 0; i < 8; i++) {
             stringBuilder.append(String.format("%02X", (int) floats[i]));
@@ -69,13 +69,13 @@ public class ImageController {
     }
 
     @PostMapping("/store2")
-    public ImageSaveResponse store2 (@RequestParam("file") MultipartFile file) throws IOException {
+    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) throws IOException {
+    public ImageSaveResponse store(@RequestParam("file") MultipartFile file) throws IOException {
         logger.info("File size: {} Name: {}", file.getSize(), file.getOriginalFilename());
         float[] floats = new float[8];
         calculate(file.getBytes(), floats);
@@ -89,15 +89,15 @@ public class ImageController {
     }
 
     @PostMapping("/match")
-    public List<ImageQueryResponse> match (@RequestParam("file") MultipartFile file, @RequestParam("pool") String pool) throws IOException {
+    public List<ImageQueryResponse> match(@RequestParam("file") MultipartFile file, @RequestParam("pool") String pool, @RequestParam("min") float min) throws IOException {
         logger.info("Not implemented pool: {}", pool);
         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 hash, (64 - l2_distance(?::FLOAT2[], vector)) / 64 AS similarity FROM image_pool.image_hash ORDER BY vector <-> ?::FLOAT2[] ASC LIMIT 2",
-                (ResultSet rs, int rowNum) -> new ImageQueryResponse(rs.getFloat("similarity"), rs.getString("hash")), binaries, binaries);
+        List<ImageQueryResponse> list = jdbcTemplate.query("SELECT * FROM (SELECT hash, (64 - l2_distance(?::FLOAT2[], vector)) / 64 AS similarity FROM image_pool.image_hash ORDER BY vector <-> ?::FLOAT2[] ASC) AS t WHERE similarity >= ? ORDER BY similarity DESC ",
+                (ResultSet rs, int rowNum) -> new ImageQueryResponse(rs.getFloat("similarity"), rs.getString("hash")), binaries, binaries, min);
         return list;
     }
 }