-
-
Notifications
You must be signed in to change notification settings - Fork 328
/
imageutil.go
134 lines (111 loc) · 2.96 KB
/
imageutil.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package utils
import (
"bytes"
"fmt"
"image"
"image/jpeg"
"image/png"
"io"
"github.com/go-rod/rod/lib/proto"
)
// ImgWithBox is a image with a box, if the box is nil, it means the whole image.
type ImgWithBox struct {
Img []byte
Box *image.Rectangle
}
// ImgOption is the option for image processing.
type ImgOption struct {
Quality int
}
// ImgProcessor is the interface for image processing.
type ImgProcessor interface {
Encode(img image.Image, opt *ImgOption) ([]byte, error)
Decode(file io.Reader) (image.Image, error)
}
type jpegProcessor struct{}
func (p jpegProcessor) Encode(img image.Image, opt *ImgOption) ([]byte, error) {
var buf bytes.Buffer
var jpegOpt *jpeg.Options
if opt != nil {
jpegOpt = &jpeg.Options{Quality: opt.Quality}
}
err := jpeg.Encode(&buf, img, jpegOpt)
return buf.Bytes(), err
}
func (p jpegProcessor) Decode(file io.Reader) (image.Image, error) {
return jpeg.Decode(file)
}
type pngProcessor struct{}
func (p pngProcessor) Encode(img image.Image, _ *ImgOption) ([]byte, error) {
var buf bytes.Buffer
err := png.Encode(&buf, img)
return buf.Bytes(), err
}
func (p pngProcessor) Decode(file io.Reader) (image.Image, error) {
return png.Decode(file)
}
// NewImgProcessor create a ImgProcessor by the format.
func NewImgProcessor(format proto.PageCaptureScreenshotFormat) (ImgProcessor, error) {
switch format {
case proto.PageCaptureScreenshotFormatJpeg:
return &jpegProcessor{}, nil
case "", proto.PageCaptureScreenshotFormatPng:
return &pngProcessor{}, nil
default:
return nil, fmt.Errorf("not support format: %v", format)
}
}
// SplicePngVertical splice png vertically, if there is only one image, it will return the image directly.
// Only support png and jpeg format yet, webP is not supported because no suitable processing
// library was found in golang.
func SplicePngVertical(files []ImgWithBox, format proto.PageCaptureScreenshotFormat, opt *ImgOption) ([]byte, error) {
if len(files) == 0 {
return nil, nil
}
if len(files) == 1 {
return files[0].Img, nil
}
var width, height int
processor, err := NewImgProcessor(format)
if err != nil {
return nil, err
}
var images []image.Image
for _, file := range files {
img, err := processor.Decode(bytes.NewReader(file.Img))
if err != nil {
return nil, err
}
images = append(images, img)
if file.Box != nil {
width = file.Box.Dx()
height += file.Box.Dy()
} else {
width = img.Bounds().Dx()
height += img.Bounds().Dy()
}
}
spliceImg := image.NewRGBA(image.Rect(0, 0, width, height))
var destY int
for i, file := range files {
img := images[i]
bounds := img.Bounds()
if file.Box != nil {
bounds = *file.Box
}
start := bounds.Min
end := bounds.Max
for y := start.Y; y < end.Y; y++ {
for x := start.X; x < end.X; x++ {
color := img.At(x, y)
spliceImg.Set(x, y-start.Y+destY, color)
}
}
destY += bounds.Dy()
}
bs, err := processor.Encode(spliceImg, opt)
if err != nil {
return nil, err
}
return bs, nil
}