Java调用Pytorch实现以图搜图

设计技术栈：
1、ElasticSearch环境；
2、Python运行环境（如果事先没有pytorch模型时，可以用python脚本创建模型）；

1、运行效果

2、创建模型（有则可以跳过）

vi script.py

import torch
import torch.nn as nn
import torchvision.models as models
 
class ImageFeatureExtractor(nn.Module):
    def __init__(self):
        super(ImageFeatureExtractor, self).__init__()
        self.resnet = models.resnet50(pretrained=True)
        #最终输出维度1024的向量，下文elastic search要设置dims为1024
        self.resnet.fc = nn.Linear(2048, 1024)
 
    def forward(self, x):
        x = self.resnet(x)
        return x
 
if __name__ == '__main__':
    model = ImageFeatureExtractor()
    model.eval()
    #根据模型随便创建一个输入
    input = torch.rand([1, 3, 224, 224])
    output = model(input)
    #以这种方式保存
    script = torch.jit.trace(model, input)
    script.save("model.pt")

2、java项目pom.xml

<dependencies>
		<dependency>
			<groupId>org.springframework.boot</groupId>
			<artifactId>spring-boot-starter-web</artifactId>
		</dependency>
		<dependency>
			<groupId>org.projectlombok</groupId>
			<artifactId>lombok</artifactId>
			<scope>provided</scope>
		</dependency>
		<dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-engine</artifactId>
            <version>0.19.0</version>
        </dependency>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-native-cpu</artifactId>
            <version>1.10.0</version>
            <scope>runtime</scope>
        </dependency>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-jni</artifactId>
            <version>1.10.0-0.19.0</version>
        </dependency>
        <dependency>
            <groupId>org.elasticsearch.client</groupId>
            <artifactId>elasticsearch-rest-high-level-client</artifactId>
        </dependency>
	</dependencies>

3、ES创建文档

PUT /isi
{
  "mappings": {
    "properties": {
      "vector": {
        "type": "dense_vector",
        "dims": 1024
      },
      "url" : {
        "type" : "keyword"
      },
      "user_id": {
          "type": "keyword"
      }
    }
  }
}

4、编写java代码调用模型

ORCUtil.java

package com.topprismcloud.rtm;

import ai.djl.Device;
import ai.djl.Model;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.transform.Normalize;
import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.transform.ToTensor;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.translate.Transform;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.ScriptQueryBuilder;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.xcontent.XContentType;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URL;
import java.nio.file.Paths;
import java.util.*;

public class ORCUtil {

