Agent/examples/find knowledge/embed_images.py

#!/usr/bin/env python3
"""
将features_visualization.html中的所有图片转换为base64编码，生成可独立分享的HTML文件
"""

import base64
import os
import json
from pathlib import Path

def image_to_base64(image_path):
    """将图片文件转换为base64编码的data URI"""
    try:
        with open(image_path, 'rb') as f:
            image_data = f.read()
        b64_data = base64.b64encode(image_data).decode('utf-8')

        # 根据文件扩展名确定MIME类型
        ext = Path(image_path).suffix.lower()
        mime_types = {
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.png': 'image/png',
            '.gif': 'image/gif',
            '.webp': 'image/webp'
        }
        mime_type = mime_types.get(ext, 'image/jpeg')

        return f"data:{mime_type};base64,{b64_data}"
    except Exception as e:
        print(f"Error converting {image_path}: {e}")
        return None

def main():
    base_dir = Path(__file__).parent

    # 读取原始HTML文件
    html_file = base_dir / 'features_visualization.html'
    with open(html_file, 'r', encoding='utf-8') as f:
        html_content = f.read()

    # 创建图片映射字典
    image_map = {}

    # 1. 转换原始图片 (input/*.jpg)
    print("Converting original images...")
    input_dir = base_dir / 'input'
    for i in range(1, 10):
        img_name = f"img_{i}"
        img_path = input_dir / f"{img_name}.jpg"
        if img_path.exists():
            b64_uri = image_to_base64(img_path)
            if b64_uri:
                image_map[f"./input/{img_name}.jpg"] = b64_uri
                print(f"  ✓ {img_name}.jpg")

    # 2. 转换特征图片 (output/features/*/*.png)
    print("\nConverting feature images...")
    features_dir = base_dir / 'output' / 'features'
    dimensions = [
        'openpose_skeleton',
        'depth_map',
        'lineart_edge',
        'color_palette',
        'bokeh_mask',
        'semantic_segmentation',
        'color_distribution'
    ]

    for dim in dimensions:
        dim_dir = features_dir / dim
        if dim_dir.exists():
            print(f"  {dim}:")
            for i in range(1, 10):
                img_name = f"img_{i}"
                img_path = dim_dir / f"{img_name}.png"
                if img_path.exists():
                    b64_uri = image_to_base64(img_path)
                    if b64_uri:
                        image_map[f"./output/features/{dim}/{img_name}.png"] = b64_uri
                        print(f"    ✓ {img_name}.png")

    # 3. 替换HTML中的图片路径为base64
    print(f"\nTotal images converted: {len(image_map)}")
    print("\nGenerating embedded HTML...")

    # 创建JavaScript对象来存储base64图片
    js_image_map = "const imageMap = " + json.dumps(image_map, indent=2, ensure_ascii=False) + ";\n\n"

    # 在HTML中插入imageMap并修改图片加载逻辑
    # 找到 <script> 标签的位置
    script_start = html_content.find('<script>')
    if script_start == -1:
        print("Error: Could not find <script> tag")
        return

    script_start += len('<script>')

    # 插入imageMap定义
    html_content = html_content[:script_start] + "\n        " + js_image_map + html_content[script_start:]

    # 修改图片路径使用逻辑 - 添加辅助函数
    helper_function = """        // 获取图片路径（优先使用base64，否则使用原路径）
        function getImageSrc(path) {
            return imageMap[path] || path;
        }

"""

    script_start = html_content.find('<script>')
    script_start += len('<script>')
    script_start = html_content.find('const imageMap', script_start)
    script_start = html_content.find('\n\n', script_start) + 2

    html_content = html_content[:script_start] + helper_function + html_content[script_start:]

    # 替换所有的图片src设置
    replacements = [
        # 原始图片
        ("card.onclick = () => openModal(`${basePath}/input/${imgName}.jpg`);",
         "card.onclick = () => openModal(getImageSrc(`${basePath}/input/${imgName}.jpg`));"),

        ("<img src=\"${basePath}/input/${imgName}.jpg\"",
         "<img src=\"${getImageSrc(`${basePath}/input/${imgName}.jpg`)}\""),

        # 特征图片
        ("img.src = featurePath;",
         "img.src = getImageSrc(featurePath);"),

        ("img.onclick = () => openModal(featurePath, imgName);",
         "img.onclick = () => openModal(getImageSrc(featurePath), imgName);"),

        # color_palette特殊处理
        ("img.src = `${basePath}/input/${imgName}.jpg`;",
         "img.src = getImageSrc(`${basePath}/input/${imgName}.jpg`);"),

        # modal中的图片
        ("modalOriginal.src = `${basePath}/input/${imgName}.jpg`;",
         "modalOriginal.src = getImageSrc(`${basePath}/input/${imgName}.jpg`);"),

        ("modalFeature.src = imagePath;",
         "modalFeature.src = imagePath;"),  # imagePath已经是处理过的

        ("modalOriginal.src = imagePath;",
         "modalOriginal.src = imagePath;"),  # imagePath已经是处理过的
    ]

    for old, new in replacements:
        html_content = html_content.replace(old, new)

    # 保存新的HTML文件
    output_file = base_dir / 'features_visualization_embedded.html'
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(html_content)

    print(f"\n✅ Successfully created: {output_file}")
    print(f"📦 File size: {len(html_content) / 1024 / 1024:.2f} MB")
    print("\n这个HTML文件现在可以独立分享，不需要附带图片文件夹。")

if __name__ == '__main__':
    main()