SpringBoot集成Milvus和deeplearning4j實現圖搜圖功能

更新時間：2024年10月30日 08:32:56 作者：HBLOG

Milvus?是一種高性能、高擴展性的向量數據庫,可在從筆記本電腦到大型分布式系統(tǒng)等各種環(huán)境中高效運行,Deeplearning4j（DL4J）是一個開源的深度學習框架,專門為Java和Scala開發(fā),本文給大家介紹了SpringBoot集成Milvus和deeplearning4j實現圖搜圖功能

1.什么是Milvus？

Milvus 是一種高性能、高擴展性的向量數據庫，可在從筆記本電腦到大型分布式系統(tǒng)等各種環(huán)境中高效運行。它既可以開源軟件的形式提供，也可以云服務的形式提供。 Milvus 是 LF AI & Data Foundation 下的一個開源項目，以 Apache 2.0 許可發(fā)布。大多數貢獻者都是高性能計算（HPC）領域的專家，擅長構建大型系統(tǒng)和優(yōu)化硬件感知代碼。核心貢獻者包括來自 Zilliz、ARM、NVIDIA、AMD、英特爾、Meta、IBM、Salesforce、阿里巴巴和微軟的專業(yè)人士

2.什么是deeplearning4j？

Deeplearning4j（DL4J）是一個開源的深度學習框架，專門為Java和Scala開發(fā)。它支持分布式計算，適合在大數據環(huán)境中運行，比如與Hadoop或Spark集成。DL4J的特點包括：

多種網絡架構：支持多種深度學習模型，包括卷積神經網絡（CNN）、循環(huán)神經網絡（RNN）和深度信念網絡（DBN）。
集成與可擴展性：能夠與大數據處理框架（如Apache Spark）和數據處理庫（如ND4J）緊密集成，方便處理大規(guī)模數據集。
易于使用：提供高層API，簡化模型構建和訓練過程，同時也允許用戶對底層實現進行細致的控制。
模型導入與導出：支持從其他框架（如Keras和TensorFlow）導入模型，并將訓練好的模型導出為多種格式，以便于部署。
性能優(yōu)化：支持多種硬件加速，包括GPU加速，能夠提高訓練和推理的效率。
支持多種應用場景：廣泛應用于計算機視覺、自然語言處理、推薦系統(tǒng)等多個領域。

Deeplearning4j是企業(yè)和開發(fā)者進行深度學習開發(fā)和研究的強大工具，特別適合于需要與Java生態(tài)系統(tǒng)兼容的場景。

3.環(huán)境搭建

First, we’ll need an instance of Milvus DB. The easiest and quickest way is to get a fully managed free Milvus DB instance provided by Zilliz Cloud: zilliz.com/
For this, we’ll need to register for a Zilliz cloud account and follow the documentation for creating a free DB cluster.

4.代碼工程

實驗目標

利用Milvus和deeplearning4j實現圖搜圖功能

pom.xml

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>3.2.1</version>
    </parent>
    <modelVersion>4.0.0</modelVersion>

    <artifactId>Milvus</artifactId>

    <properties>
        <maven.compiler.source>17</maven.compiler.source>
        <maven.compiler.target>17</maven.compiler.target>
        <deeplearning4j.version>1.0.0-M2.1</deeplearning4j.version>
        <nd4j.version>1.0.0-M2.1</nd4j.version>
    </properties>
    <dependencies>


        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-autoconfigure</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>

        <dependency>
            <groupId>io.milvus</groupId>
            <artifactId>milvus-sdk-java</artifactId>
            <version>2.4.0</version>
        </dependency>
        <dependency>
            <groupId>org.deeplearning4j</groupId>
            <artifactId>deeplearning4j-zoo</artifactId>
            <version>${deeplearning4j.version}</version>
        </dependency>
        <dependency>
            <groupId>org.nd4j</groupId>
            <artifactId>nd4j-native-platform</artifactId>
            <version>${nd4j.version}</version>
        </dependency>
        <dependency>
            <groupId>org.datavec</groupId>
            <artifactId>datavec-data-image</artifactId>
            <version>${deeplearning4j.version}</version>
        </dependency>
        <dependency>
            <groupId>org.deeplearning4j</groupId>
            <artifactId>deeplearning4j-core</artifactId>
            <version>${deeplearning4j.version}</version>
        </dependency>
        <dependency>
            <groupId>org.deeplearning4j</groupId>
            <artifactId>deeplearning4j-modelimport</artifactId>
            <version>${deeplearning4j.version}</version>
        </dependency>

    </dependencies>

    <build>
        <pluginManagement>
            <plugins>
                <plugin>
                    <groupId>org.apache.maven.plugins</groupId>
                    <artifactId>maven-compiler-plugin</artifactId>
                    <version>3.8.1</version>
                    <configuration>
                        <fork>true</fork>
                        <failOnError>false</failOnError>
                    </configuration>
                </plugin>

                <plugin>
                    <groupId>org.apache.maven.plugins</groupId>
                    <artifactId>maven-surefire-plugin</artifactId>
                    <version>2.22.2</version>
                    <configuration>
                        <forkCount>0</forkCount>
                        <failIfNoTests>false</failIfNoTests>
                    </configuration>
                </plugin>
            </plugins>
        </pluginManagement>
    </build>