	private static final String INDEX = "isi";
	private static final int IMAGE_SIZE = 224;
	private static Model model; // 模型
	private static Predictor<Image, float[]> predictor; // predictor.predict(input)相当于python中model(input)
	static {
		try {
			model = Model.newInstance("model");
			// 这里的model.pt是上面代码展示的那种方式保存的
			model.load(ORCUtil.class.getClassLoader().getResourceAsStream("model.pt"));
			Transform resize = new Resize(IMAGE_SIZE);
			Transform toTensor = new ToTensor();
			Transform normalize = new Normalize(new float[] { 0.485f, 0.456f, 0.406f },
					new float[] { 0.229f, 0.224f, 0.225f });
			// Translator处理输入Image转为tensor、输出转为float[]
			Translator<Image, float[]> translator = new Translator<Image, float[]>() {
				@Override
				public NDList processInput(TranslatorContext ctx, Image input) throws Exception {
					NDManager ndManager = ctx.getNDManager();
					System.out.println("input: " + input.getWidth() + ", " + input.getHeight());
					NDArray transform = normalize
							.transform(toTensor.transform(resize.transform(input.toNDArray(ndManager))));
					System.out.println(transform.getShape());
					NDList list = new NDList();
					list.add(transform);
					return list;
				}

				@Override
				public float[] processOutput(TranslatorContext ctx, NDList ndList) throws Exception {
					return ndList.get(0).toFloatArray();
				}
			};
			predictor = new Predictor<>(model, translator, Device.cpu(), true);
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static void upload() throws Exception {
		HttpHost host=new HttpHost("14.20.30.16", 9200, HttpHost.DEFAULT_SCHEME_NAME);
		RestClientBuilder builder=RestClient.builder(host);
		CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
		credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials("elastic", "123456"));
		builder.setHttpClientConfigCallback(f -> f.setDefaultCredentialsProvider(credentialsProvider));
		RestHighLevelClient client = new RestHighLevelClient( builder);
		// 批量上传请求
		BulkRequest bulkRequest = new BulkRequest(INDEX);
		File file = new File("D:\\001ENV\\nginx-1.24.0\\html\\resource\\new");
		for (File listFile : file.listFiles()) {
//			float[] vector = predictor.predict(ImageFactory.getInstance()
//					.fromInputStream(Test.class.getClassLoader().getResourceAsStream("new/" + listFile.getName())));
			
			float[] vector = predictor.predict(ImageFactory.getInstance()
					.fromInputStream(new FileInputStream(listFile)));
			// 构建文档
			Map<String, Object> jsonMap = new HashMap<>();
			jsonMap.put("url", "/resource/"+listFile.getName());
			jsonMap.put("vector", vector);
			jsonMap.put("user_id", "user123");
			IndexRequest request = new IndexRequest(INDEX).source(jsonMap, XContentType.JSON);
			bulkRequest.add(request);
		}
		client.bulk(bulkRequest, RequestOptions.DEFAULT);
		client.close();
	}

	// 接收待搜索图片的inputstream，搜索与其相似的图片
	public static List<SearchResult> search(InputStream input) throws Throwable {
		float[] vector = predictor.predict(ImageFactory.getInstance().fromInputStream(input));
		System.out.println(Arrays.toString(vector));

		// 展示k个结果
		int k = 100;
		// 连接Elasticsearch服务器
		RestHighLevelClient client = new RestHighLevelClient(
				RestClient.builder(new HttpHost("14.20.30.16", 9200, "http")));

		SearchRequest searchRequest = new SearchRequest(INDEX);
		Script script = new Script(ScriptType.INLINE, "painless", "cosineSimilarity(params.queryVector, doc['vector'])",
				Collections.singletonMap("queryVector", vector));

		FunctionScoreQueryBuilder functionScoreQueryBuilder = QueryBuilders
				.functionScoreQuery(QueryBuilders.matchAllQuery(), ScoreFunctionBuilders.scriptFunction(script));

		SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
		searchSourceBuilder.query(functionScoreQueryBuilder).fetchSource(null, "vector") // 不返回vector字段，太多了没用还耗时
				.size(k);

		searchRequest.source(searchSourceBuilder);

		SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);

		SearchHits hits = searchResponse.getHits();

		List<SearchResult> list = new ArrayList<>();
		for (SearchHit hit : hits) {
			// 处理搜索结果
			System.out.println(hit.toString());
			SearchResult result = new SearchResult((String) hit.getSourceAsMap().get("url"), hit.getScore());
			list.add(result);
		}

		client.close();
		return list;
	}

	public static void main(String[] args) throws Throwable {
		ORCUtil.upload();
		System.out.println("hao");
	}
}

SearchController.java

package com.topprismcloud.rtm;

import java.util.List;

import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

@RestController
@CrossOrigin
public class SearchController {

	@PostMapping("search")
	public ResponseEntity search(MultipartFile file) {
		try {
			List<SearchResult> list = ORCUtil.search(file.getInputStream());
			return ResponseEntity.ok(list);
		} catch (Throwable e) {
			return ResponseEntity.status(400).body(null);
		}
	}
}

SearchResult.java

package com.topprismcloud.rtm;

import lombok.AllArgsConstructor;
import lombok.Data;

@Data
@AllArgsConstructor
public class SearchResult {
    private String url;
    private Float score;
}

5、前端

index.html

<!DOCTYPE html>
<html lang="zh">

