Java+ElasticSearch+Pytorch實現(xiàn)以圖搜圖功能

更新時間：2023年06月12日 15:02:02 作者：歐內(nèi)的手好汗

這篇文章主要為大家詳細介紹了Java如何利用ElasticSearch和Pytorch實現(xiàn)以圖搜圖功能，文中的示例代碼講解詳細，具有一定的學習價值，感興趣的小伙伴可以了解一下

一、準備模型

創(chuàng)建demo.py，輸入代碼，借助resnet提取圖像特征

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要設(shè)置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()
    #根據(jù)模型隨便創(chuàng)建一個輸入
    input = torch.rand([1, 3, 224, 224])
    output = model(input)
    #以這種方式保存
    script = torch.jit.trace(model, input)
    script.save("model.pt")

保存好的model.pt文件放入java項目的resources中，可以在java中引入Deep-Java-Library來調(diào)用

二、創(chuàng)建Java項目

創(chuàng)建項目，引入djl和elasticsearch的依賴

<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>
            <version>7.6.2</version>
        </dependency>

然后隨便從網(wǎng)上下載點圖片，比如貓5張圖狗5張圖什么的，放到項目"resources/隨便" 路徑下，一會要提取他們的特征向量并上傳至elasticsearch

三、es創(chuàng)建文檔

需要在elastic search中創(chuàng)建一個新文檔。localhost:9200/isi (img search img)

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

完成下文上傳操作后測試搜索（params中queryVector為隨便選了個圖像提取的特征向量）：

