yml/hfb_sys
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修复 ocr 错误

Browse Source
This commit is contained in:
yml2213
2026-06-28 15:00:54 +08:00
parent 18eba36ef9
commit b124c74a46
8 changed files with 331 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/modules/chat/qrcode_ocr.go
+201 -23
View File
@@ -15,18 +15,30 @@ import (
"strings"
"time"
"unicode/utf8"
"github.com/google/uuid"
)
const (
maxQrCodeOCRFileSize = 10 * 1024 * 1024
qrCodeOCRPollInterval = 2 * time.Second
qrCodeOCRMaxPollTimes = 45
maxQrCodeOCRFileSize = 10 * 1024 * 1024
ocrJobRedisKeyPrefix = "ocr:job:"
ocrJobRedisTTL = 10 * time.Minute
ocrJobStatusSubmitted = "submitted"
ocrJobStatusRunning = "running"
ocrJobStatusDone = "done"
ocrJobStatusFailed = "failed"
)
var (
ErrQrCodeOCRNotConfigured = errors.New("二维码 OCR Token 未配置")
ErrQrCodeOCRInvalidFile = errors.New("无效的二维码图片")
ErrQrCodeOCRUnavailable = errors.New("二维码 OCR 服务不可用")
ErrQrCodeOCRJobNotFound = errors.New("OCR 任务不存在或已过期")
ErrQrCodeOCRRedisDisabled = errors.New("OCR 异步模式需要 Redis")
ocrHTTPClient = &http.Client{Timeout: 30 * time.Second}
)
type QrCodeOCRResult struct {
@@ -35,6 +47,13 @@ type QrCodeOCRResult struct {
RawText string `json:"raw_text"`
}
type OCRJobRecord struct {
PaddleJobID string `json:"paddle_job_id"`
Status string `json:"status"`
Result *QrCodeOCRResult `json:"result,omitempty"`
Error string `json:"error,omitempty"`
}
type paddleOCRJobResponse struct {
Data struct {
JobID string `json:"jobId"`
@@ -55,7 +74,51 @@ type paddleOCRJobStatusResponse struct {
Error string `json:"error"`
}
// RecognizeQrCodeGroupName 调用 PaddleOCR 异步 API 识别二维码图片中的群名。
func readAndValidateOCRFile(reader io.Reader, contentType string) ([]byte, string, error) {
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
}
return data, contentType, nil
}
func ocrJobKey(jobID string) string {
return ocrJobRedisKeyPrefix + jobID
}
func (r *Repository) saveOCRJobRecord(ctx context.Context, jobID string, record *OCRJobRecord) error {
if r.redis == nil {
return ErrQrCodeOCRRedisDisabled
}
raw, err := json.Marshal(record)
if err != nil {
return err
}
return r.redis.Set(ctx, ocrJobKey(jobID), string(raw), ocrJobRedisTTL).Err()
}
func (r *Repository) getOCRJobRecord(ctx context.Context, jobID string) (*OCRJobRecord, error) {
if r.redis == nil {
return nil, ErrQrCodeOCRRedisDisabled
}
raw, err := r.redis.Get(ctx, ocrJobKey(jobID)).Result()
if err != nil {
return nil, ErrQrCodeOCRJobNotFound
}
var record OCRJobRecord
if err := json.Unmarshal([]byte(raw), &record); err != nil {
return nil, err
}
return &record, nil
}
// RecognizeQrCodeGroupName 同步调用 PaddleOCR 识别二维码图片中的群名（保留向后兼容）。
func (r *Repository) RecognizeQrCodeGroupName(ctx context.Context, filename, contentType string, reader io.Reader) (*QrCodeOCRResult, error) {
config, err := r.GetQrCodeOCRConfig(ctx)
if err != nil {
@@ -64,22 +127,16 @@ func (r *Repository) RecognizeQrCodeGroupName(ctx context.Context, filename, con
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)
data, contentType, err := readAndValidateOCRFile(reader, contentType)
if err != nil {
return nil, err
}
jsonURL, err := waitPaddleOCRJob(ctx, config, jobID)
paddleJobID, err := submitPaddleOCRJob(ctx, config, filename, contentType, data)
if err != nil {
return nil, err
}
jsonURL, err := waitPaddleOCRJob(ctx, config, paddleJobID)
if err != nil {
return nil, err
}
@@ -99,6 +156,127 @@ func (r *Repository) RecognizeQrCodeGroupName(ctx context.Context, filename, con
}, nil
}
// SubmitOCRJob 提交 OCR 任务到 PaddleOCR，立即返回任务 ID 供前端轮询。
func (r *Repository) SubmitOCRJob(ctx context.Context, filename, contentType string, reader io.Reader) (string, error) {
if r.redis == nil {
return "", ErrQrCodeOCRRedisDisabled
}
config, err := r.GetQrCodeOCRConfig(ctx)
if err != nil {
return "", err
}
if strings.TrimSpace(config.Token) == "" {
return "", ErrQrCodeOCRNotConfigured
}
data, contentType, err := readAndValidateOCRFile(reader, contentType)
if err != nil {
return "", err
}
paddleJobID, err := submitPaddleOCRJob(ctx, config, filename, contentType, data)
if err != nil {
return "", err
}
jobID := uuid.NewString()
record := &OCRJobRecord{
PaddleJobID: paddleJobID,
Status: ocrJobStatusSubmitted,
}
if err := r.saveOCRJobRecord(ctx, jobID, record); err != nil {
return "", err
}
return jobID, nil
}
// PollOCRJobResult 轮询 OCR 任务状态，每次调用向 PaddleOCR 查询一次并更新缓存。
func (r *Repository) PollOCRJobResult(ctx context.Context, jobID string) (*OCRJobRecord, error) {
record, err := r.getOCRJobRecord(ctx, jobID)
if err != nil {
return nil, err
}
if record.Status == ocrJobStatusDone || record.Status == ocrJobStatusFailed {
return record, nil
}
config, err := r.GetQrCodeOCRConfig(ctx)
if err != nil {
return nil, err
}
statusURL := strings.TrimRight(config.JobURL, "/") + "/" + record.PaddleJobID
req, err := http.NewRequestWithContext(ctx, http.MethodGet, statusURL, nil)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
}
req.Header.Set("Authorization", "bearer "+config.Token)
resp, err := ocrHTTPClient.Do(req)
if err != nil {
// 网络错误不更新 Redis，保留当前状态让前端重试
return record, nil
}
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return record, nil
}
var statusPayload paddleOCRJobStatusResponse
if err := json.Unmarshal(respBody, &statusPayload); err != nil {
return record, nil
}
switch statusPayload.Data.State {
case "pending", "running":
record.Status = ocrJobStatusRunning
_ = r.saveOCRJobRecord(ctx, jobID, record)
return record, nil
case "done":
if statusPayload.Data.ResultURL.JSONURL == "" {
record.Status = ocrJobStatusFailed
record.Error = "未返回识别结果地址"
_ = r.saveOCRJobRecord(ctx, jobID, record)
return record, nil
}
rawText, err := fetchPaddleOCRText(ctx, statusPayload.Data.ResultURL.JSONURL)
if err != nil {
record.Status = ocrJobStatusFailed
record.Error = err.Error()
_ = r.saveOCRJobRecord(ctx, jobID, record)
return record, nil
}
candidates := parseQrCodeGroupNameCandidates(rawText)
groupName := ""
if len(candidates) > 0 {
groupName = candidates[0]
}
record.Status = ocrJobStatusDone
record.Result = &QrCodeOCRResult{
GroupName: groupName,
Candidates: candidates,
RawText: rawText,
}
_ = r.saveOCRJobRecord(ctx, jobID, record)
return record, nil
case "failed":
record.Status = ocrJobStatusFailed
if statusPayload.Data.ErrorMsg != "" {
record.Error = statusPayload.Data.ErrorMsg
} else {
record.Error = "PaddleOCR 识别失败"
}
_ = r.saveOCRJobRecord(ctx, jobID, record)
return record, nil
}
return record, nil
}
func submitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, filename, contentType string, data []byte) (string, error) {
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
@@ -130,7 +308,7 @@ func submitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, filename,
req.Header.Set("Authorization", "bearer "+config.Token)
req.Header.Set("Content-Type", writer.FormDataContentType())
resp, err := http.DefaultClient.Do(req)
resp, err := ocrHTTPClient.Do(req)
if err != nil {
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
}
@@ -150,13 +328,13 @@ func submitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, filename,
}
func waitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, jobID string) (string, error) {
for i := 0; i < qrCodeOCRMaxPollTimes; i++ {
for i := 0; i < 12; 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)
resp, err := ocrHTTPClient.Do(req)
if err != nil {
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
}
@@ -183,8 +361,8 @@ func waitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, jobID string
}
select {
case <-ctx.Done():
return "", ctx.Err()
case <-time.After(qrCodeOCRPollInterval):
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, ctx.Err())
case <-time.After(2 * time.Second):
}
}
return "", fmt.Errorf("%w: 识别超时", ErrQrCodeOCRUnavailable)
@@ -195,7 +373,7 @@ func fetchPaddleOCRText(ctx context.Context, jsonURL string) (string, error) {
if err != nil {
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
}
resp, err := http.DefaultClient.Do(req)
resp, err := ocrHTTPClient.Do(req)
if err != nil {
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
}