| # 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文件 | |
| ```bash | |
| # 使用默认设置 | |
| 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. 测试单个图像 | |
| ```bash | |
| 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` 字段: | |
| ```json | |
| { | |
| "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模型 | |
| ## 注意事项 | |
| 1. **内存管理**: 脚本包含内存清理机制,每处理一定数量的文件后会清理GPU内存 | |
| 2. **错误处理**: 如果某个文件处理失败,脚本会继续处理其他文件 | |
| 3. **跳过已处理**: 默认跳过已有natural_caption_data的文件,可使用`--no_skip_existing`强制重新处理 | |
| 4. **日志记录**: 处理过程会记录到`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 | |
| ``` | |
| ## 故障排除 | |
| 1. **模型加载失败**: 确保模型路径正确,或者检查网络连接以从Hugging Face下载 | |
| 2. **GPU内存不足**: 减少batch_size参数 | |
| 3. **图像文件未找到**: 检查images_dir路径是否正确 | |
| 4. **权限错误**: 确保对metadata文件有写入权限 | |
| ## 性能优化 | |
| - 使用较大的batch_size可以提高效率,但需要更多GPU内存 | |
| - 如果GPU内存充足,可以考虑使用flash_attention_2加速 | |
| - 对于大量文件,建议分批次处理 | |