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.OpenCvService;
import com.mokasz.image_pool.restful.entity.ImageQueryResponse;
import com.mokasz.image_pool.restful.entity.ImageSaveResponse;
import com.mokasz.image_pool.restful.service.DownloadBytes;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

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

@Slf4j
@RestController
@RequestMapping("/service/image")
@RequiredArgsConstructor
public class ImageController {

  private final OpenCvService service;

  private final JdbcTemplate jdbcTemplate;

  private byte[] getFileBytes (MultipartFile file, String url) throws IOException {
    if (url != null) {
      return DownloadBytes.downloadAsBytesArray(url);
    } else if (file != null) {
      return file.getBytes();
    } else {
      throw new RuntimeException("缺少file或url");
    }
  }

  @PostMapping("/store")
  public ImageSaveResponse store(
      @RequestParam(value = "file", required = false) MultipartFile file,
      @RequestParam(value = "url", required = false) String url,
      @RequestParam("pool") int pool) throws IOException {
    float[] floats = new float[8];
    service.calculatePHash(this.getFileBytes(file, url), floats);
    log.info("size:",floats);
    float[] binaries = new float[64];
    service.transformBinaryArray(floats, binaries);
    String hash = service.getReadableHash(floats);
    jdbcTemplate.update(
            "INSERT INTO image_pool.image_hash VALUES (?, NOW(), ?) ON CONFLICT (hash) DO NOTHING",
            hash,
            binaries);
    jdbcTemplate.update(
        "INSERT INTO image_pool.image_to_pool VALUES (?, ?, NOW()) ON CONFLICT (hash, pool_id) DO NOTHING",
        hash,
        pool);
    return new ImageSaveResponse(hash, pool);
  }

  @PostMapping("/match")
  public List<ImageQueryResponse> match(
      @RequestParam(value = "file", required = false) MultipartFile file,
      @RequestParam(value = "url", required = false) String url,
      @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];
    service.calculatePHash(this.getFileBytes(file, url), floats);
    float[] binaries = new float[64];
    service.transformBinaryArray(floats, binaries);
    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;
  }
}