lsmpp's picture
Add files using upload-large-folder tool
3f9fa87 verified
# 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加速
- 对于大量文件,建议分批次处理