mirror of
https://github.com/httprunner/httprunner.git
synced 2026-09-05 23:51:25 +08:00
feat: get ocr position by given recognition area
This commit is contained in:
+94
-23
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"image"
|
||||
"io/ioutil"
|
||||
"math"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"os"
|
||||
@@ -109,7 +110,7 @@ func getLogID(header http.Header) string {
|
||||
return logID[0]
|
||||
}
|
||||
|
||||
func (s *veDEMOCRService) FindText(text string, imageBuf []byte, index ...int) (rect image.Rectangle, err error) {
|
||||
func (s *veDEMOCRService) FindText(text string, imageBuf []byte, recAbsArea []int, index ...int) (rect image.Rectangle, err error) {
|
||||
if len(index) == 0 {
|
||||
index = []int{0} // index not specified
|
||||
}
|
||||
@@ -120,16 +121,25 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, index ...int) (
|
||||
return
|
||||
}
|
||||
|
||||
if len(recAbsArea) != 4 {
|
||||
recAbsArea = []int{0, 0, math.MaxInt64, math.MaxInt64}
|
||||
}
|
||||
|
||||
var minX, minY, maxX, maxY int
|
||||
if recAbsArea[0] < recAbsArea[2] {
|
||||
minX, maxX = recAbsArea[0], recAbsArea[2]
|
||||
} else {
|
||||
minX, maxX = recAbsArea[2], recAbsArea[0]
|
||||
}
|
||||
if recAbsArea[1] < recAbsArea[3] {
|
||||
minY, maxY = recAbsArea[1], recAbsArea[3]
|
||||
} else {
|
||||
minY, maxY = recAbsArea[3], recAbsArea[1]
|
||||
}
|
||||
|
||||
var rects []image.Rectangle
|
||||
var ocrTexts []string
|
||||
for _, ocrResult := range ocrResults {
|
||||
ocrTexts = append(ocrTexts, ocrResult.Text)
|
||||
|
||||
// not contains text
|
||||
if !strings.Contains(ocrResult.Text, text) {
|
||||
continue
|
||||
}
|
||||
|
||||
rect = image.Rectangle{
|
||||
// ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下
|
||||
Min: image.Point{
|
||||
@@ -141,7 +151,16 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, index ...int) (
|
||||
Y: int(ocrResult.Points[2].Y),
|
||||
},
|
||||
}
|
||||
rects = append(rects, rect)
|
||||
if rect.Min.X > minX && rect.Max.X < maxX && rect.Min.Y < maxY && rect.Max.Y > minY {
|
||||
ocrTexts = append(ocrTexts, ocrResult.Text)
|
||||
|
||||
// not contains text
|
||||
if !strings.Contains(ocrResult.Text, text) {
|
||||
continue
|
||||
}
|
||||
|
||||
rects = append(rects, rect)
|
||||
}
|
||||
|
||||
// contains text while not match exactly
|
||||
if ocrResult.Text != text {
|
||||
@@ -177,23 +196,36 @@ func (s *veDEMOCRService) FindText(text string, imageBuf []byte, index ...int) (
|
||||
return rects[idx], nil
|
||||
}
|
||||
|
||||
func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte) (rects []image.Rectangle, err error) {
|
||||
func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte, recAbsArea []int) (rects []image.Rectangle, err error) {
|
||||
ocrResults, err := s.getOCRResult(imageBuf)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Msg("getOCRResult failed")
|
||||
return
|
||||
}
|
||||
|
||||
if len(recAbsArea) != 4 {
|
||||
recAbsArea = []int{0, 0, math.MaxInt64, math.MaxInt64}
|
||||
}
|
||||
|
||||
var minX, minY, maxX, maxY int
|
||||
if recAbsArea[0] < recAbsArea[2] {
|
||||
minX, maxX = recAbsArea[0], recAbsArea[2]
|
||||
} else {
|
||||
minX, maxX = recAbsArea[2], recAbsArea[0]
|
||||
}
|
||||
if recAbsArea[1] < recAbsArea[3] {
|
||||
minY, maxY = recAbsArea[1], recAbsArea[3]
|
||||
} else {
|
||||
minY, maxY = recAbsArea[3], recAbsArea[1]
|
||||
}
|
||||
|
||||
var success bool
|
||||
var rect image.Rectangle
|
||||
var ocrTexts []string
|
||||
for _, text := range texts {
|
||||
var found bool
|
||||
for _, ocrResult := range ocrResults {
|
||||
// not contains text
|
||||
if !strings.Contains(ocrResult.Text, text) {
|
||||
continue
|
||||
}
|
||||
|
||||
found = true
|
||||
rect := image.Rectangle{
|
||||
rect = image.Rectangle{
|
||||
// ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下
|
||||
Min: image.Point{
|
||||
X: int(ocrResult.Points[0].X),
|
||||
@@ -204,12 +236,29 @@ func (s *veDEMOCRService) FindTexts(texts []string, imageBuf []byte) (rects []im
|
||||
Y: int(ocrResult.Points[2].Y),
|
||||
},
|
||||
}
|
||||
rects = append(rects, rect)
|
||||
break
|
||||
|
||||
if rect.Min.X > minX && rect.Max.X < maxX && rect.Min.Y < maxY && rect.Max.Y > minY {
|
||||
ocrTexts = append(ocrTexts, ocrResult.Text)
|
||||
|
||||
// not contains text
|
||||
if !strings.Contains(ocrResult.Text, text) {
|
||||
continue
|
||||
}
|
||||
|
||||
found = true
|
||||
rects = append(rects, rect)
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
rects = append(rects, image.Rectangle{})
|
||||
}
|
||||
success = found || success
|
||||
}
|
||||
|
||||
if !success {
|
||||
return rects,
|
||||
fmt.Errorf("texts %s not found in %v", texts, ocrTexts)
|
||||
}
|
||||
|
||||
return rects, nil
|
||||
@@ -219,15 +268,26 @@ type OCRService interface {
|
||||
FindText(text string, imageBuf []byte, index ...int) (rect image.Rectangle, err error)
|
||||
}
|
||||
|
||||
func (dExt *DriverExt) FindTextByOCR(ocrText string, index ...int) (x, y, width, height float64, err error) {
|
||||
func (dExt *DriverExt) FindTextByOCR(ocrText string, recognitionArea []float64, index ...int) (x, y, width, height float64, err error) {
|
||||
var bufSource *bytes.Buffer
|
||||
if bufSource, err = dExt.takeScreenShot(); err != nil {
|
||||
err = fmt.Errorf("takeScreenShot error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if len(recognitionArea) != 4 {
|
||||
recognitionArea = []float64{0, 0, 1, 1}
|
||||
}
|
||||
|
||||
absArea := []int{
|
||||
int(recognitionArea[0] * float64(dExt.windowSize.Width) * dExt.scale),
|
||||
int(recognitionArea[1] * float64(dExt.windowSize.Height) * dExt.scale),
|
||||
int(recognitionArea[2] * float64(dExt.windowSize.Width) * dExt.scale),
|
||||
int(recognitionArea[3] * float64(dExt.windowSize.Height) * dExt.scale),
|
||||
}
|
||||
|
||||
service := &veDEMOCRService{}
|
||||
rect, err := service.FindText(ocrText, bufSource.Bytes(), index...)
|
||||
rect, err := service.FindText(ocrText, bufSource.Bytes(), absArea, index...)
|
||||
if err != nil {
|
||||
log.Warn().Msgf("FindText failed: %s", err.Error())
|
||||
err = fmt.Errorf("FindText failed: %v", err)
|
||||
@@ -240,15 +300,26 @@ func (dExt *DriverExt) FindTextByOCR(ocrText string, index ...int) (x, y, width,
|
||||
return
|
||||
}
|
||||
|
||||
func (dExt *DriverExt) FindTextsByOCR(ocrTexts []string) (points [][]float64, err error) {
|
||||
func (dExt *DriverExt) FindTextsByOCR(ocrTexts []string, recognitionArea []float64) (points [][]float64, err error) {
|
||||
var bufSource *bytes.Buffer
|
||||
if bufSource, err = dExt.takeScreenShot(); err != nil {
|
||||
err = fmt.Errorf("takeScreenShot error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if len(recognitionArea) != 4 {
|
||||
recognitionArea = []float64{0, 0, 1, 1}
|
||||
}
|
||||
|
||||
absArea := []int{
|
||||
int(recognitionArea[0] * float64(dExt.windowSize.Width) * dExt.scale),
|
||||
int(recognitionArea[1] * float64(dExt.windowSize.Height) * dExt.scale),
|
||||
int(recognitionArea[2] * float64(dExt.windowSize.Width) * dExt.scale),
|
||||
int(recognitionArea[3] * float64(dExt.windowSize.Height) * dExt.scale),
|
||||
}
|
||||
|
||||
service := &veDEMOCRService{}
|
||||
rects, err := service.FindTexts(ocrTexts, bufSource.Bytes())
|
||||
rects, err := service.FindTexts(ocrTexts, bufSource.Bytes(), absArea)
|
||||
if err != nil {
|
||||
log.Warn().Msgf("FindTexts failed: %s", err.Error())
|
||||
err = fmt.Errorf("FindTexts failed: %v", err)
|
||||
|
||||
Reference in New Issue
Block a user