feat: cache screenshot ocr texts

This commit is contained in:
lilong.129
2023-04-28 14:12:31 +08:00
parent 2323015755
commit 1d41d276ab
3 changed files with 28 additions and 19 deletions
+24 -17
View File
@@ -45,23 +45,34 @@ func WithThreshold(threshold float64) CVOption {
} }
} }
type cacheStepData struct {
// cache step screenshot paths
ScreenShots []string
// cache step screenshot ocr results, key is image path, value is dumped OCRTexts
OcrResults map[string]string
}
type DriverExt struct { type DriverExt struct {
CVArgs
Device Device Device Device
Driver WebDriver Driver WebDriver
windowSize Size windowSize Size
frame *bytes.Buffer frame *bytes.Buffer
doneMjpegStream chan bool doneMjpegStream chan bool
OCRService IOCRService // used to get texts from image OCRService IOCRService // used to get texts from image
stepScreenShots map[string]string // cache screenshot ocr results, key is image path, value is dumped OCRTexts
CVArgs // cache step data
cacheStepData cacheStepData
} }
func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error) { func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error) {
dExt = &DriverExt{ dExt = &DriverExt{
Device: device, Device: device,
Driver: driver, Driver: driver,
stepScreenShots: make(map[string]string), cacheStepData: cacheStepData{
ScreenShots: make([]string, 0),
OcrResults: make(map[string]string),
},
} }
dExt.doneMjpegStream = make(chan bool, 1) dExt.doneMjpegStream = make(chan bool, 1)
@@ -160,21 +171,17 @@ func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (strin
return "", errors.Wrap(err, "encode screenshot image failed") return "", errors.Wrap(err, "encode screenshot image failed")
} }
dExt.stepScreenShots[screenshotPath] = "" dExt.cacheStepData.ScreenShots = append(dExt.cacheStepData.ScreenShots, screenshotPath)
log.Info().Str("path", screenshotPath).Msg("save screenshot file success") log.Info().Str("path", screenshotPath).Msg("save screenshot file success")
return screenshotPath, nil return screenshotPath, nil
} }
func (dExt *DriverExt) GetScreenShots() map[string]string { func (dExt *DriverExt) GetStepCacheData() cacheStepData {
defer func() { copied := dExt.cacheStepData
for key := range dExt.stepScreenShots { // clear cache
delete(dExt.stepScreenShots, key) dExt.cacheStepData = cacheStepData{
} ScreenShots: []string{},
}() OcrResults: make(map[string]string),
copied := make(map[string]string)
for key, value := range dExt.stepScreenShots {
copied[key] = value
} }
return copied return copied
} }
+1 -1
View File
@@ -292,7 +292,7 @@ func (dExt *DriverExt) GetScreenTextsByOCR() (texts OCRTexts, err error) {
} }
o, _ := json.Marshal(ocrTexts) o, _ := json.Marshal(ocrTexts)
dExt.stepScreenShots[imagePath] = string(o) dExt.cacheStepData.OcrResults[imagePath] = string(o)
return ocrTexts, nil return ocrTexts, nil
} }
+3 -1
View File
@@ -616,7 +616,9 @@ func runStepMobileUI(s *SessionRunner, step *TStep) (stepResult *StepResult, err
} }
// save attachments // save attachments
attachments["screenshots"] = uiDriver.GetScreenShots() cacheData := uiDriver.GetStepCacheData()
attachments["screenshots"] = cacheData.ScreenShots
attachments["ocr_results"] = cacheData.OcrResults
stepResult.Attachments = attachments stepResult.Attachments = attachments
}() }()