feat: support action options for ScreenShot, WithScreenShotOCR/WithScreenShotUpload/WithScreenShotLiveType/WithScreenShotUIType

This commit is contained in:
lilong.129
2023-08-22 21:03:36 +08:00
parent 12d9b92782
commit e523c4ecc9
8 changed files with 135 additions and 51 deletions

View File

@@ -111,6 +111,12 @@ type ActionOptions struct {
// set custiom options such as textview, id, description
Custom map[string]interface{} `json:"custom,omitempty" yaml:"custom,omitempty"`
// screenshot related
ScreenShotWithOCR bool `json:"screenshot_with_ocr,omitempty" yaml:"screenshot_with_ocr,omitempty"`
ScreenShotWithUpload bool `json:"screenshot_with_upload,omitempty" yaml:"screenshot_with_upload,omitempty"`
ScreenShotWithLiveType bool `json:"screenshot_with_live_type,omitempty" yaml:"screenshot_with_live_type,omitempty"`
ScreenShotWithUIType bool `json:"screenshot_with_ui_type,omitempty" yaml:"screenshot_with_ui_type,omitempty"`
}
func (o *ActionOptions) Options() []ActionOption {
@@ -185,6 +191,37 @@ func (o *ActionOptions) Options() []ActionOption {
}
}
// screenshot options
if o.ScreenShotWithOCR {
options = append(options, WithScreenShotOCR(true))
}
if o.ScreenShotWithUpload {
options = append(options, WithScreenShotUpload(true))
}
if o.ScreenShotWithLiveType {
options = append(options, WithScreenShotLiveType(true))
}
if o.ScreenShotWithUIType {
options = append(options, WithScreenShotUIType(true))
}
return options
}
func (o *ActionOptions) screenshotOptions() screenshotActionOptions {
options := screenshotActionOptions{}
if o.ScreenShotWithOCR {
options = append(options, "ocr")
}
if o.ScreenShotWithUpload {
options = append(options, "upload")
}
if o.ScreenShotWithLiveType {
options = append(options, "liveType")
}
if o.ScreenShotWithUIType {
options = append(options, "ui")
}
return options
}
@@ -363,6 +400,30 @@ func WithIgnoreNotFoundError(ignoreError bool) ActionOption {
}
}
func WithScreenShotOCR(ocrOn bool) ActionOption {
return func(o *ActionOptions) {
o.ScreenShotWithOCR = ocrOn
}
}
func WithScreenShotUpload(uploadOn bool) ActionOption {
return func(o *ActionOptions) {
o.ScreenShotWithUpload = uploadOn
}
}
func WithScreenShotLiveType(liveTypeOn bool) ActionOption {
return func(o *ActionOptions) {
o.ScreenShotWithLiveType = liveTypeOn
}
}
func WithScreenShotUIType(uiTypeOn bool) ActionOption {
return func(o *ActionOptions) {
o.ScreenShotWithUIType = uiTypeOn
}
}
func (dExt *DriverExt) ParseActionOptions(options ...ActionOption) []ActionOption {
actionOptions := NewActionOptions(options...)
@@ -547,7 +608,7 @@ func (dExt *DriverExt) DoAction(action MobileAction) (err error) {
case ACTION_ScreenShot:
// take screenshot
log.Info().Msg("take screenshot for current screen")
_, _, err := dExt.takeScreenShot(builtin.GenNameWithTimestamp("%d_screenshot"))
_, err := dExt.GetScreenResult(action.GetOptions()...)
return err
case ACTION_StartCamera:
return dExt.Driver.StartCamera()

View File

@@ -51,6 +51,9 @@ func WithThreshold(threshold float64) CVOption {
}
type ScreenResult struct {
bufSource *bytes.Buffer // raw image buffer bytes
imagePath string // image file path
Texts OCRTexts `json:"texts"` // dumped raw OCRTexts
Tags []string `json:"tags"` // tags for image, e.g. ["feed", "ad", "live"]
VideoType string `json:"video_type,omitempty"` // video type: feed, live-preview or live

View File

@@ -14,8 +14,6 @@ import (
"github.com/pkg/errors"
"github.com/rs/zerolog/log"
"gocv.io/x/gocv"
"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
)
const (
@@ -101,9 +99,11 @@ func (dExt *DriverExt) FindAllImageRect(search string) (rects []image.Rectangle,
if bufSearch, err = getBufFromDisk(search); err != nil {
return nil, err
}
if bufSource, _, err = dExt.takeScreenShot(builtin.GenNameWithTimestamp("%d_cv")); err != nil {
screenResult, err := dExt.GetScreenResult()
if err != nil {
return nil, err
}
bufSource = screenResult.bufSource
if rects, err = FindAllImageRectsFromRaw(bufSource, bufSearch, float32(dExt.threshold), TemplateMatchMode(dExt.matchMode)); err != nil {
return nil, err

View File

@@ -62,7 +62,8 @@ func (dExt *DriverExt) AutoPopupHandler() error {
// TODO: check popup by activity type
// check popup by screenshot
screenResult, err := dExt.GetScreenResult()
screenResult, err := dExt.GetScreenResult(
WithScreenShotOCR(true), WithScreenShotUpload(true))
if err != nil {
return errors.Wrap(err, "get screen result failed for popup handler")
}

View File

@@ -53,7 +53,6 @@ func (o OCRResults) ToOCRTexts() (ocrTexts OCRTexts) {
}
type ImageResult struct {
imagePath string
URL string `json:"url"` // image uploaded url
OCRResult OCRResults `json:"ocrResult"` // OCR texts
LiveType string `json:"liveType"` // 直播间类型
@@ -170,16 +169,16 @@ func newVEDEMImageService() (*veDEMImageService, error) {
// veDEMImageService implements IImageService interface
// actions:
//
// ocr - get ocr texts
// upload - get image uploaded url
// liveType - get live type
// popup - get popup windows
// close - get close popup
// ui - get ui position by type(s)
// ocr - get ocr texts
// upload - get image uploaded url
// liveType - get live type
// popup - get popup windows
// close - get close popup
// ui - get ui position by type(s)
type veDEMImageService struct{}
type (
actionOptions []string
uiTypeOptions []string
screenshotActionOptions []string
screenshotUITypeOptions []string
)
func (s *veDEMImageService) GetImage(imageBuf *bytes.Buffer, options ...interface{}) (imageResult ImageResult, err error) {
@@ -187,11 +186,11 @@ func (s *veDEMImageService) GetImage(imageBuf *bytes.Buffer, options ...interfac
bodyWriter := multipart.NewWriter(bodyBuf)
for _, option := range options {
switch ov := option.(type) {
case actionOptions:
case screenshotActionOptions:
for _, action := range ov {
bodyWriter.WriteField("actions", action)
}
case uiTypeOptions:
case screenshotUITypeOptions:
for _, uiType := range ov {
bodyWriter.WriteField("uiTypes", uiType)
}
@@ -342,51 +341,66 @@ type IImageService interface {
}
// GetScreenResult takes a screenshot, returns the image recognization result
func (dExt *DriverExt) GetScreenResult() (screenResult *ScreenResult, err error) {
func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *ScreenResult, err error) {
screenshotOptions := NewActionOptions(options...).screenshotOptions()
var fileName string
if len(screenshotOptions) == 0 {
fileName = builtin.GenNameWithTimestamp("%d_screenshot")
} else {
fileName = builtin.GenNameWithTimestamp("%d_cv")
}
startTime := time.Now()
var bufSource *bytes.Buffer
var imagePath string
if bufSource, imagePath, err = dExt.takeScreenShot(
builtin.GenNameWithTimestamp("%d_cv")); err != nil {
bufSource, imagePath, err := dExt.takeScreenShot(fileName)
if err != nil {
return
}
screenshotTakeElapsed := time.Since(startTime).Milliseconds()
imageResult, err := dExt.ImageService.GetImage(bufSource, actionOptions{"ocr", "upload", "liveType"})
if err != nil {
log.Error().Err(err).Msg("GetImage from ImageService failed")
return
}
imageResult.imagePath = imagePath
imageUrl := imageResult.URL
if imageUrl != "" {
dExt.cacheStepData.screenShotsUrls[imagePath] = imageUrl
log.Debug().Str("imagePath", imagePath).Str("imageUrl", imageUrl).Msg("log screenshot")
}
screenResult = &ScreenResult{
Texts: imageResult.OCRResult.ToOCRTexts(),
bufSource: bufSource,
imagePath: imagePath,
Tags: nil,
ScreenshotTakeElapsed: screenshotTakeElapsed,
ScreenshotCVElapsed: time.Since(startTime).Milliseconds() - screenshotTakeElapsed,
}
if imageResult.LiveType != "" && imageResult.LiveType != "NoLive" {
screenResult.Live = &LiveRoom{
LiveType: imageResult.LiveType,
var imageUrl string
if len(screenshotOptions) > 0 {
imageResult, err := dExt.ImageService.GetImage(bufSource, screenshotOptions)
if err != nil {
log.Error().Err(err).Msg("GetImage from ImageService failed")
return nil, err
}
screenResult.ScreenshotCVElapsed = time.Since(startTime).Milliseconds() - screenshotTakeElapsed
screenResult.Texts = imageResult.OCRResult.ToOCRTexts()
imageUrl = imageResult.URL
if imageUrl != "" {
dExt.cacheStepData.screenShotsUrls[imagePath] = imageUrl
}
if imageResult.LiveType != "" && imageResult.LiveType != "NoLive" {
screenResult.Live = &LiveRoom{
LiveType: imageResult.LiveType,
}
}
}
dExt.cacheStepData.screenResults[imagePath] = screenResult
log.Debug().
Str("imagePath", imagePath).
Str("imageUrl", imageUrl).
Int64("screenshot_take_elapsed(ms)", screenResult.ScreenshotTakeElapsed).
Int64("screenshot_cv_elapsed(ms)", screenResult.ScreenshotCVElapsed).
Msg("get screenshot result success")
Msg("log screenshot")
return screenResult, nil
}
func (dExt *DriverExt) GetScreenTexts() (ocrTexts OCRTexts, err error) {
screenResult, err := dExt.GetScreenResult()
screenResult, err := dExt.GetScreenResult(
WithScreenShotOCR(true), WithScreenShotUpload(true))
if err != nil {
return
}
@@ -505,14 +519,17 @@ func (u UIResults) GetUIResult(options ...ActionOption) (UIResult, error) {
}
func (dExt *DriverExt) GetUIResultMap(uiTypes []string) (uiResultMap UIResultMap, err error) {
var bufSource *bytes.Buffer
var imagePath string
if bufSource, imagePath, err = dExt.takeScreenShot(
builtin.GenNameWithTimestamp("%d_cv")); err != nil {
return
screenResult, err := dExt.GetScreenResult()
if err != nil {
return nil, err
}
bufSource := screenResult.bufSource
imagePath := screenResult.imagePath
imageResult, err := dExt.ImageService.GetImage(bufSource, actionOptions{"ui"}, uiTypeOptions(uiTypes))
imageResult, err := dExt.ImageService.GetImage(bufSource,
screenshotActionOptions{"ui"}, // turn on UI type detection
screenshotUITypeOptions(uiTypes),
)
if err != nil {
log.Error().Err(err).Msg("GetImage from ImageService failed")
return

View File

@@ -189,7 +189,8 @@ func (vc *VideoCrawler) startLiveCrawler(enterPoint PointF) error {
}
// take screenshot and get screen texts by OCR
screenResult, err := vc.driverExt.GetScreenResult()
screenResult, err := vc.driverExt.GetScreenResult(
WithScreenShotOCR(true), WithScreenShotUpload(true))
if err != nil {
log.Error().Err(err).Msg("OCR GetTexts failed")
time.Sleep(3 * time.Second)