feat: cache screenshot ocr texts

examples/uitest/demo_android_live_swipe.json

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

examples/uitest/demo_android_live_swipe_test.go

@@ -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(5).ScreenShot(). // 上划 1 次，等待 5s，截图保存
-				SwipeUp(uixt.WithIdentifier("第二次上划")).Sleep(10).ScreenShot(), // 再上划 1 次，等待 10s，截图保存
+				SwipeUp(uixt.WithIdentifier("第二次上划")).Sleep(5).ScreenShot(), // 再上划 1 次，等待 5s，截图保存
 			hrp.NewStep("exit").
 				Android().
 				AppTerminate("com.ss.android.ugc.aweme").

hrp/pkg/uixt/action.go

@@ -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()

hrp/pkg/uixt/ext.go

@@ -51,16 +51,17 @@ type DriverExt struct {
 	windowSize      Size
 	frame           *bytes.Buffer
 	doneMjpegStream chan bool
-	OCRService      IOCRService // used to get text from image
+	OCRService      IOCRService         // used to get texts from image
-	screenShots     []string    // cache screenshot paths
+	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,
+		Device:          device,
-		Driver: driver,
+		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)
+		return compressed, path, nil
-		log.Info().Str("path", path).Msg("save screenshot file success")
 	}
 
-	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

hrp/pkg/uixt/ocr_vedem.go

@@ -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
 }
 

hrp/pkg/uixt/opencv.go

@@ -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
 	}
 

hrp/pkg/uixt/video_crawler.go

@@ -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
+}

hrp/pkg/uixt/video_crawler_test.go

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

hrp/step_mobile_ui.go

@@ -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")