完善二维码池 OCR 配置与识别

2026-06-18 15:22:46 +08:00
parent e8bdf574f4
commit 35f2d93fdf
16 changed files with 1025 additions and 48 deletions
@@ -0,0 +1,327 @@
+package chat
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"net/textproto"
+	"regexp"
+	"sort"
+	"strings"
+	"time"
+	"unicode/utf8"
+)
+
+const (
+	maxQrCodeOCRFileSize  = 10 * 1024 * 1024
+	qrCodeOCRPollInterval = 2 * time.Second
+	qrCodeOCRMaxPollTimes = 45
+)
+
+var (
+	ErrQrCodeOCRNotConfigured = errors.New("二维码 OCR Token 未配置")
+	ErrQrCodeOCRInvalidFile   = errors.New("无效的二维码图片")
+	ErrQrCodeOCRUnavailable   = errors.New("二维码 OCR 服务不可用")
+)
+
+type QrCodeOCRResult struct {
+	GroupName  string   `json:"group_name"`
+	Candidates []string `json:"candidates"`
+	RawText    string   `json:"raw_text"`
+}
+
+type paddleOCRJobResponse struct {
+	Data struct {
+		JobID string `json:"jobId"`
+	} `json:"data"`
+	Message string `json:"message"`
+	Error   string `json:"error"`
+}
+
+type paddleOCRJobStatusResponse struct {
+	Data struct {
+		State     string `json:"state"`
+		ErrorMsg  string `json:"errorMsg"`
+		ResultURL struct {
+			JSONURL string `json:"jsonUrl"`
+		} `json:"resultUrl"`
+	} `json:"data"`
+	Message string `json:"message"`
+	Error   string `json:"error"`
+}
+
+// RecognizeQrCodeGroupName 调用 PaddleOCR 异步 API 识别二维码图片中的群名。
+func (r *Repository) RecognizeQrCodeGroupName(ctx context.Context, filename, contentType string, reader io.Reader) (*QrCodeOCRResult, error) {
+	config, err := r.GetQrCodeOCRConfig(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if strings.TrimSpace(config.Token) == "" {
+		return nil, ErrQrCodeOCRNotConfigured
+	}
+	data, err := io.ReadAll(io.LimitReader(reader, maxQrCodeOCRFileSize+1))
+	if err != nil || len(data) == 0 || len(data) > maxQrCodeOCRFileSize {
+		return nil, ErrQrCodeOCRInvalidFile
+	}
+	if contentType == "" {
+		contentType = http.DetectContentType(data)
+	}
+	if !strings.HasPrefix(contentType, "image/") {
+		return nil, ErrQrCodeOCRInvalidFile
+	}
+
+	jobID, err := submitPaddleOCRJob(ctx, config, filename, contentType, data)
+	if err != nil {
+		return nil, err
+	}
+	jsonURL, err := waitPaddleOCRJob(ctx, config, jobID)
+	if err != nil {
+		return nil, err
+	}
+	rawText, err := fetchPaddleOCRText(ctx, jsonURL)
+	if err != nil {
+		return nil, err
+	}
+	candidates := parseQrCodeGroupNameCandidates(rawText)
+	groupName := ""
+	if len(candidates) > 0 {
+		groupName = candidates[0]
+	}
+	return &QrCodeOCRResult{
+		GroupName:  groupName,
+		Candidates: candidates,
+		RawText:    rawText,
+	}, nil
+}
+
+func submitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, filename, contentType string, data []byte) (string, error) {
+	body := &bytes.Buffer{}
+	writer := multipart.NewWriter(body)
+	_ = writer.WriteField("model", config.Model)
+	optionalPayload, _ := json.Marshal(map[string]bool{
+		"useDocOrientationClassify": false,
+		"useDocUnwarping":           false,
+		"useChartRecognition":       false,
+	})
+	_ = writer.WriteField("optionalPayload", string(optionalPayload))
+	partHeader := make(textproto.MIMEHeader)
+	partHeader.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, escapeQuotes(filename)))
+	partHeader.Set("Content-Type", contentType)
+	fileWriter, err := writer.CreatePart(partHeader)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	if _, err := fileWriter.Write(data); err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	if err := writer.Close(); err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, config.JobURL, body)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	req.Header.Set("Authorization", "bearer "+config.Token)
+	req.Header.Set("Content-Type", writer.FormDataContentType())
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	defer resp.Body.Close()
+	respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("%w: status=%d body=%s", ErrQrCodeOCRUnavailable, resp.StatusCode, string(respBody))
+	}
+	var payload paddleOCRJobResponse
+	if err := json.Unmarshal(respBody, &payload); err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	if payload.Data.JobID == "" {
+		return "", fmt.Errorf("%w: 未返回 jobId", ErrQrCodeOCRUnavailable)
+	}
+	return payload.Data.JobID, nil
+}
+
+func waitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, jobID string) (string, error) {
+	for i := 0; i < qrCodeOCRMaxPollTimes; i++ {
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(config.JobURL, "/")+"/"+jobID, nil)
+		if err != nil {
+			return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+		}
+		req.Header.Set("Authorization", "bearer "+config.Token)
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+		}
+		respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
+		_ = resp.Body.Close()
+		if resp.StatusCode != http.StatusOK {
+			return "", fmt.Errorf("%w: status=%d body=%s", ErrQrCodeOCRUnavailable, resp.StatusCode, string(respBody))
+		}
+		var payload paddleOCRJobStatusResponse
+		if err := json.Unmarshal(respBody, &payload); err != nil {
+			return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+		}
+		switch payload.Data.State {
+		case "done":
+			if payload.Data.ResultURL.JSONURL == "" {
+				return "", fmt.Errorf("%w: 未返回识别结果地址", ErrQrCodeOCRUnavailable)
+			}
+			return payload.Data.ResultURL.JSONURL, nil
+		case "failed":
+			if payload.Data.ErrorMsg != "" {
+				return "", fmt.Errorf("%w: %s", ErrQrCodeOCRUnavailable, payload.Data.ErrorMsg)
+			}
+			return "", ErrQrCodeOCRUnavailable
+		}
+		select {
+		case <-ctx.Done():
+			return "", ctx.Err()
+		case <-time.After(qrCodeOCRPollInterval):
+		}
+	}
+	return "", fmt.Errorf("%w: 识别超时", ErrQrCodeOCRUnavailable)
+}
+
+func fetchPaddleOCRText(ctx context.Context, jsonURL string) (string, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, jsonURL, nil)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
+		return "", fmt.Errorf("%w: status=%d body=%s", ErrQrCodeOCRUnavailable, resp.StatusCode, string(respBody))
+	}
+	body, err := io.ReadAll(io.LimitReader(resp.Body, 10*1024*1024))
+	if err != nil {
+		return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
+	}
+	return collectPaddleOCRJSONLText(string(body)), nil
+}
+
+func collectPaddleOCRJSONLText(raw string) string {
+	var texts []string
+	for _, line := range strings.Split(raw, "\n") {
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+		var value any
+		if err := json.Unmarshal([]byte(line), &value); err != nil {
+			texts = append(texts, line)
+			continue
+		}
+		collectPaddleOCRText(value, &texts)
+	}
+	return strings.Join(texts, "\n")
+}
+
+func collectPaddleOCRText(value any, texts *[]string) {
+	switch v := value.(type) {
+	case string:
+		if strings.TrimSpace(v) != "" {
+			*texts = append(*texts, v)
+		}
+	case []any:
+		for _, item := range v {
+			collectPaddleOCRText(item, texts)
+		}
+	case map[string]any:
+		if markdown, ok := v["markdown"].(map[string]any); ok {
+			if text, ok := markdown["text"].(string); ok && strings.TrimSpace(text) != "" {
+				*texts = append(*texts, text)
+			}
+		}
+		for key, item := range v {
+			if key == "images" || key == "outputImages" {
+				continue
+			}
+			collectPaddleOCRText(item, texts)
+		}
+	}
+}
+
+func parseQrCodeGroupNameCandidates(rawText string) []string {
+	excluded := []*regexp.Regexp{
+		regexp.MustCompile(`使用.*微信`),
+		regexp.MustCompile(`企业微信`),
+		regexp.MustCompile(`扫码`),
+		regexp.MustCompile(`加入`),
+		regexp.MustCompile(`二维码`),
+		regexp.MustCompile(`有效`),
+		regexp.MustCompile(`天内`),
+		regexp.MustCompile(`月\d+日`),
+		regexp.MustCompile(`前`),
+		regexp.MustCompile(`^\d+$`),
+	}
+	scores := map[string]int{}
+	for _, line := range strings.Split(rawText, "\n") {
+		line = normalizeQrCodeOCRLine(line)
+		if line == "" || utf8.RuneCountInString(line) < 3 || utf8.RuneCountInString(line) > 64 {
+			continue
+		}
+		if !regexp.MustCompile(`[\p{Han}]`).MatchString(line) {
+			continue
+		}
+		skip := false
+		for _, pattern := range excluded {
+			if pattern.MatchString(line) {
+				skip = true
+				break
+			}
+		}
+		if skip {
+			continue
+		}
+		scores[line] = scoreQrCodeGroupNameCandidate(line)
+	}
+	candidates := make([]string, 0, len(scores))
+	for candidate := range scores {
+		candidates = append(candidates, candidate)
+	}
+	sort.SliceStable(candidates, func(i, j int) bool {
+		return scores[candidates[i]] > scores[candidates[j]]
+	})
+	return candidates
+}
+
+func normalizeQrCodeOCRLine(line string) string {
+	replacer := strings.NewReplacer(" ", "", "\t", "", "|", "", "｜", "", "#", "", "*", "", "`", "", ">", "")
+	line = replacer.Replace(line)
+	allowed := regexp.MustCompile(`[^\p{Han}a-zA-Z0-9_（）()《》\-—·]+`)
+	return allowed.ReplaceAllString(line, "")
+}
+
+func scoreQrCodeGroupNameCandidate(line string) int {
+	length := utf8.RuneCountInString(line)
+	score := length
+	if score > 20 {
+		score = 20
+	}
+	if strings.Contains(line, "群") {
+		score += 100
+	}
+	if strings.Contains(line, "-") || strings.Contains(line, "—") {
+		score += 30
+	}
+	if length >= 6 && length <= 24 {
+		score += 20
+	}
+	return score
+}
+
+func escapeQuotes(value string) string {
+	return strings.NewReplacer("\\", "\\\\", `"`, "\\\"").Replace(value)
+}