change: add options for FindTexts

This commit is contained in:
debugtalk
2022-11-23 23:48:55 +08:00
parent 9888cea3e5
commit 5cb7218146
+14 -20
View File
@@ -170,6 +170,8 @@ func (s *veDEMOCRService) GetTexts(imageBuf []byte, options ...DataOption) (
return return
} }
dataOptions := NewDataOptions(options...)
for _, ocrResult := range ocrResults { for _, ocrResult := range ocrResults {
rect := image.Rectangle{ rect := image.Rectangle{
// ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下 // ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下
@@ -182,6 +184,14 @@ func (s *veDEMOCRService) GetTexts(imageBuf []byte, options ...DataOption) (
Y: int(ocrResult.Points[2].Y), Y: int(ocrResult.Points[2].Y),
}, },
} }
// check if text in scope
if rect.Min.X < dataOptions.Scope[0] || rect.Max.X > dataOptions.Scope[2] ||
rect.Min.Y < dataOptions.Scope[1] || rect.Max.Y > dataOptions.Scope[3] {
// not in scope
continue
}
ocrTexts = append(ocrTexts, OCRText{ ocrTexts = append(ocrTexts, OCRText{
Text: ocrResult.Text, Text: ocrResult.Text,
Rect: rect, Rect: rect,
@@ -193,7 +203,7 @@ func (s *veDEMOCRService) GetTexts(imageBuf []byte, options ...DataOption) (
func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...DataOption) ( func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...DataOption) (
rect image.Rectangle, err error) { rect image.Rectangle, err error) {
ocrTexts, err := s.GetTexts(imageBuf) ocrTexts, err := s.GetTexts(imageBuf, options...)
if err != nil { if err != nil {
log.Error().Err(err).Msg("GetTexts failed") log.Error().Err(err).Msg("GetTexts failed")
return return
@@ -205,13 +215,6 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...Data
for _, ocrText := range ocrTexts { for _, ocrText := range ocrTexts {
rect = ocrText.Rect rect = ocrText.Rect
// check if text in scope
if rect.Min.X < dataOptions.Scope[0] || rect.Max.X > dataOptions.Scope[2] ||
rect.Min.Y < dataOptions.Scope[1] || rect.Max.Y > dataOptions.Scope[3] {
// not in scope
continue
}
// not contains text // not contains text
if !strings.Contains(ocrText.Text, text) { if !strings.Contains(ocrText.Text, text) {
continue continue
@@ -256,27 +259,18 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...Data
func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ...DataOption) ( func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ...DataOption) (
rects []image.Rectangle, err error) { rects []image.Rectangle, err error) {
ocrTexts, err := s.GetTexts(imageBuf) ocrTexts, err := s.GetTexts(imageBuf, options...)
if err != nil { if err != nil {
log.Error().Err(err).Msg("GetTexts failed") log.Error().Err(err).Msg("GetTexts failed")
return return
} }
dataOptions := NewDataOptions(options...)
var success bool var success bool
for _, text := range texts { for _, text := range texts {
var found bool var found bool
for _, ocrText := range ocrTexts { for _, ocrText := range ocrTexts {
rect := ocrText.Rect rect := ocrText.Rect
// check if text in scope
if rect.Min.X < dataOptions.Scope[0] || rect.Max.X > dataOptions.Scope[2] ||
rect.Min.Y < dataOptions.Scope[1] || rect.Max.Y > dataOptions.Scope[3] {
// not in scope
continue
}
// not contains text // not contains text
if !strings.Contains(ocrText.Text, text) { if !strings.Contains(ocrText.Text, text) {
continue continue
@@ -306,14 +300,14 @@ type OCRService interface {
FindTexts(texts []string, imageBuf []byte, options ...DataOption) (rects []image.Rectangle, err error) FindTexts(texts []string, imageBuf []byte, options ...DataOption) (rects []image.Rectangle, err error)
} }
func (dExt *DriverExt) GetTextsByOCR() (texts OCRTexts, err error) { func (dExt *DriverExt) GetTextsByOCR(options ...DataOption) (texts OCRTexts, err error) {
var bufSource *bytes.Buffer var bufSource *bytes.Buffer
if bufSource, err = dExt.takeScreenShot(); err != nil { if bufSource, err = dExt.takeScreenShot(); err != nil {
err = fmt.Errorf("takeScreenShot error: %v", err) err = fmt.Errorf("takeScreenShot error: %v", err)
return return
} }
ocrTexts, err := dExt.ocrService.GetTexts(bufSource.Bytes()) ocrTexts, err := dExt.ocrService.GetTexts(bufSource.Bytes(), options...)
if err != nil { if err != nil {
log.Error().Err(err).Msg("GetTexts failed") log.Error().Err(err).Msg("GetTexts failed")
return return