Skip to content

AI45Lab/AIGC-Identification-Toolkit

Repository files navigation

UniMark: AI生成内容标识技术开发套件

UniMark: AI生成内容标识技术开发套件

一站式开源标识技术开发套件,支持文本、图像、音频和视频内容的显式标识,隐式标识和隐水印功能
快速开始 »

English | 简体中文

关于项目

本项目提供一站式开源标识技术开发套件。支持文本、图像、音频和视频四大模态,具备显式标识和隐水印功能,覆盖GB 45438-2025《网络安全技术 人工智能生成合成内容标识方法》标准规定的标识范围。

为什么选择我们?

  • 全面覆盖:支持GB 45438-2025标准要求的标识方法
  • 多模态支持:统一处理文本、图像、音频和视频内容
  • 双模式操作:既支持AI内容生成,也支持现有文件处理

构建技术

  • Python PyTorch Flask Transformers Diffusers

(返回顶部)


📑 目录


开始使用

安装

🔧 传统安装

  1. 克隆仓库

    git clone --recurse-submodules https://github.com/MillionMillionLi/AIGC-Identification-Toolkit.git
    cd AIGC-Identification-Toolkit
  2. 安装核心依赖

    pip install -r requirements.txt
  3. 安装系统依赖

    sudo apt install ffmpeg
  4. (可选)下载 AI 生成模型

    仅当需要使用 AI 生成内容并添加水印功能时才需要此步骤。如果只处理已有文件(上传模式添加水印),可跳过此步骤。

    # 图像生成 + 水印(Stable Diffusion 2.1)
    python scripts/download_sd_model.py
    
    # 视频生成 + 水印(Wan2.1)
    python scripts/download_wan_model.py
    
    # 文本生成 + 水印(PostMark + Mistral)
    python scripts/download_postmark_deps.py
    
    # 音频生成 + 水印(Bark)
    python scripts/download_bark_model.py
    pip install git+https://github.com/suno-ai/bark.git
  5. 配置环境

    export TRANSFORMERS_OFFLINE=1
    export HF_HUB_OFFLINE=1
    export HF_ENDPOINT=https://hf-mirror.com

🐳 Docker 安装

前置要求:需要 NVIDIA GPU 和 NVIDIA Container Toolkit

  1. 拉取镜像

    docker pull millionmillionli/aigc-identification-toolkit:latest
  2. 运行容器

    docker run -d --name aigc-watermark-toolkit \
      --gpus all \
      -v /path/to/your/.cache/huggingface:/cache/huggingface \ ## 将这一行的路径改为你的实际模型缓存路径
      -v $(pwd):/app \
      -e HF_HOME=/cache/huggingface \
      -e PYTHONPATH=/app \
      -e CUDA_VISIBLE_DEVICES=0 \
      --restart unless-stopped \
      millionmillionli/aigc-identification-toolkit:latest \
      tail -f /dev/null

    (可选)准备AI模型,仅当需要使用 AI 生成内容并添加水印功能时才需要此步骤。

    需要下载的模型

    • 图像生成:Stable Diffusion 2.1 (stabilityai/stable-diffusion-2-1-base)
    • 视频生成:Wan2.1 (Wan-AI/Wan2.1-T2V-1.3B-Diffusers)
    • 文本生成:Mistral 7B + PostMark词嵌入 (mistralai/Mistral-7B-Instruct-v0.2)
    • 音频生成:Bark (suno/bark)
  3. 进入容器

    docker exec -it aigc-watermark-toolkit bash

进阶使用:如需完整配置(benchmark挂载、只读配置等),可使用 docker-compose up -d或查看docker-compose.yml 文件。

使用方法

WatermarkTool是推荐的主要入口点,提供统一的接口支持所有模态的水印和显式标识操作。

初始化

from src.unified.watermark_tool import WatermarkTool

# 使用默认配置初始化
tool = WatermarkTool()

