diff --git a/bin/main.go b/bin/main.go index 8eb4bdc..991f356 100644 --- a/bin/main.go +++ b/bin/main.go @@ -84,9 +84,10 @@ func (image *Image) GetSimilarImagesIdList(collection_name string) (ids []int64) // TODO: 处理图片不存在的情况(404) // 用向量查询相似图片 + topk := 1000 sp, _ := entity.NewIndexIvfFlatSearchParam(64) vectors := []entity.Vector{entity.FloatVector(embedding)} - resultx, err := milvusConnection.Client.Search(ctx, collection_name, nil, "", []string{"id", "article_id"}, vectors, "embedding", entity.L2, 10, sp) + resultx, err := milvusConnection.Client.Search(ctx, collection_name, nil, "", []string{"id", "article_id"}, vectors, "embedding", entity.L2, topk, sp) if err != nil { log.Println("Milvus search failed:", err) return @@ -148,17 +149,13 @@ func main() { } var ids []int64 if similar := QueryConditions("similar"); len(similar) > 0 { - // 避免报错: strconv.Atoi failed: strconv.Atoi: parsing "'8888'": invalid syntax id, err := strconv.Atoi(strings.Trim(similar[0], "'")) if err != nil { log.Println("strconv.Atoi failed:", err) return } - fmt.Println("id:", id) - // 如果指定以某个图片为基准的相似图片列表范围, 获取相似图片ID的列表 ids = (&Image{Id: id}).GetSimilarImagesIdList("default") - fmt.Println("ids:", ids) idsStr := make([]string, len(ids)) for i, v := range ids { idsStr[i] = strconv.FormatInt(v, 10) @@ -170,12 +167,8 @@ func main() { if conditions != "" { conditions = strings.Replace(conditions, " AND", "", 1) // 去掉第一个 AND conditions = " WHERE" + conditions // 拼接 WHERE - fmt.Println(conditions) // 打印查询条件 } - // 打印查询语句 - fmt.Println("SELECT id, width, height, content, update_time, create_time FROM web_images" + conditions + " LIMIT ?, ?") - // 获取图片列表 var images ListView var image_list []Image @@ -198,7 +191,6 @@ func main() { // 如果使用了相似图片列表范围, 需要按照相似图片ID原本的顺序重新排序, 需要注意的是, 相似图片ID列表中可能会包含不在当前页的图片ID if similar := QueryConditions("similar"); len(similar) > 0 { - // 用于存储排序后的图片列表 var image_list_sorted []Image for _, id := range ids { for _, image := range image_list { @@ -239,111 +231,6 @@ func main() { }) - // 获取相似图片列表 - http.HandleFunc("/similar", func(w http.ResponseWriter, r *http.Request) { - defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志 - - var collection_name = "default" // 图片集合名称 - var image Image = Image{Id: 8888} // 图片对象 - - ids := image.GetSimilarImagesIdList(collection_name) - - var listview ListView - - // 是否有下一页 - listview.Total = 10 - listview.Next = false - listview.Page = 1 - listview.PageSize = 10 - - // 获取一组ID对应的图片数据 - ids_str := strings.Trim(strings.Replace(fmt.Sprint(ids), " ", ",", -1), "[]") - println(ids_str) - rows, err := mysqlConnection.Database.Query("SELECT id, content, update_time, create_time, width, height FROM web_images WHERE id in (" + ids_str + ")") - if err != nil { - log.Println("SQL查询失败:", err.Error()) - return - } - for rows.Next() { - var image Image - err := rows.Scan(&image.Id, &image.Content, &image.UpdateTime, &image.CreateTime, &image.Width, &image.Height) - if err != nil { - log.Println("SQL查询失败:", err.Error()) - return - } - listview.List = append(listview.List, image) - } - - // 将对象转换为有缩进的JSON输出 - json, err := json.MarshalIndent(listview, "", " ") - if err != nil { - log.Println(err) - return - } - - // 输出JSON - w.Header().Set("Content-Type", "application/json") - w.Write(json) - - //result, err := milvusConnection.Client.Query( - // context.Background(), // ctx - // collection_name, // CollectionName - // nil, // PartitionName - // fmt.Sprintf("id in [%s]", id), // expr - // []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.NewIndexIvfFlatSearchParam(64) - //vectors := []entity.Vector{ - // entity.FloatVector(similar.Embedding), - //} - //resultx, err := milvusConnection.Client.Search( - // context.Background(), // ctx - // collection_name, // CollectionName - // nil, // PartitionNames - // "", // expr - // []string{"id", "article_id"}, // OutputFields - // vectors, // vectors - // "embedding", // vectorField - // entity.L2, // entity.MetricType - // 10, // topK - // sp, // searchParam - //) - //if err != nil { - // log.Println(err) - // return - //} - //// 输出结果 - //for _, item := range resultx { - // fmt.Println(item.Scores) - // fmt.Println(item.IDs) - // fmt.Println(item.ResultCount) - // fmt.Println(item.Fields) - //} - - }) - // 获取标签列表 http.HandleFunc("/tags", func(w http.ResponseWriter, r *http.Request) { defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志