mirror of
https://github.com/langhuihui/monibuca.git
synced 2025-09-27 03:25:56 +08:00
325 lines
8.5 KiB
Go
325 lines
8.5 KiB
Go
package snap
|
||
|
||
import (
|
||
"bytes"
|
||
"fmt"
|
||
"image/color"
|
||
"os"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
task "github.com/langhuihui/gotask"
|
||
"m7s.live/v5/pkg"
|
||
"m7s.live/v5/pkg/config"
|
||
"m7s.live/v5/pkg/format"
|
||
|
||
m7s "m7s.live/v5"
|
||
)
|
||
|
||
const (
|
||
SnapModeTimeInterval = iota
|
||
SnapModeIFrameInterval
|
||
SnapModeManual
|
||
)
|
||
|
||
// parseRGBA 解析rgba格式的颜色字符串
|
||
func parseRGBA(rgbaStr string) (color.RGBA, error) {
|
||
rgba := strings.TrimPrefix(rgbaStr, "rgba(")
|
||
rgba = strings.TrimSuffix(rgba, ")")
|
||
parts := strings.Split(rgba, ",")
|
||
if len(parts) != 4 {
|
||
return color.RGBA{}, fmt.Errorf("invalid rgba format")
|
||
}
|
||
|
||
r, err := strconv.Atoi(strings.TrimSpace(parts[0]))
|
||
if err != nil {
|
||
return color.RGBA{}, err
|
||
}
|
||
|
||
g, err := strconv.Atoi(strings.TrimSpace(parts[1]))
|
||
if err != nil {
|
||
return color.RGBA{}, err
|
||
}
|
||
|
||
b, err := strconv.Atoi(strings.TrimSpace(parts[2]))
|
||
if err != nil {
|
||
return color.RGBA{}, err
|
||
}
|
||
|
||
a, err := strconv.ParseFloat(strings.TrimSpace(parts[3]), 64)
|
||
if err != nil {
|
||
return color.RGBA{}, err
|
||
}
|
||
|
||
return color.RGBA{
|
||
R: uint8(r),
|
||
G: uint8(g),
|
||
B: uint8(b),
|
||
A: uint8(a * 255),
|
||
}, nil
|
||
}
|
||
|
||
// 保存截图到文件
|
||
func saveSnapshot(annexb []*format.AnnexB, savePath string, plugin *m7s.Plugin, streamPath string, snapMode int, watermarkConfig *WatermarkConfig) error {
|
||
var buf bytes.Buffer
|
||
if err := ProcessWithFFmpeg(annexb, &buf); err != nil {
|
||
return fmt.Errorf("process with ffmpeg error: %w", err)
|
||
}
|
||
|
||
// 如果配置了水印，添加水印
|
||
if watermarkConfig != nil && watermarkConfig.Text != "" {
|
||
imgData, err := AddWatermark(buf.Bytes(), *watermarkConfig)
|
||
if err != nil {
|
||
return fmt.Errorf("add watermark error: %w", err)
|
||
}
|
||
err = os.WriteFile(savePath, imgData, 0644)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
} else {
|
||
err := os.WriteFile(savePath, buf.Bytes(), 0644)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
// 保存记录到数据库
|
||
if plugin != nil && plugin.DB != nil {
|
||
record := SnapRecord{
|
||
StreamName: streamPath,
|
||
SnapMode: snapMode,
|
||
SnapTime: time.Now(),
|
||
SnapPath: savePath,
|
||
}
|
||
if err := plugin.DB.Create(&record).Error; err != nil {
|
||
return fmt.Errorf("save snapshot record failed: %w", err)
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// SnapConfig 截图配置
|
||
type SnapConfig struct {
|
||
TimeInterval time.Duration `json:"timeInterval" desc:"截图时间间隔，大于0时使用时间间隔模式"`
|
||
IFrameInterval int `json:"iFrameInterval" desc:"间隔多少帧截图，大于0时使用关键帧间隔模式"`
|
||
SavePath string `json:"savePath" desc:"截图保存路径"`
|
||
Watermark struct {
|
||
Text string `json:"text" default:"" desc:"水印文字内容"`
|
||
FontPath string `json:"fontPath" default:"" desc:"水印字体文件路径"`
|
||
FontColor string `json:"fontColor" default:"rgba(255,165,0,1)" desc:"水印字体颜色，支持rgba格式"`
|
||
FontSize float64 `json:"fontSize" default:"36" desc:"水印字体大小"`
|
||
FontSpacing float64 `json:"fontSpacing" default:"2" desc:"水印字体间距"`
|
||
OffsetX int `json:"offsetX" default:"0" desc:"水印位置X"`
|
||
OffsetY int `json:"offsetY" default:"0" desc:"水印位置Y"`
|
||
} `json:"watermark" desc:"水印配置"`
|
||
}
|
||
|
||
// SnapTask 基础截图任务结构
|
||
type SnapTask struct {
|
||
config SnapConfig
|
||
job *m7s.TransformJob
|
||
watermarkConfig *WatermarkConfig
|
||
}
|
||
|
||
// saveSnap 保存截图
|
||
func (t *SnapTask) saveSnap(annexb []*format.AnnexB, snapMode int) error {
|
||
// 生成文件名
|
||
now := time.Now()
|
||
filename := fmt.Sprintf("%s_%s.jpg", t.job.StreamPath, now.Format("20060102150405.000"))
|
||
filename = strings.ReplaceAll(filename, "/", "_")
|
||
savePath := filepath.Join(t.config.SavePath, filename)
|
||
|
||
// 处理视频帧
|
||
var buf bytes.Buffer
|
||
if err := ProcessWithFFmpeg(annexb, &buf); err != nil {
|
||
return fmt.Errorf("process with ffmpeg error: %w", err)
|
||
}
|
||
|
||
// 如果配置了水印，添加水印
|
||
if t.watermarkConfig != nil && t.watermarkConfig.Text != "" {
|
||
imgData, err := AddWatermark(buf.Bytes(), *t.watermarkConfig)
|
||
if err != nil {
|
||
return fmt.Errorf("add watermark error: %w", err)
|
||
}
|
||
err = os.WriteFile(savePath, imgData, 0644)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
} else {
|
||
err := os.WriteFile(savePath, buf.Bytes(), 0644)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
// 保存记录到数据库
|
||
if t.job.Plugin != nil && t.job.Plugin.DB != nil {
|
||
record := SnapRecord{
|
||
StreamName: t.job.StreamPath,
|
||
SnapMode: snapMode,
|
||
SnapTime: time.Now(),
|
||
SnapPath: savePath,
|
||
}
|
||
if err := t.job.Plugin.DB.Create(&record).Error; err != nil {
|
||
return fmt.Errorf("save snapshot record failed: %w", err)
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// TimeSnapTask 定时截图任务
|
||
type TimeSnapTask struct {
|
||
task.TickTask
|
||
SnapTask
|
||
}
|
||
|
||
func (t *TimeSnapTask) GetTickInterval() time.Duration {
|
||
return t.config.TimeInterval
|
||
}
|
||
|
||
// Tick 执行定时截图操作
|
||
func (t *TimeSnapTask) Tick(any) {
|
||
// 获取视频帧
|
||
annexb, err := GetVideoFrame(t.job.OriginPublisher, t.job.Plugin.Server)
|
||
if err != nil {
|
||
t.Error("get video frame failed", "error", err.Error())
|
||
return
|
||
}
|
||
|
||
if err := t.saveSnap(annexb, SnapModeTimeInterval); err != nil {
|
||
t.Error("save snapshot failed", "error", err.Error())
|
||
}
|
||
}
|
||
|
||
// IFrameSnapTask 关键帧截图任务
|
||
type IFrameSnapTask struct {
|
||
task.Task
|
||
SnapTask
|
||
subscriber *m7s.Subscriber
|
||
}
|
||
|
||
func (t *IFrameSnapTask) Start() (err error) {
|
||
subConfig := t.job.Plugin.GetCommonConf().Subscribe
|
||
subConfig.SubType = m7s.SubscribeTypeTransform
|
||
subConfig.IFrameOnly = true
|
||
t.subscriber, err = t.job.Plugin.SubscribeWithConfig(t, t.job.StreamPath, subConfig)
|
||
return
|
||
}
|
||
|
||
func (t *IFrameSnapTask) Go() (err error) {
|
||
iframeCount := 0
|
||
err = m7s.PlayBlock(t.subscriber, (func(audio *pkg.AVFrame) error)(nil), func(video *format.AnnexB) error {
|
||
iframeCount++
|
||
if iframeCount%t.config.IFrameInterval == 0 {
|
||
if err := t.saveSnap([]*format.AnnexB{video}, SnapModeIFrameInterval); err != nil {
|
||
t.Error("save snapshot failed", "error", err.Error())
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
t.Error("iframe interval snap error", "error", err.Error())
|
||
}
|
||
return
|
||
}
|
||
|
||
type Transformer struct {
|
||
task.Job
|
||
TransformJob m7s.TransformJob
|
||
}
|
||
|
||
func (r *Transformer) GetTransformJob() *m7s.TransformJob {
|
||
return &r.TransformJob
|
||
}
|
||
|
||
func NewTransform() m7s.ITransformer {
|
||
return &Transformer{}
|
||
}
|
||
|
||
func (t *Transformer) Start() (err error) {
|
||
// 为每个输出配置创建一个截图任务
|
||
for _, output := range t.TransformJob.Config.Output {
|
||
var task task.ITask
|
||
var snapConfig SnapConfig
|
||
if output.Conf != nil {
|
||
switch v := output.Conf.(type) {
|
||
case SnapConfig:
|
||
snapConfig = v
|
||
case map[string]any:
|
||
config.Parse(&snapConfig, v)
|
||
}
|
||
}
|
||
|
||
// 初始化水印配置
|
||
var watermarkConfig *WatermarkConfig
|
||
if snapConfig.Watermark.Text != "" {
|
||
watermarkConfig = &WatermarkConfig{
|
||
Text: snapConfig.Watermark.Text,
|
||
FontPath: snapConfig.Watermark.FontPath,
|
||
FontSize: snapConfig.Watermark.FontSize,
|
||
FontSpacing: snapConfig.Watermark.FontSpacing,
|
||
OffsetX: snapConfig.Watermark.OffsetX,
|
||
OffsetY: snapConfig.Watermark.OffsetY,
|
||
}
|
||
|
||
// 判断字体是否存在
|
||
if _, err := os.Stat(watermarkConfig.FontPath); os.IsNotExist(err) {
|
||
return fmt.Errorf("watermark font file not found: %w", err)
|
||
}
|
||
// 解析颜色
|
||
if snapConfig.Watermark.FontColor != "" {
|
||
fontColor, err := parseRGBA(snapConfig.Watermark.FontColor)
|
||
if err == nil {
|
||
watermarkConfig.FontColor = fontColor
|
||
} else {
|
||
t.Error("parse color failed", "error", err.Error())
|
||
watermarkConfig.FontColor = color.RGBA{uint8(255), uint8(255), uint8(255), uint8(255)}
|
||
}
|
||
}
|
||
|
||
// 预加载字体
|
||
if err := watermarkConfig.LoadFont(); err != nil {
|
||
return fmt.Errorf("load watermark font failed: %w", err)
|
||
}
|
||
t.Info("watermark config loaded",
|
||
"text", watermarkConfig.Text,
|
||
"font", watermarkConfig.FontPath,
|
||
"size", watermarkConfig.FontSize,
|
||
)
|
||
}
|
||
// 创建保存目录
|
||
if err := os.MkdirAll(snapConfig.SavePath, 0755); err != nil {
|
||
return fmt.Errorf("create save directory failed: %w", err)
|
||
}
|
||
// 根据配置创建对应的任务
|
||
if snapConfig.TimeInterval > 0 {
|
||
timeTask := &TimeSnapTask{
|
||
SnapTask: SnapTask{
|
||
config: snapConfig,
|
||
job: &t.TransformJob,
|
||
watermarkConfig: watermarkConfig,
|
||
},
|
||
}
|
||
task = timeTask
|
||
} else if snapConfig.IFrameInterval > 0 {
|
||
iframeTask := &IFrameSnapTask{
|
||
SnapTask: SnapTask{
|
||
config: snapConfig,
|
||
job: &t.TransformJob,
|
||
watermarkConfig: watermarkConfig,
|
||
},
|
||
}
|
||
task = iframeTask
|
||
}
|
||
|
||
if task != nil {
|
||
t.AddTask(task)
|
||
}
|
||
}
|
||
return nil
|
||
}
|