feat: add screenshot OCR extraction
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/mmmy/snapgo/internal/infrastructure/clipboard"
|
||||
)
|
||||
|
||||
// OCRRecognizer describes a provider capable of extracting text from a PNG
|
||||
// screenshot.
|
||||
type OCRRecognizer interface {
|
||||
RecognizeText(ctx context.Context, pngBytes []byte) (string, error)
|
||||
}
|
||||
|
||||
// CaptureOCRService wires screenshot bytes -> OCR provider -> clipboard. It
|
||||
// leaves progress reporting to the caller so native and web overlays can share
|
||||
// the same status HUD.
|
||||
type CaptureOCRService struct {
|
||||
Recognizer OCRRecognizer
|
||||
Clipboard clipboard.Writer
|
||||
}
|
||||
|
||||
// Recognize asks the configured OCR provider to extract text from the PNG.
|
||||
func (s *CaptureOCRService) Recognize(ctx context.Context, pngBytes []byte) (string, error) {
|
||||
if s.Recognizer == nil {
|
||||
return "", fmt.Errorf("ocr is not configured")
|
||||
}
|
||||
if len(pngBytes) == 0 {
|
||||
return "", fmt.Errorf("empty screenshot")
|
||||
}
|
||||
text, err := s.Recognizer.RecognizeText(ctx, pngBytes)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return "", fmt.Errorf("ocr result is empty")
|
||||
}
|
||||
return text, nil
|
||||
}
|
||||
|
||||
// CopyText writes the extracted text to the clipboard.
|
||||
func (s *CaptureOCRService) CopyText(_ context.Context, text string) error {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return fmt.Errorf("empty ocr text")
|
||||
}
|
||||
if s.Clipboard == nil {
|
||||
return fmt.Errorf("clipboard is not configured")
|
||||
}
|
||||
if err := s.Clipboard.WriteText(text); err != nil {
|
||||
return fmt.Errorf("clipboard write failed: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type fakeOCRRecognizer struct {
|
||||
image []byte
|
||||
text string
|
||||
}
|
||||
|
||||
func (f *fakeOCRRecognizer) RecognizeText(_ context.Context, pngBytes []byte) (string, error) {
|
||||
f.image = append([]byte(nil), pngBytes...)
|
||||
if f.text == "" {
|
||||
return "extracted text", nil
|
||||
}
|
||||
return f.text, nil
|
||||
}
|
||||
|
||||
func TestCaptureOCRServiceRecognizesAndCopies(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
recognizer := &fakeOCRRecognizer{}
|
||||
clip := &fakeClipboard{}
|
||||
svc := &CaptureOCRService{Recognizer: recognizer, Clipboard: clip}
|
||||
|
||||
text, err := svc.Recognize(ctx, []byte("png"))
|
||||
if err != nil {
|
||||
t.Fatalf("recognize: %v", err)
|
||||
}
|
||||
if string(recognizer.image) != "png" {
|
||||
t.Fatalf("expected screenshot bytes forwarded, got %q", string(recognizer.image))
|
||||
}
|
||||
if text != "extracted text" {
|
||||
t.Fatalf("expected extracted text, got %q", text)
|
||||
}
|
||||
if err := svc.CopyText(ctx, text); err != nil {
|
||||
t.Fatalf("copy text: %v", err)
|
||||
}
|
||||
if clip.text != "extracted text" {
|
||||
t.Fatalf("expected ocr text copied, got %q", clip.text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCaptureOCRServiceRejectsEmptyResult(t *testing.T) {
|
||||
svc := &CaptureOCRService{Recognizer: &fakeOCRRecognizer{text: " \n "}}
|
||||
if _, err := svc.Recognize(context.Background(), []byte("png")); err == nil {
|
||||
t.Fatalf("expected empty OCR result to fail")
|
||||
}
|
||||
}
|
||||
@@ -107,6 +107,29 @@ type LLMConfig struct {
|
||||
Providers map[string]LLMProviderConfig `json:"providers"`
|
||||
}
|
||||
|
||||
// OCRProviderConfig stores one cloud OCR endpoint that accepts a screenshot
|
||||
// image and returns extracted text. Aliyun and Volcengine both authenticate
|
||||
// with an AccessKey pair, but their signing schemes use slightly different
|
||||
// endpoint metadata; keeping Region/Service/Action/Version configurable lets
|
||||
// the presets track official API changes without a schema rewrite.
|
||||
type OCRProviderConfig struct {
|
||||
Label string `json:"label"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
AccessKeyID string `json:"accessKeyId"`
|
||||
AccessKeySecret string `json:"accessKeySecret"`
|
||||
Region string `json:"region"`
|
||||
Service string `json:"service"`
|
||||
Action string `json:"action"`
|
||||
Version string `json:"version"`
|
||||
TimeoutSecs int `json:"timeoutSecs"`
|
||||
}
|
||||
|
||||
// OCRConfig controls screenshot text extraction.
|
||||
type OCRConfig struct {
|
||||
ActiveProvider string `json:"activeProvider"`
|
||||
Providers map[string]OCRProviderConfig `json:"providers"`
|
||||
}
|
||||
|
||||
// SSH authentication method identifiers stored in SSHConfig.AuthMethod.
|
||||
const (
|
||||
// SSHAuthBuiltin is the legacy combined method (password → agent → key
|
||||
@@ -129,6 +152,12 @@ const (
|
||||
LLMProviderOpenAI = "openai"
|
||||
)
|
||||
|
||||
// Built-in OCR provider identifiers.
|
||||
const (
|
||||
OCRProviderAliyun = "aliyun"
|
||||
OCRProviderVolcengine = "volcengine"
|
||||
)
|
||||
|
||||
// Theme preference identifiers stored in AppConfig.Theme.
|
||||
const (
|
||||
ThemeAuto = "auto"
|
||||
@@ -172,6 +201,9 @@ type AppConfig struct {
|
||||
|
||||
// LLM holds provider settings for the "copy summary" screenshot action.
|
||||
LLM LLMConfig `json:"llm"`
|
||||
|
||||
// OCR holds provider settings for the "extract text" screenshot action.
|
||||
OCR OCRConfig `json:"ocr"`
|
||||
}
|
||||
|
||||
// DefaultAppConfig returns sane zero-value defaults used on first launch.
|
||||
@@ -190,6 +222,7 @@ func DefaultAppConfig() AppConfig {
|
||||
StrictHostKey: false,
|
||||
},
|
||||
LLM: DefaultLLMConfig(),
|
||||
OCR: DefaultOCRConfig(),
|
||||
}
|
||||
cfg.Normalize()
|
||||
return cfg
|
||||
@@ -234,6 +267,33 @@ func DefaultLLMConfig() LLMConfig {
|
||||
}
|
||||
}
|
||||
|
||||
// DefaultOCRConfig returns the built-in OCR provider presets. Users only need
|
||||
// to fill AccessKey credentials for the common regions/actions, while advanced
|
||||
// accounts can still override endpoint metadata from the settings UI.
|
||||
func DefaultOCRConfig() OCRConfig {
|
||||
return OCRConfig{
|
||||
ActiveProvider: OCRProviderAliyun,
|
||||
Providers: map[string]OCRProviderConfig{
|
||||
OCRProviderAliyun: {
|
||||
Label: "阿里云文字识别",
|
||||
Endpoint: "https://ocr-api.cn-hangzhou.aliyuncs.com",
|
||||
Action: "RecognizeGeneral",
|
||||
Version: "2021-07-07",
|
||||
TimeoutSecs: 30,
|
||||
},
|
||||
OCRProviderVolcengine: {
|
||||
Label: "火山引擎文字识别",
|
||||
Endpoint: "https://visual.volcengineapi.com",
|
||||
Region: "cn-north-1",
|
||||
Service: "cv",
|
||||
Action: "OCRNormal",
|
||||
Version: "2020-08-26",
|
||||
TimeoutSecs: 30,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize fills defaults into configs written by older app versions while
|
||||
// preserving any user-provided provider fields.
|
||||
func (c *AppConfig) Normalize() {
|
||||
@@ -295,6 +355,43 @@ func (c *AppConfig) Normalize() {
|
||||
}
|
||||
c.LLM.Providers[id] = current
|
||||
}
|
||||
|
||||
defaultOCR := DefaultOCRConfig()
|
||||
if c.OCR.ActiveProvider == "" {
|
||||
c.OCR.ActiveProvider = defaultOCR.ActiveProvider
|
||||
}
|
||||
if c.OCR.Providers == nil {
|
||||
c.OCR.Providers = map[string]OCRProviderConfig{}
|
||||
}
|
||||
for id, def := range defaultOCR.Providers {
|
||||
current, ok := c.OCR.Providers[id]
|
||||
if !ok {
|
||||
c.OCR.Providers[id] = def
|
||||
continue
|
||||
}
|
||||
if current.Label == "" {
|
||||
current.Label = def.Label
|
||||
}
|
||||
if current.Endpoint == "" {
|
||||
current.Endpoint = def.Endpoint
|
||||
}
|
||||
if current.Region == "" {
|
||||
current.Region = def.Region
|
||||
}
|
||||
if current.Service == "" {
|
||||
current.Service = def.Service
|
||||
}
|
||||
if current.Action == "" {
|
||||
current.Action = def.Action
|
||||
}
|
||||
if current.Version == "" {
|
||||
current.Version = def.Version
|
||||
}
|
||||
if current.TimeoutSecs == 0 {
|
||||
current.TimeoutSecs = def.TimeoutSecs
|
||||
}
|
||||
c.OCR.Providers[id] = current
|
||||
}
|
||||
}
|
||||
|
||||
// IsS3Configured reports whether the user has filled the mandatory S3 fields.
|
||||
@@ -329,3 +426,32 @@ func (c AppConfig) IsLLMConfigured() bool {
|
||||
_, provider, ok := c.ActiveLLMProvider()
|
||||
return ok && provider.BaseURL != "" && provider.APIKey != "" && provider.Model != ""
|
||||
}
|
||||
|
||||
// ActiveOCRProvider returns the selected OCR provider config plus a boolean
|
||||
// indicating whether the selection exists.
|
||||
func (c AppConfig) ActiveOCRProvider() (string, OCRProviderConfig, bool) {
|
||||
id := c.OCR.ActiveProvider
|
||||
if id == "" {
|
||||
id = OCRProviderAliyun
|
||||
}
|
||||
provider, ok := c.OCR.Providers[id]
|
||||
return id, provider, ok
|
||||
}
|
||||
|
||||
// IsOCRConfigured reports whether the selected OCR provider has the fields
|
||||
// required to sign and send a request.
|
||||
func (c AppConfig) IsOCRConfigured() bool {
|
||||
id, provider, ok := c.ActiveOCRProvider()
|
||||
if !ok ||
|
||||
provider.Endpoint == "" ||
|
||||
provider.AccessKeyID == "" ||
|
||||
provider.AccessKeySecret == "" ||
|
||||
provider.Action == "" ||
|
||||
provider.Version == "" {
|
||||
return false
|
||||
}
|
||||
if id == OCRProviderVolcengine && (provider.Region == "" || provider.Service == "") {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package ocr
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mmmy/snapgo/internal/domain"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func TestAliyunRecognizeTextSignsRawPNGRequest(t *testing.T) {
|
||||
var requestBody []byte
|
||||
transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if r.URL.Host != "ocr-api.cn-hangzhou.aliyuncs.com" {
|
||||
t.Fatalf("unexpected host %s", r.URL.Host)
|
||||
}
|
||||
if got := r.Header.Get("x-acs-action"); got != "RecognizeGeneral" {
|
||||
t.Fatalf("unexpected action %q", got)
|
||||
}
|
||||
if got := r.Header.Get("x-acs-version"); got != "2021-07-07" {
|
||||
t.Fatalf("unexpected version %q", got)
|
||||
}
|
||||
if got := r.Header.Get("x-acs-date"); got != "2026-07-08T01:02:03Z" {
|
||||
t.Fatalf("unexpected date %q", got)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); !strings.Contains(got, "ACS3-HMAC-SHA256 Credential=ak") {
|
||||
t.Fatalf("unexpected auth header %q", got)
|
||||
}
|
||||
var err error
|
||||
requestBody, err = io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read body: %v", err)
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(
|
||||
`{"Data":"{\"content\":\"第一行\\n第二行\"}"}`,
|
||||
)),
|
||||
}, nil
|
||||
})
|
||||
|
||||
client, err := NewClient(domain.OCRProviderAliyun, domain.OCRProviderConfig{
|
||||
Endpoint: "https://ocr-api.cn-hangzhou.aliyuncs.com",
|
||||
AccessKeyID: "ak",
|
||||
AccessKeySecret: "sk",
|
||||
Action: "RecognizeGeneral",
|
||||
Version: "2021-07-07",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("new client: %v", err)
|
||||
}
|
||||
client.httpClient.Transport = transport
|
||||
client.now = func() time.Time { return time.Date(2026, 7, 8, 1, 2, 3, 0, time.UTC) }
|
||||
client.nonce = func() string { return "nonce" }
|
||||
|
||||
text, err := client.RecognizeText(context.Background(), []byte("png"))
|
||||
if err != nil {
|
||||
t.Fatalf("recognize text: %v", err)
|
||||
}
|
||||
if string(requestBody) != "png" {
|
||||
t.Fatalf("expected raw png body, got %q", string(requestBody))
|
||||
}
|
||||
if text != "第一行\n第二行" {
|
||||
t.Fatalf("unexpected OCR text %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcengineRecognizeTextSignsBase64JSONRequest(t *testing.T) {
|
||||
var requestBody map[string]string
|
||||
transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if got := r.URL.Query().Get("Action"); got != "OCRNormal" {
|
||||
t.Fatalf("unexpected action %q", got)
|
||||
}
|
||||
if got := r.URL.Query().Get("Version"); got != "2020-08-26" {
|
||||
t.Fatalf("unexpected version %q", got)
|
||||
}
|
||||
if got := r.Header.Get("X-Date"); got != "20260708T010203Z" {
|
||||
t.Fatalf("unexpected x-date %q", got)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); !strings.Contains(got, "HMAC-SHA256 Credential=ak/20260708/cn-north-1/cv/request") {
|
||||
t.Fatalf("unexpected auth header %q", got)
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
|
||||
t.Fatalf("decode request: %v", err)
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(
|
||||
`{"ResponseMetadata":{"RequestId":"r"},"Result":{"LineTexts":["你好","世界"]}}`,
|
||||
)),
|
||||
}, nil
|
||||
})
|
||||
|
||||
client, err := NewClient(domain.OCRProviderVolcengine, domain.OCRProviderConfig{
|
||||
Endpoint: "https://visual.volcengineapi.com",
|
||||
AccessKeyID: "ak",
|
||||
AccessKeySecret: "sk",
|
||||
Region: "cn-north-1",
|
||||
Service: "cv",
|
||||
Action: "OCRNormal",
|
||||
Version: "2020-08-26",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("new client: %v", err)
|
||||
}
|
||||
client.httpClient.Transport = transport
|
||||
client.now = func() time.Time { return time.Date(2026, 7, 8, 1, 2, 3, 0, time.UTC) }
|
||||
|
||||
text, err := client.RecognizeText(context.Background(), []byte("png"))
|
||||
if err != nil {
|
||||
t.Fatalf("recognize text: %v", err)
|
||||
}
|
||||
if requestBody["image_base64"] != base64.StdEncoding.EncodeToString([]byte("png")) {
|
||||
t.Fatalf("expected base64 image body, got %#v", requestBody)
|
||||
}
|
||||
if text != "你好\n世界" {
|
||||
t.Fatalf("unexpected OCR text %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseOCRTextSupportsWordsResult(t *testing.T) {
|
||||
text, err := parseOCRText([]byte(`{"data":{"words_result":[{"words":"foo"},{"words":"bar"}]}}`))
|
||||
if err != nil {
|
||||
t.Fatalf("parse ocr text: %v", err)
|
||||
}
|
||||
if text != "foo\nbar" {
|
||||
t.Fatalf("unexpected OCR text %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseOCRTextReturnsProviderError(t *testing.T) {
|
||||
_, err := parseOCRText([]byte(`{"ResponseMetadata":{"Error":{"Code":"BadRequest","Message":"bad image"}}}`))
|
||||
if err == nil || !strings.Contains(err.Error(), "bad image") {
|
||||
t.Fatalf("expected provider error, got %v", err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user