POST /isi/_search
{
  "query": {
    "function_score": {
      "query": {
        "match_all": {}
      },
      "script_score": {
        "script": {
          "source": "cosineSimilarity(params.queryVector, 'vector') ",
          "params": {
            "queryVector": [-0.21950562, 0.0979692, 0.30605257, -0.04246464, 0.3086218, 0.2133326, -0.13531154, 0.16382562, 0.2505685, 0.35654455, 0.50346404, -0.2031727, -0.4501943, 0.23117387, 0.39451313, 0.044487886, -0.11032343, 0.47252116, 0.24667346, -0.2052311, -0.10872754, 0.22328046, 0.13366169, -0.5555884, 0.23139203, 0.024292288, 0.3071902, 0.23381571, -0.14484097, -0.80570614, 0.096950606, -0.034106746, 0.3221968, 0.35980088, -0.24408965, 0.10010342, 0.34878045, 0.25403115, 0.8813986, -0.23978959, -0.101492174, -0.34241566, -0.258092, 0.38593173, 0.24993907, -0.6891467, 0.5723483, -0.4987241, -0.46613082, 0.07435644, -0.32876882, 0.1923833, 0.41619772, 0.006919967, -0.35519657, -0.2463252, -0.07216969, -0.10412077, -0.3964988, -0.43174505, 0.6576338, -0.09753291, 0.058324523, -0.366405, -0.08003934, -0.41232625, -0.59834087, 0.35432702, -0.33971205, -0.695481, -0.38738084, -0.08746443, 0.37581405, 0.5092232, 0.26168102, 0.33873072, 0.3769325, 0.5525994, -0.018578911, -0.16984223, 0.24996795, -0.33088574, 0.22646378, 0.28422385, -0.4198824, 0.23480973, -0.17118981, 0.21741581, -0.21377188, -0.21778852, 0.052144438, 0.6118544, 0.29015478, 0.38010067, -0.2526567, 0.31930774, 0.39023396, -0.18484715, 0.25706646, 0.5203727, 0.29022205, -0.21464778, 0.47365767, 0.87486994, 0.44488958, 0.35022217, 0.27183002, -0.10723972, 0.22553404, -0.15306596, 0.22945945, -0.3670853, 0.19239302, -0.44274554, -0.57217056, 0.03906954, 0.7685065, -0.9496267, 0.4024507, -0.13379225, 0.011526011, -0.09900194, -0.16814715, 0.13582, -0.30829066, -0.13147047, 0.2822718, -0.1368126, -0.027293338, 0.49085197, -0.5331921, 0.49128994, 0.07755277, 0.0067159105, -0.42123628, 0.4408007, 0.58957845, -0.611145, -0.20726573, -0.14454971, 0.66820395, -0.007127879, 0.39164197, 0.69164586, 0.0024765078, -0.002838524, -0.6342908, 0.082881235, 0.298888, 0.5127087, -0.1300959, 0.16589926, 0.14517388, 0.2331577, -0.79271543, 0.20724288, -0.08834252, 0.2896087, 0.45591825, 0.0028544534, -0.51650685, 0.40780628, -0.2327805, 0.36442646, -0.6027139, -0.13944842, 0.24956803, -0.024805166, 0.24770494, -0.45614043, 0.051073316, 0.16322246, 0.28946966, -0.74529195, -0.4670576, -0.0660522, 0.2550549, -0.08014119, 0.06633573, -0.24153815, 0.36240074, 0.6568509, 0.1186171, 0.24172828, -0.47089025, -0.03256646, 0.14894187, 0.37634763, -0.24351446, 0.16164464, -0.06065242, 0.5091901, 0.5561973, -0.12233099, 0.13057524, -0.6718906, 0.35523456, 0.07152326, -0.20505619, -0.72078145, 0.044493422, -0.7440514, 0.032849424, 0.20111637, 0.43541732, 0.57285374, 0.11018203, -0.28251386, 0.2966433, -0.23974887, 0.068521775, 0.108346015, 0.30997896, 0.1311228, 0.22840464, -0.027650226, 0.73658687, -0.42126977, -0.10911406, -0.03391467, 0.035138264, -0.37075225, -0.5644764, 0.18124935, -0.16772854, 0.19432716, -0.22632165, 0.3900388, -0.69183433, 0.3003796, -0.3586413, 0.53591347, 0.6152693, -0.0038650539, -0.21756373, -0.31812942, -0.5402912, 0.2289098, -0.26077187, -0.2690417, -0.23407489, -0.3686389, -0.40951043, 0.18680233, 0.09277787, -0.013292864, -0.10835711, -0.53350085, 0.14608094, -0.15405877, -0.4699281, 0.10309839, 0.270473, -0.53506243, 0.001998501, -0.24999668, -0.1813581, -0.005529306, -0.56031805, -0.048346244, 0.16996299, -0.29963455, 0.20901474, -0.30934745, -0.21063489, 0.113431595, 0.032042094, -0.41378844, 0.17604393, -0.4631637, -0.5003293, -0.55831015, -0.15353276, -0.46837738, -0.77764505, 0.4621635, 0.28037566, 0.044108637, -0.11912877, -0.082292914, 0.6248694, -0.25019014, 0.2364985, 0.37894827, -0.1412379, -0.4572027, 0.08775234, -0.1661222, -0.09539573, 0.4805072, -0.44165152, 0.092903554, 0.28206846, 0.19792132, -0.05991053, -0.1682668, -0.5983673, 0.29130918, 0.17447682, -0.13979116, -0.5393585, 0.16808756, 0.7923586, 0.09991479, -0.22254673, -0.04165893, -0.2054404, -0.23928185, -0.36898398, -0.25248212, 0.49941728, 0.41065186, -0.23073834, -0.23773289, -0.1929644, -0.3499782, 0.1223462, -0.06807637, 0.175438, 0.14037246, 0.04750026, 0.31662372, 0.9243611, 0.12812611, -0.069829896, -0.29231697, -0.14734478, 0.1471743, -0.22162598, 0.37801683, 0.078523755, 0.022035534, 0.64812034, -0.23139268, -0.24911498, -0.33378544, -0.07344267, 0.11432794, -0.059079893, 0.31315947, -0.48160297, 0.045891166, 0.09479678, 0.4726333, -0.039052956, -0.2787302, -0.11797555, 0.25318092, -0.27943715, -0.22414759, -0.5546054, -0.106284395, 0.028235137, -0.3618798, -0.3404342, 0.31008774, -0.062293224, 0.053846367, 0.057488017, 0.09902769, 0.70674986, 0.37460673, 0.032727424, -0.4898543, -0.40673503, -0.19604088, 0.4853623, -0.20603043, -0.067109674, -0.53158236, -0.09847969, -0.14446007, -0.15639015, 0.13992839, -0.11348952, 0.15039599, 0.18439567, -0.20131497, 0.20396306, -0.3851034, -0.42031956, 0.2908249, 0.38035524, 0.15540914, 0.030931382, 0.33113614, 0.2741093, 0.18215193, -0.017978923, -0.0023425352, 0.27586395, -0.48403275, 0.023819037, -0.31364787, 0.21789135, -0.3554474, 0.0189421, 0.85861605, 0.15560226, 0.014263891, -0.16498215, -0.39344954, 0.5993788, 0.2708438, -0.29564035, 0.029406447, -0.43017676, 0.057052203, -0.04981024, -0.27520806, -0.51138026, 0.5926964, -0.40741482, 0.08575866, -0.196537, 0.3554017, -0.14750098, 0.051478647, -0.44164056, 0.13783655, 0.697562, 0.069351025, -0.14384158, 0.16349174, 0.36234668, -0.29366237, -0.28419775, -0.22636151, -0.122755915, -0.08138535, -0.7832234, 0.11401084, 0.43588766, 0.3236169, 0.14211948, 0.11028929, 0.2964773, -0.35562044, -0.3229665, -0.12864646, -0.3883256, 0.18198651, -0.45499957, 0.6918359, 0.1301559, -0.19978304, -0.1603161, -0.21330322, 0.07482076, 0.19179785, 0.45639312, 0.010576941, -0.3680949, -0.67871827, 0.14103784, 0.26047683, 0.64846706, -0.6718977, -0.5179457, 0.5580428, -0.48272127, 0.09030259, -0.47150746, 0.534373, 0.20664622, 0.5013874, -0.20477112, 0.22002026, -0.042670928, -0.047632568, -0.14199638, -0.36322978, 0.14286354, 0.35466686, 0.31752202, -0.3477305, -0.0045454763, -0.066675276, -0.2702982, 0.21498637, 0.08594364, 0.23323308, -0.6374196, 0.36372712, -0.30689493, -0.15897107, -0.16212063, -0.3653109, -0.40084177, 0.050653785, 0.13755074, -0.11666774, -0.2285859, -0.037478417, 0.20095918, 0.3487094, -0.08746929, 0.61240536, -0.26032016, 0.4687963, 0.14449233, 0.47144917, 0.039736982, 0.079499036, 0.19588508, -0.017659996, 0.5632192, -0.18003824, -0.040397547, -0.36472237, -0.025306623, 0.437257, -0.086474344, -0.7352421, 0.13874608, -0.110756345, -0.048157282, 0.19240974, 0.080513775, 0.59602517, -0.34077823, 0.35610113, -0.25455856, -0.36457375, 0.37937617, -0.400827, 0.054261968, 0.2879811, -0.11387855, 0.5244568, -0.107315816, 0.27161226, -0.1021186, -0.18614362, 0.2978657, 0.5370135, 0.36572018, -0.15257043, 0.20518257, 0.6419209, -0.3316342, -0.2725687, 0.014353298, 0.26933295, 0.26414502, 0.12794185, -0.21693292, 0.20204909, -0.21943296, -0.2593051, 0.13386369, -0.58061814, -0.63560903, 0.36805475, 0.016099993, -0.30377084, -0.06283789, -0.2682347, 0.056822743, -0.40932798, 0.36230356, -0.0491954, -0.029950788, -0.014880693, 0.20313863, 0.2214678, 0.2732247, -0.20896281, 0.017126573, -0.122736566, 0.06616182, 0.024874818, -0.7285719, -0.08402412, 0.28992975, 0.6755032, -0.071360886, 0.08794832, 0.09367153, 0.2191656, -0.28197074, -0.060419604, 0.008227522, -0.77392685, -0.37148783, 0.028109314, 0.6432903, 0.107398115, -0.031399835, -0.17004932, 0.34670952, -0.2593519, 0.3303603, 0.05096604, -0.59591603, -0.28604108, -0.37413853, -0.50223345, -0.06456756, 0.30097467, -0.344833, -0.44159552, 0.1301419, 0.4851787, -0.48891386, 0.29695496, 0.04624813, 0.44581613, -0.342841, 0.030060228, 0.41584003, -0.06151448, -0.6391304, 0.24350402, 0.2373283, -0.21993239, -0.15819496, 0.33998314, -0.104932904, 0.53482264, 0.18774611, -0.30718842, 0.09050197, 0.06582601, -0.8714315, 0.2866774, 0.10656398, 0.108911455, 0.12436204, 0.6072432, -0.263783, 0.3477571, -0.21450447, -0.20965956, 0.2725455, -0.15962526, -0.023959063, -0.16272986, 0.37898353, 0.1893706, -0.35078412, 0.018863793, 0.19243363, 0.7553659, 0.3897343, 0.16990745, -0.12922706, 0.32337534, -0.07977969, 0.09969508, 0.12787843, 0.14316258, -0.38789797, -0.18665363, 0.41474488, -0.04381171, 0.47398177, 0.20612329, -0.13801742, 0.48971528, 0.15693656, 0.10821125, -0.10725921, -0.20428485, -0.84528387, 0.18022658, 0.50938493, -0.32568434, 0.30802926, 0.33309558, 0.1919713, 0.20726888, -0.16194591, 0.17306438, -0.15405764, -0.57394016, 0.6925947, -0.31852844, -0.07849608, -0.5219136, -0.12416126, -0.20998093, 0.6197391, -0.11049731, -0.07111119, 0.22960934, -0.15123159, 0.22498849, -0.07220747, -0.28159276, -0.16107027, -0.010942766, 0.5636157, 0.4077794, 0.39782813, 0.456499, 0.22233048, 0.56309587, 0.26227084, -0.18100007, 0.06122207, 0.27089763, 0.17011975, -0.42344883, -0.063430965, 0.070528686, -0.046008512, -0.29042992, -0.07066448, -0.2578915, -0.27239347, 0.2880362, -0.056104008, -0.40367386, -0.091103815, 0.46031728, -0.36084417, -0.01598189, 0.19975084, 0.01695741, 0.3267317, 0.22532314, -0.55215025, -0.098993674, 0.36677533, 0.44303438, 0.3397658, -0.42336193, 0.002683131, 0.17797257, 0.6305417, 0.54148203, 0.17323923, 0.11428201, -0.07747766, -0.11240339, 0.11639454, 0.05241075, 0.035248175, -0.57705295, 0.45263726, -0.35879546, -0.7651455, -0.03033166, 0.47368425, -0.02433325, -0.15444314, -0.27954623, 0.30544212, 0.19804852, -0.66339266, -0.018637381, -0.3836641, 0.10387643, -0.23915236, 0.097831056, -0.18519881, 0.42123106, -0.0021492783, -0.4928366, 0.051339585, -0.50189865, -0.0325974, 0.03475754, 0.24877562, -0.50540763, 0.14656179, -0.033425312, -0.2698435, 0.1414198, 0.015859405, 0.4277053, -0.040847912, 0.032052774, 0.39479595, -0.0018053818, -0.37721512, -0.027027369, 0.44188333, 0.18346275, 0.6159405, -0.0010263352, 0.120682925, -0.5515572, 0.4246414, 0.37855124, 0.31135443, 0.255429, 0.010005429, -0.8138245, -0.26479146, -0.34098482, 0.14558652, 0.63190436, 0.1779253, 0.43572387, 0.6876498, 0.06974258, 0.007930072, -0.09172004, 0.18957798, -0.16211304, 0.18704513, -0.25963065, -0.26715553, -0.22632961, -0.3099424, 0.3464097, -0.12967771, 0.16652606, 0.2921636, 0.09758349, 0.2582998, 0.11978268, 0.42495492, 0.02736637, -0.32260302, -0.3379873, -0.23938976, -0.19942743, 0.30798694, 0.25228044, -0.033107795, 0.09772943, 0.38394168, 0.7219979, -0.5064522, -0.21723904, -0.2033075, 0.020857109, 0.13053142, -0.38791847, 0.4991684, 0.20062184, 0.49477854, -0.26213312, -0.61973774, -0.074013926, -0.12128413, -0.56617993, 0.13392372, 0.73387975, 0.5033897, 0.33373255, -0.06803796, 0.5550287, 0.26606622, -0.35267583, -0.23695293, 0.26170373, -0.12340009, 0.80251247, -0.70798254, -0.028666062, 0.6997679, -0.05996991, 0.06898104, -0.14557816, 0.054661553, 0.5187798, 0.41702572, 0.792891, -0.17265478, -0.06679568, -0.331478, 0.0694997, 0.4253223, -0.2783028, 0.23903547, -0.58266413, 0.09287575, 0.045140624, -0.10417832, 0.08257238, 0.48208177, -0.24164109, 0.81102467, -0.40342188, 0.65527093, -0.12488523, 0.078327045, -0.5329088, 0.37736076, 0.2925939, 0.20142855, 0.21402623, -0.21197478, -0.31154165, 0.45887777, 0.205758, 0.12233909, -0.26103553, 1.1294454, -0.7648704, 0.32436037, -0.06368509, 0.57072765, 0.9322751, -0.29020756, 0.44769418, -0.839836, 0.07865648, -0.2559945, 0.4581841, -0.017776983, 0.18255703, 0.2528128, -0.41778934, -0.071126916, -0.041809052, -0.53156054, 0.16023451, -0.2608511, 0.0725673, -0.15921246, -0.03191948, -0.366381, 0.53149635, -0.2550226, -0.022553788, -0.36375383, 0.40580854, -0.076502666, -0.04272891, -0.28619775, -0.26721123, 0.56044143, -0.040593743, -0.28715926, -0.0043915436, 0.11840753, -0.35239887, -0.30920973, -0.14502974, -0.36411104, 0.44530326, 0.43969297, -0.23792548, 0.30757633, 0.26880985, 0.18359815, -0.5675205, 0.19222523, 0.22303401, 0.21661428, 0.22141027, -0.10556421, 0.11646886, 0.17539617, 0.96604997, 0.055217523, -0.7456562, -0.106842384, -0.19286019, 0.17667075, 0.92509866, 0.57278365, -0.024029609, -0.8224203, -0.12689532, 0.079639494, -0.06534128, 0.7061269, -0.09063532, 0.5011331, -0.5051317, -0.054662913, -0.26086497, -0.53341925, 0.9624672, 0.08449669, -0.21910548, 0.36410314, -0.24794322, 0.16658492, 0.7944018, -0.058724128, -0.22618303, 0.5062074, -0.516353, 0.69395834, -0.23764399, -0.13169304, 0.51044196, -0.042955525, -0.42410484, -0.4293069, 0.13401544, 0.80136365, 0.30296534, -0.06788176, 0.16880289, 0.27950272, -0.37403736, 0.11813866, -0.41821468, 0.0033562258, -0.53348655, -0.22950119, 0.3889678, 0.10558852, -0.25912097, -0.03190498, 0.028149713, 0.36284888, -0.63619995, 0.8380439, 0.6589971, 0.6046954, -0.2093836, 0.08808039, 0.48332697, -0.010615652, -0.40519536, 0.011716956, 0.096273005, -0.27340046, -0.19237258, -0.2970637, -0.44011658, 0.17786184, 0.0071578454, 0.23985118, -0.040508576]
          }
        }
      }
    }
  },
  "_source": ["url"],
  "size": 100
}

