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