refactor: move LoadImage

This commit is contained in:
lilong.129
2025-04-30 16:21:01 +08:00
parent fcddcfb630
commit 6569121d5d
7 changed files with 156 additions and 160 deletions
+43
View File
@@ -6,9 +6,13 @@ import (
"context" "context"
"crypto/hmac" "crypto/hmac"
"crypto/sha256" "crypto/sha256"
"encoding/base64"
"encoding/csv" "encoding/csv"
builtinJSON "encoding/json" builtinJSON "encoding/json"
"fmt" "fmt"
"image"
"image/jpeg"
"image/png"
"math" "math"
"math/rand" "math/rand"
"net" "net"
@@ -26,6 +30,7 @@ import (
"github.com/httprunner/httprunner/v5/code" "github.com/httprunner/httprunner/v5/code"
"github.com/httprunner/httprunner/v5/internal/json" "github.com/httprunner/httprunner/v5/internal/json"
"github.com/httprunner/httprunner/v5/uixt/types"
) )
func Dump2JSON(data interface{}, path string) error { func Dump2JSON(data interface{}, path string) error {
@@ -484,3 +489,41 @@ func RunCommandWithCallback(cmdName string, args []string, callback LineCallback
return err return err
} }
} }
// LoadImage loads image file and returns base64 encoded string and image size
func LoadImage(imagePath string) (base64Str string, size types.Size, err error) {
// Read the image file
imageFile, err := os.OpenFile(imagePath, os.O_RDONLY, 0o600)
if err != nil {
return "", types.Size{}, fmt.Errorf("failed to open image file: %w", err)
}
defer imageFile.Close()
// Decode the image to get its resolution
imageData, format, err := image.Decode(imageFile)
if err != nil {
return "", types.Size{}, fmt.Errorf("failed to decode image: %w", err)
}
// Get the resolution of the image
width := imageData.Bounds().Dx()
height := imageData.Bounds().Dy()
size = types.Size{Width: width, Height: height}
// Convert image to base64
buf := new(bytes.Buffer)
if format == "jpeg" || format == "jpg" {
if err := jpeg.Encode(buf, imageData, nil); err != nil {
return "", types.Size{}, fmt.Errorf("failed to encode image to buffer: %w", err)
}
} else {
// default use png
if err := png.Encode(buf, imageData); err != nil {
return "", types.Size{}, fmt.Errorf("failed to encode image to buffer: %w", err)
}
}
base64Str = fmt.Sprintf("data:image/%s;base64,%s", format,
base64.StdEncoding.EncodeToString(buf.Bytes()))
return base64Str, size, nil
}
+1 -1
View File
@@ -1 +1 @@
v5.0.0-beta-2504301521 v5.0.0-beta-2504301621
+2 -1
View File
@@ -4,6 +4,7 @@ import (
"context" "context"
"testing" "testing"
"github.com/httprunner/httprunner/v5/internal/builtin"
"github.com/httprunner/httprunner/v5/uixt/option" "github.com/httprunner/httprunner/v5/uixt/option"
"github.com/httprunner/httprunner/v5/uixt/types" "github.com/httprunner/httprunner/v5/uixt/types"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
@@ -50,7 +51,7 @@ func TestValidAssertions(t *testing.T) {
for _, tc := range testCases { for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) { t.Run(tc.name, func(t *testing.T) {
imageBase64, size, err := loadImage(tc.imagePath) imageBase64, size, err := builtin.LoadImage(tc.imagePath)
require.NoError(t, err) require.NoError(t, err)
result, err := asserter.Assert(&AssertOptions{ result, err := asserter.Assert(&AssertOptions{
+8 -83
View File
@@ -1,18 +1,12 @@
package ai package ai
import ( import (
"bytes"
"context" "context"
"encoding/base64"
"fmt"
"image"
"image/jpeg"
"image/png"
"os"
"testing" "testing"
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
"github.com/httprunner/httprunner/v5/code" "github.com/httprunner/httprunner/v5/code"
"github.com/httprunner/httprunner/v5/internal/builtin"
"github.com/httprunner/httprunner/v5/uixt/option" "github.com/httprunner/httprunner/v5/uixt/option"
"github.com/httprunner/httprunner/v5/uixt/types" "github.com/httprunner/httprunner/v5/uixt/types"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
@@ -20,7 +14,7 @@ import (
) )
func TestVLMPlanning(t *testing.T) { func TestVLMPlanning(t *testing.T) {
imageBase64, size, err := loadImage("testdata/llk_1.png") imageBase64, size, err := builtin.LoadImage("testdata/llk_1.png")
require.NoError(t, err) require.NoError(t, err)
userInstruction := `连连看是一款经典的益智消除类小游戏,通常以图案或图标为主要元素。以下是连连看的基本规则说明: userInstruction := `连连看是一款经典的益智消除类小游戏,通常以图案或图标为主要元素。以下是连连看的基本规则说明:
@@ -104,7 +98,7 @@ func TestVLMPlanning(t *testing.T) {
} }
func TestXHSPlanning(t *testing.T) { func TestXHSPlanning(t *testing.T) {
imageBase64, size, err := loadImage("testdata/xhs-feed.jpeg") imageBase64, size, err := builtin.LoadImage("testdata/xhs-feed.jpeg")
require.NoError(t, err) require.NoError(t, err)
userInstruction := "点击第二个帖子的作者头像" userInstruction := "点击第二个帖子的作者头像"
@@ -177,7 +171,7 @@ func TestXHSPlanning(t *testing.T) {
} }
func TestChatList(t *testing.T) { func TestChatList(t *testing.T) {
imageBase64, size, err := loadImage("testdata/chat_list.jpeg") imageBase64, size, err := builtin.LoadImage("testdata/chat_list.jpeg")
require.NoError(t, err) require.NoError(t, err)
userInstruction := "请结合图片的文字信息,请告诉我一共有多少个群聊,哪些群聊右下角有绿点" userInstruction := "请结合图片的文字信息,请告诉我一共有多少个群聊,哪些群聊右下角有绿点"
@@ -232,7 +226,7 @@ func TestHandleSwitch(t *testing.T) {
} }
for _, tc := range testCases { for _, tc := range testCases {
imageBase64, size, err := loadImage(tc.imageFile) imageBase64, size, err := builtin.LoadImage(tc.imageFile)
require.NoError(t, err) require.NoError(t, err)
opts := &PlanningOptions{ opts := &PlanningOptions{
@@ -263,7 +257,7 @@ func TestHandleSwitch(t *testing.T) {
} }
func TestValidateInput(t *testing.T) { func TestValidateInput(t *testing.T) {
imageBase64, size, err := loadImage("testdata/popup_risk_warning.png") imageBase64, size, err := builtin.LoadImage("testdata/popup_risk_warning.png")
require.NoError(t, err) require.NoError(t, err)
tests := []struct { tests := []struct {
@@ -388,87 +382,18 @@ func TestProcessVLMResponse(t *testing.T) {
} }
} }
func TestSavePositionImg(t *testing.T) {
imageBase64, _, err := loadImage("testdata/popup_risk_warning.png")
require.NoError(t, err)
params := struct {
InputImgBase64 string
Rect struct {
X float64
Y float64
}
OutputPath string
}{
InputImgBase64: imageBase64,
Rect: struct {
X float64
Y float64
}{
X: 100,
Y: 100,
},
OutputPath: "testdata/output.png",
}
err = SavePositionImg(params)
require.NoError(t, err)
// cleanup
defer os.Remove(params.OutputPath)
}
func TestLoadImage(t *testing.T) { func TestLoadImage(t *testing.T) {
// Test PNG image // Test PNG image
pngBase64, pngSize, err := loadImage("testdata/llk_1.png") pngBase64, pngSize, err := builtin.LoadImage("testdata/llk_1.png")
require.NoError(t, err) require.NoError(t, err)
assert.NotEmpty(t, pngBase64) assert.NotEmpty(t, pngBase64)
assert.Greater(t, pngSize.Width, 0) assert.Greater(t, pngSize.Width, 0)
assert.Greater(t, pngSize.Height, 0) assert.Greater(t, pngSize.Height, 0)
// Test JPEG image // Test JPEG image
jpegBase64, jpegSize, err := loadImage("testdata/xhs-feed.jpeg") jpegBase64, jpegSize, err := builtin.LoadImage("testdata/xhs-feed.jpeg")
require.NoError(t, err) require.NoError(t, err)
assert.NotEmpty(t, jpegBase64) assert.NotEmpty(t, jpegBase64)
assert.Greater(t, jpegSize.Width, 0) assert.Greater(t, jpegSize.Width, 0)
assert.Greater(t, jpegSize.Height, 0) assert.Greater(t, jpegSize.Height, 0)
} }
// loadImage loads image and returns base64 encoded string
func loadImage(imagePath string) (base64Str string, size types.Size, err error) {
// Read the image file
imageFile, err := os.OpenFile(imagePath, os.O_RDONLY, 0o600)
if err != nil {
return "", types.Size{}, fmt.Errorf("failed to open image file: %w", err)
}
defer imageFile.Close()
// Decode the image to get its resolution
imageData, format, err := image.Decode(imageFile)
if err != nil {
return "", types.Size{}, fmt.Errorf("failed to decode image: %w", err)
}
// Get the resolution of the image
width := imageData.Bounds().Dx()
height := imageData.Bounds().Dy()
size = types.Size{Width: width, Height: height}
// Convert image to base64
buf := new(bytes.Buffer)
// 根据图像格式选择正确的编码器
if format == "jpeg" || format == "jpg" {
if err := jpeg.Encode(buf, imageData, nil); err != nil {
return "", types.Size{}, fmt.Errorf("failed to encode image to buffer: %w", err)
}
} else {
// 默认使用 PNG 编码
if err := png.Encode(buf, imageData); err != nil {
return "", types.Size{}, fmt.Errorf("failed to encode image to buffer: %w", err)
}
}
base64Str = fmt.Sprintf("data:image/%s;base64,%s", format,
base64.StdEncoding.EncodeToString(buf.Bytes()))
return base64Str, size, nil
}
-75
View File
@@ -1,14 +1,6 @@
package ai package ai
import ( import (
"bytes"
"encoding/base64"
"fmt"
"image"
"image/color"
"image/draw"
"image/png"
"os"
"strings" "strings"
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
@@ -109,70 +101,3 @@ func logResponse(resp *schema.Message) {
} }
logger.Msg("log response message") logger.Msg("log response message")
} }
// SavePositionImg saves an image with position markers
func SavePositionImg(params struct {
InputImgBase64 string
Rect struct {
X float64
Y float64
}
OutputPath string
}) error {
// 解码Base64图像
imgData := params.InputImgBase64
// 如果包含了数据URL前缀,去掉它
if strings.HasPrefix(imgData, "data:image/") {
parts := strings.Split(imgData, ",")
if len(parts) > 1 {
imgData = parts[1]
}
}
// 解码Base64
unbased, err := base64.StdEncoding.DecodeString(imgData)
if err != nil {
return fmt.Errorf("无法解码Base64图像: %w", err)
}
// 解码图像
reader := bytes.NewReader(unbased)
img, _, err := image.Decode(reader)
if err != nil {
return fmt.Errorf("无法解码图像数据: %w", err)
}
// 创建一个可以在其上绘制的图像
bounds := img.Bounds()
rgba := image.NewRGBA(bounds)
draw.Draw(rgba, bounds, img, bounds.Min, draw.Src)
// 在点击/拖动位置绘制标记
markRadius := 10
x, y := int(params.Rect.X), int(params.Rect.Y)
// 绘制红色圆圈
for i := -markRadius; i <= markRadius; i++ {
for j := -markRadius; j <= markRadius; j++ {
if i*i+j*j <= markRadius*markRadius {
if x+i >= 0 && x+i < bounds.Max.X && y+j >= 0 && y+j < bounds.Max.Y {
rgba.Set(x+i, y+j, color.RGBA{255, 0, 0, 255})
}
}
}
}
// 保存图像
outFile, err := os.Create(params.OutputPath)
if err != nil {
return fmt.Errorf("无法创建输出文件: %w", err)
}
defer outFile.Close()
// 编码为PNG并保存
if err := png.Encode(outFile, rgba); err != nil {
return fmt.Errorf("无法编码和保存图像: %w", err)
}
return nil
}
+70
View File
@@ -2,8 +2,11 @@ package uixt
import ( import (
"bytes" "bytes"
"encoding/base64"
"fmt" "fmt"
"image" "image"
"image/color"
"image/draw"
"image/gif" "image/gif"
"image/jpeg" "image/jpeg"
"image/png" "image/png"
@@ -294,3 +297,70 @@ func compressImageBuffer(raw *bytes.Buffer) (compressed *bytes.Buffer, err error
// return compressed image buffer // return compressed image buffer
return &buf, nil return &buf, nil
} }
// SavePositionImg saves an image with position markers
func SavePositionImg(params struct {
InputImgBase64 string
Rect struct {
X float64
Y float64
}
OutputPath string
}) error {
// 解码Base64图像
imgData := params.InputImgBase64
// 如果包含了数据URL前缀,去掉它
if strings.HasPrefix(imgData, "data:image/") {
parts := strings.Split(imgData, ",")
if len(parts) > 1 {
imgData = parts[1]
}
}
// 解码Base64
unbased, err := base64.StdEncoding.DecodeString(imgData)
if err != nil {
return fmt.Errorf("无法解码Base64图像: %w", err)
}
// 解码图像
reader := bytes.NewReader(unbased)
img, _, err := image.Decode(reader)
if err != nil {
return fmt.Errorf("无法解码图像数据: %w", err)
}
// 创建一个可以在其上绘制的图像
bounds := img.Bounds()
rgba := image.NewRGBA(bounds)
draw.Draw(rgba, bounds, img, bounds.Min, draw.Src)
// 在点击/拖动位置绘制标记
markRadius := 30
x, y := int(params.Rect.X), int(params.Rect.Y)
// 绘制红色圆圈
for i := -markRadius; i <= markRadius; i++ {
for j := -markRadius; j <= markRadius; j++ {
if i*i+j*j <= markRadius*markRadius {
if x+i >= 0 && x+i < bounds.Max.X && y+j >= 0 && y+j < bounds.Max.Y {
rgba.Set(x+i, y+j, color.RGBA{255, 0, 0, 255})
}
}
}
}
// 保存图像
outFile, err := os.Create(params.OutputPath)
if err != nil {
return fmt.Errorf("无法创建输出文件: %w", err)
}
defer outFile.Close()
// 编码为PNG并保存
if err := png.Encode(outFile, rgba); err != nil {
return fmt.Errorf("无法编码和保存图像: %w", err)
}
return nil
}
+32
View File
@@ -3,9 +3,11 @@
package uixt package uixt
import ( import (
"os"
"testing" "testing"
"time" "time"
"github.com/httprunner/httprunner/v5/internal/builtin"
"github.com/httprunner/httprunner/v5/uixt/option" "github.com/httprunner/httprunner/v5/uixt/option"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
@@ -251,3 +253,33 @@ func TestDriverExt_Action_Offset(t *testing.T) {
option.WithTapRandomRect(true)) option.WithTapRandomRect(true))
assert.Nil(t, err) assert.Nil(t, err)
} }
func TestSavePositionImg(t *testing.T) {
imageBase64, _, err := builtin.LoadImage("ai/testdata/popup_risk_warning.png")
require.NoError(t, err)
params := struct {
InputImgBase64 string
Rect struct {
X float64
Y float64
}
OutputPath string
}{
InputImgBase64: imageBase64,
Rect: struct {
X float64
Y float64
}{
X: 500,
Y: 500,
},
OutputPath: "ai/testdata/output.png",
}
err = SavePositionImg(params)
require.NoError(t, err)
// cleanup
defer os.Remove(params.OutputPath)
}