diff --git a/hrp/pkg/uixt/ext.go b/hrp/pkg/uixt/ext.go index b774ffb7..70de478c 100644 --- a/hrp/pkg/uixt/ext.go +++ b/hrp/pkg/uixt/ext.go @@ -86,7 +86,7 @@ type DriverExt struct { windowSize Size frame *bytes.Buffer doneMjpegStream chan bool - OCRService IOCRService // used to get texts from image + ImageService IImageService // used to extract image data interruptSignal chan os.Signal // cache step data @@ -110,7 +110,7 @@ func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error) return nil, err } - if dExt.OCRService, err = newVEDEMOCRService(); err != nil { + if dExt.ImageService, err = newVEDEMImageService(); err != nil { return nil, err } diff --git a/hrp/pkg/uixt/ocr_vedem.go b/hrp/pkg/uixt/ocr_vedem.go index afc97a76..e9b81816 100644 --- a/hrp/pkg/uixt/ocr_vedem.go +++ b/hrp/pkg/uixt/ocr_vedem.go @@ -28,6 +28,25 @@ type OCRResult struct { Points []PointF `json:"points"` } +func (o OCRResult) ToOCRText() OCRText { + rect := image.Rectangle{ + // ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下 + Min: image.Point{ + X: int(o.Points[0].X), + Y: int(o.Points[0].Y), + }, + Max: image.Point{ + X: int(o.Points[2].X), + Y: int(o.Points[2].Y), + }, + } + + return OCRText{ + Text: o.Text, + Rect: rect, + } +} + type ImageResult struct { URL string `json:"url"` // image uploaded url OCRResult []OCRResult `json:"ocrResult"` // OCR texts @@ -137,15 +156,15 @@ func (t OCRTexts) FindTexts(texts []string, options ...ActionOption) ( return results, nil } -func newVEDEMOCRService() (*veDEMOCRService, error) { +func newVEDEMImageService() (*veDEMImageService, error) { if err := checkEnv(); err != nil { return nil, err } - return &veDEMOCRService{}, nil + return &veDEMImageService{}, nil } -// veDEMOCRService implements IOCRService interface -type veDEMOCRService struct{} +// veDEMImageService implements IImageService interface +type veDEMImageService struct{} var actions = []string{ "ocr", // get ocr texts @@ -155,7 +174,7 @@ var actions = []string{ // "close", } -func (s *veDEMOCRService) getImageResult(imageBuf *bytes.Buffer) ( +func (s *veDEMImageService) GetImage(imageBuf *bytes.Buffer) ( imageResult ImageResult, err error) { bodyBuf := &bytes.Buffer{} @@ -265,40 +284,9 @@ func (s *veDEMOCRService) getImageResult(imageBuf *bytes.Buffer) ( Msg("request veDEM OCR service failed") } - return imageResponse.Result, nil -} - -func (s *veDEMOCRService) GetTexts(imageBuf *bytes.Buffer) ( - ocrTexts OCRTexts, url string, err error) { - - imageResult, err := s.getImageResult(imageBuf) - if err != nil { - log.Error().Err(err).Msg("getOCRResult failed") - return - } - - for _, ocrResult := range imageResult.OCRResult { - rect := image.Rectangle{ - // ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下 - Min: image.Point{ - X: int(ocrResult.Points[0].X), - Y: int(ocrResult.Points[0].Y), - }, - Max: image.Point{ - X: int(ocrResult.Points[2].X), - Y: int(ocrResult.Points[2].Y), - }, - } - - ocrTexts = append(ocrTexts, OCRText{ - Text: ocrResult.Text, - Rect: rect, - }) - } - url = imageResult.URL - - log.Debug().Interface("texts", ocrTexts).Msg("get screen texts by veDEM OCR") - return + imageResult = imageResponse.Result + log.Debug().Interface("imageResult", imageResult).Msg("get image data by veDEM") + return imageResult, nil } func checkEnv() error { @@ -327,9 +315,9 @@ func getLogID(header http.Header) string { return logID[0] } -type IOCRService interface { - // GetTexts returns ocr texts and uploaded image url - GetTexts(imageBuf *bytes.Buffer) (texts OCRTexts, url string, err error) +type IImageService interface { + // GetImage returns image result including ocr texts, uploaded image url, etc + GetImage(imageBuf *bytes.Buffer) (imageResult ImageResult, err error) } // GetScreenTextsByOCR takes a screenshot, returns the image path and OCR texts. @@ -340,13 +328,17 @@ func (dExt *DriverExt) GetScreenTextsByOCR() (imagePath string, ocrTexts OCRText return } - var imageUrl string - ocrTexts, imageUrl, err = dExt.OCRService.GetTexts(bufSource) + imageResult, err := dExt.ImageService.GetImage(bufSource) if err != nil { log.Error().Err(err).Msg("GetScreenTextsByOCR failed") return } + for _, ocrResult := range imageResult.OCRResult { + ocrTexts = append(ocrTexts, ocrResult.ToOCRText()) + } + + imageUrl := imageResult.URL if imageUrl != "" { dExt.cacheStepData.screenShotsUrls[imagePath] = imageUrl log.Debug().Str("imagePath", imagePath).Str("imageUrl", imageUrl).Msg("log screenshot") diff --git a/hrp/pkg/uixt/ocr_vedem_test.go b/hrp/pkg/uixt/ocr_vedem_test.go index ba0a4e74..ec7b5a4e 100644 --- a/hrp/pkg/uixt/ocr_vedem_test.go +++ b/hrp/pkg/uixt/ocr_vedem_test.go @@ -10,11 +10,11 @@ import ( ) func checkOCR(buff *bytes.Buffer) error { - service, err := newVEDEMOCRService() + service, err := newVEDEMImageService() if err != nil { return err } - imageResult, err := service.getImageResult(buff) + imageResult, err := service.GetImage(buff) if err != nil { return err }