textin/scan.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
	}
}