feat: analysis admin

2026-11-02 10:06:20 +08:00
parent 5d20c150f4
commit f77913f8a7
34 changed files with 2302 additions and 139 deletions
@@ -0,0 +1,82 @@
+package face
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"lone-services/pkg/utils"
+
+	"github.com/google/uuid"
+)
+
+// PrepareImagePath 将上传字节或已有路径解析为 insightface 可读的本地路径。
+// 返回 cleanup 用于删除临时文件（共享路径不删）。
+func PrepareImagePath(file []byte, filePath string) (string, func(), error) {
+	noop := func() {}
+	if len(file) > utils.NumberZero {
+		dir := utils.GetConfigString("file.upload_path")
+		if dir == utils.StringEmpty {
+			dir = os.TempDir()
+		}
+		if err := os.MkdirAll(dir, 0o755); err != nil {
+			return "", noop, err
+		}
+		path := filepath.Join(dir, "face-login-"+uuid.NewString()+".jpg")
+		if err := os.WriteFile(path, file, 0o644); err != nil {
+			return "", noop, err
+		}
+		return path, func() { _ = os.Remove(path) }, nil
+	}
+
+	path := strings.TrimSpace(filePath)
+	if path == utils.StringEmpty {
+		return "", noop, fmt.Errorf("file or file_path required")
+	}
+	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
+		local, err := downloadToTemp(path)
+		if err != nil {
+			return "", noop, err
+		}
+		return local, func() { _ = os.Remove(local) }, nil
+	}
+	if filepath.IsAbs(path) {
+		return path, noop, nil
+	}
+	base := utils.GetConfigString("file.upload_path")
+	return filepath.Join(base, strings.TrimPrefix(path, "/")), noop, nil
+}
+
+func downloadToTemp(url string) (string, error) {
+	client := &http.Client{Timeout: 15 * time.Second}
+	resp, err := client.Get(url)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("download face image status %d", resp.StatusCode)
+	}
+	dir := utils.GetConfigString("file.upload_path")
+	if dir == utils.StringEmpty {
+		dir = os.TempDir()
+	}
+	if err = os.MkdirAll(dir, 0o755); err != nil {
+		return "", err
+	}
+	path := filepath.Join(dir, "face-dl-"+uuid.NewString()+".jpg")
+	f, err := os.Create(path)
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+	if _, err = io.Copy(f, resp.Body); err != nil {
+		_ = os.Remove(path)
+		return "", err
+	}
+	return path, nil
+}
@@ -0,0 +1,65 @@
+package face
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+
+	"lone-services/pkg/utils"
+)
+
+const faceEmbeddingDim = 512
+
+type insightFaceResp struct {
+	Embedding []float32 `json:"embedding"`
+	Error     string    `json:"error"`
+}
+
+func GetFaceVector(imagePath string) ([]float32, error) {
+	host := utils.GetConfigString("services.skin_py_face_host")
+	apiPath := utils.GetConfigString("services.skin_py_face_path")
+	if host == utils.StringEmpty || apiPath == utils.StringEmpty {
+		return nil, fmt.Errorf("insightface config empty")
+	}
+	url := strings.TrimRight(host, "/") + apiPath
+
+	body, err := json.Marshal(map[string]string{"path": imagePath})
+	if err != nil {
+		return nil, fmt.Errorf("marshal request: %w", err)
+	}
+	req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	client := &http.Client{Timeout: 10 * time.Second}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	respBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("read response: %w", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		var errResp insightFaceResp
+		_ = json.Unmarshal(respBody, &errResp)
+		return nil, fmt.Errorf("face extraction failed (status %d): %s", resp.StatusCode, errResp.Error)
+	}
+
+	var result insightFaceResp
+	if err = json.Unmarshal(respBody, &result); err != nil {
+		return nil, fmt.Errorf("parse response: %w", err)
+	}
+	if len(result.Embedding) != faceEmbeddingDim {
+		return nil, fmt.Errorf("invalid embedding dimension: %d", len(result.Embedding))
+	}
+	return result.Embedding, nil
+}
@@ -0,0 +1,191 @@
+package face
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+
+	"lone-services/pkg/utils"
+
+	"github.com/milvus-io/milvus-sdk-go/v2/client"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+const (
+	collectionName   = "face_vectors"
+	faceVectorField  = "face_vector"
+	faceUserIDField  = "user_id"
+	searchTopK       = 1
+	similarityMin    = float32(0.55)
+	hnswM            = 16
+	hnswEfConstruct  = 200
+	hnswEf           = 64
+	collectionShards = 2
+)
+
+var (
+	milvusMu     sync.RWMutex
+	milvusClient client.Client
+)
+
+func InitMilvus() error {
+	ctx := context.Background()
+	host := utils.GetConfigString("milvus.milvus_host")
+	port := utils.GetConfigString("milvus.milvus_grpc_port")
+	if host == utils.StringEmpty || port == utils.StringEmpty {
+		return fmt.Errorf("milvus config empty")
+	}
+	address := fmt.Sprintf("%s:%s", host, port)
+	c, err := client.NewGrpcClient(ctx, address)
+	if err != nil {
+		return fmt.Errorf("connect milvus %s: %w", address, err)
+	}
+	if err = ensureCollection(ctx, c); err != nil {
+		_ = c.Close()
+		return err
+	}
+	if err = ensureIndex(ctx, c); err != nil {
+		_ = c.Close()
+		return err
+	}
+	if err = c.LoadCollection(ctx, collectionName, false); err != nil {
+		_ = c.Close()
+		return err
+	}
+
+	milvusMu.Lock()
+	if milvusClient != nil {
+		_ = milvusClient.Close()
+	}
+	milvusClient = c
+	milvusMu.Unlock()
+	logx.Infof("milvus connected: %s", address)
+	return nil
+}
+
+func CloseMilvus() {
+	milvusMu.Lock()
+	defer milvusMu.Unlock()
+	if milvusClient != nil {
+		_ = milvusClient.Close()
+		milvusClient = nil
+	}
+}
+
+func getMilvus() client.Client {
+	milvusMu.RLock()
+	defer milvusMu.RUnlock()
+	return milvusClient
+}
+
+func ensureCollection(ctx context.Context, c client.Client) error {
+	has, err := c.HasCollection(ctx, collectionName)
+	if err != nil {
+		return err
+	}
+	if has {
+		return nil
+	}
+	schema := &entity.Schema{
+		CollectionName: collectionName,
+		Description:    "user face vectors (ArcFace 512d)",
+		AutoID:         false,
+		Fields: []*entity.Field{
+			{
+				Name:       faceUserIDField,
+				DataType:   entity.FieldTypeInt64,
+				PrimaryKey: true,
+			},
+			{
+				Name:     faceVectorField,
+				DataType: entity.FieldTypeFloatVector,
+				TypeParams: map[string]string{
+					"dim": fmt.Sprintf("%d", faceEmbeddingDim),
+				},
+			},
+		},
+	}
+	return c.CreateCollection(ctx, schema, int32(collectionShards))
+}
+
+func ensureIndex(ctx context.Context, c client.Client) error {
+	indexes, err := c.DescribeIndex(ctx, collectionName, faceVectorField)
+	if err != nil {
+		if !strings.Contains(err.Error(), "index not found") {
+			return err
+		}
+	} else if len(indexes) > utils.NumberZero {
+		return nil
+	}
+	idx, err := entity.NewIndexHNSW(entity.IP, hnswM, hnswEfConstruct)
+	if err != nil {
+		return err
+	}
+	return c.CreateIndex(ctx, collectionName, faceVectorField, idx, false)
+}
+
+func UpsertFace(userID int64, vector []float32) error {
+	c := getMilvus()
+	if c == nil {
+		return errors.New("milvus client is nil")
+	}
+	if userID < utils.NumberOne || len(vector) != faceEmbeddingDim {
+		return errors.New("invalid face upsert args")
+	}
+	ctx := context.Background()
+	_ = c.Delete(ctx, collectionName, "", fmt.Sprintf("%s == %d", faceUserIDField, userID))
+	_, err := c.Insert(ctx, collectionName, "",
+		entity.NewColumnInt64(faceUserIDField, []int64{userID}),
+		entity.NewColumnFloatVector(faceVectorField, faceEmbeddingDim, [][]float32{vector}),
+	)
+	return err
+}
+
+func SearchUser(vector []float32) (int64, float32, error) {
+	c := getMilvus()
+	if c == nil {
+		return 0, 0, errors.New("milvus client is nil")
+	}
+	if len(vector) != faceEmbeddingDim {
+		return 0, 0, errors.New("invalid face vector")
+	}
+	ctx := context.Background()
+	searchParam, err := entity.NewIndexHNSWSearchParam(hnswEf)
+	if err != nil {
+		return 0, 0, err
+	}
+	result, err := c.Search(
+		ctx,
+		collectionName,
+		nil,
+		"",
+		[]string{faceUserIDField},
+		[]entity.Vector{entity.FloatVector(vector)},
+		faceVectorField,
+		entity.IP,
+		searchTopK,
+		searchParam,
+	)
+	if err != nil {
+		return 0, 0, err
+	}
+	if len(result) == utils.NumberZero || result[0].IDs.Len() == utils.NumberZero {
+		return 0, 0, errors.New("face not found")
+	}
+	ids, ok := result[0].IDs.(*entity.ColumnInt64)
+	if !ok {
+		return 0, 0, errors.New("invalid milvus id column")
+	}
+	userID, err := ids.ValueByIdx(0)
+	if err != nil {
+		return 0, 0, err
+	}
+	score := result[0].Scores[0]
+	if score < similarityMin {
+		return 0, score, errors.New("face similarity too low")
+	}
+	return userID, score, nil
+}