File size: 4,019 Bytes
3f9fa87
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# 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加速
- 对于大量文件,建议分批次处理