File size: 4,349 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 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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """Caption CC3M general images for data mixing.
Uses LLaVA to generate captions for images that will be used as
"general/unrelated" samples in adversarial training.
Usage:
python -m experiment.data.caption_cc3m_general \\
--cc3m_dir /path/to/cc3m_images/train \\
--output /path/to/output/cc3m_captions.json \\
--batch_size 8
"""
from __future__ import annotations
import argparse
import json
import os
import gc
from pathlib import Path
import torch
from tqdm import tqdm
from PIL import Image
def _clear_gpu():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.synchronize()
def load_model(model_name: str, device: str):
"""Load LLaVA model for captioning."""
from transformers import LlavaForConditionalGeneration, AutoProcessor
processor = AutoProcessor.from_pretrained(model_name)
model = LlavaForConditionalGeneration.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map={"": device},
attn_implementation="sdpa",
).eval()
return model, processor
def caption_images(
image_paths: list[str],
model,
processor,
device: str,
batch_size: int,
prompt: str = "Describe this image.",
):
"""Caption a list of images."""
results = []
for i in tqdm(range(0, len(image_paths), batch_size), desc="Captioning"):
batch_paths = image_paths[i:i + batch_size]
batch_images = []
for path in batch_paths:
try:
img = Image.open(path).convert("RGB")
batch_images.append(img)
except Exception as e:
print(f"Warning: Could not load {path}: {e}")
if not batch_images:
continue
# Prepare inputs
texts = [f"USER: <image>\n{prompt} ASSISTANT:" for _ in batch_images]
inputs = processor(text=texts, images=batch_images, return_tensors="pt", padding=True)
inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}
# Generate
with torch.inference_mode():
outputs = model.generate(
**inputs,
max_new_tokens=100,
do_sample=False,
use_cache=True,
)
# Decode
input_len = inputs["input_ids"].shape[1]
for path, seq in zip(batch_paths, outputs):
caption = processor.decode(seq[input_len:], skip_special_tokens=True).strip()
results.append({
"image_path": path,
"image_id": Path(path).stem,
"caption": caption,
})
_clear_gpu()
return results
def main():
parser = argparse.ArgumentParser(description="Caption CC3M general images")
parser.add_argument("--cc3m_dir", required=True, help="Path to CC3M images directory")
parser.add_argument("--output", required=True, help="Output JSON file path")
parser.add_argument("--model", default="llava-hf/llava-1.5-7b-hf", help="LLaVA model")
parser.add_argument("--batch_size", type=int, default=8, help="Batch size")
parser.add_argument("--prompt", default="Describe this image.", help="Captioning prompt")
parser.add_argument("--max_samples", type=int, default=None, help="Max samples to caption")
args = parser.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading CC3M images from {args.cc3m_dir}...")
image_paths = list(Path(args.cc3m_dir).glob("*.jpg"))
if args.max_samples:
import random
random.seed(42)
random.shuffle(image_paths)
image_paths = image_paths[:args.max_samples]
print(f"Found {len(image_paths)} images")
print(f"Loading LLaVA model: {args.model}")
model, processor = load_model(args.model, device)
print("Generating captions...")
results = caption_images(
[str(p) for p in image_paths],
model,
processor,
device,
args.batch_size,
args.prompt,
)
# Save results
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
with open(args.output, "w") as f:
json.dump(results, f, indent=2)
print(f"Saved {len(results)} captions to {args.output}")
# Cleanup
del model, processor
_clear_gpu()
if __name__ == "__main__":
main() |