From 87f9c6808e92e9cc84470450a3b0b50bbc203e74 Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Sat, 14 Sep 2024 11:21:43 +0800
Subject: [PATCH] fix: compress screenshot image file

---
 hrp/pkg/uixt/screenshot.go | 73 +++++++++++++++-----------------------
 1 file changed, 29 insertions(+), 44 deletions(-)

diff --git a/hrp/pkg/uixt/screenshot.go b/hrp/pkg/uixt/screenshot.go
index d7ebb69f..607cfc2f 100644
--- a/hrp/pkg/uixt/screenshot.go
+++ b/hrp/pkg/uixt/screenshot.go
@@ -4,9 +4,9 @@ import (
 	"bytes"
 	"fmt"
 	"image"
-	_ "image/gif"
+	"image/gif"
 	"image/jpeg"
-	_ "image/png"
+	"image/png"
 	"os"
 	"path/filepath"
 	"strings"
@@ -180,52 +180,18 @@ func (dExt *DriverExt) GetScreenShot(fileName string) (raw *bytes.Buffer, path s
 		return nil, "", errors.Wrap(code.DeviceScreenShotError, err.Error())
 	}
 
-	// compress image data
-	compressed, err := compressImageBuffer(raw)
-	if err != nil {
-		log.Error().Err(err).Msg("compress screenshot data failed")
-		return nil, "", errors.Wrap(code.DeviceScreenShotError,
-			fmt.Sprintf("compress screenshot data failed: %s", err.Error()))
-	}
-
 	// save screenshot to file
 	path = filepath.Join(env.ScreenShotsPath, fileName)
-	path, err = saveScreenShot(compressed, path)
+	path, err = saveScreenShot(raw, path)
 	if err != nil {
 		log.Error().Err(err).Msg("save screenshot file failed")
 		return nil, "", errors.Wrap(code.DeviceScreenShotError,
 			fmt.Sprintf("save screenshot file failed: %s", err.Error()))
 	}
-	return compressed, path, nil
+	return raw, path, 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: 95}
-		err = jpeg.Encode(&buf, img, jpegOptions)
-		if err != nil {
-			return nil, err
-		}
-	default:
-		return nil, fmt.Errorf("unsupported image format: %s", format)
-	}
-
-	// 返回压缩后的图像数据
-	return &buf, nil
-}
-
-// saveScreenShot saves image file with file name
+// saveScreenShot saves compressed image file with file name
 func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
 	// notice: screenshot data is a stream, so we need to copy it to a new buffer
 	copiedBuffer := &bytes.Buffer{}
@@ -248,17 +214,36 @@ func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
 		_ = file.Close()
 	}()
 
+	// compress image and save to file
 	switch format {
-	case "jpeg", "png":
-		jpegOptions := &jpeg.Options{}
+	case "jpeg":
+		jpegOptions := &jpeg.Options{Quality: 90}
 		err = jpeg.Encode(file, img, jpegOptions)
+	case "png":
+		encoder := png.Encoder{
+			CompressionLevel: png.BestCompression,
+		}
+		err = encoder.Encode(file, img)
+	case "gif":
+		gifOptions := &gif.Options{
+			NumColors: 256,
+		}
+		err = gif.Encode(file, img, gifOptions)
 	default:
-		return "", fmt.Errorf("unsupported image format: %s", format)
+		return "", fmt.Errorf("unsupported image format %s", format)
 	}
 	if err != nil {
-		return "", errors.Wrap(err, "encode screenshot image failed")
+		return "", errors.Wrap(err, "save image file failed")
 	}
 
-	log.Info().Str("path", screenshotPath).Msg("save screenshot file success")
+	var fileSize int64
+	fileInfo, err := file.Stat()
+	if err == nil {
+		fileSize = fileInfo.Size()
+	}
+	log.Info().Str("path", screenshotPath).
+		Int("rawBytes", raw.Len()).Int64("saveBytes", fileSize).
+		Msg("save screenshot file success")
+
 	return screenshotPath, nil
 }