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

2025-12-21

AI Summary

Existing image rectification and dewarping methods suffer from poor generalization, high task coupling, and difficulty in unifying diverse lens distortions. To address these challenges, this paper proposes UniRect, a unified distortion correction framework. Methodologically, it formulates geometric rectification and quality restoration as a joint inverse problem, introducing two novel components: the Residual Progressive Thin-Plate Spline (RP-TPS) module for geometric deformation modeling and Residual Mamba Blocks (RMBs) for semantic-aware restoration. To mitigate task interference across distortion types, UniRect incorporates a Sparse Mixture-of-Experts (MoE) architecture and pioneers the integration of the Mamba sequence model for geometry–semantics co-modeling. Leveraging unified distortion representation and prompt-driven multi-task training, UniRect achieves state-of-the-art performance on multiple rectification and rectangling benchmarks. It significantly improves geometric accuracy and texture fidelity under complex distortions and enables cross-task generalization for practical deployment.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Calibration & Uncertainty Quantification
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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.

Problem

• Unified model for image correction and rectangling tasks
• Addresses diverse distortions with a single task-agnostic framework
• Improves generalization across different image rectification problems

Innovation

• Unified framework with dual-component structure for distortion rectification
• Residual Progressive Thin-Plate Spline model for geometric deformations
• Residual Mamba Blocks and Sparse Mixture-of-Experts for restoration