视频水印添加

根据合规要求，需要将AI生成的视频或者图片添加水印，提示用户这是AI生成的。

add_watermark_to_frame

 函数的功能是将指定的水印文本添加到输入的图像帧中，并返回带有水印的图像帧。具体步骤如下：

1. 转换图像格式：将OpenCV格式的图像转换为Pillow格式，以便进行绘图操作。
2. 设置水印参数：定义水印字体、大小、颜色和位置。
3. 绘制水印：使用Pillow库在图像上绘制水印文本。
4. 转换回OpenCV格式：将处理后的Pillow图像转换回OpenCV格式，以便后续处理或保存。

add_watermark_to_video

 函数的功能是将水印添加到整个视频的所有帧中，并保存为新的视频文件。具体步骤如下：

1. 打开输入视频：读取输入视频文件。
2. 获取视频属性：获取视频的宽度、高度和帧率。
3. 定义输出视频编码器：设置输出视频的编码器和文件路径。
4. 逐帧处理：循环读取每一帧，调用 `add_watermark_to_frame` 添加水印，并写入输出视频。
5. 释放资源：关闭视频文件并释放资源。
通过这两个函数的结合，可以实现对视频每一帧添加水印并保存为新视频的功能。

import json
import os
import time
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from loguru import logger


# 加载json文件
def load_json(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
        # 安全地访问 "watermark" 键
    watermark = data.get("watermark")
    if watermark is None:
        return
    return watermark

def add_watermark_to_frame(frame):
    config_path = r"/root/autodl-fs/config/config.json"
    for idx in range(5):
        try:
            config_path = f"{os.environ['DG_CONFIG_PATH']}/config.json"
            break
        except Exception as e:
            time.sleep(1)
            logger.error(f"环境变量获取异常：{e}, 第{idx + 1}次重试")

    watermark_text = load_json(config_path)

    if not watermark_text:
        return frame

    # start_time = time.time()
    # 将OpenCV图像转换为Pillow图像
    pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pil_image)

    frame_width, frame_height = pil_image.size
    # 设置水印字体和参数
    area = frame_width * frame_height
    font_size = int(area ** 0.5 * 0.03)  # 字体大小
    # 获取当前绝对路径
    current_dir = os.path.dirname(os.path.dirname(__file__))
    # 获取当前绝对路径底下的configs下的simhei.ttf
    font_path = os.path.join(current_dir, "configs", "simhei.ttf") # 指定中文字体文件路径，如黑体字体

    try:
        # Try to load a system font (you can specify any available TTF font file)
        font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
    except IOError:
        # Fallback to default font if specified font is not found
        font = ImageFont.load_default()

    # 加载水印字体宽度，需要将水印定位到右下角
    font_width = font.getlength(watermark_text)
    # print(f"水印宽度：{font_width}")
    font_color = (255, 255, 255)  # 白色
    position = (frame_width - font_width - font_size*0.1, frame_height - font_size - font_size*0.1)  # 右下角位置

    # 加载字体并绘制中文水印
    draw.text(position, watermark_text, font=font, fill=font_color)

    # 将Pillow图像转换回OpenCV图像
    frame_with_watermark = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    # print(f"水印添加完成，耗时：{time.time() - start_time}秒")

    return frame_with_watermark


def add_watermark_to_video(input_video_path, output_video_path):
    # 打开输入视频
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        print("Error: Could not open video.")
        return

    # 获取视频的宽度、高度和帧率
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # 定义输出视频的编码器和文件
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        frame_with_watermark = add_watermark_to_frame(frame)

        # 写入输出视频
        out.write(frame_with_watermark)

    # 释放资源
    cap.release()
    out.release()
    print("Watermark added successfully. Output video saved as:", output_video_path)


if __name__ == "__main__":
    input_video_path = 'E:\\a.mp4'  # Path to your input video file
    output_video_path = 'output.mp4'  # Path where you want to save the output video file

    add_watermark_to_video(input_video_path, output_video_path)