# 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加速
- 对于大量文件，建议分批次处理