Include clean PRIMA code and evaluator; document metrics, initialization, and MIT contributions
91c1e9e verified | """Run PRIMA on ordered image groups and text prompts.""" | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| from tools.config import parse_with_config, validate_paths | |
| def parse_args(): | |
| from tools.model_loader import DEFAULT_HUB_MODEL, add_model_arguments | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| add_model_arguments(parser, merged_checkpoint=DEFAULT_HUB_MODEL) | |
| parser.add_argument('--image_paths', nargs='+', help='One image group for CLI inference.') | |
| parser.add_argument('--prompts', nargs='+', help='Prompts paired with image groups in the config.') | |
| parser.add_argument('--output_save_path', default='outputs/inference') | |
| parser.add_argument('--max_new_tokens', type=int, default=1024) | |
| return parse_with_config(parser) | |
| def main(): | |
| args = parse_args() | |
| validate_paths(args) | |
| if not args.image_paths or not args.prompts: | |
| raise ValueError('Provide image_paths and prompts in the config or CLI.') | |
| groups = args.image_paths | |
| if all(isinstance(item, str) for item in groups): | |
| groups = [groups] | |
| if len(groups) != len(args.prompts): | |
| raise ValueError('Each image group must have exactly one prompt.') | |
| if any(not isinstance(group, list) or not 1 <= len(group) <= args.max_num_images for group in groups): | |
| raise ValueError('Each image group must be a nonempty list within max_num_images.') | |
| for group in groups: | |
| for path in group: | |
| if not Path(path).is_file(): | |
| raise ValueError(f'Image does not exist: {path}') | |
| import numpy as np | |
| import torch | |
| from model.llava import conversation | |
| from tools.chat import Chat | |
| from tools.model_loader import load_bundle | |
| from tools.outputs import image_mask_slots | |
| model, tokenizer, qformer_tokenizer = load_bundle(args) | |
| chat = Chat(model, tokenizer, qformer_tokenizer, args, device=f'cuda:{args.local_rank}') | |
| output = Path(args.output_save_path) | |
| output.mkdir(parents=True, exist_ok=True) | |
| for index, (images, question) in enumerate(zip(groups, args.prompts)): | |
| turn = output / f'turn{index}' | |
| turn.mkdir(exist_ok=True) | |
| conv = conversation.conv_templates[args.conv_type].copy() | |
| encoded = [] | |
| for image in images: | |
| chat.upload_img(image, conv, encoded) | |
| chat.ask(question, conv, num_images=len(images)) | |
| with torch.inference_mode(): | |
| answer, masks, _ = chat.answer(conv, encoded, question, | |
| max_new_tokens=args.max_new_tokens, | |
| max_length=args.model_max_length + args.max_new_tokens) | |
| mask_paths = [] | |
| for image_index, predicted in enumerate(image_mask_slots(masks, len(images))): | |
| paths = [] | |
| if predicted is not None: | |
| for mask_index, mask in enumerate(predicted): | |
| name = f'image{image_index + 1}_mask{mask_index + 1}.npy' | |
| np.save(turn / name, (mask > 0).cpu().numpy()) | |
| paths.append(name) | |
| mask_paths.append(paths) | |
| record = {'images': images, 'prompt': question, 'response': answer, 'masks': mask_paths} | |
| (turn / 'dialogue.json').write_text(json.dumps(record, indent=2) + '\n') | |
| print(answer) | |
| if __name__ == '__main__': | |
| try: | |
| main() | |
| except ValueError as error: | |
| raise SystemExit(str(error)) | |