</project>

特征抽取

package com.et.imagesearch;

import org.deeplearning4j.zoo.model.ResNet50;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.datavec.image.loader.NativeImageLoader;

import java.io.File;
import java.io.IOException;

public class FeatureExtractor {
    private ComputationGraph model;

   public FeatureExtractor() throws IOException {
      try {
         ZooModel<ComputationGraph> zooModel = ResNet50.builder().build();
         model = (ComputationGraph) zooModel.initPretrained();
      } catch (Exception e) {
         throw new IOException("Failed to initialize the pre-trained model: " + e.getMessage(), e);
      }
   }

    public INDArray extractFeatures(File imageFile) throws IOException {
        NativeImageLoader loader = new NativeImageLoader(224, 224, 3);
        INDArray image = loader.asMatrix(imageFile);
        ImagePreProcessingScaler scaler = new ImagePreProcessingScaler(0, 1);
        scaler.transform(image);

        return model.outputSingle(image);
    }
}

加載圖像: 使用 NativeImageLoader 將圖像加載為一個 INDArray，并將圖像的大小調整為 224x224 像素，通道數為 3（即 RGB 圖像）。
預處理圖像: 使用 ImagePreProcessingScaler 將圖像數據縮放到 [0, 1] 的范圍，以便模型可以更好地處理。
特征提取: 使用模型的 outputSingle() 方法將預處理后的圖像輸入模型，返回提取的特征向量。

Milvus數據庫操作

package com.et.imagesearch;

import io.milvus.client.*;
import io.milvus.param.*;
import io.milvus.param.collection.*;
import io.milvus.param.dml.*;
import io.milvus.grpc.*;
import io.milvus.param.index.CreateIndexParam;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

public class MilvusManager {
    private  MilvusServiceClient milvusClient;

    public MilvusManager() {
      milvusClient = new MilvusServiceClient(
            ConnectParam.newBuilder()
                  .withUri("https://xxx.gcp-us-west1.cloud.zilliz.com")
                  .withToken("xxx")
                  .build());
    }

    public void createCollection() {
        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .build();

        FieldType vectorField = FieldType.newBuilder()
                .withName("embedding")
                .withDataType(DataType.FloatVector)
                .withDimension(1000)
                .build();

        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName("image_collection")
                .withDescription("Image collection")
                .withShardsNum(2)
                .addFieldType(idField)
                .addFieldType(vectorField)
                .build();

        milvusClient.createCollection(createCollectionParam);
    }

    public void insertData(long id, INDArray features) {
        List<Long> ids = Collections.singletonList(id);
        float[] floatArray = features.toFloatVector();

        List<Float> floatList = new ArrayList<>();
        for (float f : floatArray) {
            floatList.add(f); 
        }

        List<List<Float>> vectors = Collections.singletonList(floatList);

        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("id",ids));
        fields.add(new InsertParam.Field("embedding", vectors));
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName("image_collection")
                .withFields(fields)
                .build();

        milvusClient.insert(insertParam);

    }
   public void flush() {
      milvusClient.flush(FlushParam.newBuilder()
            .withCollectionNames(Collections.singletonList("image_collection"))
            .withSyncFlush(true)
            .withSyncFlushWaitingInterval(50L)
            .withSyncFlushWaitingTimeout(30L)
            .build());
   }

   public void buildindex() {
      // build index
      System.out.println("Building AutoIndex...");
      final IndexType INDEX_TYPE = IndexType.AUTOINDEX;   // IndexType
      long startIndexTime = System.currentTimeMillis();
      R<RpcStatus> indexR = milvusClient.createIndex(
            CreateIndexParam.newBuilder()
                  .withCollectionName("image_collection")
                  .withFieldName("embedding")
                  .withIndexType(INDEX_TYPE)
                  .withMetricType(MetricType.L2)
                  .withSyncMode(Boolean.TRUE)
                  .withSyncWaitingInterval(500L)
                  .withSyncWaitingTimeout(30L)
                  .build());
      long endIndexTime = System.currentTimeMillis();
      System.out.println("Succeed in " + (endIndexTime - startIndexTime) / 1000.00 + " seconds!");
   }
}

