diff --git a/example/default/main.go b/example/default/main.go index 497601c..90a0eb2 100644 --- a/example/default/main.go +++ b/example/default/main.go @@ -16,6 +16,7 @@ import ( _ "m7s.live/v5/plugin/rtmp" _ "m7s.live/v5/plugin/rtsp" _ "m7s.live/v5/plugin/sei" + _ "m7s.live/v5/plugin/snap" _ "m7s.live/v5/plugin/srt" _ "m7s.live/v5/plugin/stress" _ "m7s.live/v5/plugin/transcode" diff --git a/go.mod b/go.mod index 456da7b..f44941f 100644 --- a/go.mod +++ b/go.mod @@ -11,9 +11,11 @@ require ( github.com/cilium/ebpf v0.15.0 github.com/cloudwego/goref v0.0.0-20240724113447-685d2a9523c8 github.com/deepch/vdk v0.0.27 + github.com/disintegration/imaging v1.6.2 github.com/emiago/sipgo v0.22.0 github.com/go-delve/delve v1.23.0 github.com/gobwas/ws v1.3.2 + github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 github.com/google/gopacket v1.1.19 github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1 github.com/husanpao/ip v0.0.0-20220711082147-73160bb611a8 @@ -34,7 +36,8 @@ require ( github.com/shirou/gopsutil/v4 v4.24.8 github.com/vishvananda/netlink v1.1.0 github.com/yapingcat/gomedia v0.0.0-20240601043430-920523f8e5c7 - golang.org/x/text v0.17.0 + golang.org/x/image v0.22.0 + golang.org/x/text v0.20.0 google.golang.org/genproto/googleapis/api v0.0.0-20240711142825-46eb208f015d google.golang.org/grpc v1.65.0 google.golang.org/protobuf v1.34.2 @@ -114,7 +117,7 @@ require ( github.com/vishvananda/netns v0.0.0-20191106174202-0a2b9b5464df // indirect github.com/yusufpapurcu/wmi v1.2.4 // indirect golang.org/x/arch v0.6.0 // indirect - golang.org/x/sync v0.8.0 // indirect + golang.org/x/sync v0.9.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20240711142825-46eb208f015d // indirect ) @@ -132,7 +135,7 @@ require ( github.com/quangngotan95/go-m3u8 v0.1.0 go.uber.org/mock v0.4.0 // indirect golang.org/x/crypto v0.26.0 // indirect - golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7 // indirect + golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7 golang.org/x/mod v0.19.0 // indirect golang.org/x/net v0.27.0 golang.org/x/sys v0.25.0 diff --git a/go.sum b/go.sum index a5e03f6..c000afe 100644 --- a/go.sum +++ b/go.sum @@ -59,6 +59,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/deepch/vdk v0.0.27 h1:j/SHaTiZhA47wRpaue8NRp7P9xwOOO/lunxrDJBwcao= github.com/deepch/vdk v0.0.27/go.mod h1:JlgGyR2ld6+xOIHa7XAxJh+stSDBAkdNvIPkUIdIywk= +github.com/disintegration/imaging v1.6.2 h1:w1LecBlG2Lnp8B3jk5zSuNqd7b4DXhcjwek1ei82L+c= +github.com/disintegration/imaging v1.6.2/go.mod h1:44/5580QXChDfwIclfc/PCwrr44amcmDAg8hxG0Ewe4= github.com/emiago/sipgo v0.22.0 h1:GaQ51m26M9QnVBVY2aDJ/mXqq/BDfZ1A+nW7XgU/4Ts= github.com/emiago/sipgo v0.22.0/go.mod h1:a77FgPEEjJvfYWYfP3p53u+dNhWEMb/VGVS6guvBzx0= github.com/frankban/quicktest v1.11.3/go.mod h1:wRf/ReqHper53s+kmmSZizM8NamnL3IM0I9ntUbOk+k= @@ -86,6 +88,8 @@ github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6Wezm github.com/gobwas/ws v1.3.2 h1:zlnbNHxumkRvfPWgfXu8RBwyNR1x8wh9cf5PTOCqs9Q= github.com/gobwas/ws v1.3.2/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY= github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= +github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 h1:DACJavvAHhabrF08vX0COfcOBJRhZ8lUbR+ZWIs0Y5g= +github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= @@ -365,6 +369,9 @@ golang.org/x/crypto v0.26.0 h1:RrRspgV4mU+YwB4FYnuBoKsUapNIL5cohGAmSH3azsw= golang.org/x/crypto v0.26.0/go.mod h1:GY7jblb9wI+FOo5y8/S2oY4zWP07AkOJ4+jxCqdqn54= golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7 h1:wDLEX9a7YQoKdKNQt88rtydkqDxeGaBUTnIYc3iG/mA= golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7/go.mod h1:M4RDyNAINzryxdtnbRXRL/OHtkFuWGRjvuhBJpk2IlY= +golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= +golang.org/x/image v0.22.0 h1:UtK5yLUzilVrkjMAZAZ34DXGpASN8i8pj8g+O+yd10g= +golang.org/x/image v0.22.0/go.mod h1:9hPFhljd4zZ1GNSIZJ49sqbp45GKK9t6w+iXvGqZUz4= golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= @@ -400,8 +407,8 @@ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ= -golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/sync v0.9.0 h1:fEo0HyrW1GIgZdpbhCRO0PkJajUS5H9IFUztCgEo2jQ= +golang.org/x/sync v0.9.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -468,8 +475,8 @@ golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.10.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/text v0.12.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= -golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc= -golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= +golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug= +golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4= golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk= golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/plugin/snap/api.go b/plugin/snap/api.go new file mode 100755 index 0000000..03f22c5 --- /dev/null +++ b/plugin/snap/api.go @@ -0,0 +1,84 @@ +package plugin_snap + +import ( + "io" + "net/http" + + "m7s.live/v5/pkg" + "m7s.live/v5/pkg/config" + snap "m7s.live/v5/plugin/snap/pkg" +) + +func (t *SnapPlugin) doSnap(rw http.ResponseWriter, r *http.Request) { + streamPath := r.PathValue("streamPath") + targetStreamPath := streamPath + + ok := t.Server.Streams.Has(streamPath) + if !ok { + http.Error(rw, pkg.ErrNotFound.Error(), http.StatusNotFound) + return + } + + pr, pw := io.Pipe() + + flusher, ok := rw.(http.Flusher) + if !ok { + http.Error(rw, "Streaming unsupported", http.StatusInternalServerError) + return + } + + rw.Header().Set("Content-Type", "image/jpeg") + rw.Header().Set("Transfer-Encoding", "chunked") + rw.Header().Del("Content-Length") + + done := make(chan error, 1) + + go func() { + defer pw.Close() + defer pr.Close() + defer close(done) + + buf := make([]byte, 32*1024) + _, err := io.CopyBuffer(rw, pr, buf) + if err != nil { + t.Error("write response error", err) + done <- err + return + } + flusher.Flush() + done <- nil + }() + + var transformer *snap.Transformer + if tm, ok := t.Server.Transforms.Get(targetStreamPath); ok { + transformer, ok = tm.TransformJob.Transformer.(*snap.Transformer) + if !ok { + http.Error(rw, "not a snap transformer", http.StatusInternalServerError) + return + } + } else { + transformer = snap.NewTransform().(*snap.Transformer) + transformer.TransformJob.Init(transformer, &t.Plugin, streamPath, config.Transform{ + Output: []config.TransfromOutput{ + { + Target: targetStreamPath, + StreamPath: targetStreamPath, + Conf: pw, + }, + }, + }).WaitStarted() + } + + transformer.TriggerSnap() + + if err := <-done; err != nil { + t.Error("snapshot failed", err) + return + } +} + +func (config *SnapPlugin) RegisterHandler() map[string]http.HandlerFunc { + return map[string]http.HandlerFunc{ + "/{streamPath...}": config.doSnap, + } +} diff --git a/plugin/snap/index.go b/plugin/snap/index.go new file mode 100755 index 0000000..cba2b83 --- /dev/null +++ b/plugin/snap/index.go @@ -0,0 +1,13 @@ +package plugin_snap + +import ( + m7s "m7s.live/v5" +) + +var _ = m7s.InstallPlugin[SnapPlugin]() + +type SnapPlugin struct { + //pb.UnimplementedApiServer + m7s.Plugin + LogToFile string +} diff --git a/plugin/snap/pkg/transform.go b/plugin/snap/pkg/transform.go new file mode 100644 index 0000000..baceec9 --- /dev/null +++ b/plugin/snap/pkg/transform.go @@ -0,0 +1,153 @@ +package snap + +import ( + "fmt" + "io" + "os/exec" + + "m7s.live/v5/pkg" + "m7s.live/v5/pkg/filerotate" + + m7s "m7s.live/v5" + "m7s.live/v5/pkg/task" +) + +// 定义传输模式的常量 +const ( + SNAP_MODE_PIPE SnapMode = "pipe" + SNAP_MODE_REMOTE SnapMode = "remote" +) + +type ( + SnapMode string + SnapConfig struct { + Mode SnapMode `default:"pipe" json:"mode" desc:"截图模式"` //截图模式 + Remote string `json:"remote" desc:"远程地址"` + } + SnapRule struct { + From SnapConfig `json:"from"` + LogToFile string `json:"logtofile" desc:"截图是否写入日志"` //截图日志写入文件 + } + + Config struct { + Input interface{} `json:"input"` + Output []Output `json:"output"` + } + + Output struct { + Target string `json:"target"` + Conf interface{} `json:"conf"` + } +) + +func NewTransform() m7s.ITransformer { + ret := &Transformer{ + snapChan: make(chan struct{}, 1), + } + ret.SetDescription(task.OwnerTypeKey, "Snap") + return ret +} + +type Transformer struct { + m7s.DefaultTransformer + SnapRule + logFile *filerotate.File + ffmpeg *exec.Cmd + snapChan chan struct{} +} + +func (t *Transformer) TriggerSnap() { + select { + case t.snapChan <- struct{}{}: + default: + // 如果通道已满，移除旧的请求 + <-t.snapChan + t.snapChan <- struct{}{} + } +} + +func (t *Transformer) Run() (err error) { + + // 等待截图触发信号 + <-t.snapChan + s := t.GetTransformJob().Plugin.Server + publisher, ok := s.Streams.Get(t.TransformJob.StreamPath) + if !ok || publisher.VideoTrack.AVTrack == nil { + return pkg.ErrNotFound + } + + err = publisher.VideoTrack.WaitReady() + if err != nil { + return err + } + + reader := pkg.NewAVRingReader(publisher.VideoTrack.AVTrack, "Origin") + err = reader.StartRead(publisher.VideoTrack.GetIDR()) + if err != nil { + return err + } + defer reader.StopRead() + + if reader.Value.Raw == nil { + if err = reader.Value.Demux(publisher.VideoTrack.ICodecCtx); err != nil { + return err + } + } + var annexb pkg.AnnexB + var track pkg.AVTrack + + track.ICodecCtx, track.SequenceFrame, err = annexb.ConvertCtx(publisher.VideoTrack.ICodecCtx) + if track.ICodecCtx == nil { + return fmt.Errorf("unsupported codec") + } + annexb.Mux(track.ICodecCtx, &reader.Value) + + // 创建ffmpeg命令 + cmd := exec.Command("ffmpeg", "-hide_banner", "-i", "pipe:0", "-vframes", "1", "-f", "mjpeg", "pipe:1") + + // 获取输入和输出pipe + stdin, err := cmd.StdinPipe() + if err != nil { + return err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + return err + } + + // 启动ffmpeg进程 + if err = cmd.Start(); err != nil { + return err + } + + // 将annexb数据写入到ffmpeg的stdin + _, err = annexb.WriteTo(stdin) + stdin.Close() + if err != nil { + return err + } + + // 从ffmpeg的stdout读取图片数据并写入到输出配置中 + _, err = io.Copy(t.TransformJob.Config.Output[0].Conf.(io.Writer), stdout) + if err != nil { + return err + } + + // 等待ffmpeg进程结束 + if err = cmd.Wait(); err != nil { + return err + } + + return nil +} + +func (t *Transformer) Dispose() { + close(t.snapChan) + if t.ffmpeg != nil { + err := t.ffmpeg.Process.Kill() + t.Error("kill ffmpeg", "err", err) + } + if t.logFile != nil { + _ = t.logFile.Close() + } +} diff --git a/plugin/snap/pkg/watermark/watermark.go b/plugin/snap/pkg/watermark/watermark.go new file mode 100755 index 0000000..2db01be --- /dev/null +++ b/plugin/snap/pkg/watermark/watermark.go @@ -0,0 +1,334 @@ +package watermark + +import ( + "fmt" + "image" + "image/color" + "image/draw" + "math" + + "github.com/disintegration/imaging" + "github.com/golang/freetype" + "github.com/golang/freetype/truetype" + "golang.org/x/image/font" + "golang.org/x/image/math/fixed" +) + +// TextConfig 文字配置结构体 +type TextConfig struct { + Text string // 水印文字 + Font *truetype.Font // 字体 + FontSize float64 // 字体大小 + Spacing float64 // 字符间距 + RowSpacing float64 // 行间距 + ColSpacing float64 // 列间距 + Rows int // 行数 + Cols int // 列数 + DPI float64 // 分辨率 + Color color.RGBA // 文字颜色 + Angle float64 // 旋转角度 + IsGrid bool + OffsetX int + OffsetY int +} + +// CalculateTextDimensions 计算文字尺寸的公共函数 +func CalculateTextDimensions(config TextConfig, face font.Face) (width, height float64) { + // 计算文字宽度 + var textWidth float64 + prevC := rune(-1) + for i, c := range config.Text { + if prevC >= 0 { + advance := face.Kern(prevC, c) + textWidth += float64(advance) / 64 + } + advance, _ := face.GlyphAdvance(c) + textWidth += float64(advance) / 64 + if i < len(config.Text)-1 { + textWidth += config.Spacing + } + prevC = c + } + + // 计算文字高度 + metrics := face.Metrics() + textHeight := float64(metrics.Height) / 64 + + return textWidth, textHeight +} + +func DrawWatermark(baseImg image.Image, config TextConfig, isDebug bool) (image.Image, error) { + switch { + case config.IsGrid: + return DrawWatermarkGrid(baseImg, config, isDebug) + default: + return DrawWatermarkSingle(baseImg, config, isDebug) + } +} + +func DrawWatermarkSingle(baseImg image.Image, config TextConfig, isDebug bool) (image.Image, error) { + bounds := baseImg.Bounds() + // width, height := bounds.Dx(), bounds.Dy() + + // 设置字体选项 + opts := truetype.Options{ + Size: config.FontSize, + DPI: config.DPI, + Hinting: font.HintingFull, + } + face := truetype.NewFace(config.Font, &opts) + defer face.Close() + + // 计算文字尺寸 + textWidth, textHeight := CalculateTextDimensions(config, face) + + // 创建一个与文字大小相同的透明图层 + textImg := image.NewRGBA(image.Rect(0, 0, int(math.Ceil(textWidth)), int(math.Ceil(textHeight)))) + + // 设置字体上下文 + c := freetype.NewContext() + c.SetDPI(config.DPI) + c.SetFont(config.Font) + c.SetFontSize(config.FontSize) + c.SetClip(textImg.Bounds()) + c.SetDst(textImg) + c.SetSrc(image.NewUniform(config.Color)) + c.SetHinting(font.HintingFull) + + // 绘制文字，从左上角开始 + pt := freetype.Pt(0, int(textHeight)) + + prevC := rune(-1) + for i, char := range config.Text { + if prevC >= 0 { + kern := face.Kern(prevC, char) + pt.X += kern + } + + _, err := c.DrawString(string(char), pt) + if err != nil { + return nil, err + } + + advance, _ := face.GlyphAdvance(char) + pt.X += advance + if i < len(config.Text)-1 { + pt.X += fixed.Int26_6(config.Spacing * 64) + } + + prevC = char + } + + if isDebug { + imaging.Save(textImg, "watermark.png") + } + + // 创建最终图像 + finalImg := image.NewRGBA(bounds) + + // 复制原图 + draw.Draw(finalImg, bounds, baseImg, image.Point{}, draw.Src) + + // 确保偏移量在图像范围内 + x := config.OffsetX + y := config.OffsetY + + // 如果需要旋转 + if config.Angle != 0 { + // 旋转文字，以文字左上角为圆心 + rotated := imaging.Rotate(textImg, config.Angle, color.Transparent) + if isDebug { + imaging.Save(rotated, "rotated_watermark.png") + } + + rotatedBounds := rotated.Bounds() + + // 归化角度到-180到180之间 + normalizedAngle := normalizeAngle(config.Angle) + + // 根据角度范围决定对齐方式 + var rect image.Rectangle + switch { + case normalizedAngle > 0 && normalizedAngle <= 90: + // 左下角对齐 + rect = image.Rect(x, y-rotatedBounds.Dy(), x+rotatedBounds.Dx(), y) + case normalizedAngle > 90 && normalizedAngle <= 180: + // 右下角对齐 + rect = image.Rect(x-rotatedBounds.Dx(), y-rotatedBounds.Dy(), x, y) + case normalizedAngle > -90 && normalizedAngle <= 0: + // 左上角对齐 + rect = image.Rect(x, y, x+rotatedBounds.Dx(), y+rotatedBounds.Dy()) + default: // -180 到 -90 + // 右上角对齐 + rect = image.Rect(x-rotatedBounds.Dx(), y, x, y+rotatedBounds.Dy()) + } + + // 混合旋转后的文字图层 + draw.Draw(finalImg, + rect, + rotated, + rotatedBounds.Min, + draw.Over) + } else { + // 不旋转时直接绘制 + draw.Draw(finalImg, + image.Rect(x, y, x+textImg.Bounds().Dx(), y+textImg.Bounds().Dy()), + textImg, + textImg.Bounds().Min, + draw.Over) + } + + return finalImg, nil +} + +// normalizeAngle 将角度归化到-180到180之间 +func normalizeAngle(angle float64) float64 { + // 先归化到0-360 + angle = math.Mod(angle, 360) + if angle > 180 { + angle -= 360 + } else if angle <= -180 { + angle += 360 + } + return angle +} + +// DrawWatermark 绘制水印网格并旋转 +func DrawWatermarkGrid(baseImg image.Image, config TextConfig, isDebug bool) (image.Image, error) { + bounds := baseImg.Bounds() + width, height := bounds.Dx(), bounds.Dy() + + opts := truetype.Options{ + Size: config.FontSize, + DPI: config.DPI, + Hinting: font.HintingFull, + } + face := truetype.NewFace(config.Font, &opts) + defer face.Close() + + // 计算矩形的外接圆半径 + radius := math.Sqrt(float64(width*width+height*height)) / 2 + + // 计算外接圆的外切正方形边长 + squareSize := int(math.Ceil(radius * 2)) + + // 如果行数或列数为0，自动计算 + if config.Rows == 0 || config.Cols == 0 { + textWidth, textHeight := CalculateTextDimensions(config, face) + + // 计算单个水印占用的空间 + cellWidth := textWidth + config.ColSpacing + cellHeight := textHeight + config.RowSpacing + + // 自动计算行数和列数 + if config.Cols == 0 { + config.Cols = int(math.Floor(float64(squareSize) / cellWidth)) + } + if config.Rows == 0 { + config.Rows = int(math.Floor(float64(squareSize) / cellHeight)) + } + + // 确保至少有一行一列 + if config.Cols < 1 { + config.Cols = 1 + } + if config.Rows < 1 { + config.Rows = 1 + } + } + + // 创建一个正方形的透明图层 + textImg := image.NewRGBA(image.Rect(0, 0, squareSize, squareSize)) + + // 设置字体上下文 + c := freetype.NewContext() + c.SetDPI(config.DPI) + c.SetFont(config.Font) + c.SetFontSize(config.FontSize) + c.SetClip(textImg.Bounds()) + c.SetDst(textImg) + c.SetSrc(image.NewUniform(config.Color)) + c.SetHinting(font.HintingFull) + + textWidth, textHeight := CalculateTextDimensions(config, face) + + // 计算网格的总宽度和高度 + gridWidth := textWidth + config.ColSpacing + gridHeight := textHeight + config.RowSpacing + + // 计算起始位置，使整个网格居中 + startX := (float64(squareSize) - (float64(config.Cols)*gridWidth - config.ColSpacing)) / 2 + startY := (float64(squareSize) - (float64(config.Rows)*gridHeight - config.RowSpacing)) / 2 + + // 绘制文字网格 + for row := 0; row < config.Rows; row++ { + for col := 0; col < config.Cols; col++ { + x := startX + float64(col)*gridWidth + y := startY + float64(row)*gridHeight + + pt := freetype.Pt( + int(x), + int(y+textHeight), + ) + + prevC := rune(-1) + for i, char := range config.Text { + if prevC >= 0 { + kern := face.Kern(prevC, char) + pt.X += kern + } + + _, err := c.DrawString(string(char), pt) + if err != nil { + return nil, err + } + + advance, _ := face.GlyphAdvance(char) + pt.X += advance + if i < len(config.Text)-1 { + pt.X += fixed.Int26_6(config.Spacing * 64) + } + + prevC = char + } + } + } + + if isDebug { + // 保存文字模板 + imaging.Save(textImg, "watermark.png") + } + // 旋转整个文字网格 + rotated := imaging.Rotate(textImg, config.Angle, color.Transparent) + if isDebug { + // 保存旋转后的文字模板 + imaging.Save(rotated, "rotated_watermark.png") + } + + // 创建最终图像 + finalImg := image.NewRGBA(bounds) + + // 复制原图 + draw.Draw(finalImg, bounds, baseImg, image.Point{}, draw.Src) + + // 计算旋转后图片的位置（确保居中） + rotatedBounds := rotated.Bounds() + x := (width - rotatedBounds.Dx()) / 2 + y := (height - rotatedBounds.Dy()) / 2 + if isDebug { + fmt.Printf("width: %d, height: %d, x: %d, y: %d\n", width, height, x, y) + } + // 混合旋转后的文字图层 + draw.Draw(finalImg, + image.Rect(x, y, x+rotatedBounds.Dx(), y+rotatedBounds.Dy()), + rotated, + rotatedBounds.Min, + draw.Over) + + return finalImg, nil +} + +// CreateTextColor 创建文字颜色 +func CreateTextColor(r, g, b uint8, a uint8) *image.Uniform { + return image.NewUniform(color.RGBA{r, g, b, a}) +} diff --git a/plugin/snap/pkg/watermark/watermark_test.go b/plugin/snap/pkg/watermark/watermark_test.go new file mode 100644 index 0000000..381f77b --- /dev/null +++ b/plugin/snap/pkg/watermark/watermark_test.go @@ -0,0 +1,205 @@ +package watermark + +import ( + "image/color" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/disintegration/imaging" + "github.com/golang/freetype/truetype" + "github.com/stretchr/testify/assert" + xfont "golang.org/x/image/font" +) + +func loadTestFont(t *testing.T) *truetype.Font { + var fontPath string + switch runtime.GOOS { + case "windows": + fontPath = "C:\\Windows\\Fonts\\simsun.ttc" + case "darwin": + fontPath = "/System/Library/Fonts/STHeiti Light.ttc" + case "linux": + fontPath = "/usr/share/fonts/truetype/winfonts/simsun.ttc" + default: + t.Fatal("不支持的操作系统") + } + + fontBytes, err := os.ReadFile(fontPath) + if err != nil { + t.Fatalf("读取字体文件失败: %v", err) + } + + font, err := truetype.Parse(fontBytes) + if err != nil { + t.Fatalf("解析字体失败: %v", err) + } + + return font +} + +func TestDrawWatermarkSingle(t *testing.T) { + // 准备测试图片 + tmpDir := t.TempDir() + testImagePath := filepath.Join(tmpDir, "test_input.png") + + // 创建一个测试图片 + img := imaging.New(800, 600, color.White) + err := imaging.Save(img, testImagePath) + if err != nil { + t.Fatalf("创建测试图片失败: %v", err) + } + + // 加载测试图片 + baseImg, err := imaging.Open(testImagePath) + if err != nil { + t.Fatalf("加载测试图片失败: %v", err) + } + + // 加载字体 + font := loadTestFont(t) + + // 创建水印配置 + config := TextConfig{ + Text: "测试水印", + Font: font, + FontSize: 36, + Spacing: 10, + RowSpacing: 10, + ColSpacing: 20, + DPI: 72, + Color: color.RGBA{R: 255, G: 0, B: 0, A: 255}, // 红色 + IsGrid: false, + Angle: 45, + OffsetX: 100, + OffsetY: 100, + } + + // 测试单个水印 + result, err := DrawWatermarkSingle(baseImg, config, true) + if err != nil { + t.Fatalf("绘制单个水印失败: %v", err) + } + + // 保存结果用于目视检查 + outputPath := filepath.Join(tmpDir, "test_output_single.png") + err = imaging.Save(result, outputPath) + if err != nil { + t.Fatalf("保存结果图片失败: %v", err) + } + + // 验证结果图片存在且大小正确 + stat, err := os.Stat(outputPath) + assert.NoError(t, err) + assert.True(t, stat.Size() > 0) +} + +func TestDrawWatermarkGrid(t *testing.T) { + // 准备测试图片 + tmpDir := t.TempDir() + testImagePath := filepath.Join(tmpDir, "test_input.png") + + // 创建一个测试图片 + img := imaging.New(800, 600, color.White) + err := imaging.Save(img, testImagePath) + if err != nil { + t.Fatalf("创建测试图片失败: %v", err) + } + + // 加载测试图片 + baseImg, err := imaging.Open(testImagePath) + if err != nil { + t.Fatalf("加载测试图片失败: %v", err) + } + + // 加载字体 + font := loadTestFont(t) + + // 创建水印配置 + config := TextConfig{ + Text: "测试水印", + Font: font, + FontSize: 36, + Spacing: 10, + RowSpacing: 10, + ColSpacing: 20, + Rows: 3, + Cols: 4, + DPI: 72, + Color: color.RGBA{R: 0, G: 0, B: 255, A: 255}, // 蓝色 + IsGrid: true, + Angle: 30, + } + + // 测试网格水印 + result, err := DrawWatermarkGrid(baseImg, config, true) + if err != nil { + t.Fatalf("绘制网格水印失败: %v", err) + } + + // 保存结果用于目视检查 + outputPath := filepath.Join(tmpDir, "test_output_grid.png") + err = imaging.Save(result, outputPath) + if err != nil { + t.Fatalf("保存结果图片失败: %v", err) + } + + // 验证结果图片存在且大小正确 + stat, err := os.Stat(outputPath) + assert.NoError(t, err) + assert.True(t, stat.Size() > 0) +} + +func TestCalculateTextDimensions(t *testing.T) { + // 加载字体 + font := loadTestFont(t) + + // 创建配置 + config := TextConfig{ + Text: "测试文本", + Font: font, + FontSize: 36, + Spacing: 10, + DPI: 72, + } + + // 设置字体选项 + opts := truetype.Options{ + Size: config.FontSize, + DPI: config.DPI, + Hinting: xfont.HintingFull, + } + face := truetype.NewFace(font, &opts) + defer face.Close() + + // 计算文字尺寸 + width, height := CalculateTextDimensions(config, face) + + // 验证结果 + assert.True(t, width > 0, "文字宽度应该大于0") + assert.True(t, height > 0, "文字高度应该大于0") +} + +func TestNormalizeAngle(t *testing.T) { + tests := []struct { + input float64 + expected float64 + }{ + {0, 0}, + {180, 180}, + {-180, -180}, + {360, 0}, + {-360, 0}, + {540, 180}, + {-540, -180}, + {270, -90}, + {-270, 90}, + } + + for _, test := range tests { + result := normalizeAngle(test.input) + assert.Equal(t, test.expected, result, "输入角度 %f 应该归一化为 %f，但得到 %f", + test.input, test.expected, result) + } +}