GithubBackup/httprunner
1
0
Fork 0
mirror of https://github.com/httprunner/httprunner.git synced 2026-06-01 22:09:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: get all ocr texts

Browse Source
This commit is contained in:
debugtalk
2022-11-23 20:49:44 +08:00
parent 992b1a29ce
commit 08fc2fbdd2
1 changed files with 81 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
hrp/pkg/uixt/ocr_vedem.go
View File

@@ -147,8 +147,23 @@ func getLogID(header http.Header) string {
return logID[0] return logID[0]
} }
func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...DataOption) (rect image.Rectangle, err error) { type OCRText struct {
data := NewData(map[string]interface{}{}, options...) Text string
Rect image.Rectangle
}
type OCRTexts []OCRText
func (t OCRTexts) Texts() []string {
var texts []string
for _, text := range t {
texts = append(texts, text.Text)
}
return texts
}
func (s *veDEMOCRService) GetAllTexts(imageBuf []byte) (
ocrTexts OCRTexts, err error) {
ocrResults, err := s.getOCRResult(imageBuf) ocrResults, err := s.getOCRResult(imageBuf)
if err != nil { if err != nil {
@@ -156,10 +171,8 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...Data
return return
} }
var rects []image.Rectangle
var ocrTexts []string
for _, ocrResult := range ocrResults { for _, ocrResult := range ocrResults {
rect = image.Rectangle{ rect := image.Rectangle{
// ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下 // ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下
Min: image.Point{ Min: image.Point{
X: int(ocrResult.Points[0].X), X: int(ocrResult.Points[0].X),
@@ -170,31 +183,57 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...Data
Y: int(ocrResult.Points[2].Y), Y: int(ocrResult.Points[2].Y),
}, },
} }
if rect.Min.X >= data.Scope[0] && rect.Max.X <= data.Scope[2] && rect.Min.Y >= data.Scope[1] && rect.Max.Y <= data.Scope[3] { ocrTexts = append(ocrTexts, OCRText{
ocrTexts = append(ocrTexts, ocrResult.Text) Text: ocrResult.Text,
Rect: rect,
})
}
return
}
// not contains text func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...DataOption) (
if !strings.Contains(ocrResult.Text, text) { rect image.Rectangle, err error) {
continue
}
rects = append(rects, rect) ocrTexts, err := s.GetAllTexts(imageBuf)
if err != nil {
log.Error().Err(err).Msg("GetAllTexts failed")
return
}
// contains text while not match exactly data := NewData(map[string]interface{}{}, options...)
if ocrResult.Text != text {
continue
}
// match exactly, and not specify index, return the first one var rects []image.Rectangle
if data.Index == 0 { for _, ocrText := range ocrTexts {
return rect, nil rect = ocrText.Rect
}
// check if text in scope
if rect.Min.X < data.Scope[0] || rect.Max.X > data.Scope[2] ||
rect.Min.Y < data.Scope[1] || rect.Max.Y > data.Scope[3] {
// not in scope
continue
}
// not contains text
if !strings.Contains(ocrText.Text, text) {
continue
}
rects = append(rects, rect)
// contains text while not match exactly
if ocrText.Text != text {
continue
}
// match exactly, and not specify index, return the first one
if data.Index == 0 {
return rect, nil
} }
} }
if len(rects) == 0 { if len(rects) == 0 {
return image.Rectangle{}, errors.Wrap(code.OCRTextNotFoundError, return image.Rectangle{}, errors.Wrap(code.OCRTextNotFoundError,
fmt.Sprintf("text %s not found in %v", text, ocrTexts)) fmt.Sprintf("text %s not found in %v", text, ocrTexts.Texts()))
} }
// get index // get index
@@ -215,45 +254,38 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, options ...Data
return rects[idx], nil return rects[idx], nil
} }
func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ...DataOption) (rects []image.Rectangle, err error) { func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ...DataOption) (
ocrResults, err := s.getOCRResult(imageBuf) rects []image.Rectangle, err error) {
ocrTexts, err := s.GetAllTexts(imageBuf)
if err != nil { if err != nil {
log.Error().Err(err).Msg("getOCRResult failed") log.Error().Err(err).Msg("GetAllTexts failed")
return return
} }
data := NewData(map[string]interface{}{}, options...) data := NewData(map[string]interface{}{}, options...)
ocrTexts := map[string]bool{}
var success bool var success bool
var rect image.Rectangle
for _, text := range texts { for _, text := range texts {
var found bool var found bool
for _, ocrResult := range ocrResults { for _, ocrText := range ocrTexts {
rect = image.Rectangle{ rect := ocrText.Rect
// ocrResult.Points 顺序：左上 -> 右上 -> 右下 -> 左下
Min: image.Point{ // check if text in scope
X: int(ocrResult.Points[0].X), if rect.Min.X < data.Scope[0] || rect.Max.X > data.Scope[2] ||
Y: int(ocrResult.Points[0].Y), rect.Min.Y < data.Scope[1] || rect.Max.Y > data.Scope[3] {
}, // not in scope
Max: image.Point{ continue
X: int(ocrResult.Points[2].X),
Y: int(ocrResult.Points[2].Y),
},
} }
if rect.Min.X >= data.Scope[0] && rect.Max.X <= data.Scope[2] && rect.Min.Y >= data.Scope[1] && rect.Max.Y <= data.Scope[3] { // not contains text
ocrTexts[ocrResult.Text] = true if !strings.Contains(ocrText.Text, text) {
continue
// not contains text
if !strings.Contains(ocrResult.Text, text) {
continue
}
found = true
rects = append(rects, rect)
break
} }
found = true
rects = append(rects, rect)
break
} }
if !found { if !found {
rects = append(rects, image.Rectangle{}) rects = append(rects, image.Rectangle{})
@@ -263,7 +295,7 @@ func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, options ...
if !success { if !success {
return rects, errors.Wrap(code.OCRTextNotFoundError, return rects, errors.Wrap(code.OCRTextNotFoundError,
fmt.Sprintf("texts %s not found in %v", texts, ocrTexts)) fmt.Sprintf("texts %s not found in %v", texts, ocrTexts.Texts()))
} }
return rects, nil return rects, nil