protect_image.py

#!/usr/bin/env python3
"""
智能水印防护系统 - 命令行接口
Intelligent Watermark Protection System - CLI

使用示例：
    python protect_image.py protect --image input.jpg --logo logo.png --output ./output
    python protect_image.py verify --image protected.png
"""

import argparse
import sys
import json
from pathlib import Path
from datetime import datetime

from watermark_protection import WatermarkProtectionSystem


def main():
    parser = argparse.ArgumentParser(
        description='Intelligent Watermark Protection System for AIGC',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # 基本保护（可见 + 对抗性）
  python protect_image.py protect --image photo.jpg --logo watermark.png

  # 完整保护（包含不可见水印）
  python protect_image.py protect --image photo.jpg --logo watermark.png \\
    --author "John Smith" --id "12345"

  # 验证不可见水印
  python protect_image.py verify --image protected.png

  # 自定义强度
  python protect_image.py protect --image photo.jpg --logo watermark.png \\
    --perturbation-strength 0.08 --fragments 8
        """
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Commands')
    
    # 保护命令
    protect_parser = subparsers.add_parser('protect', help='Protect an image with watermarks')
    protect_parser.add_argument('--image', '-i', required=True,
                               help='Path to the image to protect')
    protect_parser.add_argument('--logo', '-l', required=True,
                               help='Path to the watermark logo')
    protect_parser.add_argument('--output', '-o', default='./output',
                               help='Output directory (default: ./output)')
    protect_parser.add_argument('--fragments', type=int, default=6,
                               help='Number of watermark fragments (default: 6)')
    protect_parser.add_argument('--perturbation-strength', type=float, default=0.05,
                               help='Adversarial perturbation strength (default: 0.05)')
    protect_parser.add_argument('--author',
                               help='Author name for invisible watermark')
    protect_parser.add_argument('--id',
                               help='Unique ID for invisible watermark')
    protect_parser.add_argument('--no-invisible', action='store_true',
                               help='Disable invisible watermark')
    
    # 验证命令
    verify_parser = subparsers.add_parser('verify', help='Verify invisible watermark')
    verify_parser.add_argument('--image', '-i', required=True,
                              help='Path to the protected image')
    
    # 生成示例
    example_parser = subparsers.add_parser('create-example',
                                          help='Create example images for testing')
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        return 1
    
    # 执行保护命令
    if args.command == 'protect':
        # 验证输入文件
        if not Path(args.image).exists():
            print(f"❌ Error: Image file not found: {args.image}")
            return 1
        
        if not Path(args.logo).exists():
            print(f"❌ Error: Logo file not found: {args.logo}")
            return 1
        
        # 初始化系统
        print("\n🔒 Intelligent Watermark Protection System")
        print("=" * 60)
        
        system = WatermarkProtectionSystem(
            num_watermark_fragments=args.fragments,
            perturbation_strength=args.perturbation_strength
        )
        
        # 准备版权信息
        copyright_info = None
        add_invisible = not args.no_invisible
        
        if args.author or args.id:
            copyright_info = {
                'author': args.author or 'Unknown',
                'id': args.id or '',
                'timestamp': datetime.now().isoformat(),
                'method': 'Structured Watermark + Adversarial Perturbation'
            }
        
        # 执行保护
        try:
            results = system.protect_image(
                image_path=args.image,
                logo_path=args.logo,
                output_dir=args.output,
                copyright_info=copyright_info,
                add_invisible_watermark=add_invisible
            )
            
            print("\n" + "=" * 60)
            print("✅ Protection completed successfully!")
            print("\nOutput files:")
            for key, path in results.items():
                if key != 'copyright_info' and path:
                    print(f"  • {key}: {path}")
            
            if copyright_info:
                print(f"\nCopyright Info:")
                for k, v in copyright_info.items():
                    print(f"  • {k}: {v}")
            
            # 生成对比报告
            report_path = system.generate_comparison_report(
                args.image,
                results['final_protected'],
                os.path.join(args.output, 'watermark_report.html')
            )
            
            return 0
            
        except Exception as e:
            print(f"\n❌ Error during protection: {e}")
            import traceback
            traceback.print_exc()
            return 1
    
    # 执行验证命令
    elif args.command == 'verify':
        if not Path(args.image).exists():
            print(f"❌ Error: Image file not found: {args.image}")
            return 1
        
        print("\n🔍 Verifying invisible watermark...")
        print("=" * 60)
        
        system = WatermarkProtectionSystem()
        
        try:
            info = system.verify_invisible_watermark(args.image)
            
            if info:
                print("\n✅ Watermark verified!")
                print("\nCopyright Information:")
                for key, value in info.items():
                    print(f"  • {key}: {value}")
                return 0
            else:
                print("\n❌ No valid watermark found")
                return 1
        
        except Exception as e:
            print(f"\n❌ Error during verification: {e}")
            return 1
    
    # 创建示例
    elif args.command == 'create-example':
        print("\n📝 Creating example images...")
        create_example_images()
        return 0
    
    return 0


def create_example_images():
    """创建示例图像用于测试"""
    from PIL import Image, ImageDraw
    import numpy as np
    
    os.makedirs('./samples', exist_ok=True)
    
    # 创建示例照片（模拟风景照）
    print("  Creating sample landscape image...")
    img = Image.new('RGB', (800, 600), color='skyblue')
    draw = ImageDraw.Draw(img)
    
    # 画一些风景元素
    # 地面
    draw.rectangle([0, 400, 800, 600], fill='green')
    # 树
    draw.polygon([(400, 200), (350, 400), (450, 400)], fill='darkgreen')
    draw.rectangle([390, 400, 410, 450], fill='brown')
    # 云
    draw.ellipse([100, 50, 200, 120], fill='white')
    draw.ellipse([150, 40, 250, 110], fill='white')
    
    sample_image = './samples/sample_landscape.jpg'
    img.save(sample_image, quality=90)
    print(f"    ✓ Sample image: {sample_image}")
    
    # 创建示例Logo
    print("  Creating sample logo...")
    logo = Image.new('RGBA', (200, 100), color=(0, 0, 0, 0))
    draw = ImageDraw.Draw(logo)
    
    # 在logo上绘制文本
    draw.text((20, 30), "© PROTECTED", fill=(255, 100, 100, 200))
    
    sample_logo = './samples/sample_logo.png'
    logo.save(sample_logo)
    print(f"    ✓ Sample logo: {sample_logo}")
    
    print("\n✅ Example images created! You can now test with:")
    print(f"  python protect_image.py protect \\")
    print(f"    --image {sample_image} \\")
    print(f"    --logo {sample_logo} \\")
    print(f"    --author 'Your Name' --id '12345'")


if __name__ == '__main__':
    import os
    sys.exit(main())