569 lines
14 KiB
Go
569 lines
14 KiB
Go
// 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
|
|
}
|