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