sstm-flow 0.1.4

Python package for learning-based SSTM optical flow inference

SSTM-PyPI

Package Description

SSTM-PyPI is a python package for dense optical flow estimation from image sequences based on our deep-learning method, SSTM: Spatiotemporal recurrent transformers for multi-frame optical flow estimation.
It supports both two-frame and three-frame inference.

For a two-frame input, the package predicts a single flow field:

• f12: optical flow from frame 1 → frame 2

For a three-frame input, the package predicts two flow fields:

• f12: optical flow from frame 1 → frame 2
• f23: optical flow from frame 2 → frame 3

The package also supports:

• input image paths
• already loaded OpenCV / NumPy images
• optional saving of flow fields as .flo
• optional color-wheel visualization of predicted flow

Installation

conda create -n sstm_env python=3.10
conda activate sstm_env
pip install torch torchvision
pip install sstm-flow

Usage

Two-frame optical flow

from sstm import estimate_flow

flow = estimate_flow(
    "checkpoints/ft-sintel.pth",
    "frame_0013.png",
    "frame_0014.png"
)

Three-frame optical flow

from sstm import estimate_flow

f12, f23 = estimate_flow(
    "checkpoints/ft-sintel.pth",
    "frame_0013.png",
    "frame_0014.png",
    "frame_0015.png"
)

Using already loaded images

The API also supports images already loaded with OpenCV or NumPy.

import cv2
from sstm import estimate_flow

img1 = cv2.imread("frame_0013.png")
img2 = cv2.imread("frame_0014.png")

flow = estimate_flow(
    "checkpoints/ft-sintel.pth",
    img1,
    img2
)

Faster inference (recommended)

If running inference on multiple frame pairs, load the model once.

import cv2
from sstm import load_model, infer_flow_pair

model = load_model("checkpoints/ft-sintel.pth")

img1 = cv2.imread("frame_0013.png")
img2 = cv2.imread("frame_0014.png")

flow = infer_flow_pair(model, img1, img2)

Saving flow outputs using .flo and standard flow color-wheel visualization

from sstm import estimate_flow

estimate_flow(
    "checkpoints/ft-sintel.pth",
    "frame_0013.png",
    "frame_0014.png",
    output_dir="results",
    save_flo=True,
    save_vis_outputs=True
)

This will output:

results/
├── flow0001.flo
├── flow0001.png

Model weight description and task specific recommendation

Under checkpoints/. there are three different fine-tuned versions of SSTM

Model	Description	Recommended Use
ft-speckle-sintel	Fine-tuned for motion patterns that exhibit local spatial deformation.	Biomechanics, elastography, biological tissue motion, speckle tracking
ft-sintel	General-purpose optical flow model trained on non-rigid motion with diverse dynamics.	Natural scenes, animation, and general research applications
ft-kitti	Fine-tuned for rigid scene motion typical in driving environments.	Autonomous driving, vehicle motion, rigid object tracking

Citation

If you find this work useful please cite:

@article{ferede2023sstm,
  title={SSTM: Spatiotemporal recurrent transformers for multi-frame optical flow estimation},
  author={Ferede, Fisseha Admasu and Balasubramanian, Madhusudhanan},
  journal={Neurocomputing},
  volume={558},
  pages={126705},
  year={2023},
  publisher={Elsevier}
}