PRIMA / code /README.md

Use one paper evaluator and clarify metric coverage

1445825 verified about 13 hours ago

preview code

Raw History Blame Contribute Delete

12.9 kB

PRIMA: Multi-Image Vision-Language Models for Reasoning Segmentation [WACV 2027]

Kiet A. Nguyen¹, Muntasir Wahed¹, Adheesh Sunil Juvekar¹, Xinzhuo Li¹, Xiaona Zhou¹, Vedant Shah¹, Tianjiao Yu¹, Pinar Yanardag², Ismini Lourentzou¹

¹PLAN Lab, University of Illinois Urbana-Champaign ²Virginia Tech

📢 Latest Updates

Oct. 29, 2026: Inference, training, and evaluation code is now available, together with the PRIMA checkpoint and M⁴Seg dataset.
🎉 PRIMA has been accepted to WACV 2027!

PRIMA Overview

PRIMA is a vision-language model for multi-image pixel-grounded reasoning. Given a set of images and a question, it compares objects and parts across the images and generates an explanation with segmentation masks for the regions it refers to.

M4Seg examples of functional, spatial, open-ended, and numerical reasoning, with matching colors linking answer phrases and masks.

Matching colors connect phrases in each answer to the regions they describe. The examples cover object functions, spatial relationships, scene context, and counting.

Abstract

Despite significant advancements in Large Vision-Language Models (LVLMs)’ capabilities, existing pixel-grounding models primarily operate in single-image settings, limiting their ability to perform detailed, fine-grained comparisons across multiple images. Conversely, current multi-image understanding models lack pixel-level grounding. We introduce the task of multi-image pixel-grounded reasoning alongside PRIMA, an LVLM that integrates reasoning segmentation with robust multi-image reasoning to produce contextually rich, pixel-grounded explanations. Central to PRIMA is SQuARE, a vision module that injects cross-image relational context into compact query-based visual tokens before fusing them with the language backbone. To support training and evaluation, we curate M⁴Seg, a new multi-image reasoning segmentation benchmark consisting of ~750K question–answer pairs that require fine-grained visual understanding across multiple images. PRIMA outperforms state-of-the-art baselines, yielding 11.8% and 12.2% gains in Recall and S-IoU, respectively. Ablations further demonstrate SQuARE's effectiveness in capturing cross-image relationships.

🏆 Contributions

Multi-image pixel-grounded reasoning. We introduce a task that combines reasoning across images with object- and part-level segmentation.
PRIMA and SQuARE. Our visual encoder gathers shared relational context across the image set and adds it to compact, instruction-guided visual queries before language fusion.
M⁴Seg benchmark. We introduce multi-image question-answer pairs with segmentation annotations covering functional, spatial, numerical, and open-ended reasoning.

🚀 Dive Deeper: Code, Data, and Checkpoints

Installation: Environment setup and dependencies.
Data Preparation: M⁴Seg image sources, annotation conventions, and configurable data paths.
Inference: Running PRIMA on your own images and interpreting its text and mask outputs.
Training: Single-GPU and distributed training, trainable components, and checkpoint resume.
Evaluation: M⁴Seg generation, scoring, sharding, and metric definitions.
PRIMA Checkpoint: Complete merged model for inference and evaluation.
M⁴Seg Dataset: Training and test annotations, RLE masks, and source-image preparation tools.

The PRIMA checkpoint and M⁴Seg annotations are available on Hugging Face. The checkpoint includes merged weights and both tokenizers. The dataset contains train/test annotations, relative image paths, and RLE masks; obtain the original images separately. Data paths in the example configurations remain placeholders.

🩰 PRIMA Architecture

PRIMA combines multi-image reasoning with pixel-level segmentation. Its SQuARE encoder gathers context across the images, the language model generates an explanation with [SEG] tokens, and a SAM-based decoder turns the corresponding token embeddings into masks.

Key components:

EVA/Q-Former visual interface: extracts compact visual tokens guided by the question.
SQuARE relational attention: uses learnable queries to gather shared context from all input images and enrich the global visual queries.
Language model with LoRA: generates the explanation and image-specific grounding tokens.
Text-to-mask projection and SAM decoder: produce masks for the objects and parts named in the answer.

PRIMA architecture: SQuARE, a language model with LoRA, and a SAM-based segmentation decoder.

SQuARE module: relational queries attend across image features and enrich the instruction-guided global query pathway.

The released checkpoint includes all model components and both tokenizers. Load it with this repository's scripts or directly with Hugging Face Transformers; see the inference guide.

🖼️ M⁴Seg Dataset

M⁴Seg pairs image sets with reasoning questions, natural-language answers, and object or part masks. Each grounded phrase is linked to its source image, allowing evaluation of both the explanation and its segmentation masks.

Reasoning type	Example question
Functional	Which vehicle is better suited to carrying larger cargo?
Spatial	Which bird is facing left across the images?
Numerical	Which image has fewer suitcases, and by how many?
Open-ended	How do the tables' functions relate to their surroundings?

Training and evaluation use a single M⁴Seg adapter. Its shared utilities resolve the constituent image sources and decode their masks. Prepare the data with the data guide.

Split	Examples	Grounding annotations
Train	743,463	2,595,754
Test	1,000	3,620

⚡ Quick Start

Install the environment with the installation guide, then set your image or dataset paths. Inference and evaluation load the complete checkpoint from PLAN-Lab/PRIMA; a downloaded local model directory works too.

Run inference on your images

Pair each image group with a question in the inference configuration:

merged_ckpt_path: PLAN-Lab/PRIMA
image_paths:
  - - /path/to/first_image.jpg
    - /path/to/second_image.jpg
prompts:
  - Identify and segment the objects shared by these images.
output_save_path: outputs/inference

python run.py --config run_config.yaml

Each image group gets its own output directory containing the question, response, and mask paths in JSON, plus Boolean NumPy masks at the original image resolutions. Masks are grouped by source image and follow the generated [SEG] (IMAGE1), [SEG] (IMAGE2), and subsequent image routes. See inference for details.

Prepare M⁴Seg

Download the annotations and source-image preparation files:

huggingface-cli download PLAN-Lab/M4Seg \
  --repo-type dataset --local-dir ./data/m4seg \
  --include "annotations/*" "metadata/*" "prepare_images.py" "PREPARATION.md"

Follow the data guide to obtain the original images and arrange them under your dataset root. The downloaded bundle includes native JSON annotations for the training and evaluation scripts.

Train on M⁴Seg

Set the training annotations and model assets in the training configuration. For one GPU:

python train.py --config configs/train.yaml

For four GPUs:

deepspeed --num_gpus 4 train.py --config configs/train.yaml

The example uses a per-device batch size of four and ten accumulation steps, giving an effective batch size of 160 on four GPUs. Use --resume /path/to/recoverable_checkpoint or --auto_resume to continue training. See training for the initialization and export formats.

Evaluate and compute metrics

Set the dataset root and M⁴Seg test annotations in the evaluation configuration. Generate predictions and compute all seven mask/text metrics with:

python evaluate.py --config configs/evaluate.yaml \
  --compute_metrics --text_model sentence-transformers/all-MiniLM-L6-v2

Omit --compute_metrics to generate predictions only. Scores are saved alongside the predictions. New prediction records include native image dimensions, so scoring does not require a separate image-size cache. METEOR requires Java.

To score existing predictions separately:

python -m eval.m4seg /path/to/predictions \
  --gt-json /path/to/evaluation_annotations.json \
  --penalize-missing \
  --with-text --text-model sentence-transformers/all-MiniLM-L6-v2 --with-meteor \
  --output outputs/metrics.json

Missing or failed predictions receive zero for mask and phrase scores under the default policy. METEOR is computed over annotated images with predicted masks. See evaluation for metric definitions and multi-GPU sharding.

License

Original PRIMA code and contributions are released under the MIT License. New M⁴Seg annotations are released under CC BY 4.0. Inherited code, pretrained weights, source masks, and images retain their upstream terms.

📜 Citation

If you use PRIMA or M⁴Seg in your research, please cite our paper:

@inproceedings{nguyen2027prima,
  title     = {PRIMA: Multi-Image Vision-Language Models for Reasoning Segmentation},
  author    = {Nguyen, Kiet A. and Wahed, Muntasir and Juvekar, Adheesh Sunil and
               Li, Xinzhuo and Zhou, Xiaona and Shah, Vedant and Yu, Tianjiao and
               Yanardag, Pinar and Lourentzou, Ismini},
  booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision},
  year      = {2027}
}

🙏 Acknowledgments

We thank the authors of the projects that PRIMA builds on:

GLaMM and LLaVA for the vision-language model integration.
Segment Anything for the grounding module.
EVA and LAVIS/BLIP for the visual encoder, Q-Former, and image processors.

Original copyright and license headers are retained. Copies of the applicable Apache 2.0, LAVIS BSD 3-Clause, and EVA MIT licenses accompany the source. The MIT grant for PRIMA's original contributions does not replace these inherited terms or the terms of pretrained model components.