"""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))