相似图片

bin/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -15,6 +16,7 @@ import (
 
 	"git.satori.love/gameui/webp/models"
 	_ "github.com/go-sql-driver/mysql"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 )
 
 // string 转换为 int, 如果转换失败则返回默认值
@@ -186,25 +188,70 @@ func main() {
 		**/
 	})
 
-	//// 获取相似图片列表
+	type Similar struct {
-	//http.HandleFunc("/similar", func(w http.ResponseWriter, r *http.Request) {
+		Id        int64     `json:"id"`
-	//	defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
+		ArticleId int64     `json:"article_id"`
-	//	id := "8888"
+		Embedding []float32 `json:"embedding"`
-	//	// 先查询图片的向量在 mulvis 中是否存在
+	}
-	//	var collection_name = "default" // 图片集合名称
+
-	//	result, err := milvusConnection.Client.Query(
+	// 获取相似图片列表
-	//		context.Background(),              // ctx
+	http.HandleFunc("/similar", func(w http.ResponseWriter, r *http.Request) {
-	//		collection_name,                   // CollectionName
+		defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
-	//		[]string{},                        // PartitionName
+		id := "8888"
-	//		fmt.Sprintf("id in [%s]", id),     // expr
+		// 先查询图片的向量在 mulvis 中是否存在
-	//		[]string{"book_id", "book_intro"}, // OutputFields
+		var collection_name = "default" // 图片集合名称
-	//	)
+		result, err := milvusConnection.Client.Query(
-	//	if err != nil {
+			context.Background(),                      // ctx
-	//		log.Println(err)
+			collection_name,                           // CollectionName
-	//		return
+			[]string{},                                // PartitionName
-	//	}
+			fmt.Sprintf("id in [%s]", id),             // expr
-	//	log.Println(result)
+			[]string{"id", "embedding", "article_id"}, // OutputFields
-	//})
+		)
+		if err != nil {
+			log.Println(err)
+			return
+		}
+
+		// TODO: 不存在则重建向量
+		var similar Similar
+		for _, item := range result {
+			if item.Name() == "id" {
+				similar.Id = item.FieldData().GetScalars().GetLongData().GetData()[0]
+				continue
+			}
+			if item.Name() == "article_id" {
+				similar.ArticleId = item.FieldData().GetScalars().GetLongData().GetData()[0]
+				continue
+			}
+			if item.Name() == "embedding" {
+				similar.Embedding = item.FieldData().GetVectors().GetFloatVector().Data
+				continue
+			}
+		}
+
+		// 用向量查询相似图片
+		sp, _ := entity.NewIndexFlatSearchParam()
+		vectors := []entity.Vector{
+			entity.FloatVector(similar.Embedding),
+		}
+		resultx, err := milvusConnection.Client.Search(
+			context.Background(),         // ctx
+			collection_name,              // CollectionName
+			[]string{},                   // PartitionNames
+			"",                           // expr
+			[]string{"id", "article_id"}, // OutputFields
+			vectors,                      // vectors
+			"embedding",                  // vectorField
+			entity.L2,                    // entity.MetricType
+			10,                           // topK
+			sp,
+		)
+		if err != nil {
+			log.Println(err)
+			return
+		}
+		println(resultx)
+	})
 
 	// 获取图片信息列表(分页)
 	http.HandleFunc("/images", func(w http.ResponseWriter, r *http.Request) {