# 使用自定义配置初始化
tool = WatermarkTool(config_path="path/to/config.yaml")

核心方法

embed() - 嵌入水印或显式标识
def embed(self,
          content: Union[str, Path],
          message: str,
          modality: str,
          operation: str = 'watermark',
          **kwargs) -> Union[str, PIL.Image, torch.Tensor, Path]:
    """
    嵌入水印或添加显式标识

    Args:
        content: 输入内容
            - 文本模态: 提示文本(AI生成)或文本文件路径(上传模式)
            - 图像模态: 提示文本(AI生成)或图像文件路径(上传模式)
            - 音频模态: 提示文本(AI生成)或音频文件路径(上传模式)
            - 视频模态: 提示文本(AI生成)或视频文件路径(上传模式)
        message: 要嵌入的水印信息或显式标识文本
        modality: 模态类型 ('text', 'image', 'audio', 'video')
        operation: 操作类型 ('watermark' 或 'visible_mark')
        **kwargs: 模态特定参数

    Returns:
        处理后的内容(格式根据模态而定)
    """

使用示例:

# 隐式水印(默认operation='watermark')
img_wm = tool.embed("a cat under the sun", "img_msg", 'image')


# 上传文件模式
img_wm = tool.embed("", "file_msg", 'image', image_input="/path/to/image.jpg")


# 显式标识
marked_img = tool.embed("/path/to/image.jpg", "AI标识", 'image',
                       operation='visible_mark', position='bottom_right')
extract() - 提取水印或检测显式标识
def extract(self,
           content: Union[str, PIL.Image, torch.Tensor, Path],
           modality: str,
           operation: str = 'watermark',
           **kwargs) -> Dict[str, Any]:
    """
    提取水印或检测显式标识

    Args:
        content: 待检测的内容
        modality: 模态类型
        operation: 操作类型 ('watermark' 或 'visible_mark')
        **kwargs: 检测参数

    Returns:
        检测结果字典:
        {
            'detected': bool,      # 是否检测到水印/标识
            'message': str,        # 提取的消息内容
            'confidence': float,   # 置信度 (0.0-1.0)
            'metadata': dict       # 额外的元数据信息
        }
    """

使用示例:

# 提取隐式水印
img_result = tool.extract(watermarked_image, 'image')


# 检测显式标识
mark_result = tool.extract(marked_content, 'text', operation='visible_mark')

Benchmarks

评估各模态水印算法的性能表现,提供标准化的测试数据集、攻击方式和评估指标,帮助用户选择最适合应用场景的算法。


Image-Bench

评估图像水印算法对传统失真攻击(亮度、对比度、模糊、噪声、JPEG压缩)的鲁棒性。

核心特性:

  • 数据集: W-Bench DISTORTION_1K(1000张图像)

  • 评估指标:: PSNR, SSIM, LPIPS, TPR,Bit accuracy

快速使用:

python benchmarks/Image-Bench/run_benchmark.py