createCollection():
- 創(chuàng)建一個名為 image_collection 的集合，包含兩個字段：
  - id: 主鍵，類型為 Int64。
  - embedding: 特征向量，類型為 FloatVector，維度為 1000。
- 使用 CreateCollectionParam 指定集合的名稱、描述和分片數量，并調用 createCollection 方法執(zhí)行創(chuàng)建操作。
insertData(long id, INDArray features):
- 插入一條新數據到 image_collection 集合中。
- 將 INDArray 類型的特征向量轉換為 List<List<Float>> 格式，以滿足 Milvus 的插入要求。
- 創(chuàng)建一個 InsertParam 實例，包含 ID 和特征向量，并調用 insert 方法執(zhí)行插入操作。
flush():
- 刷新 image_collection 集合，確保所有待處理的插入操作都被寫入數據庫。
- 使用 FlushParam 配置同步刷新模式和等待參數，確保操作的可靠性。
buildindex():
- 構建 image_collection 集合中 embedding 字段的索引，以加快后續(xù)的相似性搜索。
- 使用 CreateIndexParam 指定集合名稱、字段名稱、索引類型（自動索引）和度量類型（L2距離）。
- 調用 createIndex 方法執(zhí)行索引創(chuàng)建，并輸出所用時間。

圖片搜索功能

package com.et.imagesearch;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

public class ImageSearcher {
   private  MilvusServiceClient milvusClient;

   public ImageSearcher() {
      milvusClient = new MilvusServiceClient(
            ConnectParam.newBuilder()
                  .withUri("https://ixxxxx.gcp-us-west1.cloud.zilliz.com")
                  .withToken("xxx")
                  .build());
   }

   public void search(INDArray queryFeatures) {
      float[] floatArray = queryFeatures.toFloatVector();
      List<Float> floatList = new ArrayList<>();
      for (float f : floatArray) {
         floatList.add(f);
      }
      List<List<Float>> vectors = Collections.singletonList(floatList);


      SearchParam searchParam = SearchParam.newBuilder()
            .withCollectionName("image_collection")
            .withMetricType(MetricType.L2)
            .withTopK(5)
            .withVectors(vectors)
            .withVectorFieldName("embedding")
            .build();

      R<SearchResults> searchResults = milvusClient.search(searchParam);


      System.out.println("Searching vector: " + queryFeatures.toFloatVector());
      System.out.println("Result: " + searchResults.getData().getResults().getFieldsDataList());
   }
}

特征轉換: 將 INDArray 轉換為 float[] 數組，然后將其轉換為 List<Float>。這是因為 Milvus 需要特定格式的向量輸入。
構建搜索參數: 創(chuàng)建一個 SearchParam 對象，指定要搜索的集合名稱、度量類型（例如 L2 距離）、返回的最相似的前 K 個結果、向量字段名稱以及搜索的向量數據。
執(zhí)行搜索: 使用 milvusClient 的 search 方法執(zhí)行搜索，并將結果存儲在 searchResults 中。
結果輸出: 打印出搜索的特征向量和搜索結果。

Main主類

package com.et.imagesearch;

import org.nd4j.linalg.api.ndarray.INDArray;

import java.io.File;
import java.io.IOException;

public class Main {
    public static void main(String[] args) throws IOException {
        FeatureExtractor extractor = new FeatureExtractor();
        MilvusManager milvusManager = new MilvusManager();
        ImageSearcher searcher = new ImageSearcher();

        milvusManager.createCollection();

        // images extract
        File[] imageFiles = new File("/Users/liuhaihua/ai/ut-zap50k-images-square/Boots/Ankle/Columbia").listFiles();
        if (imageFiles != null) {
            for (int i = 0; i < imageFiles.length; i++) {
                INDArray features = extractor.extractFeatures(imageFiles[i]);
                milvusManager.insertData(i, features);
            }
        }
      milvusManager.flush();
      milvusManager.buildindex();


        // query
        File queryImage = new File("/Users/liuhaihua/ai/ut-zap50k-images-square/Boots/Ankle/Columbia/7247580.16952.jpg");
        INDArray queryFeatures = extractor.extractFeatures(queryImage);
        searcher.search(queryFeatures);
    }
}

以上只是一些關鍵代碼，所有代碼請參見下面代碼倉庫