#!/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"