AI 内容水印技术：生成内容的溯源与检测

import torch
import torch.nn.functional as F
from collections import defaultdict
import hashlib

class StatisticalWatermark:
    """统计水印——Green List/Red List 方法"""
    
    def __init__(self, model, green_ratio: float = 0.5, 
                 green_bias: float = 2.0, hash_key: int = 42):
        self.model = model
        self.green_ratio = green_ratio
        self.green_bias = green_bias  # 绿名单token的logit偏置
        self.hash_key = hash_key
    
    def _get_green_list(self, prev_token: int, vocab_size: int) -> set:
        """根据前一个token确定绿名单"""
        # 使用哈希函数确定性地生成绿名单
        rng = torch.Generator()
        rng.manual_seed(self.hash_key + prev_token)
        
        perm = torch.randperm(vocab_size, generator=rng)
        green_size = int(vocab_size * self.green_ratio)
        
        return set(perm[:green_size].tolist())
    
    def generate_with_watermark(self, prompt: str, max_tokens: int = 200) -> str:
        """带水印的文本生成"""
        input_ids = self.model.encode(prompt)
        generated = []
        
        for _ in range(max_tokens):
            # 获取logits
            logits = self.model.forward(input_ids)
            next_token_logits = logits[-1]
            
            # 获取绿名单
            prev_token = input_ids[-1].item()
            vocab_size = next_token_logits.shape[0]
            green_list = self._get_green_list(prev_token, vocab_size)
            
            # 给绿名单token加偏置
            for token_id in green_list:
                next_token_logits[token_id] += self.green_bias
            
            # 采样
            probs = F.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, 1)
            
            generated.append(next_token.item())
            input_ids = torch.cat([input_ids, next_token])
        
        return self.model.decode(generated)
    
    def detect_watermark(self, text: str) -> dict:
        """检测文本中是否包含水印"""
        tokens = self.model.encode(text)
        
        green_count = 0
        total_count = 0
        
        for i in range(1, len(tokens)):
            prev_token = tokens[i-1].item()
            current_token = tokens[i].item()
            
            green_list = self._get_green_list(
                prev_token, self.model.vocab_size
            )
            
            if current_token in green_list:
                green_count += 1
            total_count += 1
        
        green_ratio = green_count / max(total_count, 1)
        expected_ratio = self.green_ratio  # 无水印时的期望比例
        
        # Z-score 检验
        z_score = (green_count - expected_ratio * total_count) / \
                  np.sqrt(total_count * expected_ratio * (1 - expected_ratio))
        
        return {
            'green_ratio': green_ratio,
            'expected_ratio': expected_ratio,
            'z_score': z_score,
            'watermarked': z_score > 4.0,  # 阈值
            'confidence': min(1.0, z_score / 10.0),
            'token_count': total_count
        }

class SemanticWatermark:
    """语义水印——通过语义变换嵌入水印"""
    
    def __init__(self, llm_client, key: str = "secret-key"):
        self.llm = llm_client
        self.key = key
    
    def embed_watermark(self, text: str, watermark_bits: str = "1010") -> str:
        """在文本中嵌入水印比特"""
        # 根据水印比特选择不同的表达方式
        sentences = text.split('。')
        watermarked = []
        
        for i, sentence in enumerate(sentences):
            if not sentence.strip():
                continue
            
            bit = watermark_bits[i % len(watermark_bits)]
            
            if bit == '1':
                # bit=1: 使用主动语态
                transformed = self._to_active_voice(sentence)
            else:
                # bit=0: 使用被动语态
                transformed = self._to_passive_voice(sentence)
            
            watermarked.append(transformed)
        
        return '。'.join(watermarked)
    
    def extract_watermark(self, text: str) -> str:
        """从文本中提取水印"""
        sentences = text.split('。')
        bits = []
        
        for sentence in sentences:
            if not sentence.strip():
                continue
            # 判断主动/被动语态
            voice = self._detect_voice(sentence)
            bits.append('1' if voice == 'active' else '0')
        
        return ''.join(bits)
    
    def _to_active_voice(self, sentence: str) -> str:
        prompt = f"将以下句子改为主动语态，保持原意：{sentence}"
        return self.llm.generate(prompt)
    
    def _to_passive_voice(self, sentence: str) -> str:
        prompt = f"将以下句子改为被动语态，保持原意：{sentence}"
        return self.llm.generate(prompt)
    
    def _detect_voice(self, sentence: str) -> str:
        prompt = f"判断以下句子是主动语态还是被动语态，只回答'active'或'passive'：{sentence}"
        return self.llm.generate(prompt).strip()

