From ce0959022bcc90ca5f6a451cfc52b3e6703a2c72 Mon Sep 17 00:00:00 2001 From: debugtalk Date: Thu, 24 Nov 2022 20:23:09 +0800 Subject: [PATCH] feat: add option WithScreenShot --- hrp/pkg/uixt/demo/main_test.go | 28 ---------------------------- hrp/pkg/uixt/ext.go | 4 ++-- hrp/pkg/uixt/interface.go | 11 +++++++++++ hrp/pkg/uixt/ocr_vedem.go | 30 +++++++++++++++++++----------- hrp/pkg/uixt/ocr_vedem_test.go | 11 ++++++++--- 5 files changed, 40 insertions(+), 44 deletions(-) diff --git a/hrp/pkg/uixt/demo/main_test.go b/hrp/pkg/uixt/demo/main_test.go index de169814..1ff036ce 100644 --- a/hrp/pkg/uixt/demo/main_test.go +++ b/hrp/pkg/uixt/demo/main_test.go @@ -3,7 +3,6 @@ package demo import ( - "fmt" "testing" "time" @@ -45,30 +44,3 @@ func TestIOSDemo(t *testing.T) { } } } - -func TestIOSGetOCRTexts(t *testing.T) { - device, err := uixt.NewIOSDevice( - uixt.WithWDAPort(8700), uixt.WithWDAMjpegPort(8800), - uixt.WithResetHomeOnStartup(false), // not reset home on startup - ) - if err != nil { - t.Fatal(err) - } - - capabilities := uixt.NewCapabilities() - capabilities.WithDefaultAlertAction(uixt.AlertActionAccept) // or uixt.AlertActionDismiss - driverExt, err := device.NewDriver(capabilities) - if err != nil { - t.Fatal(err) - } - - for { - texts, err := driverExt.GetTextsByOCR() - if err != nil { - t.Fatal(err) - } - - fmt.Println(texts) - time.Sleep(3 * time.Second) - } -} diff --git a/hrp/pkg/uixt/ext.go b/hrp/pkg/uixt/ext.go index 88f6ab50..b544cd19 100644 --- a/hrp/pkg/uixt/ext.go +++ b/hrp/pkg/uixt/ext.go @@ -260,7 +260,7 @@ func (dExt *DriverExt) takeScreenShot() (raw *bytes.Buffer, err error) { } // saveScreenShot saves image file to $CWD/screenshots/ folder -func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) { +func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) { img, format, err := image.Decode(raw) if err != nil { return "", errors.Wrap(err, "decode screenshot image failed") @@ -304,7 +304,7 @@ func (dExt *DriverExt) ScreenShot(fileName string) (string, error) { return "", errors.Wrap(err, "screenshot failed") } - path, err := dExt.saveScreenShot(raw, fileName) + path, err := saveScreenShot(raw, fileName) if err != nil { return "", errors.Wrap(err, "save screenshot failed") } diff --git a/hrp/pkg/uixt/interface.go b/hrp/pkg/uixt/interface.go index 5dfba32a..cd837ff1 100644 --- a/hrp/pkg/uixt/interface.go +++ b/hrp/pkg/uixt/interface.go @@ -784,6 +784,7 @@ type DataOptions struct { IgnoreNotFoundError bool // ignore error if target element not found MaxRetryTimes int // max retry times if target element not found Interval float64 // interval between retries in seconds + ScreenShotFilename string // turn on screenshot and specify file name } type DataOption func(data *DataOptions) @@ -860,6 +861,16 @@ func WithDataWaitTime(sec float64) DataOption { } } +func WithScreenShot(fileName ...string) DataOption { + return func(data *DataOptions) { + if len(fileName) > 0 { + data.ScreenShotFilename = fileName[0] + } else { + data.ScreenShotFilename = fmt.Sprintf("screenshot_%d", time.Now().Unix()) + } + } +} + func NewDataOptions(options ...DataOption) *DataOptions { dataOptions := &DataOptions{} for _, option := range options { diff --git a/hrp/pkg/uixt/ocr_vedem.go b/hrp/pkg/uixt/ocr_vedem.go index d4233b51..9d1bd229 100644 --- a/hrp/pkg/uixt/ocr_vedem.go +++ b/hrp/pkg/uixt/ocr_vedem.go @@ -56,7 +56,7 @@ func checkEnv() error { return nil } -func (s *veDEMOCRService) getOCRResult(imageBuf []byte) ([]OCRResult, error) { +func (s *veDEMOCRService) getOCRResult(imageBuf *bytes.Buffer) ([]OCRResult, error) { bodyBuf := &bytes.Buffer{} bodyWriter := multipart.NewWriter(bodyBuf) bodyWriter.WriteField("withDet", "true") @@ -67,7 +67,7 @@ func (s *veDEMOCRService) getOCRResult(imageBuf []byte) ([]OCRResult, error) { return nil, errors.Wrap(code.OCRRequestError, fmt.Sprintf("create form file error: %v", err)) } - size, err := formWriter.Write(imageBuf) + size, err := formWriter.Write(imageBuf.Bytes()) if err != nil { return nil, errors.Wrap(code.OCRRequestError, fmt.Sprintf("write form error: %v", err)) @@ -161,7 +161,7 @@ func (t OCRTexts) Texts() (texts []string) { return texts } -func (s *veDEMOCRService) GetTexts(imageBuf []byte, options ...DataOption) ( +func (s *veDEMOCRService) GetTexts(imageBuf *bytes.Buffer, options ...DataOption) ( ocrTexts OCRTexts, err error) { ocrResults, err := s.getOCRResult(imageBuf) @@ -172,6 +172,14 @@ func (s *veDEMOCRService) GetTexts(imageBuf []byte, options ...DataOption) ( dataOptions := NewDataOptions(options...) + if dataOptions.ScreenShotFilename != "" { + path, err := saveScreenShot(imageBuf, dataOptions.ScreenShotFilename) + if err != nil { + return nil, errors.Wrap(err, "save screenshot failed") + } + log.Info().Str("path", path).Msg("save screenshot") + } + for _, ocrResult := range ocrResults { rect := image.Rectangle{ // ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下 @@ -200,7 +208,7 @@ func (s *veDEMOCRService) GetTexts(imageBuf []byte, options ...DataOption) ( return } -func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...DataOption) ( +func (s *veDEMOCRService) FindText(text string, imageBuf *bytes.Buffer, options ...DataOption) ( rect image.Rectangle, err error) { ocrTexts, err := s.GetTexts(imageBuf, options...) @@ -256,7 +264,7 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...Data return rects[idx], nil } -func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ...DataOption) ( +func (s *veDEMOCRService) FindTexts(texts []string, imageBuf *bytes.Buffer, options ...DataOption) ( rects []image.Rectangle, err error) { ocrTexts, err := s.GetTexts(imageBuf, options...) @@ -295,9 +303,9 @@ func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ... } type OCRService interface { - GetTexts(imageBuf []byte, options ...DataOption) (ocrTexts OCRTexts, err error) - FindText(text string, imageBuf []byte, options ...DataOption) (rect image.Rectangle, err error) - FindTexts(texts []string, imageBuf []byte, options ...DataOption) (rects []image.Rectangle, err error) + GetTexts(imageBuf *bytes.Buffer, options ...DataOption) (ocrTexts OCRTexts, err error) + FindText(text string, imageBuf *bytes.Buffer, options ...DataOption) (rect image.Rectangle, err error) + FindTexts(texts []string, imageBuf *bytes.Buffer, options ...DataOption) (rects []image.Rectangle, err error) } func (dExt *DriverExt) GetTextsByOCR(options ...DataOption) (texts OCRTexts, err error) { @@ -307,7 +315,7 @@ func (dExt *DriverExt) GetTextsByOCR(options ...DataOption) (texts OCRTexts, err return } - ocrTexts, err := dExt.ocrService.GetTexts(bufSource.Bytes(), options...) + ocrTexts, err := dExt.ocrService.GetTexts(bufSource, options...) if err != nil { log.Error().Err(err).Msg("GetTexts failed") return @@ -323,7 +331,7 @@ func (dExt *DriverExt) FindTextByOCR(ocrText string, options ...DataOption) (x, return } - rect, err := dExt.ocrService.FindText(ocrText, bufSource.Bytes(), options...) + rect, err := dExt.ocrService.FindText(ocrText, bufSource, options...) if err != nil { log.Warn().Msgf("FindText failed: %s", err.Error()) return @@ -342,7 +350,7 @@ func (dExt *DriverExt) FindTextsByOCR(ocrTexts []string, options ...DataOption) return } - rects, err := dExt.ocrService.FindTexts(ocrTexts, bufSource.Bytes(), options...) + rects, err := dExt.ocrService.FindTexts(ocrTexts, bufSource, options...) if err != nil { log.Warn().Msgf("FindTexts failed: %s", err.Error()) return diff --git a/hrp/pkg/uixt/ocr_vedem_test.go b/hrp/pkg/uixt/ocr_vedem_test.go index 2bb650b7..f3fe1d32 100644 --- a/hrp/pkg/uixt/ocr_vedem_test.go +++ b/hrp/pkg/uixt/ocr_vedem_test.go @@ -3,12 +3,13 @@ package uixt import ( + "bytes" "fmt" "os" "testing" ) -func checkOCR(buff []byte) error { +func checkOCR(buff *bytes.Buffer) error { service, err := newVEDEMOCRService() if err != nil { return err @@ -33,19 +34,23 @@ func TestOCRWithScreenshot(t *testing.T) { t.Fatal(err) } - if err := checkOCR(raw.Bytes()); err != nil { + if err := checkOCR(raw); err != nil { t.Fatal(err) } } func TestOCRWithLocalFile(t *testing.T) { imagePath := "/Users/debugtalk/Downloads/s1.png" + file, err := os.ReadFile(imagePath) if err != nil { t.Fatal(err) } - if err := checkOCR(file); err != nil { + buf := new(bytes.Buffer) + buf.Read(file) + + if err := checkOCR(buf); err != nil { t.Fatal(err) } }