#!/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()