Files
webp/bin/main.go
2023-04-12 15:52:02 +08:00

371 lines
12 KiB
Go

package main
import (
"context"
"fmt"
"log"
"net/http"
"runtime"
"strings"
"time"
"regexp"
"strconv"
"encoding/json"
"git.satori.love/gameui/webp/models"
_ "github.com/go-sql-driver/mysql"
"github.com/milvus-io/milvus-sdk-go/v2/entity"
"github.com/spf13/viper"
)
// string 转换为 int, 如果转换失败则返回默认值
func stringToInt(str string, defaultValue int) int {
if str == "" {
return defaultValue
}
value, err := strconv.Atoi(str)
if err != nil {
return defaultValue
}
return value
}
func LogComponent(startTime int64, r *http.Request) {
endTime := fmt.Sprintf("%dms", (time.Now().UnixNano()-startTime)/1000000)
log.Println(r.Method, r.URL.Path, endTime)
}
type Image struct {
Id int `json:"id"`
Content string `json:"content"`
CreateTime time.Time `json:"create_time"`
UpdateTime time.Time `json:"update_time"`
}
type Tag struct {
Id int `json:"id"`
Name string `json:"name"`
CreateTime time.Time `json:"create_time"`
UpdateTime time.Time `json:"update_time"`
}
type ListView struct {
Page int `json:"page"`
PageSize int `json:"page_size"`
Total int `json:"total"`
Next bool `json:"next"`
List []interface{} `json:"list"`
}
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 {
log.Println("Milvus load collection failed:", err)
return
}
// 获取标签列表
http.HandleFunc("/tags", func(w http.ResponseWriter, r *http.Request) {
defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
// 标签的原理
// 1. 通过文章的 tag 字段, 获取所有的标签
// 2. 通过标签的 id, 获取标签的名称
// 热门权重指数的标签排序
// 1. 标签的权重指数 = (标签的文章数 * 标签的文章数) * 近期增幅
// 2. 标签的近期增幅 = (标签的文章数 - 标签的文章数) / 标签的文章数
// 标签是一个虚拟表, ORC 提取的数据都带有多个维度的比重概率(分布概率, 对比度概率, 文字大小, 文字重量, 词频概率, 词性概率, 词长概率, 词序概率)
// 经过规则过滤后, 用动态调参的指数计算乘积作为权重, 权重仍达到某个阈值的数据才会被视为标签
// 获取查询条件(忽略空值), 超级简洁写法
QueryConditions := func(key string) (list []string) {
if r.FormValue(key) != "" {
list = strings.Split(r.FormValue(key), ",")
}
return
}
// 拼接查询条件, 超级简洁写法
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, ","))
}
// 获取标签列表
var tags ListView
tags.Page, tags.PageSize = stringToInt(r.FormValue("page"), 1), stringToInt(r.FormValue("pageSize"), 10)
rows, err := mysqlConnection.Database.Query("SELECT id, name, update_time, create_time FROM web_tags"+conditions+" ORDER BY id DESC LIMIT ?, ?", (tags.Page-1)*tags.PageSize, tags.PageSize)
if err != nil {
log.Println(err)
return
}
defer rows.Close()
for rows.Next() {
var tag Tag
if err := rows.Scan(&tag.Id, &tag.Name, &tag.UpdateTime, &tag.CreateTime); err != nil {
log.Println(err)
continue
}
tags.List = append(tags.List, tag)
}
if err := rows.Err(); err != nil {
log.Println(err)
return
}
// 获取总数
if err := mysqlConnection.Database.QueryRow("SELECT COUNT(*) FROM web_tags" + conditions).Scan(&tags.Total); err != nil {
log.Println(err)
return
}
// 是否有下一页
tags.Next = tags.Total > tags.Page*tags.PageSize
// 将对象转换为有缩进的JSON输出
json, err := json.MarshalIndent(tags, "", " ")
if err != nil {
log.Println(err)
return
}
// 输出JSON
w.Header().Set("Content-Type", "application/json")
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.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) {
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) // 最后打印日志
reg := regexp.MustCompile(`^/img/([0-9a-zA-Z]+)-([0-9a-zA-Z]+).(jpg|jpeg|png|webp)$`)
matches := reg.FindStringSubmatch(r.URL.Path)
if len(matches) != 4 {
http.Error(w, "URL 格式错误", http.StatusNotFound)
return
}
group, id, format, width, height, fit := matches[1], matches[2], matches[3], stringToInt(r.URL.Query().Get("width"), 0), stringToInt(r.URL.Query().Get("height"), 0), r.URL.Query().Get("fit")
content, err := mysqlConnection.GetImageContent(group, id)
if err != nil {
log.Println("获取图片失败", format, err)
http.Error(w, err.Error(), http.StatusNotFound)
return
}
var img models.Image
img.Init(content)
data, err := img.ToWebP(width, height, fit)
if err != nil {
log.Println("转换图片失败", err)
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
w.Header().Set("Content-Type", "image/webp")
w.Header().Set("Cache-Control", "max-age=31536000")
w.Write(data)
})
// URL 格式: /webp/{type}-{id}-{version}-{width}-{height}-{fit}.{format}
http.HandleFunc("/webp/", func(w http.ResponseWriter, r *http.Request) {
defer LogComponent(time.Now().UnixNano(), r) // 最后打印日志
reg := regexp.MustCompile(`^/webp/([0-9a-zA-Z]+)-([0-9a-zA-Z]+)-([0-9a-zA-Z]+)-([0-9]+)-([0-9]+)-([a-zA-Z]+).(jpg|jpeg|png|webp)$`)
matches := reg.FindStringSubmatch(r.URL.Path)
if len(matches) != 8 {
log.Println("URL 格式错误", matches)
w.WriteHeader(http.StatusNotFound)
return
}
group, id, version, width, height, fit, format := matches[1], matches[2], matches[3], stringToInt(matches[4], 0), stringToInt(matches[5], 0), matches[6], matches[7]
content, err := mysqlConnection.GetImageContent(group, id)
if err != nil {
log.Println("获取图片失败", version, format, err)
w.WriteHeader(http.StatusNotFound)
return
}
var img models.Image
img.Init(content)
data, err := img.ToWebP(width, height, fit)
if err != nil {
log.Println("转换图片失败", err)
w.WriteHeader(http.StatusBadRequest)
return
}
w.Header().Set("Content-Type", "image/webp")
w.Header().Set("Cache-Control", "max-age=31536000")
w.Write(data)
})
// 从Viper中读取配置
port := viper.GetString("server.port")
log.Println("Server is running at http://localhost:" + port)
http.ListenAndServe(":"+port, nil)
}