feat: add screenshot OCR extraction

2026-07-08 00:42:49 +08:00
parent f90612976f
commit baec01610c
14 changed files with 1420 additions and 46 deletions
@@ -0,0 +1,568 @@
+// Package ocr contains cloud OCR adapters.
+package ocr
+
+import (
+	"bytes"
+	"context"
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"sort"
+	"strings"
+	"time"
+
+	"github.com/mmmy/snapgo/internal/application"
+	"github.com/mmmy/snapgo/internal/domain"
+)
+
+// Client calls one configured OCR provider.
+type Client struct {
+	providerID string
+	cfg        domain.OCRProviderConfig
+	httpClient *http.Client
+	now        func() time.Time
+	nonce      func() string
+}
+
+var _ application.OCRRecognizer = (*Client)(nil)
+
+// NewClient validates cfg and returns a reusable OCR client.
+func NewClient(providerID string, cfg domain.OCRProviderConfig) (*Client, error) {
+	if strings.TrimSpace(cfg.Endpoint) == "" {
+		return nil, fmt.Errorf("ocr endpoint is required")
+	}
+	if strings.TrimSpace(cfg.AccessKeyID) == "" {
+		return nil, fmt.Errorf("ocr access key id is required")
+	}
+	if strings.TrimSpace(cfg.AccessKeySecret) == "" {
+		return nil, fmt.Errorf("ocr access key secret is required")
+	}
+	if strings.TrimSpace(cfg.Action) == "" {
+		return nil, fmt.Errorf("ocr action is required")
+	}
+	if strings.TrimSpace(cfg.Version) == "" {
+		return nil, fmt.Errorf("ocr version is required")
+	}
+	switch providerID {
+	case domain.OCRProviderAliyun:
+	case domain.OCRProviderVolcengine:
+		if strings.TrimSpace(cfg.Region) == "" {
+			return nil, fmt.Errorf("volcengine ocr region is required")
+		}
+		if strings.TrimSpace(cfg.Service) == "" {
+			return nil, fmt.Errorf("volcengine ocr service is required")
+		}
+	default:
+		return nil, fmt.Errorf("unsupported ocr provider %q", providerID)
+	}
+	timeout := cfg.TimeoutSecs
+	if timeout <= 0 {
+		timeout = 30
+	}
+	return &Client{
+		providerID: providerID,
+		cfg:        cfg,
+		httpClient: &http.Client{Timeout: time.Duration(timeout) * time.Second},
+		now:        time.Now,
+		nonce:      randomNonce,
+	}, nil
+}
+
+// RecognizeText extracts visible text from a PNG screenshot.
+func (c *Client) RecognizeText(ctx context.Context, pngBytes []byte) (string, error) {
+	if len(pngBytes) == 0 {
+		return "", fmt.Errorf("empty screenshot")
+	}
+	switch c.providerID {
+	case domain.OCRProviderAliyun:
+		return c.recognizeAliyun(ctx, pngBytes)
+	case domain.OCRProviderVolcengine:
+		return c.recognizeVolcengine(ctx, pngBytes)
+	default:
+		return "", fmt.Errorf("unsupported ocr provider %q", c.providerID)
+	}
+}
+
+func (c *Client) recognizeAliyun(ctx context.Context, pngBytes []byte) (string, error) {
+	endpoint, err := parseEndpoint(c.cfg.Endpoint)
+	if err != nil {
+		return "", err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint.String(), bytes.NewReader(pngBytes))
+	if err != nil {
+		return "", fmt.Errorf("build aliyun ocr request: %w", err)
+	}
+	req.Header.Set("Accept", "application/json")
+	req.Header.Set("Content-Type", "application/octet-stream")
+	c.signAliyun(req, pngBytes)
+	return c.doOCRRequest(req, "aliyun ocr")
+}
+
+func (c *Client) recognizeVolcengine(ctx context.Context, pngBytes []byte) (string, error) {
+	endpoint, err := parseEndpoint(c.cfg.Endpoint)
+	if err != nil {
+		return "", err
+	}
+	query := endpoint.Query()
+	query.Set("Action", c.cfg.Action)
+	query.Set("Version", c.cfg.Version)
+	endpoint.RawQuery = query.Encode()
+
+	body := map[string]string{
+		"image_base64": base64.StdEncoding.EncodeToString(pngBytes),
+	}
+	payload, err := json.Marshal(body)
+	if err != nil {
+		return "", fmt.Errorf("marshal volcengine ocr request: %w", err)
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint.String(), bytes.NewReader(payload))
+	if err != nil {
+		return "", fmt.Errorf("build volcengine ocr request: %w", err)
+	}
+	req.Header.Set("Accept", "application/json")
+	req.Header.Set("Content-Type", "application/json")
+	c.signVolcengine(req, payload)
+	return c.doOCRRequest(req, "volcengine ocr")
+}
+
+func (c *Client) doOCRRequest(req *http.Request, label string) (string, error) {
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("%s request: %w", label, err)
+	}
+	defer resp.Body.Close()
+	data, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return "", fmt.Errorf("%s status %d: %s", label, resp.StatusCode, compactBody(data))
+	}
+	text, err := parseOCRText(data)
+	if err != nil {
+		return "", fmt.Errorf("%s response: %w", label, err)
+	}
+	return text, nil
+}
+
+func (c *Client) signAliyun(req *http.Request, payload []byte) {
+	now := c.now().UTC()
+	payloadHash := sha256Hex(payload)
+	nonce := c.nonce()
+
+	req.Header.Set("x-acs-action", c.cfg.Action)
+	req.Header.Set("x-acs-version", c.cfg.Version)
+	req.Header.Set("x-acs-date", now.Format("2006-01-02T15:04:05Z"))
+	req.Header.Set("x-acs-signature-nonce", nonce)
+	req.Header.Set("x-acs-content-sha256", payloadHash)
+
+	signedHeaders := []string{
+		"content-type",
+		"host",
+		"x-acs-action",
+		"x-acs-content-sha256",
+		"x-acs-date",
+		"x-acs-signature-nonce",
+		"x-acs-version",
+	}
+	canonicalHeaders := canonicalHeaders(req, signedHeaders)
+	canonicalRequest := strings.Join([]string{
+		req.Method,
+		canonicalURI(req.URL),
+		canonicalQuery(req.URL),
+		canonicalHeaders,
+		strings.Join(signedHeaders, ";"),
+		payloadHash,
+	}, "\n")
+	stringToSign := "ACS3-HMAC-SHA256\n" + sha256Hex([]byte(canonicalRequest))
+	signature := hmacSHA256Hex([]byte(c.cfg.AccessKeySecret), []byte(stringToSign))
+	req.Header.Set(
+		"Authorization",
+		fmt.Sprintf("ACS3-HMAC-SHA256 Credential=%s,SignedHeaders=%s,Signature=%s",
+			c.cfg.AccessKeyID,
+			strings.Join(signedHeaders, ";"),
+			signature,
+		),
+	)
+}
+
+func (c *Client) signVolcengine(req *http.Request, payload []byte) {
+	now := c.now().UTC()
+	xDate := now.Format("20060102T150405Z")
+	shortDate := now.Format("20060102")
+	payloadHash := sha256Hex(payload)
+
+	req.Header.Set("X-Date", xDate)
+	req.Header.Set("X-Content-Sha256", payloadHash)
+
+	signedHeaders := []string{
+		"content-type",
+		"host",
+		"x-content-sha256",
+		"x-date",
+	}
+	canonicalHeaders := canonicalHeaders(req, signedHeaders)
+	canonicalRequest := strings.Join([]string{
+		req.Method,
+		canonicalURI(req.URL),
+		canonicalQuery(req.URL),
+		canonicalHeaders,
+		strings.Join(signedHeaders, ";"),
+		payloadHash,
+	}, "\n")
+	scope := strings.Join([]string{shortDate, c.cfg.Region, c.cfg.Service, "request"}, "/")
+	stringToSign := strings.Join([]string{
+		"HMAC-SHA256",
+		xDate,
+		scope,
+		sha256Hex([]byte(canonicalRequest)),
+	}, "\n")
+	signingKey := volcengineSigningKey(c.cfg.AccessKeySecret, shortDate, c.cfg.Region, c.cfg.Service)
+	signature := hmacSHA256Hex(signingKey, []byte(stringToSign))
+	req.Header.Set(
+		"Authorization",
+		fmt.Sprintf("HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
+			c.cfg.AccessKeyID,
+			scope,
+			strings.Join(signedHeaders, ";"),
+			signature,
+		),
+	)
+}
+
+func parseEndpoint(raw string) (*url.URL, error) {
+	endpoint := strings.TrimSpace(raw)
+	if endpoint == "" {
+		return nil, fmt.Errorf("ocr endpoint is required")
+	}
+	if !strings.Contains(endpoint, "://") {
+		endpoint = "https://" + endpoint
+	}
+	parsed, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, fmt.Errorf("parse ocr endpoint: %w", err)
+	}
+	if parsed.Scheme == "" || parsed.Host == "" {
+		return nil, fmt.Errorf("invalid ocr endpoint %q", raw)
+	}
+	if parsed.Path == "" {
+		parsed.Path = "/"
+	}
+	return parsed, nil
+}
+
+func parseOCRText(data []byte) (string, error) {
+	var root any
+	if err := json.Unmarshal(data, &root); err != nil {
+		return "", fmt.Errorf("decode json: %w", err)
+	}
+	if msg := responseError(root); msg != "" {
+		return "", errors.New(msg)
+	}
+	parts := collectOCRParts(root, "")
+	if len(parts) == 0 {
+		return "", fmt.Errorf("no text found")
+	}
+	return strings.Join(dedupeNonEmpty(parts), "\n"), nil
+}
+
+func responseError(v any) string {
+	obj, ok := v.(map[string]any)
+	if !ok {
+		return ""
+	}
+	if meta, ok := obj["ResponseMetadata"].(map[string]any); ok {
+		if errObj, ok := meta["Error"].(map[string]any); ok {
+			if msg := stringValue(errObj["Message"]); msg != "" {
+				return msg
+			}
+			if code := stringValue(errObj["Code"]); code != "" {
+				return code
+			}
+		}
+	}
+	if errObj, ok := obj["Error"].(map[string]any); ok {
+		if msg := stringValue(errObj["Message"]); msg != "" {
+			return msg
+		}
+		if msg := stringValue(errObj["message"]); msg != "" {
+			return msg
+		}
+	}
+	if code, ok := numericCode(obj["code"]); ok && code != 0 && code != 10000 {
+		if msg := stringValue(obj["message"]); msg != "" {
+			return msg
+		}
+		if msg := stringValue(obj["msg"]); msg != "" {
+			return msg
+		}
+		return fmt.Sprintf("provider code %.0f", code)
+	}
+	if code := stringValue(obj["Code"]); code != "" && !isSuccessCode(code) {
+		if msg := stringValue(obj["Message"]); msg != "" {
+			return msg
+		}
+		return code
+	}
+	return ""
+}
+
+func collectOCRParts(v any, parentKey string) []string {
+	switch value := v.(type) {
+	case string:
+		text := strings.TrimSpace(value)
+		if text == "" {
+			return nil
+		}
+		if looksLikeJSON(text) {
+			var nested any
+			if err := json.Unmarshal([]byte(text), &nested); err == nil {
+				return collectOCRParts(nested, parentKey)
+			}
+		}
+		if isTextKey(parentKey) || isContainerTextKey(parentKey) {
+			return []string{text}
+		}
+		return nil
+	case []any:
+		parts := make([]string, 0, len(value))
+		for _, item := range value {
+			parts = append(parts, collectOCRParts(item, parentKey)...)
+		}
+		return parts
+	case map[string]any:
+		parts := make([]string, 0)
+		for _, key := range priorityOCRKeys() {
+			if child, ok := value[key]; ok {
+				parts = append(parts, collectOCRParts(child, key)...)
+			}
+		}
+		keys := make([]string, 0, len(value))
+		for key := range value {
+			keys = append(keys, key)
+		}
+		sort.Strings(keys)
+		for _, key := range keys {
+			if isPriorityOCRKey(key) || isMetadataKey(key) {
+				continue
+			}
+			parts = append(parts, collectOCRParts(value[key], key)...)
+		}
+		return parts
+	default:
+		return nil
+	}
+}
+
+func priorityOCRKeys() []string {
+	return []string{
+		"Data",
+		"data",
+		"Result",
+		"result",
+		"content",
+		"Content",
+		"text",
+		"Text",
+		"DetectedText",
+		"detected_text",
+		"line_text",
+		"LineText",
+		"line_texts",
+		"LineTexts",
+		"word",
+		"Word",
+		"words",
+		"Words",
+		"words_result",
+		"WordsResult",
+		"prism_wordsInfo",
+		"prism_words_info",
+		"ocr_infos",
+		"OCRInfos",
+		"items",
+		"Items",
+		"blocks",
+		"Blocks",
+		"regions",
+		"Regions",
+	}
+}
+
+func isPriorityOCRKey(key string) bool {
+	for _, item := range priorityOCRKeys() {
+		if item == key {
+			return true
+		}
+	}
+	return false
+}
+
+func isTextKey(key string) bool {
+	switch normalizeKey(key) {
+	case "content", "text", "detectedtext", "linetext", "word", "words", "value":
+		return true
+	default:
+		return false
+	}
+}
+
+func isContainerTextKey(key string) bool {
+	switch normalizeKey(key) {
+	case "data", "result", "linetexts", "texts":
+		return true
+	default:
+		return false
+	}
+}
+
+func isMetadataKey(key string) bool {
+	switch normalizeKey(key) {
+	case "requestid", "request", "code", "status", "statuscode", "success",
+		"error", "errors", "message", "msg", "cost", "angle", "probability",
+		"confidence", "height", "width", "left", "top", "right", "bottom",
+		"x", "y", "responsemetadata":
+		return true
+	default:
+		return false
+	}
+}
+
+func normalizeKey(key string) string {
+	key = strings.ToLower(key)
+	key = strings.ReplaceAll(key, "_", "")
+	key = strings.ReplaceAll(key, "-", "")
+	return key
+}
+
+func dedupeNonEmpty(parts []string) []string {
+	seen := map[string]struct{}{}
+	out := make([]string, 0, len(parts))
+	for _, part := range parts {
+		for _, line := range strings.Split(part, "\n") {
+			line = strings.TrimSpace(line)
+			if line == "" {
+				continue
+			}
+			if _, ok := seen[line]; ok {
+				continue
+			}
+			seen[line] = struct{}{}
+			out = append(out, line)
+		}
+	}
+	return out
+}
+
+func looksLikeJSON(text string) bool {
+	return strings.HasPrefix(text, "{") || strings.HasPrefix(text, "[")
+}
+
+func stringValue(v any) string {
+	if s, ok := v.(string); ok {
+		return strings.TrimSpace(s)
+	}
+	return ""
+}
+
+func numericCode(v any) (float64, bool) {
+	switch n := v.(type) {
+	case float64:
+		return n, true
+	case int:
+		return float64(n), true
+	case json.Number:
+		f, err := n.Float64()
+		return f, err == nil
+	default:
+		return 0, false
+	}
+}
+
+func isSuccessCode(code string) bool {
+	switch strings.ToLower(strings.TrimSpace(code)) {
+	case "", "ok", "success", "200", "10000":
+		return true
+	default:
+		return false
+	}
+}
+
+func canonicalURI(u *url.URL) string {
+	if u == nil || u.EscapedPath() == "" {
+		return "/"
+	}
+	return u.EscapedPath()
+}
+
+func canonicalQuery(u *url.URL) string {
+	if u == nil || u.RawQuery == "" {
+		return ""
+	}
+	values, _ := url.ParseQuery(u.RawQuery)
+	return values.Encode()
+}
+
+func canonicalHeaders(req *http.Request, signedHeaders []string) string {
+	lines := make([]string, 0, len(signedHeaders))
+	for _, key := range signedHeaders {
+		var value string
+		if key == "host" {
+			value = req.URL.Host
+		} else {
+			value = req.Header.Get(key)
+		}
+		lines = append(lines, key+":"+normalizeHeaderValue(value)+"\n")
+	}
+	return strings.Join(lines, "")
+}
+
+func normalizeHeaderValue(value string) string {
+	return strings.Join(strings.Fields(strings.TrimSpace(value)), " ")
+}
+
+func volcengineSigningKey(secret, date, region, service string) []byte {
+	kDate := hmacSHA256([]byte(secret), []byte(date))
+	kRegion := hmacSHA256(kDate, []byte(region))
+	kService := hmacSHA256(kRegion, []byte(service))
+	return hmacSHA256(kService, []byte("request"))
+}
+
+func hmacSHA256(key, data []byte) []byte {
+	mac := hmac.New(sha256.New, key)
+	mac.Write(data)
+	return mac.Sum(nil)
+}
+
+func hmacSHA256Hex(key, data []byte) string {
+	return hex.EncodeToString(hmacSHA256(key, data))
+}
+
+func sha256Hex(data []byte) string {
+	sum := sha256.Sum256(data)
+	return hex.EncodeToString(sum[:])
+}
+
+func randomNonce() string {
+	buf := make([]byte, 16)
+	if _, err := rand.Read(buf); err != nil {
+		return fmt.Sprintf("%d", time.Now().UnixNano())
+	}
+	return hex.EncodeToString(buf)
+}
+
+func compactBody(data []byte) string {
+	text := strings.TrimSpace(string(data))
+	if text == "" {
+		return "<empty body>"
+	}
+	if len(text) > 800 {
+		return text[:800] + "..."
+	}
+	return text
+}