diff --git a/bin/main.go b/bin/main.go index 8d47e80..8eb4bdc 100644 --- a/bin/main.go +++ b/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) // 最后打印日志