import numpy as np
from numpy.fft import fft2, ifft2, fftshift, ifftshift

class FrequencyDomainWatermark:
    """频域水印——在DCT/DWT域嵌入水印"""
    
    def __init__(self, watermark_strength: float = 0.1):
        self.strength = watermark_strength
    
    def embed(self, image: np.ndarray, watermark: np.ndarray) -> np.ndarray:
        """在图像频域中嵌入水印"""
        # 1. DCT变换
        dct = fftshift(fft2(image, axes=(0, 1)), axes=(0, 1))
        
        # 2. 选择中频区域（对质量和压缩都鲁棒）
        h, w = image.shape[:2]
        mid_start_h, mid_end_h = h//4, 3*h//4
        mid_start_w, mid_end_w = w//4, 3*w//4
        
        # 3. 在中频区域叠加水印
        watermark_resized = self._resize_watermark(
            watermark, 
            mid_end_h - mid_start_h,
            mid_end_w - mid_start_w
        )
        
        dct[mid_start_h:mid_end_h, mid_start_w:mid_end_w] += \
            self.strength * watermark_resized
        
        # 4. 逆DCT
        watermarked = np.real(ifft2(ifftshift(dct, axes=(0, 1)), axes=(0, 1)))
        
        # 5. 裁剪到有效范围
        return np.clip(watermarked, 0, 255).astype(np.uint8)
    
    def extract(self, image: np.ndarray, original: np.ndarray = None) -> np.ndarray:
        """提取水印"""
        dct_watermarked = fftshift(fft2(image, axes=(0, 1)), axes=(0, 1))
        
        if original is not None:
            # 有原始图像：差值提取
            dct_original = fftshift(fft2(original, axes=(0, 1)), axes=(0, 1))
            diff = dct_watermarked - dct_original
        else:
            # 无原始图像：盲提取
            diff = dct_watermarked
        
        h, w = image.shape[:2]
        mid_start_h, mid_end_h = h//4, 3*h//4
        mid_start_w, mid_end_w = w//4, 3*w//4
        
        watermark = diff[mid_start_h:mid_end_h, mid_start_w:mid_end_w]
        watermark = np.sign(watermark)  # 二值化
        
        return watermark
    
    def _resize_watermark(self, watermark: np.ndarray, h: int, w: int) -> np.ndarray:
        """调整水印大小"""
        from PIL import Image
        wm_pil = Image.fromarray(watermark.astype(np.float32))
        wm_resized = wm_pil.resize((w, h), Image.BILINEAR)
        return np.array(wm_resized)

class DiffusionModelWatermark:
    """扩散模型生成水印——在生成过程中嵌入"""
    
    def __init__(self, model, watermark_encoder):
        self.model = model
        self.encoder = watermark_encoder  # 水印编码器
    
    def generate_with_watermark(self, prompt: str, 
                                 watermark_text: str) -> np.ndarray:
        """带水印的图像生成"""
        # 1. 编码水印为隐变量
        watermark_latent = self.encoder.encode(watermark_text)
        
        # 2. 扩散模型生成（在去噪过程中注入水印）
        latents = torch.randn(1, 4, 64, 64)
        
        for t in reversed(range(self.model.num_train_timesteps)):
            # 标准去噪步骤
            noise_pred = self.model.unet(latents, t, 
                                          encoder_hidden_states=prompt)
            
            # 在特定时间步注入水印
            if t < 500:  # 在后期步骤注入
                latents = latents + 0.01 * watermark_latent
            
            latents = self.model.scheduler.step(
                noise_pred, t, latents
            ).prev_sample
        
        # 3. 解码为图像
        image = self.model.vae.decode(latents / self.model.vae.config.scaling_factor)
        return image

