Rectification Reimagined: A Unified Mamba Model for Image Correction and Rectangling with Prompts

TL;DR

This paper introduces UniRect, a unified deep learning framework for image correction and rectangling that generalizes across tasks using a novel distortion model, residual modules, and a mixture-of-experts approach.

Contribution

The paper presents a comprehensive, task-agnostic rectification framework with novel modules and a mixture-of-experts design, enabling effective multi-task image correction and rectangling.

Findings

Achieves state-of-the-art performance on multiple image correction tasks.

Effectively handles diverse geometric distortions with a unified model.

Demonstrates strong generalization across different practical photography scenarios.

Abstract

Image correction and rectangling are valuable tasks in practical photography systems such as smartphones. Recent remarkable advancements in deep learning have undeniably brought about substantial performance improvements in these fields. Nevertheless, existing methods mainly rely on task-specific architectures. This significantly restricts their generalization ability and effective application across a wide range of different tasks. In this paper, we introduce the Unified Rectification Framework (UniRect), a comprehensive approach that addresses these practical tasks from a consistent distortion rectification perspective. Our approach incorporates various task-specific inverse problems into a general distortion model by simulating different types of lenses. To handle diverse distortions, UniRect adopts one task-agnostic rectification framework with a dual-component structure: a {Deformation Module}, which utilizes a novel Residual Progressive Thin-Plate Spline (RP-TPS) model to address complex geometric deformations, and a subsequent Restoration Module, which employs Residual Mamba Blocks (RMBs) to counteract the degradation caused by the deformation process and enhance the fidelity of the output image. Moreover, a Sparse Mixture-of-Experts (SMoEs) structure is designed to circumvent heavy task competition in multi-task learning due to varying distortions. Extensive experiments demonstrate that our models have achieved state-of-the-art performance compared with other up-to-date methods.