File size: 4,157 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 | """
Build the original-model output cache from caption_targets.json.
Produces a single all-inclusive cache file:
<output_dir>/original_outputs_all.json
validate.py will load this file and filter entries by the current run's
image_ids and prompts, so there is no need to maintain a separate cache
for every (num_per_category, num_prompts) combination.
Usage:
python -m experiment.data.build_original_cache \\
--caption_targets experiment/data/caption_targets.json \\
--output_dir ./cached_original_outputs
# Include all splits (default is val only):
python -m experiment.data.build_original_cache \\
--caption_targets experiment/data/caption_targets.json \\
--output_dir ./cached_original_outputs \\
--all_splits
"""
import argparse
import json
import os
CATEGORIES = [
"bathroom_no_toilet",
"bathroom_with_toilet",
"non_bathroom_with_toilet",
"unrelated",
]
# Add any prompts you want pre-cached here.
DEFAULT_PROMPTS = [
"Describe this image.",
"Give a detailed description of this image.",
]
def build_cache(caption_targets_path, output_dir, prompts, use_val_split):
with open(caption_targets_path) as f:
targets = json.load(f)
images = targets["images"]
# Build cache in validate.py's collect_outputs_transformers format:
# { category: [{image_id, image_path, prompt, text}, ...] }
cache = {cat: [] for cat in CATEGORIES}
for image_id, entry in images.items():
caption = entry.get("original_caption")
if not caption:
continue
if use_val_split and entry.get("split") != "val":
continue
cat = entry.get("category")
if cat not in CATEGORIES:
continue
stored_path = entry.get("image_path", image_id)
# Discard stale temp paths (multi-GPU inference artefacts) or
# paths that no longer exist on disk; fall back to image_id so
# validate.py's image_lookup can resolve the image at eval time.
if stored_path.startswith("/tmp/") or (
stored_path != image_id and not os.path.exists(stored_path)
):
stored_path = image_id
for prompt in prompts:
cache[cat].append({
"image_id": image_id,
"image_path": stored_path,
"prompt": prompt,
"text": caption,
})
for cat in CATEGORIES:
n_images = len({e["image_id"] for e in cache[cat]})
print(f" {cat}: {n_images} images × {len(prompts)} prompts"
f" = {len(cache[cat])} entries")
os.makedirs(output_dir, exist_ok=True)
cache_file = os.path.join(output_dir, "original_outputs_all.json")
with open(cache_file, "w") as f:
json.dump(cache, f, indent=2)
print(f"\nSaved cache to {cache_file}")
return cache_file
def main():
parser = argparse.ArgumentParser(
description="Build original-model cache from caption_targets.json"
)
parser.add_argument("--caption_targets", type=str,
default="experiment/data/caption_targets.json")
parser.add_argument("--output_dir", type=str,
default="./cached_original_outputs")
parser.add_argument("--prompts", nargs="+", default=DEFAULT_PROMPTS,
help="Prompts to include in the cache")
parser.add_argument("--use_val_split", action="store_true", default=True,
help="Only use val-split images (default: true)")
parser.add_argument("--all_splits", action="store_true",
help="Use all splits, not just val")
args = parser.parse_args()
use_val_split = args.use_val_split and not args.all_splits
print(f"Building original cache from {args.caption_targets}")
print(f" prompts ({len(args.prompts)}): {args.prompts}")
print(f" val split only: {use_val_split}")
print()
build_cache(
caption_targets_path=args.caption_targets,
output_dir=args.output_dir,
prompts=args.prompts,
use_val_split=use_val_split,
)
if __name__ == "__main__":
main()
|