golang 为图片加水印

前段时间一个需求涉及到给图片加水印，考虑图片安全性，决定放在后端加水印。记录一下代码。

思路

思路是先为水印文字生成一个仅包含水印文字的图片，把这个图片倾斜一定角度 (一般水印都是倾斜的)，之后把倾斜的水印文字图片贴在原图上，得到最终的水印图片。

代码

// watermark.go
package main

import (
	"image"
	"image/color"
	"image/draw"
	"github.com/disintegration/imaging"
	"github.com/golang/freetype"
	"github.com/golang/freetype/truetype"
	"golang.org/x/image/font"
	"golang.org/x/image/font/gofont/gomono"
	"github.com/pkg/errors"
)

func AddWatermarkForImage(oriImage image.Image, uid string) (*image.RGBA, error) {
	watermarkedImage := image.NewRGBA(oriImage.Bounds())
	draw.Draw(watermarkedImage, oriImage.Bounds(), oriImage, image.Point{}, draw.Src)

	watermark, err := MakeImageByText(uid, color.Transparent)
	if err != nil {
		return nil, err
	}
	rotatedWatermark := imaging.Rotate(watermark, 30, color.Transparent)

	x, y := 0, 0
	for y <= watermarkedImage.Bounds().Max.Y {
		for x <= watermarkedImage.Bounds().Max.X {
			offset := image.Pt(x, y)
			draw.Draw(watermarkedImage, rotatedWatermark.Bounds().Add(offset), rotatedWatermark, image.Point{}, draw.Over)
			// 稀疏一点, 稍微提升点速度
			x += rotatedWatermark.Bounds().Dx() * 2
		}
		y += rotatedWatermark.Bounds().Dy()
		x = 0
	}
	return watermarkedImage, nil
}

// MakeImageByText 根据文本内容制作一个仅包含该文本内容的图片
func MakeImageByText(text string, bgColor color.Color) (image.Image, error) {
	fontSize := float64(72)
	freetypeCtx := MakeFreetypeCtx(fontSize)

	width, height := int(fontSize)*len(text), int(fontSize)*2
	rgbaRect := image.NewRGBA(image.Rect(0, 0, width, height))

	// 仅当非透明时才做一次额外的渲染
	if bgColor != color.Transparent {
		bgUniform := image.NewUniform(bgColor)
		draw.Draw(rgbaRect, rgbaRect.Bounds(), bgUniform, image.Pt(0, 0), draw.Src)
	}

	freetypeCtx.SetClip(rgbaRect.Rect)
	freetypeCtx.SetDst(rgbaRect)
	pt := freetype.Pt(0, int(freetypeCtx.PointToFixed(fontSize)>>6))
	_, err := freetypeCtx.DrawString(text, pt)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return rgbaRect, nil
}

// MustParseFont 通过单测来保证该方法必不会 panic
func MustParseFont() *truetype.Font {
	ft, err := freetype.ParseFont(gomono.TTF)
	if err != nil {
		panic(err)
	}
	return ft
}

func MakeFreetypeCtx(fontSize float64) *freetype.Context {
	fontColor := color.RGBA{R: 0, G: 0, B: 0, A: 50}
	fontColorUniform := image.NewUniform(fontColor)

	freetypeCtx := freetype.NewContext()
	freetypeCtx.SetDPI(100)
	freetypeCtx.SetFont(MustParseFont())
	freetypeCtx.SetFontSize(fontSize)
	freetypeCtx.SetSrc(fontColorUniform)
	freetypeCtx.SetHinting(font.HintingNone)
	return freetypeCtx
}

// watermark_test.go
package main

import (
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	"image/png"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestMakeImageByText(t *testing.T) {
	t.Run("bg white", func(t *testing.T) {
		img, err := MakeImageByText("hello", color.White)
		assert.NoError(t, err)

		helloPng, err := os.Create("hello_white.png")
		assert.NoError(t, err)
		defer helloPng.Close()

		err = png.Encode(helloPng, img)
		assert.NoError(t, err)

		helloJpeg, err := os.Create("hello_white.jpeg")
		assert.NoError(t, err)
		defer helloJpeg.Close()

		err = jpeg.Encode(helloJpeg, img, nil)
		assert.NoError(t, err)
	})
	t.Run("bg transparent", func(t *testing.T) {
		img, err := MakeImageByText("hello", color.Transparent)
		assert.NoError(t, err)

		helloPng, err := os.Create("hello_transparent.png")
		assert.NoError(t, err)
		defer helloPng.Close()

		err = png.Encode(helloPng, img)
		assert.NoError(t, err)

		helloJpeg, err := os.Create("hello_transparent.jpeg")
		assert.NoError(t, err)
		defer helloJpeg.Close()

		// jpeg 没有 alpha 通道, 所以会是全黑的
		err = jpeg.Encode(helloJpeg, img, nil)
		assert.NoError(t, err)
	})
}

func TestMustParseFont(t *testing.T) {
	ft := MustParseFont()
	assert.NotNil(t, ft)
}

func BaseImageForTest() image.Image {
	rgbaRect := image.NewRGBA(image.Rect(0, 0, 3000, 2000))
	bgColor := color.RGBA{R: 0xff, G: 0, B: 0, A: 0xff}
	bg := image.NewUniform(bgColor)
	draw.Draw(rgbaRect, rgbaRect.Bounds(), bg, image.Pt(0, 0), draw.Src)
	return rgbaRect
}

func TestWassObject_AddWatermark(t *testing.T) {
	baseImg := BaseImageForTest()

	watermarkedImg, err := AddWatermarkForImage(baseImg, "hello.world")
	assert.NoError(t, err)

	helloWatermarkedJpeg, err := os.Create("hello_watermarked.jpeg")
	assert.NoError(t, err)
	defer helloWatermarkedJpeg.Close()

	err = jpeg.Encode(helloWatermarkedJpeg, watermarkedImg, nil)
	assert.NoError(t, err)
}

效果

坑

1. 最开始的思路是，计算出原图的对角线长度，制作出一个长宽均为对角线长度的 mask，把水印文字填充在 mask 上，旋转 mask，再把旋转后的 mask 盖在原图上。后发现因为旋转的是一个大图，所以旋转的耗时比较久，对于一些比较大的原始图片，旋转可能花个五六秒 (2核4G的机器)。因此改为了只旋转那个比较小的水印文字图。
2. 对图片的处理本质是一个矩阵运算，因此计算量还是比较大的，cpu 和内存的占用量会比较大，功能上线前最好做一下压测

参考链接

https://johnpili.com/generate-text-to-image-in-go/

https://blog.csdn.net/diandianxiyu_geek/article/details/119546482