Selaa lähdekoodia

Feature: simple pool storage and matching

wudi 5 vuotta sitten
vanhempi
commit
a93397a357

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

@@ -21,7 +21,6 @@ import org.springframework.web.multipart.MultipartFile;
 
 import java.io.IOException;
 import java.sql.ResultSet;
-import java.util.ArrayList;
 import java.util.List;
 
 @RestController
@@ -37,21 +36,22 @@ 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, @RequestParam("pool") int pool) throws IOException {
         logger.info("File size: {} Name: {}", file.getSize(), file.getOriginalFilename());
         float[] floats = new float[8];
         calculate(file.getBytes(), floats);
@@ -85,19 +85,19 @@ public class ImageController {
         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") String pool, @RequestParam("min") float min) throws IOException {
-        logger.info("Not implemented pool: {}", pool);
+    public List<ImageQueryResponse> match (@RequestParam("file") MultipartFile file, @RequestParam("pool") int pool, @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 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);
+        List<ImageQueryResponse> list = jdbcTemplate.query("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 LIMIT ?",
+                (ResultSet rs, int rowNum) -> new ImageQueryResponse(rs.getFloat("similarity"), rs.getString("hash")), binaries, pool, binaries, limit);
         return list;
     }
 }