File size: 3,789 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
124
125
126
127
128
129
130
131
132
#!/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