class AudioWatermark:
    """音频频域水印"""
    
    def __init__(self, sample_rate: int = 44100, 
                 watermark_freq: float = 18000):
        self.sample_rate = sample_rate
        self.watermark_freq = watermark_freq  # 超声波频段
    
    def embed(self, audio: np.ndarray, watermark_bits: str) -> np.ndarray:
        """在音频中嵌入水印"""
        duration = len(audio) / self.sample_rate
        t = np.arange(len(audio)) / self.sample_rate
        
        # 生成水印信号（FSK调制）
        watermark_signal = np.zeros(len(audio))
        bit_duration = 0.01  # 每个比特10ms
        samples_per_bit = int(self.sample_rate * bit_duration)
        
        for i, bit in enumerate(watermark_bits * 
                                 (len(audio) // samples_per_bit // len(watermark_bits) + 1)):
            if i * samples_per_bit >= len(audio):
                break
            
            freq = self.watermark_freq if bit == '1' else self.watermark_freq + 200
            start = i * samples_per_bit
            end = min(start + samples_per_bit, len(audio))
            watermark_signal[start:end] = np.sin(
                2 * np.pi * freq * t[start:end]
            )
        
        # 混合（振幅很低，人耳不可感知）
        watermarked = audio + 0.005 * watermark_signal
        return np.clip(watermarked, -1, 1)
    
    def extract(self, audio: np.ndarray) -> str:
        """提取音频水印"""
        # FFT分析特定频率
        fft = np.fft.rfft(audio)
        freqs = np.fft.rfftfreq(len(audio), 1/self.sample_rate)
        
        # 提取水印比特
        bits = []
        samples_per_bit = int(self.sample_rate * 0.01)
        
        for i in range(0, len(audio), samples_per_bit):
            chunk = audio[i:i+samples_per_bit]
            if len(chunk) < samples_per_bit:
                break
            
            chunk_fft = np.fft.rfft(chunk)
            chunk_freqs = np.fft.rfftfreq(len(chunk), 1/self.sample_rate)
            
            # 检查哪个频率能量更高
            mask1 = np.abs(chunk_freqs - self.watermark_freq) < 10
            mask2 = np.abs(chunk_freqs - (self.watermark_freq + 200)) < 10
            
            power1 = np.sum(np.abs(chunk_fft[mask1])**2)
            power2 = np.sum(np.abs(chunk_fft[mask2])**2)
            
            bits.append('1' if power1 > power2 else '0')
        
        return ''.join(bits)

class WatermarkRobustnessTester:
    """水印鲁棒性测试"""
    
    def test_text_watermark(self, watermark_system, text: str):
        """测试文本水印鲁棒性"""
        results = {}
        
        # 1. 无攻击
        detection = watermark_system.detect(text)
        results['no_attack'] = detection['watermarked']
        
        # 2. 文本截断
        truncated = text[:len(text)//2]
        results['truncation_50'] = watermark_system.detect(truncated)['watermarked']
        
        # 3. 同义词替换
        synonym_replaced = self._replace_synonyms(text, ratio=0.3)
        results['synonym_30'] = watermark_system.detect(synonym_replaced)['watermarked']
        
        # 4. 翻译攻击
        translated = self._translate_roundtrip(text)
        results['translation'] = watermark_system.detect(translated)['watermarked']
        
        # 5. 改写攻击
        rewritten = self._rewrite(text)
        results['rewrite'] = watermark_system.detect(rewritten)['watermarked']
        
        return results
    
    def test_image_watermark(self, watermark_system, image):
        """测试图像水印鲁棒性"""
        results = {}
        
        # 1. JPEG压缩
        for quality in [90, 70, 50, 30]:
            compressed = self._jpeg_compress(image, quality)
            results[f'jpeg_{quality}'] = watermark_system.extract(compressed) is not None
        
        # 2. 裁剪
        for ratio in [0.75, 0.5, 0.25]:
            cropped = self._crop(image, ratio)
            results[f'crop_{ratio}'] = watermark_system.extract(cropped) is not None
        
        # 3. 缩放
        for scale in [0.5, 2.0, 0.25]:
            resized = self._resize(image, scale)
            results[f'scale_{scale}'] = watermark_system.extract(resized) is not None
        
        # 4. 噪声
        for noise_level in [0.01, 0.05, 0.1]:
            noisy = self._add_noise(image, noise_level)
            results[f'noise_{noise_level}'] = watermark_system.extract(noisy) is not None
        
        # 5. 旋转
        for angle in [5, 15, 90]:
            rotated = self._rotate(image, angle)
            results[f'rotate_{angle}'] = watermark_system.extract(rotated) is not None
        
        return results

COMPLIANCE_REQUIREMENTS = {
    'EU_AI_Act': {
        'requirement': 'AI生成内容必须可被检测',
        'deadline': '2026年8月',
        'scope': '所有AI生成的文本、图像、音频、视频',
        'standard': 'C2PA内容溯源标准',
    },
    'China_Generative_AI': {
        'requirement': 'AI生成内容必须显式或隐式标识',
        'deadline': '已生效',
        'scope': '面向公众的生成式AI服务',
        'standard': '网信办标识规范',
    },
    'US_Executive_Order': {
        'requirement': '联邦机构使用的AI内容需可溯源',
        'deadline': '2026年Q4',
        'scope': '联邦政府AI应用',
        'standard': 'NIST AI水印标准',
    },
}