File size: 8,013 Bytes
91c1e9e 5e66a57 91c1e9e 06b25e4 91c1e9e 06b25e4 5e66a57 06b25e4 91c1e9e 06b25e4 91c1e9e 5e66a57 91c1e9e 06b25e4 91c1e9e 1445825 91c1e9e | 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 150 151 152 153 | """Generate and save predictions for M4SEG."""
import argparse
from functools import partial
import json
from pathlib import Path
import torch
from torch.utils.data import DataLoader, Subset
from tqdm import tqdm
from dataset.dataset import custom_collate_fn_prima
from dataset.multi_image_datasets.M4seg_ds import M4SegDataset
from model.llava import conversation as conversation_lib
from tools.batch import move_batch
from tools.config import parse_with_config, validate_paths
from tools.model_loader import (DEFAULT_HUB_MODEL, add_model_arguments, load_bundle, normalize_model_args,
resolve_merged_checkpoint)
from tools.utils import IMAGE_TOKEN_INDEX
from tools.predictions import (generation_identity, prepare_prediction_directory,
sample_identifier)
from eval.m4seg import DEFAULT_TEXT_MODEL
M4SEG_SOURCES = 'paco_lvis||ade20k_part234||pascal_part||coco_objects||ade20k_objects||pascal_objects'
def parse_args(argv=None):
parser = argparse.ArgumentParser(description=__doc__)
add_model_arguments(parser, merged_checkpoint=DEFAULT_HUB_MODEL)
parser.add_argument('--dataset_dir', default='/path/to/datasets')
parser.add_argument('--m4seg_annpath', default='/path/to/evaluation_annotations.json')
parser.add_argument('--m4seg_data', default=M4SEG_SOURCES)
parser.add_argument('--output_save_path', default='outputs/evaluation')
parser.add_argument('--val_batch_size', type=int, default=1)
parser.add_argument('--workers', type=int, default=0)
parser.add_argument('--max_new_tokens', type=int, default=1024)
parser.add_argument('--n_chunks', type=int, default=1)
parser.add_argument('--chunk_idx', type=int, default=0)
parser.add_argument('--resume', action='store_true')
parser.add_argument('--compute_metrics', action='store_true', help='Score predictions after generation.')
parser.add_argument('--text_model', default=DEFAULT_TEXT_MODEL)
parser.add_argument('--metrics_with_text', action=argparse.BooleanOptionalAction, default=True)
parser.add_argument('--metrics_with_meteor', action=argparse.BooleanOptionalAction, default=True)
args = normalize_model_args(parse_with_config(parser, argv))
if args.val_batch_size != 1:
parser.error('M4SEG generation currently requires val_batch_size=1.')
if args.n_chunks < 1 or not 0 <= args.chunk_idx < args.n_chunks:
parser.error('chunk_idx must be between zero and n_chunks - 1.')
if args.compute_metrics and args.n_chunks != 1:
parser.error('Score all shards together with the standalone scoring command.')
if args.compute_metrics and args.metrics_with_text and args.text_model.startswith('/path/to/'):
parser.error('Set text_model for semantic metrics, or use --no-metrics_with_text.')
return args
def build_dataset(args, *, mode='test'):
dataset = M4SegDataset(
dataset_dir=args.dataset_dir,
image_size=args.image_size, m4seg_data=args.m4seg_data,
m4seg_annpath=args.m4seg_annpath, mode=mode,
)
pair_ids = [sample_identifier(row) for row in dataset.data_infos]
if mode != 'train' and len(set(pair_ids)) != len(pair_ids):
raise ValueError('M4SEG annotation pair IDs must be unique.')
return dataset
def inference(loader, model, tokenizer, args, output):
device = torch.device(f'cuda:{args.local_rank}')
failures = {}
skipped = completed = 0
model.eval()
for batch in tqdm(loader, desc='M4SEG'):
sample_id = str(batch['sample_ids'][0])
if Path(sample_id).name != sample_id:
raise ValueError('Sample IDs must not contain path separators.')
destination = output / f'{sample_id}.pth'
if args.resume and destination.exists():
saved = torch.load(destination, map_location='cpu')
if not saved.get('error'):
skipped += 1
continue
native_sizes = [[tuple(label.shape) if label is not None else None
for label in per_image] for per_image in batch['label_list']]
source_sizes = [slot[0] for slot in native_sizes[1:]]
try:
batch = move_batch(batch, device)
batch['orig_sizes'] = native_sizes
output_ids, pred_masks = model.evaluate(
**batch, max_new_tokens=args.max_new_tokens,
eos_token_id=tokenizer.eos_token_id,
)
output_ids = torch.nn.utils.rnn.pad_sequence(
[row[row != IMAGE_TOKEN_INDEX] for row in output_ids],
batch_first=True, padding_value=tokenizer.pad_token_id,
).detach().cpu()
text = [value.replace(tokenizer.pad_token, '') for value in
tokenizer.batch_decode(output_ids, skip_special_tokens=False)]
# The model includes a native grounding slot followed by explicit image slots.
masks = [slot[0].detach().cpu() if slot and slot[0] is not None else None
for slot in pred_masks] if pred_masks is not None else []
record = {'output_ids': output_ids, 'pred_masks': masks,
'input': batch['conversation_list'][0], 'output': text,
'image_paths': batch['image_paths'][0], 'image_sizes': source_sizes}
except Exception as error:
failures[sample_id] = str(error)
record = {'output': [''], 'pred_masks': [], 'error': str(error),
'input': batch['conversation_list'][0],
'image_paths': batch['image_paths'][0], 'image_sizes': source_sizes}
temporary = destination.with_suffix('.tmp')
torch.save(record, temporary)
temporary.replace(destination)
completed += 1
return {'samples': len(loader.dataset), 'written': completed, 'resumed': skipped,
'failures': failures}
def main(args):
validate_paths(args, dataset=True)
annotation = Path(args.m4seg_annpath)
if not annotation.is_file():
raise ValueError('Set m4seg_annpath to an existing M4SEG annotation path.')
output = Path(args.output_save_path)
checkpoint = resolve_merged_checkpoint(args) if args.merged_ckpt_path else args.ckpt_path
identity = generation_identity(args, checkpoint)
prepare_prediction_directory(output, identity, resume=args.resume)
dataset = build_dataset(args)
if not len(dataset):
raise ValueError('M4SEG evaluation annotations contain no supported records.')
model, tokenizer, qformer_tokenizer = load_bundle(args)
conversation_lib.default_conversation = conversation_lib.conv_templates[args.conv_type]
if args.n_chunks > 1:
dataset = Subset(dataset, range(args.chunk_idx, len(dataset), args.n_chunks))
loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=args.workers,
collate_fn=partial(custom_collate_fn_prima, tokenizer=tokenizer,
qformer_tokenizer=qformer_tokenizer, inference=True,
use_mm_start_end=args.use_mm_start_end,
local_rank=args.local_rank, img_emb_len=args.img_emb_len,
padding_side='left'))
summary = inference(loader, model, tokenizer, args, output)
(output / 'summary.json').write_text(json.dumps(summary, indent=2) + '\n')
print(json.dumps(summary, indent=2))
if args.compute_metrics:
from eval.m4seg import score_predictions
del model
torch.cuda.empty_cache()
report = score_predictions(
[output], dataset.data_infos, penalize_missing=True,
with_text=args.metrics_with_text, with_meteor=args.metrics_with_meteor,
text_model=args.text_model, ground_truth=annotation,
)
(output / 'metrics.json').write_text(json.dumps(report, indent=2) + '\n')
print(json.dumps(report, indent=2))
|