feat: cache screenshot ocr texts

This commit is contained in:
lilong.129
2023-04-28 11:27:18 +08:00
parent c0ea52d02a
commit 9981dad1f8
9 changed files with 56 additions and 35 deletions

View File

@@ -97,7 +97,7 @@
},
{
"method": "sleep",
"params": 10
"params": 5
},
{
"method": "screenshot"
@@ -109,7 +109,7 @@
},
{
"method": "sleep",
"params": 10
"params": 5
},
{
"method": "screenshot"

View File

@@ -37,8 +37,8 @@ func TestAndroidLiveSwipeTest(t *testing.T) {
SleepRandom(0, 5, 0.6, 5, 10, 0.4),
hrp.NewStep("向上滑动,等待 10s").
Android().
SwipeUp(uixt.WithIdentifier("第一次上划")).Sleep(10).ScreenShot(). // 上划 1 次,等待 10s截图保存
SwipeUp(uixt.WithIdentifier("第二次上划")).Sleep(10).ScreenShot(), // 再上划 1 次,等待 10s截图保存
SwipeUp(uixt.WithIdentifier("第一次上划")).Sleep(5).ScreenShot(). // 上划 1 次,等待 5s截图保存
SwipeUp(uixt.WithIdentifier("第二次上划")).Sleep(5).ScreenShot(), // 再上划 1 次,等待 5s截图保存
hrp.NewStep("exit").
Android().
AppTerminate("com.ss.android.ugc.aweme").

View File

@@ -340,7 +340,7 @@ func (dExt *DriverExt) DoAction(action MobileAction) error {
case ACTION_ScreenShot:
// take screenshot
log.Info().Msg("take screenshot for current screen")
_, err := dExt.TakeScreenShot(builtin.GenNameWithTimestamp("%d_screenshot"))
_, _, err := dExt.TakeScreenShot(builtin.GenNameWithTimestamp("%d_screenshot"))
return err
case ACTION_StartCamera:
return dExt.Driver.StartCamera()

View File

@@ -51,16 +51,17 @@ type DriverExt struct {
windowSize Size
frame *bytes.Buffer
doneMjpegStream chan bool
OCRService IOCRService // used to get text from image
screenShots []string // cache screenshot paths
OCRService IOCRService // used to get texts from image
stepScreenShots map[string]OCRTexts // cache screenshot ocr results, key is image path
CVArgs
}
func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error) {
dExt = &DriverExt{
Device: device,
Driver: driver,
Device: device,
Driver: driver,
stepScreenShots: make(map[string]OCRTexts),
}
dExt.doneMjpegStream = make(chan bool, 1)
@@ -86,22 +87,22 @@ func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error)
// TakeScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
// if fileName is empty, it will not save image file and only return raw image data
func (dExt *DriverExt) TakeScreenShot(fileName ...string) (raw *bytes.Buffer, err error) {
func (dExt *DriverExt) TakeScreenShot(fileName ...string) (raw *bytes.Buffer, path string, err error) {
// iOS 优先使用 MJPEG 流进行截图,性能最优
// 如果 MJPEG 流未开启,则使用 WebDriver 的截图接口
if dExt.frame != nil {
return dExt.frame, nil
return dExt.frame, "", nil
}
if raw, err = dExt.Driver.Screenshot(); err != nil {
log.Error().Err(err).Msg("capture screenshot data failed")
return nil, err
return nil, "", err
}
// compress image data
compressed, err := compressImageBuffer(raw)
if err != nil {
log.Error().Err(err).Msg("compress screenshot data failed")
return nil, err
return nil, "", err
}
// save screenshot to file
@@ -110,13 +111,12 @@ func (dExt *DriverExt) TakeScreenShot(fileName ...string) (raw *bytes.Buffer, er
path, err := dExt.saveScreenShot(compressed, path)
if err != nil {
log.Error().Err(err).Msg("save screenshot file failed")
return nil, err
return nil, "", err
}
dExt.screenShots = append(dExt.screenShots, path)
log.Info().Str("path", path).Msg("save screenshot file success")
return compressed, path, nil
}
return compressed, nil
return compressed, "", nil
}
func compressImageBuffer(raw *bytes.Buffer) (compressed *bytes.Buffer, err error) {
@@ -160,14 +160,23 @@ func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (strin
return "", errors.Wrap(err, "encode screenshot image failed")
}
dExt.stepScreenShots[screenshotPath] = nil
log.Info().Str("path", screenshotPath).Msg("save screenshot file success")
return screenshotPath, nil
}
func (dExt *DriverExt) GetScreenShots() []string {
func (dExt *DriverExt) GetScreenShots() map[string]OCRTexts {
defer func() {
dExt.screenShots = nil
for key := range dExt.stepScreenShots {
delete(dExt.stepScreenShots, key)
}
}()
return dExt.screenShots
copied := make(map[string]OCRTexts)
for key, value := range dExt.stepScreenShots {
copied[key] = value
}
return copied
}
// isPathExists returns true if path exists, whether path is file or dir

View File

@@ -279,7 +279,8 @@ type IOCRService interface {
func (dExt *DriverExt) GetScreenTextsByOCR() (texts OCRTexts, err error) {
var bufSource *bytes.Buffer
if bufSource, err = dExt.TakeScreenShot(
var imagePath string
if bufSource, imagePath, err = dExt.TakeScreenShot(
builtin.GenNameWithTimestamp("%d_ocr")); err != nil {
return
}
@@ -290,6 +291,7 @@ func (dExt *DriverExt) GetScreenTextsByOCR() (texts OCRTexts, err error) {
return
}
dExt.stepScreenShots[imagePath] = ocrTexts
return ocrTexts, nil
}

View File

@@ -101,7 +101,7 @@ 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 {
if bufSource, _, err = dExt.TakeScreenShot(builtin.GenNameWithTimestamp("%d_cv")); err != nil {
return nil, err
}

View File

@@ -25,15 +25,21 @@ func (dExt *DriverExt) VideoCrawler(configs *VideoCrawlerConfigs) (err error) {
// loop until target count achieved
for {
// take screenshot and get screen texts by OCR
_, err := dExt.GetScreenTextsByOCR()
texts, err := dExt.GetScreenTextsByOCR()
if err != nil {
log.Error().Err(err).Msg("OCR GetTexts failed")
return err
}
// TODO: check if text popup exists
// check if text popup exists
if isTextPopup(texts) {
log.Info().Msg("text popup found")
}
// TODO: check if live video
// check if live video
if isLiveVideo(texts) {
log.Info().Msg("live video found")
}
// assert feed video type
@@ -48,3 +54,13 @@ func (dExt *DriverExt) VideoCrawler(configs *VideoCrawlerConfigs) (err error) {
// return nil
}
func isTextPopup(texts OCRTexts) bool {
texts.FindTexts([]string{"确定", "取消"})
return false
}
func isLiveVideo(texts OCRTexts) bool {
_, err := texts.FindTexts([]string{"点击进入直播间", "直播中"})
return err == nil
}

View File

@@ -5,18 +5,12 @@ package uixt
import "testing"
func TestVideoCrawler(t *testing.T) {
device, err := NewAndroidDevice()
if err != nil {
t.Fatal(err)
}
driver, err := device.NewDriver(nil)
if err != nil {
t.Fatal(err)
}
setupAndroid(t)
configs := &VideoCrawlerConfigs{
AppPackageName: "com.ss.android.ugc.aweme",
}
err = driver.VideoCrawler(configs)
err := driverExt.VideoCrawler(configs)
if err != nil {
t.Fatal(err)
}

View File

@@ -609,7 +609,7 @@ func runStepMobileUI(s *SessionRunner, step *TStep) (stepResult *StepResult, err
}
// take screenshot after each step
_, err := uiDriver.TakeScreenShot(
_, _, err := uiDriver.TakeScreenShot(
builtin.GenNameWithTimestamp("%d_step_") + step.Name)
if err != nil {
log.Error().Err(err).Str("step", step.Name).Msg("take screenshot failed on step finished")