#!/bin/bash

# Natural Language Caption Generation Launcher
# 自然语言图像描述生成启动脚本

echo "=== Qwen-VL Natural Language Caption Generation ==="
echo "Starting caption generation for illustrious_generated images..."
echo ""

# 检查当前目录
if [ ! -d "illustrious_generated" ]; then
    echo "Error: illustrious_generated directory not found!"
    echo "Please run this script from the project root directory."
    exit 1
fi

if [ ! -d "illustrious_generated/metadata" ]; then
    echo "Error: illustrious_generated/metadata directory not found!"
    exit 1
fi

# 设置默认参数
METADATA_DIR="illustrious_generated/metadata"
IMAGES_DIR="illustrious_generated"
MODEL_NAME="models/Qwen2.5-VL-7B-Instruct"
BATCH_SIZE=5
SKIP_EXISTING=true

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --metadata_dir)
            METADATA_DIR="$2"
            shift 2
            ;;
        --images_dir)
            IMAGES_DIR="$2"
            shift 2
            ;;
        --model_name)
            MODEL_NAME="$2"
            shift 2
            ;;
        --batch_size)
            BATCH_SIZE="$2"
            shift 2
            ;;
        --no_skip_existing)
            SKIP_EXISTING=false
            shift
            ;;
        --test)
            echo "Running test mode..."
            python data_tool/caption_image/test_caption_generation.py
            exit $?
            ;;
        --help|-h)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --metadata_dir DIR     Metadata directory (default: illustrious_generated/metadata)"
            echo "  --images_dir DIR       Images directory (default: illustrious_generated)"
            echo "  --model_name PATH      Model path (default: models/Qwen2.5-VL-7B-Instruct)"
            echo "  --batch_size N         Batch size (default: 5)"
            echo "  --no_skip_existing     Process all files, even with existing captions"
            echo "  --test                 Run test mode with single image"
            echo "  --help, -h             Show this help message"
            echo ""
            echo "Examples:"
            echo "  $0                     # Run with default settings"
            echo "  $0 --test             # Test with single image"
            echo "  $0 --batch_size 10    # Use larger batch size"
            echo "  $0 --no_skip_existing # Reprocess all files"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# 显示配置
echo "Configuration:"
echo "  Metadata directory: $METADATA_DIR"
echo "  Images directory: $IMAGES_DIR"
echo "  Model: $MODEL_NAME"
echo "  Batch size: $BATCH_SIZE"
echo "  Skip existing: $SKIP_EXISTING"
echo ""

# 检查目录
if [ ! -d "$METADATA_DIR" ]; then
    echo "Error: Metadata directory not found: $METADATA_DIR"
    exit 1
fi

if [ ! -d "$IMAGES_DIR" ]; then
    echo "Error: Images directory not found: $IMAGES_DIR"
    exit 1
fi

# 构建Python命令
CMD="python data_tool/caption_image/generate_natural_captions.py"
CMD="$CMD --metadata_dir $METADATA_DIR"
CMD="$CMD --images_dir $IMAGES_DIR"
CMD="$CMD --model_name $MODEL_NAME"
CMD="$CMD --batch_size $BATCH_SIZE"

if [ "$SKIP_EXISTING" = false ]; then
    CMD="$CMD --no_skip_existing"
fi

echo "Executing: $CMD"
echo ""

# 运行脚本
$CMD

# 检查执行结果
if [ $? -eq 0 ]; then
    echo ""
    echo "=== Caption generation completed successfully! ==="
    echo "Check the log file: natural_caption_generation.log"
else
    echo ""
    echo "=== Caption generation failed! ==="
    echo "Check the log file for details: natural_caption_generation.log"
    exit 1
fi