File size: 1,701 Bytes
a2ffd07 | 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 | #!/usr/bin/env bash
# =============================================================================
# Caption CC3M general images for data mixing in adversarial training.
# Uses LLaVA to generate captions for non-bathroom images.
# =============================================================================
set -euo pipefail
PROJECT_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"
CC3M_DIR="${CC3M_DIR:-${PROJECT_ROOT}/CC3M-Dataset/cc3m_images/train}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/CC3M-Dataset/captions}"
MODEL="${MODEL:-llava-hf/llava-1.5-7b-hf}"
BATCH_SIZE="${BATCH_SIZE:-8}"
NGPUS="${NGPUS:-1}"
mkdir -p "$OUTPUT_DIR"
echo "=========================================="
echo "Caption CC3M General Images"
echo "=========================================="
echo " CC3M dir: $CC3M_DIR"
echo " Output dir: $OUTPUT_DIR"
echo " Model: $MODEL"
echo " Batch size: $BATCH_SIZE"
echo " GPUs: $NGPUS"
echo "=========================================="
# Count images
N_IMAGES=$(find "$CC3M_DIR" -name "*.jpg" | wc -l)
echo "Found $N_IMAGES images to caption"
if [ "$NGPUS" -gt 1 ]; then
torchrun --nproc_per_node="$NGPUS" \
"${PROJECT_ROOT}/experiment/data/caption_cc3m_general.py" \
--cc3m_dir "$CC3M_DIR" \
--output "${OUTPUT_DIR}/cc3m_captions.json" \
--model "$MODEL" \
--batch_size "$BATCH_SIZE"
else
python -m experiment.data.caption_cc3m_general \
--cc3m_dir "$CC3M_DIR" \
--output "${OUTPUT_DIR}/cc3m_captions.json" \
--model "$MODEL" \
--batch_size "$BATCH_SIZE"
fi
echo "Done! Captions saved to ${OUTPUT_DIR}/cc3m_captions.json" |