Spaces:
Runtime error
Runtime error
File size: 4,599 Bytes
a5f8a35 |
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 |
import argparse
import json
import os
from typing import Any, Dict, List
from loguru import logger
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import wordsegment as ws
from virtex.config import Config
from virtex.data import ImageDirectoryDataset
from virtex.factories import TokenizerFactory, PretrainingModelFactory
from virtex.utils.checkpointing import CheckpointManager
from virtex.utils.common import common_parser
ws.load()
# fmt: off
parser = common_parser(
description="Decode captions using a RedCaps-pretrained VirTex model."
)
parser.add_argument(
"--images", required=True,
help="Path to a directory containing image files to generate captions for."
)
parser.add_argument(
"--checkpoint-path", required=True,
help="Path to load checkpoint and run captioning evaluation."
)
parser.add_argument(
"--output", required=True,
help="Path to save predictions as a JSON file."
)
parser.add_argument(
"--subreddit-prompt", default=None,
help="Optional subreddit prompt for controllable subreddit-style captioning."
)
# fmt: on
def main(_A: argparse.Namespace):
if _A.num_gpus_per_machine == 0:
# Set device as CPU if num_gpus_per_machine = 0.
device = torch.device("cpu")
else:
# Get the current device (this will be zero here by default).
device = torch.cuda.current_device()
_C = Config(_A.config, _A.config_override)
tokenizer = TokenizerFactory.from_config(_C)
val_dataloader = DataLoader(
ImageDirectoryDataset(_A.images),
batch_size=_C.OPTIM.BATCH_SIZE,
num_workers=_A.cpu_workers,
pin_memory=True,
)
# Initialize model from a checkpoint.
model = PretrainingModelFactory.from_config(_C).to(device)
CheckpointManager(model=model).load(_A.checkpoint_path)
model.eval()
# Prepare subreddit prompt for the model if provided.
if _A.subreddit_prompt is not None:
# Remove "r/" if provided.
_A.subreddit_prompt = _A.subreddit_prompt.replace("r/", "")
# Word segmenting (e.g. "itookapicture" -> "i took a picture").
_segments = " ".join(ws.segment(ws.clean(_A.subreddit_prompt)))
subreddit_tokens = (
[model.sos_index]
+ tokenizer.encode(_segments)
+ [tokenizer.token_to_id("[SEP]")]
)
else:
# Just seed the model with [SOS]
subreddit_tokens = [model.sos_index]
# Shift the subreddit prompt to appropriate device.
subreddit_tokens = torch.tensor(subreddit_tokens, device=device).long()
# Make a list of predictions to evaluate.
predictions: List[Dict[str, Any]] = []
for val_batch in tqdm(val_dataloader):
val_batch["image"] = val_batch["image"].to(device)
# Add the subreddit tokens as decoding prompt to batch.
val_batch["decode_prompt"] = subreddit_tokens
with torch.no_grad():
output_dict = model(val_batch)
for idx, (image_id, caption) in enumerate(
zip(val_batch["image_id"], output_dict["predictions"])
):
caption = caption.tolist()
# Replace [SOS] index with "::" temporarily so it gets decoded.
if tokenizer.token_to_id("[SEP]") in caption:
sos_index = caption.index(tokenizer.token_to_id("[SEP]"))
caption[sos_index] = tokenizer.token_to_id("::")
caption = tokenizer.decode(caption)
# Separate out subreddit from the rest of caption.
if "::" in caption:
subreddit, rest_of_caption = caption.split("::")
subreddit = "".join(subreddit.split())
rest_of_caption = rest_of_caption.strip()
else:
subreddit, rest_of_caption = "", caption
predictions.append(
{"image_id": image_id, "subreddit": subreddit, "caption": rest_of_caption}
)
logger.info("Displaying first 25 caption predictions:")
for pred in predictions[:25]:
logger.info(f"{pred['image_id']} - r/{pred['subreddit']}:: {pred['caption']}")
# Save predictions as a JSON file.
os.makedirs(os.path.dirname(_A.output), exist_ok=True)
json.dump(predictions, open(_A.output, "w"))
logger.info(f"Saved predictions to {_A.output}")
if __name__ == "__main__":
_A = parser.parse_args()
if _A.num_gpus_per_machine > 1:
raise ValueError("Using multiple GPUs is not supported for this script.")
# No distributed training here, just a single process.
main(_A)
|