四、調(diào)用pytorch模型代碼

創(chuàng)建Test類，copy一下我的，感興趣可以去djl的官網(wǎng)學習更多內(nèi)容。寫完后，就可以獲取“隨便”文件夾中的圖像的特征向量，上傳到es里了。

package org.gwen;
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.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.common.xcontent.XContentType;
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.gwen.entity.SearchResult;
import java.io.File;
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 Test {
    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(Test.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轉(zhuǎn)為tensor、輸出轉(zhuǎn)為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();
        }
    }
}

五、es上傳和搜索

上傳：遍歷每張圖片，獲取每張圖片的特征，上傳到es

搜索：獲取輸入圖像的特征，創(chuàng)建SearchRequest在es中通過painless腳本進行余弦相似度對比檢索。

首先創(chuàng)建SearchResult類表示es搜索的結(jié)果，包括圖像url和相關(guān)度評分score

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

然后在Test里實現(xiàn)upload和search

    public static void upload() throws Exception {
        RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(new HttpHost("localhost", 9200, "http")));
        //批量上傳請求
        BulkRequest bulkRequest = new BulkRequest(INDEX);
        File file = new File("C:\\Users\\Administrator\\IdeaProjects\\img_search_img\\src\\main\\resources\\隨便");
        for (File listFile : file.listFiles()) {
            float[] vector = predictor.predict(ImageFactory.getInstance().fromInputStream(Test.class.getClassLoader().getResourceAsStream("隨便/" + listFile.getName())));
            // 構(gòu)建文檔
            Map<String, Object> jsonMap = new HashMap<>();
            jsonMap.put("url", listFile.getAbsolutePath());
            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個結(jié)果
        int k = 100;
        // 連接Elasticsearch服務(wù)器
        RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(new HttpHost("localhost", 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) {
            // 處理搜索結(jié)果
            System.out.println(hit.toString());
            SearchResult result = new SearchResult((String) hit.getSourceAsMap().get("url"), hit.getScore());
            list.add(result);
        }
        client.close();
        return list;
    }

六、測試

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

頁面：

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>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('系統(tǒng)僅支持 JPG、PNG、JPEG 格式的圖片，請您調(diào)整格式后重新上傳')
            return
        }
        $('.result-block').empty().append($('<h1>正在識別中...</h1>'))
        $("#image").attr("src",getObjectURL(f));
        let formData = new FormData()
        formData.append('file',f)
        $.ajax({
            url: 'http://localhost:8080/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="file:///${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>

到此這篇關(guān)于Java+ElasticSearch+Pytorch實現(xiàn)以圖搜圖功能的文章就介紹到這了,更多相關(guān)Java 以圖搜圖內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: