UFM: A Simple Path towards Unified Dense Correspondence with Flow

Carnegie Mellon University

Yuchen Zhang, Nikhil Keetha, Chenwei Lyu, Bhuvan Jhamb, Yutian Chen, Yuheng Qiu, Jay Karhade, Shreyas Jha, Yaoyu Hu, Deva Ramanan, Sebastian Scherer, Wenshan Wang

UFM unifies the tasks of Optical Flow Estimation and Wide Baseline Matching and provides accurate dense correspondences for in-the-wild images at significantly fast inference speeds.

Updates

• [2025/06/10] Initial release of model checkpoint and inference code.

Stay Tuned for the Upcoming Updates!

• Training and benchmarking code for all results presented in the paper.
• UFM-Tiny for real-time applications such as robotics.

Overview

UFM (Unified Flow & Matching, UniFlowMatch) is a simple, end-to-end trained transformer model that directly regresses pixel displacement images (flow) and can be applied concurrently to both optical flow and wide-baseline matching tasks.

Quick Start

Installation

We use UniCeption, a library which contains modular, config-swappable components for assembling end-to-end networks. To install UFM, recursively clone this repository and install the package with all dependencies:

git clone --recursive https://github.com/UniFlowMatch/UFM.git
cd UFM

# In case you cloned without --recursive:
# git submodule update --init

# Create and activate conda environment
conda create -n ufm python=3.11 -y
conda activate ufm

# Install UniCeption dependency
cd UniCeption
pip install -e .
cd ..

# Install UFM with all dependencies
pip install -e .

# Optional: Install with specific extras
# pip install -e ".[dev]"     # For development
# pip install -e ".[demo]"    # For demo
# pip install -e ".[all]"     # All optional dependencies

# Optional: For development and linting
pre-commit install  # Install pre-commit hooks

Verify Installation

Verify your installation by running the basic model test:

# Test installation
ufm test

# Or run the basic model test
python uniflowmatch/models/ufm.py

Verify that ufm_output.png looks like examples/example_ufm_output.png.

Command Line Interface

UFM provides a convenient CLI for common tasks:

# Test installation
ufm test

# Launch interactive demo
ufm demo

# Launch demo with specific settings
ufm demo --port 8080 --share --model refine

# Run inference on image pair
ufm infer source.jpg target.jpg --output results/

# Run inference with refinement model
ufm infer img1.png img2.png --model refine --output ./output

Python API

import cv2
import torch

# Load the base model (for general use)
from uniflowmatch.models.ufm import UniFlowMatchConfidence
model = UniFlowMatchConfidence.from_pretrained("infinity1096/UFM-Base")

# Or load the refinement model (for higher accuracy)
from uniflowmatch.models.ufm import UniFlowMatchClassificationRefinement
model = UniFlowMatchClassificationRefinement.from_pretrained("infinity1096/UFM-Refine")

# Set the model to evaluation mode
model.eval()

# Load images using cv2 or PIL
source_image = cv2.imread("path/to/source.jpg")
target_image = cv2.imread("path/to/target.jpg")
source_rgb = cv2.cvtColor(source_image, cv2.COLOR_BGR2RGB)  # Convert to RGB
target_rgb = cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB)  # Convert to RGB

# Convert to torch tensors (uint8 or float32)
# Forward call takes care of normalizing uint8 images appropriate to the UFM model
source_image = torch.from_numpy(source_rgb)  # Shape: (H, W, 3)
target_image = torch.from_numpy(target_rgb)  # Shape: (H, W, 3)

# Predict correspondences
with torch.no_grad():
    result = model.predict_correspondences_batched(
        source_image=source_image,
        target_image=target_image,
    )

    flow = result.flow.flow_output[0].cpu().numpy()
    covisibility = result.covisibility.mask[0].cpu().numpy()

Interactive Demo

Online Demo

Try our online demo without installation: 🤗 Hugging Face Demo

Local Gradio Demo

Run the interactive Gradio demo locally to visualize UFM outputs:

# Using the CLI (recommended)
ufm demo

# Or run directly
python gradio_demo.py

# Advanced options
ufm demo --port 8080 --share --model refine

License

This code is licensed under a fully open-source BSD-3-Clause license. The pre-trained UFM model checkpoints inherit the licenses of the underlying training datasets and as result, may not be used for commercial purposes (CC BY-NC-SA 4.0). Please refer to the respective training dataset licenses for more details.

Based on community interest, we can look into releasing an Apache 2.0 licensed version of the model in the future. Please upvote the issue here if you would like to see this happen.

Acknowledgements

We thank the folowing projects for their open-source code: DUSt3R, MASt3R, RoMA, and DINOv2.

Citation

If you find our repository useful, please consider giving it a star ⭐ and citing our paper in your work:

@inproceedings{zhang2025ufm,
 title={UFM: A Simple Path towards Unified Dense Correspondence with Flow},
 author={Zhang, Yuchen and Keetha, Nikhil and Lyu, Chenwei and Jhamb, Bhuvan and Chen, Yutian and Qiu, Yuheng and Karhade, Jay and Jha, Shreyas and Hu, Yaoyu and Ramanan, Deva and Scherer, Sebastian and Wang, Wenshan},
 booktitle={arXiV},
 year={2025}
}