<head>
    <meta charset="UTF-8">
    <title>以图搜图</title>
    <style>
        body {
            background: url("/img/bg.jpg");
            background-attachment: fixed;
            background-size: 100% 100%;
        }

        body>div {
            width: 1000px;
            margin: 50px auto;
            padding: 10px 20px;
            border: 1px solid lightgray;
            border-radius: 20px;
            box-sizing: border-box;
            background: rgba(255, 255, 255, 0.7);
        }

        .upload {
            display: inline-block;
            width: 300px;
            height: 280px;
            border: 1px dashed lightcoral;
            vertical-align: top;
        }

        .upload .cover {
            width: 200px;
            height: 200px;
            margin: 10px 50px;
            border: 1px solid black;
            box-sizing: border-box;
            text-align: center;
            line-height: 200px;
            position: relative;
        }

        .upload img {
            width: 198px;
            height: 198px;
            position: absolute;
            left: 0;
            top: 0;
        }

        .upload input {
            margin-left: 50px;
        }

        .upload button {
            width: 80px;
            height: 30px;
            margin-left: 110px;
        }

        .result-block {
            display: inline-block;
            margin-left: 40px;
            border: 1px solid lightgray;
            border-radius: 10px;
            min-height: 500px;
            width: 600px;
        }

        .result-block h1 {
            text-align: center;
            margin-top: 100px;
        }

        .result {
            padding: 10px;
            cursor: pointer;
            display: inline-block;
        }

        .result:hover {
            background: rgb(240, 240, 240);
        }

        .result p {
            width: 110px;
            overflow: hidden;
            white-space: nowrap;
            text-overflow: ellipsis;
        }

        .result img {
            width: 160px;
            height: 160px;
        }

        .result .prob {
            color: rgb(37, 147, 60)
        }
    </style>
    <script src="js/jquery-3.6.0.js"></script>
</head>

<body>
    <div>
        <div class="upload">
            <div class="cover">
                请选择图片
                <img id="image" src="" />
            </div>
            <input id="file" type="file">
        </div>
        <div class="result-block">
            <h1>请选择图片</h1>
        </div>
    </div>
    <ul id="box">

    </ul>
    <script>
        var file = $('#file')
        file.change(function () {
            let f = this.files[0]
            let index = f.name.lastIndexOf('.')
            let fileText = f.name.substring(index, f.name.length)
            let ext = fileText.toLowerCase() //文件类型
            console.log(ext)
            if (ext != '.png' && ext != '.jpg' && ext != '.jpeg') {
                alert('系统仅支持 JPG、PNG、JPEG 格式的图片，请您调整格式后重新上传')
                return
            }
            $('.result-block').empty().append($('<h1>正在识别中...</h1>'))
            $("#image").attr("src", getObjectURL(f));
            let formData = new FormData()
            formData.append('file', f)
            $.ajax({
                url: 'http://10.1.2.240:8081/search',
                method: 'post',
                data: formData,
                processData: false,
                contentType: false,
                success: res => {
                    console.log('shibie', res)
                    $('.result-block').empty()
                    for (let item of res) {
                        console.log(item)
                        let html = `<div class="result">
                                    <img src="${item.url}"/>
                                    <div style="display: inline-block;vertical-align: top">
                                        <p class="prob">得分：${item.score.toFixed(4)}</p>
                                    </div>
                                </div>`
                        $('.result-block').append($(html))
                    }

                }
            })
        });
        $('#button').click(function (e) {
            var file = $('#file')[0].files[0] //单个
            console.log(file)
        })
        function getObjectURL(file) {
            var url = null;
            if (window.createObjcectURL != undefined) {
                url = window.createOjcectURL(file);
            } else if (window.URL != undefined) {
                url = window.URL.createObjectURL(file);
            } else if (window.webkitURL != undefined) {
                url = window.webkitURL.createObjectURL(file);
            }
            return url;
        }
        function detect() {

        }
    </script>
</body>

</html>

6、打包后的源代码

以图搜图Java+html源代码

相关参考文章：Java调用Pytorch模型进行图像识别