From 5be8c09402256ce93b2ed027832df2336b0f5fab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BD=99=E6=B3=93=E9=93=AE?= <yuhongzheng@bytedance.com>
Date: Tue, 25 Feb 2025 17:54:25 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E5=8E=8B=E7=BC=A9?=
 =?UTF-8?q?=E6=96=B9=E6=B3=95=E5=A4=B1=E6=95=88=E5=92=8C=E5=8E=8B=E7=BC=A9?=
 =?UTF-8?q?=E6=85=A2=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/uixt/driver_ext_screenshot.go | 46 ++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/pkg/uixt/driver_ext_screenshot.go b/pkg/uixt/driver_ext_screenshot.go
index 9ac9d31c..b8d99944 100644
--- a/pkg/uixt/driver_ext_screenshot.go
+++ b/pkg/uixt/driver_ext_screenshot.go
@@ -61,6 +61,7 @@ func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult
 	}
 
 	var bufSource *bytes.Buffer
+	var compressBufSource *bytes.Buffer
 	var imageResult *ai.CVResult
 	var imagePath string
 	var windowSize types.Size
@@ -69,10 +70,14 @@ func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult
 	// get screenshot info with retry
 	for i := 0; i < 3; i++ {
 		imagePath = filepath.Join(config.ScreenShotsPath, fileName)
-		bufSource, err = dExt.ScreenShot(option.WithScreenShotFileName(imagePath))
+		bufSource, err = dExt.ScreenShot()
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		compressBufSource, err = compressImageBuffer(bufSource)
 		if err != nil {
 			lastErr = err
-			time.Sleep(time.Second * 1)
 			continue
 		}
 
@@ -83,12 +88,12 @@ func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult
 		}
 
 		screenResult = &ScreenResult{
-			bufSource:  bufSource,
+			bufSource:  compressBufSource,
 			ImagePath:  imagePath,
 			Tags:       nil,
 			Resolution: windowSize,
 		}
-		imageResult, err = dExt.CVService.ReadFromBuffer(bufSource, opts...)
+		imageResult, err = dExt.CVService.ReadFromBuffer(compressBufSource, opts...)
 		if err != nil {
 			log.Error().Err(err).Msg("ReadFromBuffer from ImageService failed")
 			lastErr = err
@@ -222,7 +227,11 @@ func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
 		encoder := png.Encoder{
 			CompressionLevel: png.BestCompression,
 		}
+		time1 := time.Now()
+		log.Info().Str("time1", time1.String()).Msg("png encode start")
 		err = encoder.Encode(file, img)
+		time2 := time.Now()
+		log.Info().Int64("time", time2.Sub(time1).Milliseconds()).Msg("png encode time")
 	case "gif":
 		gifOptions := &gif.Options{
 			NumColors: 256,
@@ -246,3 +255,32 @@ func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
 
 	return screenshotPath, nil
 }
+
+func compressImageBuffer(raw *bytes.Buffer) (compressed *bytes.Buffer, err error) {
+	// 解码原始图像数据
+	img, format, err := image.Decode(raw)
+	if err != nil {
+		return nil, err
+	}
+
+	// 创建一个用来保存压缩后数据的buffer
+	var buf bytes.Buffer
+
+	switch format {
+	// Convert to jpeg uniformly and compress with a compression rate of 95
+	case "jpeg", "png":
+		jpegOptions := &jpeg.Options{Quality: 60}
+		time1 := time.Now()
+		err = jpeg.Encode(&buf, img, jpegOptions)
+		if err != nil {
+			return nil, err
+		}
+		time2 := time.Now()
+		log.Info().Int64("time", time2.Sub(time1).Milliseconds()).Msg("jpeg encode time")
+	default:
+		return nil, fmt.Errorf("unsupported image format: %s", format)
+	}
+
+	// 返回压缩后的图像数据
+	return &buf, nil
+}