使用自定义数据集:

  1. 准备图像数据:将PNG图像放入自定义目录(如 benchmarks/Image-Bench/dataset/my_dataset/
  2. 修改配置 configs/videoseal_distortion.yaml
    dataset:
      path: benchmarks/Image-Bench/dataset/my_dataset

评估指标

指标类别 指标 判定阈值 指标说明
质量 PSNR ≥ 35.0 dB Peak Signal-to-Noise Ratio(峰值信噪比),越高越好
质量 SSIM ≥ 0.95 Structural Similarity Index(结构相似度),越接近 1 越好
质量 LPIPS ≤ 0.015 Learned Perceptual Similarity(感知相似度),越低越好
鲁棒性 TPR ≥ 0.80 True Positive Rate(检测成功率),越高表示鲁棒性越强
鲁棒性 Bit Accuracy ≥ 0.85 水印比特准确率,决定解码结果与原始水印的接近程度

结果分析

TPR Bit Accuracy 质量评估指标
VideoSeal TPR Radar VideoSeal Bit Accuracy Radar
指标数值达到阈值
PSNR45.52 dB
SSIM0.9953
LPIPS0.0025

详细文档: benchmarks/Image-Bench/README.md


Audio-Bench

评估音频水印算法(AudioSeal)对多种音频攻击的鲁棒性,覆盖噪声、滤波、压缩等常见干扰。

核心特性:

  • 📊 数据集: AudioMark Dataset
  • 🔧 攻击类型: 高斯噪声、背景噪声、量化、滤波、平滑、回声、MP3压缩

快速使用:

python benchmarks/Audio-Bench/run_benchmark.py

使用自定义数据集:

  1. 准备音频数据:将音频文件(支持WAV/MP3/FLAC/M4A)放入自定义目录
  2. 修改配置 configs/audioseal_robustness.yaml
    dataset:
      path: benchmarks/Audio-Bench/dataset/my_audio_dataset

评估指标

指标类别 指标 判定阈值 指标说明
质量 SNR ≥ 20.0 dB Signal-to-Noise Ratio,原音频 vs 水印音频,越高越好
鲁棒性 TPR (Detection Probability) ≥ 0.80 以检测概率判定的真阳性率
鲁棒性 Bit Accuracy ≥ 0.875 图案水印比特正确率,越高越好

结果分析

TPR (Detection Probability) Bit Accuracy 质量评估指标
AudioSeal TPR Probability Radar AudioSeal Bit Accuracy Radar
指标数值达到阈值
SNR23

详细文档: benchmarks/Audio-Bench/README.md


Video-Bench

评估视频水印算法(VideoSeal)在图像级和视频级扰动下的鲁棒性,严格遵循VideoMarkBench论文方法。

核心特性:

  • 📊 数据集: VideoMarkBench Dataset
  • 🔧 攻击类型: 高斯噪声、模糊、JPEG压缩、裁剪、帧平均、帧交换、帧删除

快速使用:

python benchmarks/Video-Bench/run_benchmark.py

使用自定义数据集:

  1. 准备视频数据:将视频文件(支持MP4/AVI/MOV/MKV)放入自定义目录,支持子目录
  2. 修改配置 configs/videoseal_robustness.yaml
    dataset:
      path: benchmarks/Video-Bench/dataset/my_video_dataset

评估指标

指标类别 指标 判定阈值 指标说明
质量 PSNR ≥ 35.0 dB Peak Signal-to-Noise Ratio,越高越好
质量 SSIM ≥ 0.95 Structural Similarity Index,越接近 1 越好
质量 tLP ≤ 0.20 Temporal LPIPS,衡量跨帧感知一致性,越低越好
鲁棒性 FNR ≤ 0.01 False Negative Rate,漏检率,越低表示鲁棒性越强
鲁棒性 Bit Accuracy ≥ 0.85 解码比特准确率,越高越好

结果分析

FNR Bit Accuracy 质量评估指标
VideoSeal FNR Radar VideoSeal Bit Accuracy Radar
指标数值达到阈值
PSNR40.59
SSIM0.97
tLP0.001

详细文档: benchmarks/Video-Bench/README.md


Text-Bench(规划中)

(返回顶部)

致谢

本项目基于以下优秀的开源工作构建:

水印算法

AI 生成模型

评估与基准测试

  • VINE - W-Bench 数据集和图像失真攻击实现
  • AudioMarkBench - 音频水印评估框架

(返回顶部)

引用

如果您觉得本项目对您的研究有所帮助,请考虑引用我们的论文:

@article{li2025unimark,
  title={UniMark: Artificial Intelligence Generated Content Identification Toolkit},
  author={Meilin Li and Ji He and Yu Yi and Jia Xu and Shanzhe Lei and Yan Teng and Yingchun Wang and Xuhong Wang},
  journal={arXiv preprint arXiv:2512.12324},
  year={2025}
}

(返回顶部)

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published