textin/scan.go
2024-11-03 11:36:46 +08:00

129 lines
3.3 KiB
Go

package textin
import (
"bytes"
"errors"
"fmt"
"strings"
"time"
"apigo.cc/ai"
"github.com/ssgo/httpclient"
"github.com/ssgo/u"
)
var hc = httpclient.GetClient(time.Second * 60)
func Scan(aiConf *ai.AIConfig, image []byte, conf map[string]any) (ai.ScanResult, error) {
if bytes.HasPrefix(image, []byte("http")) && bytes.Contains(image, []byte("://")) {
if r := hc.Get(u.String(image)); r.Error != nil {
return ai.ScanResult{}, r.Error
} else {
image = r.Bytes()
}
} else if bytes.HasPrefix(image, []byte("data:")) {
// TODO 解析BLOB
}
endpoint := "https://api.textin.com/ai/service/v2"
if aiConf.Endpoint != "" {
endpoint = aiConf.Endpoint
}
typ := u.String(conf["type"])
aksk := strings.SplitN(aiConf.ApiKey, ",", 2)
if len(aksk) == 1 {
aksk = append(aksk, aksk[0])
}
character := u.Int(conf["character"])
straighten := u.Int(conf["straighten"])
if typ == "table" {
output_order := u.String(conf["output_order"])
if output_order == "" {
output_order = "perpendicular"
}
table_type_hint := u.String(conf["table_type_hint"])
if table_type_hint == "" {
table_type_hint = "automatic"
}
excel := u.Int(conf["excel"])
resMap := hc.Post(fmt.Sprintf("%s/recognize/table/multipage?character=%d&straighten=%d&output_order=%s&table_type_hint=%s&excel=%d", endpoint, character, straighten, output_order, table_type_hint, excel), image, "x-ti-app-id", aksk[0], "x-ti-secret-code", aksk[1]).Map()
r := struct {
Code int
Message string
Duration float64
Result struct {
Pages []struct {
Tables []struct {
Lines []struct {
Text string
}
Table_cells []struct {
Text string
}
}
}
}
}{}
u.Convert(resMap, &r)
if r.Code != 200 {
return ai.ScanResult{}, errors.New(r.Message)
}
results := make([]string, 0)
for _, table := range r.Result.Pages[len(r.Result.Pages)-1].Tables {
partResults := make([]string, 0)
for _, line := range table.Lines {
partResults = append(partResults, line.Text)
}
if len(partResults) > 0 {
results = append(results, strings.Join(partResults, " "))
}
partResults = make([]string, 0)
for _, line := range table.Table_cells {
partResults = append(partResults, line.Text)
}
if len(partResults) > 0 {
results = append(results, strings.Join(partResults, " "))
}
}
detail := map[string]any{}
u.Convert(resMap["result"], &detail)
return ai.ScanResult{
Result: strings.Join(results, "\n"),
Detail: detail,
UsedTime: int64(r.Duration),
}, nil
} else {
// 通用文本识别
resMap := hc.Post(fmt.Sprintf("%s/recognize/multipage?character=%d&straighten=%d", endpoint, character, straighten), image, "x-ti-app-id", aksk[0], "x-ti-secret-code", aksk[1]).Map()
r := struct {
Code int
Message string
Duration float64
Result struct {
Pages []struct {
Lines []struct {
Text string
}
}
}
}{}
u.Convert(resMap, &r)
if r.Code != 200 {
return ai.ScanResult{}, errors.New(r.Message)
}
results := make([]string, 0)
for _, line := range r.Result.Pages[len(r.Result.Pages)-1].Lines {
results = append(results, line.Text)
}
detail := map[string]any{}
u.Convert(resMap["result"], &detail)
return ai.ScanResult{
Result: strings.Join(results, " "),
Detail: detail,
UsedTime: int64(r.Duration),
}, nil
}
}