相似推荐合并接口
This commit is contained in:
458
bin/main.go
458
bin/main.go
@@ -39,6 +39,8 @@ func LogComponent(startTime int64, r *http.Request) {
|
||||
|
||||
type Image struct {
|
||||
Id int `json:"id"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
Content string `json:"content"`
|
||||
CreateTime time.Time `json:"create_time"`
|
||||
UpdateTime time.Time `json:"update_time"`
|
||||
@@ -59,13 +61,50 @@ type ListView struct {
|
||||
List []interface{} `json:"list"`
|
||||
}
|
||||
|
||||
var mysqlConnection models.MysqlConnection
|
||||
var milvusConnection models.MilvusConnection
|
||||
|
||||
func (image *Image) GetSimilarImagesIdList(collection_name string) (ids []int64) {
|
||||
ctx := context.Background()
|
||||
|
||||
// 先从milvus中查询图片的向量
|
||||
result, err := milvusConnection.Client.Query(ctx, collection_name, nil, fmt.Sprintf("id in [%d]", image.Id), []string{"embedding"})
|
||||
if err != nil {
|
||||
log.Println("Milvus query failed:", err)
|
||||
return
|
||||
}
|
||||
var embedding []float32
|
||||
for _, item := range result {
|
||||
if item.Name() == "embedding" {
|
||||
embedding = item.FieldData().GetVectors().GetFloatVector().Data
|
||||
continue
|
||||
}
|
||||
}
|
||||
// TODO: 处理向量不存在的情况(生成)
|
||||
// TODO: 处理图片不存在的情况(404)
|
||||
|
||||
// 用向量查询相似图片
|
||||
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)
|
||||
if err != nil {
|
||||
log.Println("Milvus search failed:", err)
|
||||
return
|
||||
}
|
||||
|
||||
// 输出结果
|
||||
for _, item := range resultx {
|
||||
//fmt.Println(item.Scores)
|
||||
//fmt.Println(item.IDs.FieldData().GetScalars().GetLongData().GetData())
|
||||
ids = item.IDs.FieldData().GetScalars().GetLongData().GetData()
|
||||
}
|
||||
|
||||
return ids
|
||||
}
|
||||
|
||||
func main() {
|
||||
runtime.GOMAXPROCS(runtime.NumCPU())
|
||||
|
||||
var mysqlConnection models.MysqlConnection
|
||||
mysqlConnection.Init()
|
||||
|
||||
var milvusConnection models.MilvusConnection
|
||||
milvusConnection.Init()
|
||||
err := milvusConnection.Client.LoadCollection(context.Background(), "default", false)
|
||||
if err != nil {
|
||||
@@ -73,6 +112,238 @@ func main() {
|
||||
return
|
||||
}
|
||||
|
||||
// 获取图片信息列表(分页)
|
||||
http.HandleFunc("/images", func(w http.ResponseWriter, r *http.Request) {
|
||||
defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
|
||||
|
||||
// 私域: (自己的图片, 自己的文章, 自己的精选集, 点赞收藏精选集)
|
||||
// 条件查询(模糊搜索, 时间区间, 作者, 标签, 分类, 精选集, 状态, 置顶, 模糊权重)(权重规则:权重指数)
|
||||
// 条件筛选(交集, 并集, 差集, 子集)
|
||||
// 排序
|
||||
// 分页
|
||||
|
||||
// 获取查询条件(忽略空值), 超级简洁写法
|
||||
QueryConditions := func(key string) (list []string) {
|
||||
for _, item := range strings.Split(r.URL.Query().Get(key), ",") {
|
||||
if item != "" {
|
||||
list = append(list, fmt.Sprintf("'%s'", item))
|
||||
}
|
||||
}
|
||||
return list
|
||||
}
|
||||
|
||||
// 拼接查询条件, 超级简洁写法
|
||||
conditions := ""
|
||||
if authors := QueryConditions("authors"); len(authors) > 0 {
|
||||
conditions += fmt.Sprintf(" AND author IN (%s)", strings.Join(authors, ","))
|
||||
}
|
||||
if tags := QueryConditions("tags"); len(tags) > 0 {
|
||||
conditions += fmt.Sprintf(" AND tag IN (%s)", strings.Join(tags, ","))
|
||||
}
|
||||
if categories := QueryConditions("categories"); len(categories) > 0 {
|
||||
conditions += fmt.Sprintf(" AND categorie IN (%s)", strings.Join(categories, ","))
|
||||
}
|
||||
if sets := QueryConditions("sets"); len(sets) > 0 {
|
||||
conditions += fmt.Sprintf(" AND sets IN (%s)", strings.Join(sets, ","))
|
||||
}
|
||||
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)
|
||||
}
|
||||
if len(idsStr) > 0 {
|
||||
conditions += fmt.Sprintf(" AND id IN (%s)", strings.Join(idsStr, ",")) // 拼接查询条件
|
||||
}
|
||||
}
|
||||
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
|
||||
images.Page, images.PageSize = stringToInt(r.URL.Query().Get("page"), 1), stringToInt(r.URL.Query().Get("pageSize"), 10)
|
||||
rows, err := mysqlConnection.Database.Query("SELECT id, width, height, content, update_time, create_time FROM web_images"+conditions+" LIMIT ?, ?", (images.Page-1)*images.PageSize, images.PageSize)
|
||||
if err != nil {
|
||||
log.Println("获取图片列表失败", err)
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var image Image
|
||||
rows.Scan(&image.Id, &image.Width, &image.Height, &image.Content, &image.UpdateTime, &image.CreateTime)
|
||||
image.UpdateTime = image.UpdateTime.UTC()
|
||||
image.CreateTime = image.CreateTime.UTC()
|
||||
image.Content = regexp.MustCompile(`http:`).ReplaceAllString(image.Content, "https:")
|
||||
image_list = append(image_list, image)
|
||||
}
|
||||
|
||||
// 如果使用了相似图片列表范围, 需要按照相似图片ID原本的顺序重新排序, 需要注意的是, 相似图片ID列表中可能会包含不在当前页的图片ID
|
||||
if similar := QueryConditions("similar"); len(similar) > 0 {
|
||||
// 用于存储排序后的图片列表
|
||||
var image_list_sorted []Image
|
||||
for _, id := range ids {
|
||||
for _, image := range image_list {
|
||||
if image.Id == int(id) {
|
||||
image_list_sorted = append(image_list_sorted, image)
|
||||
}
|
||||
}
|
||||
}
|
||||
image_list = image_list_sorted
|
||||
}
|
||||
|
||||
// 将 []Image 转换为 []interface{}
|
||||
images.List = make([]interface{}, len(image_list))
|
||||
for i, v := range image_list {
|
||||
images.List[i] = v
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
err = mysqlConnection.Database.QueryRow("SELECT COUNT(*) FROM web_images" + conditions).Scan(&images.Total)
|
||||
if err != nil {
|
||||
log.Println("获取图片总数失败", err)
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// 是否有下一页
|
||||
images.Next = images.Total > images.Page*images.PageSize
|
||||
|
||||
// 将对象转换为有缩进的JSON输出
|
||||
data, err := json.MarshalIndent(images, "", " ")
|
||||
if err != nil {
|
||||
log.Println("转换图片列表失败", err)
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json; charset=UTF-8")
|
||||
w.Write(data)
|
||||
|
||||
})
|
||||
|
||||
// 获取相似图片列表
|
||||
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) // 最后打印日志
|
||||
@@ -151,185 +422,6 @@ func main() {
|
||||
w.Write(json)
|
||||
})
|
||||
|
||||
type Similar struct {
|
||||
Id int64 `json:"id"`
|
||||
ArticleId int64 `json:"article_id"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
}
|
||||
|
||||
// 获取相似图片列表
|
||||
http.HandleFunc("/similar", func(w http.ResponseWriter, r *http.Request) {
|
||||
defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
|
||||
id := "8888"
|
||||
// 先查询图片的向量在 mulvis 中是否存在
|
||||
var collection_name = "default" // 图片集合名称
|
||||
result, err := milvusConnection.Client.Query(
|
||||
context.Background(), // ctx
|
||||
collection_name, // CollectionName
|
||||
[]string{}, // 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)
|
||||
}
|
||||
|
||||
//func printResult(sRet *client.SearchResult) {
|
||||
// randoms := make([]float64, 0, sRet.ResultCount)
|
||||
// scores := make([]float32, 0, sRet.ResultCount)
|
||||
//
|
||||
// var randCol *entity.ColumnDouble
|
||||
// for _, field := range sRet.Fields {
|
||||
// if field.Name() == randomCol {
|
||||
// c, ok := field.(*entity.ColumnDouble)
|
||||
// if ok {
|
||||
// randCol = c
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// for i := 0; i < sRet.ResultCount; i++ {
|
||||
// val, err := randCol.ValueByIdx(i)
|
||||
// if err != nil {
|
||||
// log.Fatal(err)
|
||||
// }
|
||||
// randoms = append(randoms, val)
|
||||
// scores = append(scores, sRet.Scores[i])
|
||||
// }
|
||||
// fmt.Printf("\trandoms: %v, scores: %v\n", randoms, scores)
|
||||
//}
|
||||
})
|
||||
|
||||
// 获取图片信息列表(分页)
|
||||
http.HandleFunc("/images", func(w http.ResponseWriter, r *http.Request) {
|
||||
defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
|
||||
|
||||
// 私域: (自己的图片, 自己的文章, 自己的精选集, 点赞收藏精选集)
|
||||
// 条件查询(模糊搜索, 时间区间, 作者, 标签, 分类, 精选集, 状态, 置顶, 模糊权重)(权重规则:权重指数)
|
||||
// 条件筛选(交集, 并集, 差集, 子集)
|
||||
// 排序
|
||||
// 分页
|
||||
|
||||
// 获取查询条件(忽略空值), 超级简洁写法
|
||||
QueryConditions := func(key string) (list []string) {
|
||||
for _, item := range strings.Split(r.URL.Query().Get(key), ",") {
|
||||
if item != "" {
|
||||
list = append(list, fmt.Sprintf("'%s'", item))
|
||||
}
|
||||
}
|
||||
return list
|
||||
}
|
||||
|
||||
// 拼接查询条件, 超级简洁写法
|
||||
conditions := ""
|
||||
if authors := QueryConditions("authors"); len(authors) > 0 {
|
||||
conditions += fmt.Sprintf(" AND author IN (%s)", strings.Join(authors, ","))
|
||||
}
|
||||
if tags := QueryConditions("tags"); len(tags) > 0 {
|
||||
conditions += fmt.Sprintf(" AND tag IN (%s)", strings.Join(tags, ","))
|
||||
}
|
||||
if categories := QueryConditions("categories"); len(categories) > 0 {
|
||||
conditions += fmt.Sprintf(" AND categorie IN (%s)", strings.Join(categories, ","))
|
||||
}
|
||||
if sets := QueryConditions("sets"); len(sets) > 0 {
|
||||
conditions += fmt.Sprintf(" AND sets IN (%s)", strings.Join(sets, ","))
|
||||
}
|
||||
if conditions != "" {
|
||||
conditions = strings.Replace(conditions, " AND", "", 1) // 去掉第一个 AND
|
||||
conditions = " WHERE" + conditions // 拼接 WHERE
|
||||
fmt.Println(conditions) // 打印查询条件
|
||||
}
|
||||
|
||||
// 获取图片列表
|
||||
var images ListView
|
||||
images.Page, images.PageSize = stringToInt(r.URL.Query().Get("page"), 1), stringToInt(r.URL.Query().Get("pageSize"), 10)
|
||||
rows, err := mysqlConnection.Database.Query("SELECT id, content, update_time, create_time FROM web_images"+conditions+" LIMIT ?, ?", (images.Page-1)*images.PageSize, images.PageSize)
|
||||
if err != nil {
|
||||
log.Println("获取图片列表失败", err)
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// 处理结果集
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var image Image
|
||||
rows.Scan(&image.Id, &image.Content, &image.UpdateTime, &image.CreateTime)
|
||||
image.UpdateTime = image.UpdateTime.UTC()
|
||||
image.CreateTime = image.CreateTime.UTC()
|
||||
image.Content = regexp.MustCompile(`http:`).ReplaceAllString(image.Content, "https:")
|
||||
images.List = append(images.List, image)
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
err = mysqlConnection.Database.QueryRow("SELECT COUNT(*) FROM web_images" + conditions).Scan(&images.Total)
|
||||
if err != nil {
|
||||
log.Println("获取图片总数失败", err)
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// 是否有下一页
|
||||
images.Next = images.Total > images.Page*images.PageSize
|
||||
|
||||
// 将对象转换为有缩进的JSON输出
|
||||
data, err := json.MarshalIndent(images, "", " ")
|
||||
if err != nil {
|
||||
log.Println("转换图片列表失败", err)
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json; charset=UTF-8")
|
||||
w.Write(data)
|
||||
|
||||
})
|
||||
|
||||
// URL 格式: /img/{type}-{id}.{format}?width=320&height=320&fit=cover
|
||||
http.HandleFunc("/img/", func(w http.ResponseWriter, r *http.Request) {
|
||||
defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
|
||||
|
Reference in New Issue
Block a user