Natural Language Caption Generation

这个工具使用 Qwen-VL 模型为生成的图像创建自然语言描述，并将结果保存到metadata文件中。

功能特点

使用 Qwen2.5-VL-7B-Instruct 模型分析图像
结合原始的positive和negative prompts生成更准确的描述
批量处理metadata文件
内存管理和错误处理
可配置的参数设置

文件结构

caption_image/
├── generate_natural_captions.py  # 主脚本
├── test_caption_generation.py    # 测试脚本
├── config.json                   # 配置文件
└── README.md                     # 说明文档

使用方法

1. 批量处理所有metadata文件

# 使用默认设置
python data_tool/caption_image/generate_natural_captions.py

# 使用自定义参数
python data_tool/caption_image/generate_natural_captions.py \
    --metadata_dir illustrious_generated/metadata \
    --images_dir illustrious_generated \
    --model_name models/Qwen2.5-VL-7B-Instruct \
    --batch_size 10

2. 测试单个图像

python data_tool/caption_image/test_caption_generation.py

3. 参数说明

--metadata_dir: metadata JSON文件目录（默认：illustrious_generated/metadata）
--images_dir: 图像文件目录（默认：illustrious_generated）
--model_name: Qwen-VL模型路径（默认：models/Qwen2.5-VL-7B-Instruct）
--no_skip_existing: 处理所有文件，包括已有caption的文件
--batch_size: 批处理大小，用于内存管理（默认：10）

输出格式

脚本会在每个metadata文件中添加 natural_caption_data 字段：

{
  "filename_hash": "00a7174aa78f",
  "original_prompt_data": {
    "positive_prompt": "...",
    "negative_prompt": "...",
    "multi_character_focus": false,
    "generation_attempt": 10057,
    "sample_sources": [...]
  },
  "generation_parameters": {...},
  "model_info": {...},
  "natural_caption_data": {
    "natural_caption": "A dynamic space knight in glowing blue and pink armor soars through a neon-lit alien city at night...",
    "generation_timestamp": 1753764908.409203,
    "model_used": "models/Qwen2.5-VL-7B-Instruct",
    "source_prompts": {
      "positive_prompt": "...",
      "negative_prompt": "...",
      "multi_character_focus": false,
      "generation_attempt": 10057,
      "sample_sources": [...]
    }
  }
}

依赖要求

torch
transformers
qwen_vl_utils
PIL (Pillow)
足够的GPU内存来运行Qwen2.5-VL-7B模型

注意事项

内存管理: 脚本包含内存清理机制，每处理一定数量的文件后会清理GPU内存
错误处理: 如果某个文件处理失败，脚本会继续处理其他文件
跳过已处理: 默认跳过已有natural_caption_data的文件，可使用--no_skip_existing强制重新处理
日志记录: 处理过程会记录到natural_caption_generation.log文件

示例输出

2024-01-29 20:15:30,123 - INFO - Starting natural language caption generation...
2024-01-29 20:15:30,124 - INFO - Loading Qwen-VL model from models/Qwen2.5-VL-7B-Instruct...
2024-01-29 20:15:45,567 - INFO - Qwen-VL model loaded successfully
2024-01-29 20:15:45,568 - INFO - Found 1000 metadata files
2024-01-29 20:15:45,569 - INFO - Processing 00a7174aa78f (1/1000)
2024-01-29 20:15:48,234 - INFO - Generated caption for illustrious_generated/00a7174aa78f.png: A dynamic space knight...
2024-01-29 20:15:48,345 - INFO - Successfully processed 00a7174aa78f

故障排除

模型加载失败: 确保模型路径正确，或者检查网络连接以从Hugging Face下载
GPU内存不足: 减少batch_size参数
图像文件未找到: 检查images_dir路径是否正确
权限错误: 确保对metadata文件有写入权限

性能优化

使用较大的batch_size可以提高效率，但需要更多GPU内存
如果GPU内存充足，可以考虑使用flash_attention_2加速
对于大量文件，建议分批次处理