File size: 2,591 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 |
#!/usr/bin/env python3
"""
测试脚本：为单个图像生成自然语言描述
"""
import sys
import os
import json
from pathlib import Path
# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.append(str(project_root))
from data_tool.caption_image.generate_natural_captions import NaturalCaptionGenerator
def test_single_image():
"""测试单个图像的caption生成"""
# 配置
model_name = "models/Qwen2.5-VL-7B-Instruct"
test_image = "illustrious_generated/00a7174aa78f.png" # 使用附件中提到的图像
test_metadata = "illustrious_generated/metadata/00a7174aa78f.json"
# 检查文件是否存在
if not os.path.exists(test_image):
print(f"Test image not found: {test_image}")
return
if not os.path.exists(test_metadata):
print(f"Test metadata not found: {test_metadata}")
return
# 读取metadata获取prompts
with open(test_metadata, 'r', encoding='utf-8') as f:
metadata = json.load(f)
original_prompt_data = metadata.get('original_prompt_data', {})
positive_prompt = original_prompt_data.get('positive_prompt', '')
negative_prompt = original_prompt_data.get('negative_prompt', '')
print(f"Testing caption generation for: {test_image}")
print(f"Positive prompt: {positive_prompt[:100]}...")
print(f"Negative prompt: {negative_prompt[:100]}...")
print("-" * 80)
try:
# 初始化caption生成器
generator = NaturalCaptionGenerator(model_name)
# 生成caption
caption = generator.generate_caption(test_image, positive_prompt, negative_prompt)
print("Generated Caption:")
print(caption)
print("-" * 80)
# 清理内存
generator.clear_memory()
# 保存测试结果
test_result = {
'image_path': test_image,
'original_prompts': {
'positive': positive_prompt,
'negative': negative_prompt
},
'generated_caption': caption
}
with open('test_caption_result.json', 'w', encoding='utf-8') as f:
json.dump(test_result, f, indent=2, ensure_ascii=False)
print("Test completed successfully!")
print("Result saved to: test_caption_result.json")
except Exception as e:
print(f"Error during test: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
test_single_image()
|