feat: request vedem ocr with uploading image

This commit is contained in:
lilong.129
2023-05-09 22:48:42 +08:00
parent b480c6058b
commit 41e4d89c95
4 changed files with 71 additions and 35 deletions
+12 -1
View File
@@ -45,8 +45,10 @@
{ {
"method": "tap_ocr", "method": "tap_ocr",
"params": "我知道了", "params": "我知道了",
"options": {
"ignore_NotFoundError": true "ignore_NotFoundError": true
} }
}
] ]
} }
}, },
@@ -57,9 +59,11 @@
{ {
"method": "swipe_to_tap_text", "method": "swipe_to_tap_text",
"params": "点击进入直播间", "params": "点击进入直播间",
"options": {
"identifier": "进入直播间", "identifier": "进入直播间",
"max_retry_times": 10 "max_retry_times": 10
} }
}
] ]
} }
}, },
@@ -69,7 +73,10 @@
"actions": [ "actions": [
{ {
"method": "swipe", "method": "swipe",
"params": "up" "params": "up",
"options": {
}
}, },
{ {
"method": "sleep_random", "method": "sleep_random",
@@ -93,7 +100,9 @@
{ {
"method": "swipe", "method": "swipe",
"params": "up", "params": "up",
"options": {
"identifier": "第一次上划" "identifier": "第一次上划"
}
}, },
{ {
"method": "sleep", "method": "sleep",
@@ -105,7 +114,9 @@
{ {
"method": "swipe", "method": "swipe",
"params": "up", "params": "up",
"options": {
"identifier": "第二次上划" "identifier": "第二次上划"
}
}, },
{ {
"method": "sleep", "method": "sleep",
+5 -4
View File
@@ -65,6 +65,7 @@ type OcrResult struct {
type cacheStepData struct { type cacheStepData struct {
// cache step screenshot paths // cache step screenshot paths
ScreenShots []string ScreenShots []string
screenShotsUrls map[string]string // map screenshot file path to uploaded url
// cache step screenshot ocr results, key is image path, value is OcrResult // cache step screenshot ocr results, key is image path, value is OcrResult
OcrResults map[string]*OcrResult OcrResults map[string]*OcrResult
// cache feed/live video stat // cache feed/live video stat
@@ -73,6 +74,7 @@ type cacheStepData struct {
func (d *cacheStepData) reset() { func (d *cacheStepData) reset() {
d.ScreenShots = make([]string, 0) d.ScreenShots = make([]string, 0)
d.screenShotsUrls = make(map[string]string)
d.OcrResults = make(map[string]*OcrResult) d.OcrResults = make(map[string]*OcrResult)
d.VideoStat = nil d.VideoStat = nil
} }
@@ -95,12 +97,10 @@ func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error)
dExt = &DriverExt{ dExt = &DriverExt{
Device: device, Device: device,
Driver: driver, Driver: driver,
cacheStepData: cacheStepData{ cacheStepData: cacheStepData{},
ScreenShots: make([]string, 0),
OcrResults: make(map[string]*OcrResult),
},
interruptSignal: make(chan os.Signal, 1), interruptSignal: make(chan os.Signal, 1),
} }
dExt.cacheStepData.reset()
signal.Notify(dExt.interruptSignal, syscall.SIGTERM, syscall.SIGINT) signal.Notify(dExt.interruptSignal, syscall.SIGTERM, syscall.SIGINT)
dExt.doneMjpegStream = make(chan bool, 1) dExt.doneMjpegStream = make(chan bool, 1)
@@ -208,6 +208,7 @@ func (dExt *DriverExt) GetStepCacheData() map[string]interface{} {
cacheData := make(map[string]interface{}) cacheData := make(map[string]interface{})
cacheData["video_stat"] = dExt.cacheStepData.VideoStat cacheData["video_stat"] = dExt.cacheStepData.VideoStat
cacheData["screenshots"] = dExt.cacheStepData.ScreenShots cacheData["screenshots"] = dExt.cacheStepData.ScreenShots
cacheData["screenshots_urls"] = dExt.cacheStepData.screenShotsUrls
ocrResults := make(map[string]interface{}) ocrResults := make(map[string]interface{})
for imagePath, ocrResult := range dExt.cacheStepData.OcrResults { for imagePath, ocrResult := range dExt.cacheStepData.OcrResults {
+42 -15
View File
@@ -31,6 +31,7 @@ type OCRResult struct {
type ResponseOCR struct { type ResponseOCR struct {
Code int `json:"code"` Code int `json:"code"`
Message string `json:"message"` Message string `json:"message"`
URL string `json:"url"` // image uploaded url
OCRResult []OCRResult `json:"ocrResult"` OCRResult []OCRResult `json:"ocrResult"`
} }
@@ -141,33 +142,40 @@ func newVEDEMOCRService() (*veDEMOCRService, error) {
// veDEMOCRService implements IOCRService interface // veDEMOCRService implements IOCRService interface
type veDEMOCRService struct{} type veDEMOCRService struct{}
func (s *veDEMOCRService) getOCRResult(imageBuf *bytes.Buffer) ([]OCRResult, error) { func (s *veDEMOCRService) getOCRResult(imageBuf *bytes.Buffer) (
ocrResutls []OCRResult, url string, err error) {
bodyBuf := &bytes.Buffer{} bodyBuf := &bytes.Buffer{}
bodyWriter := multipart.NewWriter(bodyBuf) bodyWriter := multipart.NewWriter(bodyBuf)
bodyWriter.WriteField("withDet", "true") bodyWriter.WriteField("withDet", "true")
bodyWriter.WriteField("upload", "true") // get image uploaded url
// bodyWriter.WriteField("timestampOnly", "true") // bodyWriter.WriteField("timestampOnly", "true")
formWriter, err := bodyWriter.CreateFormFile("image", "screenshot.png") formWriter, err := bodyWriter.CreateFormFile("image", "screenshot.png")
if err != nil { if err != nil {
return nil, errors.Wrap(code.OCRRequestError, err = errors.Wrap(code.OCRRequestError,
fmt.Sprintf("create form file error: %v", err)) fmt.Sprintf("create form file error: %v", err))
return
} }
size, err := formWriter.Write(imageBuf.Bytes()) size, err := formWriter.Write(imageBuf.Bytes())
if err != nil { if err != nil {
return nil, errors.Wrap(code.OCRRequestError, err = errors.Wrap(code.OCRRequestError,
fmt.Sprintf("write form error: %v", err)) fmt.Sprintf("write form error: %v", err))
return
} }
err = bodyWriter.Close() err = bodyWriter.Close()
if err != nil { if err != nil {
return nil, errors.Wrap(code.OCRRequestError, err = errors.Wrap(code.OCRRequestError,
fmt.Sprintf("close body writer error: %v", err)) fmt.Sprintf("close body writer error: %v", err))
return
} }
req, err := http.NewRequest("POST", env.VEDEM_OCR_URL, bodyBuf) req, err := http.NewRequest("POST", env.VEDEM_OCR_URL, bodyBuf)
if err != nil { if err != nil {
return nil, errors.Wrap(code.OCRRequestError, err = errors.Wrap(code.OCRRequestError,
fmt.Sprintf("construct request error: %v", err)) fmt.Sprintf("construct request error: %v", err))
return
} }
token := builtin.Sign("auth-v2", env.VEDEM_OCR_AK, env.VEDEM_OCR_SK, bodyBuf.Bytes()) token := builtin.Sign("auth-v2", env.VEDEM_OCR_AK, env.VEDEM_OCR_SK, bodyBuf.Bytes())
@@ -195,41 +203,52 @@ func (s *veDEMOCRService) getOCRResult(imageBuf *bytes.Buffer) ([]OCRResult, err
log.Error().Err(err). log.Error().Err(err).
Str("X-TT-LOGID", logID). Str("X-TT-LOGID", logID).
Int("imageBufSize", size). Int("imageBufSize", size).
Msgf("request OCR service failed, retry %d", i) Msgf("request veDEM OCR service failed, retry %d", i)
time.Sleep(1 * time.Second) time.Sleep(1 * time.Second)
} }
if resp == nil { if resp == nil {
return nil, code.OCRServiceConnectionError err = code.OCRServiceConnectionError
return
} }
defer resp.Body.Close() defer resp.Body.Close()
results, err := ioutil.ReadAll(resp.Body) results, err := ioutil.ReadAll(resp.Body)
if err != nil { if err != nil {
return nil, errors.Wrap(code.OCRResponseError, err = errors.Wrap(code.OCRResponseError,
fmt.Sprintf("read response body error: %v", err)) fmt.Sprintf("read response body error: %v", err))
return
} }
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
return nil, errors.Wrap(code.OCRResponseError, err = errors.Wrap(code.OCRResponseError,
fmt.Sprintf("unexpected response status code: %d, results: %v", fmt.Sprintf("unexpected response status code: %d, results: %v",
resp.StatusCode, string(results))) resp.StatusCode, string(results)))
return
} }
var ocrResult ResponseOCR var ocrResult ResponseOCR
err = json.Unmarshal(results, &ocrResult) err = json.Unmarshal(results, &ocrResult)
if err != nil { if err != nil {
return nil, errors.Wrap(code.OCRResponseError, err = errors.Wrap(code.OCRResponseError,
fmt.Sprintf("json unmarshal response body error: %v", err)) fmt.Sprintf("json unmarshal response body error: %v", err))
return
} }
return ocrResult.OCRResult, nil if ocrResult.Code != 0 {
log.Error().
Int("code", ocrResult.Code).
Str("message", ocrResult.Message).
Msg("request veDEM OCR service failed")
}
return ocrResult.OCRResult, ocrResult.URL, nil
} }
func (s *veDEMOCRService) GetTexts(imageBuf *bytes.Buffer) ( func (s *veDEMOCRService) GetTexts(imageBuf *bytes.Buffer) (
ocrTexts OCRTexts, err error) { ocrTexts OCRTexts, url string, err error) {
ocrResults, err := s.getOCRResult(imageBuf) ocrResults, url, err := s.getOCRResult(imageBuf)
if err != nil { if err != nil {
log.Error().Err(err).Msg("getOCRResult failed") log.Error().Err(err).Msg("getOCRResult failed")
return return
@@ -262,6 +281,7 @@ func checkEnv() error {
if env.VEDEM_OCR_URL == "" { if env.VEDEM_OCR_URL == "" {
return errors.Wrap(code.OCREnvMissedError, "VEDEM_OCR_URL missed") return errors.Wrap(code.OCREnvMissedError, "VEDEM_OCR_URL missed")
} }
log.Info().Str("VEDEM_OCR_URL", env.VEDEM_OCR_URL).Msg("get env")
if env.VEDEM_OCR_AK == "" { if env.VEDEM_OCR_AK == "" {
return errors.Wrap(code.OCREnvMissedError, "VEDEM_OCR_AK missed") return errors.Wrap(code.OCREnvMissedError, "VEDEM_OCR_AK missed")
} }
@@ -284,7 +304,8 @@ func getLogID(header http.Header) string {
} }
type IOCRService interface { type IOCRService interface {
GetTexts(imageBuf *bytes.Buffer) (texts OCRTexts, err error) // GetTexts returns ocr texts and uploaded image url
GetTexts(imageBuf *bytes.Buffer) (texts OCRTexts, url string, err error)
} }
// GetScreenTextsByOCR takes a screenshot, returns the image path and OCR texts. // GetScreenTextsByOCR takes a screenshot, returns the image path and OCR texts.
@@ -295,12 +316,18 @@ func (dExt *DriverExt) GetScreenTextsByOCR() (imagePath string, ocrTexts OCRText
return return
} }
ocrTexts, err = dExt.OCRService.GetTexts(bufSource) var imageUrl string
ocrTexts, imageUrl, err = dExt.OCRService.GetTexts(bufSource)
if err != nil { if err != nil {
log.Error().Err(err).Msg("GetScreenTextsByOCR failed") log.Error().Err(err).Msg("GetScreenTextsByOCR failed")
return return
} }
if imageUrl != "" {
dExt.cacheStepData.screenShotsUrls[imagePath] = imageUrl
log.Debug().Str("imagePath", imagePath).Str("imageUrl", imageUrl).Msg("log screenshot")
}
dExt.cacheStepData.OcrResults[imagePath] = &OcrResult{ dExt.cacheStepData.OcrResults[imagePath] = &OcrResult{
Texts: ocrTexts, Texts: ocrTexts,
} }
+4 -7
View File
@@ -14,22 +14,19 @@ func checkOCR(buff *bytes.Buffer) error {
if err != nil { if err != nil {
return err return err
} }
ocrResults, err := service.getOCRResult(buff) ocrResults, url, err := service.getOCRResult(buff)
if err != nil { if err != nil {
return err return err
} }
fmt.Println(ocrResults) fmt.Println(ocrResults)
fmt.Println(url)
return nil return nil
} }
func TestOCRWithScreenshot(t *testing.T) { func TestOCRWithScreenshot(t *testing.T) {
device, _ := NewAndroidDevice() setupAndroid(t)
driver, err := device.NewDriver(nil)
if err != nil {
t.Fatal(err)
}
raw, err := driver.Driver.Screenshot() raw, err := driverExt.Driver.Screenshot()
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }