code/README.md

# <img src="assets/prima-logo.png" height="60" alt="PRIMA logo"> PRIMA: Multi-Image Vision-Language Models for Reasoning Segmentation [WACV 2027]

#### [Kiet A. Nguyen](https://scholar.google.com/citations?user=ECznem4AAAAJ)<sup>1</sup>, [Muntasir Wahed](https://mwahed.com/)<sup>1</sup>, [Adheesh Sunil Juvekar](https://juvekaradheesh.github.io/)<sup>1</sup>, [Xinzhuo Li](https://xinzhuo20.github.io/)<sup>1</sup>, [Xiaona Zhou](https://scholar.google.com/citations?user=TgdQrLwAAAAJ)<sup>1</sup>, [Vedant Shah](https://scholar.google.com/citations?user=szV_wkoAAAAJ)<sup>1</sup>, [Tianjiao Yu](https://tianjiao-yu.github.io/)<sup>1</sup>, [Pinar Yanardag](https://pinguar.org/)<sup>2</sup>, [Ismini Lourentzou](https://isminoula.github.io/)<sup>1</sup>

<picture>
  <source media="(prefers-color-scheme: dark)" srcset="assets/plan-logo-full-white.svg">
  <source media="(prefers-color-scheme: light)" srcset="assets/plan-logo-full.svg">
  <img src="assets/plan-logo-full.svg" height="60" alt="PLAN Lab logo">
</picture>

#### <sup>1</sup>[PLAN Lab](https://plan-lab.github.io/), University of Illinois Urbana-Champaign &nbsp; <sup>2</sup>Virginia Tech

[![Paper PDF](https://img.shields.io/badge/paper-PDF-blue)](https://arxiv.org/pdf/2412.15209)
[![Paper arXiv](https://img.shields.io/badge/paper-arxiv-green)](https://arxiv.org/abs/2412.15209)
[![Model](https://img.shields.io/badge/model-HuggingFace-yellow)](https://huggingface.co/PLAN-Lab/PRIMA)
[![Dataset](https://img.shields.io/badge/dataset-M4Seg-orange)](https://huggingface.co/datasets/PLAN-Lab/M4Seg)
[![Project Website](https://img.shields.io/badge/project-website-purple)](https://plan-lab.github.io/projects/prima/)

---

## 📢 Latest Updates

- **Oct. 29, 2026:** Inference, training, and evaluation code is now available, together with the [PRIMA checkpoint](https://huggingface.co/PLAN-Lab/PRIMA) and [M⁴Seg dataset](https://huggingface.co/datasets/PLAN-Lab/M4Seg).
- 🎉 **PRIMA has been accepted to WACV 2027!**

---

## <img src="assets/prima-logo.png" height="40" alt=""> 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.

<p align="center">
  <img src="assets/m4seg-examples.png" alt="M4Seg examples of functional, spatial, open-ended, and numerical reasoning, with matching colors linking answer phrases and masks.">
</p>

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

<details>
<summary><strong>Abstract</strong></summary>

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.

</details>

---

## 🏆 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**](docs/installation.md): Environment setup and dependencies.
- [**Data Preparation**](docs/datasets.md): M⁴Seg image sources, annotation conventions, and configurable data paths.
- [**Inference**](docs/inference.md): Running PRIMA on your own images and interpreting its text and mask outputs.
- [**Training**](docs/training.md): Single-GPU and distributed training, trainable components, and checkpoint resume.
- [**Evaluation**](docs/evaluation.md): M⁴Seg generation, scoring, sharding, and metric definitions.
- [**PRIMA Checkpoint**](https://huggingface.co/PLAN-Lab/PRIMA): Complete merged model for inference and evaluation.
- [**M⁴Seg Dataset**](https://huggingface.co/datasets/PLAN-Lab/M4Seg): Training and test annotations, RLE masks, and source-image preparation tools.

The [PRIMA checkpoint](https://huggingface.co/PLAN-Lab/PRIMA) and [M⁴Seg annotations](https://huggingface.co/datasets/PLAN-Lab/M4Seg) 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.

<p align="center">
  <img src="assets/prima-architecture.png" alt="PRIMA architecture: SQuARE, a language model with LoRA, and a SAM-based segmentation decoder.">
</p>

<p align="center">
  <img src="assets/square.png" width="620" alt="SQuARE module: relational queries attend across image features and enrich the instruction-guided global query pathway.">
</p>

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](docs/inference.md).

---

<a id="m4seg"></a>

## 🖼️ 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](docs/datasets.md).

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

---

## ⚡ Quick Start

Install the environment with the [installation guide](docs/installation.md), 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](run_config.yaml):

```yaml
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
```

```bash
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](docs/inference.md) for details.

### Prepare M⁴Seg

Download the annotations and source-image preparation files:

```bash
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](docs/datasets.md) 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](configs/train.yaml). For one GPU:

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

For four GPUs:

```bash
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](docs/training.md) for the initialization and export formats.

### Evaluate and compute metrics

Set the dataset root and M⁴Seg test annotations in [the evaluation configuration](configs/evaluate.yaml). Generate predictions and compute all seven mask/text metrics with:

```bash
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:

```bash
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](docs/evaluation.md) for metric definitions and multi-GPU sharding.

---

## License

Original PRIMA code and contributions are released under the [MIT License](LICENSE). New M⁴Seg annotations are released under [CC BY 4.0](https://creativecommons.org/licenses/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:

```bibtex
@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](https://github.com/mbzuai-oryx/groundingLMM) and [LLaVA](https://github.com/haotian-liu/LLaVA) for the vision-language model integration.
- [Segment Anything](https://github.com/facebookresearch/segment-anything) for the grounding module.
- [EVA](https://github.com/baaivision/EVA) and [LAVIS/BLIP](https://github.com/salesforce/LAVIS) for the visual encoder, Q-Former, and image processors.

Original copyright and license headers are retained. Copies of the applicable [Apache 2.0](LICENSES/Apache-2.0.txt), [LAVIS BSD 3-Clause](LICENSES/BSD-3-Clause-LAVIS.txt), and [EVA MIT](LICENSES/MIT-EVA.txt) 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.