A document image rectification method and system based on space-channel sparse attention

Abstract

This invention relates to a document image correction method and system based on spatial-channel sparse attention, belonging to the field of computer vision and intelligent document processing technology. The method includes: acquiring a document image to be corrected and inputting it into a document image correction network; in the encoder-decoder of the document image correction network, the attention branch of the transformer block adopts adaptive sparse probability window attention, setting a learnable sparsity parameter for each attention head, and converting the window correlation matrix into a sparse attention distribution through the α-entmax function; a parallel spatial-channel attention module is embedded in the feedforward branch of the document image correction network, performing channel recalibration and spatial recalibration in parallel on the normalized features. The decoder fuses encoder features through skip connections, and the output layer fuses high-dimensional features with shallowly embedded feature residuals to directly reconstruct the corrected image. This invention does not require explicit estimation of the deformation field, can effectively handle curling, wrinkling, and perspective distortion, and balances correction accuracy and computational efficiency.

Description

Technical Field

[0001] This invention belongs to the field of computer vision and intelligent document processing technology, specifically relating to a document image correction method and system based on spatial-channel sparse attention. Background Technology

[0002] With the development of mobile office, electronic archiving, mobile photo scanning, and intelligent OCR technologies, more and more paper documents are entering the digital processing workflow as images captured by mobile devices. However, due to factors such as skewed shooting angles, document surface curling, paper folding, page undulations, seam distortion caused by binding, and local perspective stretching, the original captured document images often have varying degrees of geometric distortion. These distortions can cause text lines to be uneven, character sizes to be inconsistent, boundaries to be curved, and local areas to be compressed or expanded, thus affecting subsequent layout analysis, character recognition, information extraction, and visual reading effects.

[0003] Existing document image correction techniques mainly fall into two categories: traditional image processing methods and deep learning methods. Traditional image processing methods typically rely on edge, corner, contour, line, or quadrilateral boundary detection, and achieve document flattening by estimating perspective transformations, affine transformations, or geometric mapping relationships. These methods often struggle to achieve stable results when faced with complex backgrounds, blurred edges, uneven lighting, missing document boundaries, or non-planar curvature. Especially for non-linear deformation scenarios such as page bending, curling, and wavy folds, methods based on rigid transformation models are ill-suited.

[0004] Deep learning methods directly learn the mapping relationship between the input distorted image and the target corrected image through convolutional neural networks or visual transformers, offering stronger feature representation and end-to-end learning capabilities compared to traditional methods. However, existing technologies still face several challenges. First, some methods overly rely on the local receptive field of convolution, resulting in insufficient modeling of long-range geometric dependencies. Second, directly employing standard global self-attention leads to high computational complexity and memory overhead, making it difficult to balance high-resolution processing with efficiency. Third, document images contain local textures, edges, and character structures, as well as global-scale deformation propagation; using only a single type of feature enhancement method cannot simultaneously address channel importance selection and spatial key region localization. Fourth, in the multi-scale encoding and reconstruction process, the lack of an effective feature modulation mechanism can easily lead to detail loss, insufficient boundary restoration, and local structural distortion.

[0005] Therefore, there is an urgent need for a document image correction method that combines local enhancement, window-level global modeling, sparse attention modulation, multi-scale feature recovery, and spatial-channel joint recalibration capabilities to improve the correction accuracy, structure preservation, and computational efficiency of document images in complex scenarios. Summary of the Invention

[0006] The purpose of this invention is to address the problems of insufficient modeling of complex nonlinear document deformation, inadequate restoration of local details, and high computational overhead under high resolution conditions in the existing technology, and to provide a document image correction method and system based on spatial-channel sparse attention.

[0007] To achieve the above objectives, the technical solution of the present invention is: a document image correction method based on spatial-channel sparse attention, comprising:

[0008] The document image to be corrected is obtained and input into a pre-constructed document image correction network to extract multi-scale features and generate correction results.

[0009] The document image correction network includes a convolutional layer module, an input layer, an encoder, a bottleneck layer, a decoder, and an output layer connected in sequence.

[0010] The convolutional layer module is used to extract two-dimensional shallow embedding features from the document image to be corrected, and the input layer is used to enhance the receptive field of the two-dimensional shallow embedding features and rearrange them into sequence features;

[0011] The encoder includes multiple cascaded base layers and downsampling modules corresponding to the base layers. The decoder includes multiple cascaded base layers and upsampling modules corresponding to the base layers. The decoder also fuses features with the corresponding layers of the encoder through skip connections. The base layer includes a transformer block, which includes an attention branch and a feedforward branch. The attention branch is used to restore the sequence features to two-dimensional features, divide them into multiple local windows according to a preset window size, and perform adaptive sparse probability window attention calculation within each local window. The adaptive sparse probability window attention calculation is based on the query vector, key vector, and value vector to obtain a window relevance matrix, and converts the window relevance matrix into a sparse attention distribution through an α-entmax sparse probability mapping function based on the learnable sparsity parameter of each attention head. The feedforward branch includes a parallel spatial-channel attention feedforward modulation module, which is used to restore the sequence features to two-dimensional features, perform parallel modulation of spatial attention and channel attention, and rearrange the modulation results into sequence features for output.

[0012] The output layer is used to map the decoded high-dimensional features into output features with the same dimension as the two-dimensional shallow embedding features. The output features and the two-dimensional shallow embedding features are residually fused to directly reconstruct the corrected document image.

[0013] The document image correction network is a direct image-to-image end-to-end reconstruction network that does not explicitly output deformation field, control point grid, sampling coordinates, perspective matrix, or dense displacement field.

[0014] Furthermore, the input layer includes a receptive field enhancement convolutional unit, which generates extended receptive field features through depthwise convolution, and enhances the input features by combining channel attention and spatial weighting mechanisms, and then maps the enhanced features to sequence features.

[0015] Furthermore, the adaptive sparse probabilistic window attention calculation includes: performing a linear mapping on the input features to obtain a query vector, a key vector, and a value vector; calculating a relevance score based on the query vector and the key vector within a local window, and obtaining a window relevance matrix by combining the relative position bias; and setting a learnable sparsity parameter for each attention head. The window relevance matrix of the corresponding attention head is input into the sparse probability mapping function to generate the sparse attention distribution of the corresponding attention head; wherein, the attention weights of low response positions in the sparse attention distribution are compressed to zero or close to zero, while high response positions retain larger attention weights; the value vectors are weighted and aggregated based on the sparse attention distributions corresponding to each attention head to obtain the adaptive sparse probability window attention output.

[0016] Furthermore, the learnable sparsity parameter is a parameter set independently for each attention head in each transformer block, used to control the sparsity of the attention distribution of the corresponding attention head; for the h-th attention head in the k-th transformer block, a basic parameter is set. The sparsity parameters are obtained through the following constraint mapping:

[0017] = 1+

[0018] in, The Sigmoid function is defined such that the learnable sparsity parameter lies in the interval greater than 1 and less than 2. These are the basic parameters that can be learned.

[0019] Window correlation matrix of the h-th attention head in the k-th transformer block Its attention distribution is represented as:

[0020] = Entmax

[0021] in, For the sparse attention distribution corresponding to the attention head, Entmax is the sparse probability normalization function.

[0022] Furthermore, the learnable basic parameters The learnable sparsity parameter is initialized to 0, making its initial value 1.5. During training, the learnable basic parameters are updated together with the network parameters so that different layers and different attention heads can adaptively form attention distributions with different sparsity.

[0023] Furthermore, the parallel spatial-channel attention feedforward modulation module includes a normalization submodule, a channel attention submodule, a spatial attention submodule, a first projection convolution submodule, and a second projection convolution submodule; its processing includes:

[0024] The input sequence features are restored into a two-dimensional feature map according to the spatial dimensions;

[0025] The two-dimensional feature map is normalized to obtain a normalized feature map;

[0026] The channel attention submodule performs global average pooling and global max pooling on the normalized feature map respectively, and inputs the pooling results into the shared mapping branch to generate channel attention weights.

[0027] The spatial attention submodule performs average aggregation and maximum aggregation on the normalized feature map in the channel dimension, and then concatenates the aggregation results and inputs them into the convolutional mapping branch to generate spatial attention weights.

[0028] The channel attention weights are multiplied element-wise with the normalized feature map by the first projection convolution submodule, and then projected by the first 1×1 convolution to obtain the channel modulation features.

[0029] The spatial attention weights are multiplied element-wise with the normalized feature map by the second projection convolution submodule, and then projected by the second 1×1 convolution to obtain the spatial modulation features.

[0030] The channel modulation features and spatial modulation features are added and fused to obtain the parallel spatial-channel attention modulation result;

[0031] The modulation result is rearranged into sequence features and used as the output of the feedforward branch.

[0032] Furthermore, the method also includes a training step, which employs the following loss function:

[0033]

[0034] in, To predict the pixel-level reconstruction loss between the corrected image and the target corrected document image, L1 loss is used; This is the structural similarity loss, used to constrain the overall page layout structure; Edge gradient loss is used to constrain the geometry of text edges, document boundaries, and crease regions; This is used to constrain the consistency of high-level visual structures, representing a perceptual loss. , , , The corresponding loss weights.

[0035] Furthermore, =0.8, =0.1, =0.05, .

[0036] The present invention also provides a document image correction system based on spatial-channel sparse attention that implements the document image correction method described above, comprising:

[0037] The image acquisition module is used to acquire images of the document to be corrected.

[0038] A feature embedding module is used to perform convolutional layers and input projection on the document image to be corrected;

[0039] The encoding module is used to extract multi-scale features through multiple base layers and a downsampling module;

[0040] The deep modeling module is used to perform deep modeling of features through adaptive sparse probability window attention and parallel spatial-channel attention modulation.

[0041] The decoding and reconstruction module is used to gradually restore spatial resolution through multiple upsampling modules, skip connections, and a base layer;

[0042] The output module is used to output the corrected document image.

[0043] The present invention also provides an electronic device or a computer-readable storage medium, the electronic device including a processor, a memory and a computer program stored in the memory, or the computer-readable storage medium storing a computer program, which, when executed by the processor, implements the document image correction method as described above.

[0044] Compared with the prior art, the present invention has the following beneficial effects:

[0045] 1. This invention introduces a parallel spatial-channel attention module in the feedforward branch, and explicitly models the channel importance and spatial key regions, which can enhance the text edges, page outlines, crease textures and local distortion features in document images, and improve the expressive ability in complex document scenes.

[0046] 2. This invention introduces an adaptive sparse probability window attention mechanism, which directly generates the attention distribution through a sparse probability mapping function and learns the sparsity parameter independently for each attention head, enabling different attention heads to form different attention patterns, thereby more effectively suppressing low-relevance positions and highlighting key positions.

[0047] 3. This invention introduces channel recalibration and spatial weighting in the shallow stage by enhancing the receptive field of the input layer, providing a more stable and discriminative initial representation for subsequent multi-scale feature extraction.

[0048] 4. This invention achieves joint recovery of multi-scale semantic information and high-resolution detail information through encoder-decoder structure and skip connections, which is beneficial to improving the overall quality of document image geometric correction and the ability to preserve local details. Attached Figure Description

[0049] Figure 1 This is a schematic diagram of the overall architecture of the document image correction network provided in an embodiment of the present invention.

[0050] Figure 2 This is a schematic diagram of the adaptive sparse probability window attention structure provided in an embodiment of the present invention.

[0051] Figure 3 This is a schematic diagram of the parallel spatial-channel attention module structure provided in an embodiment of the present invention.

[0052] Figure 4 This is a schematic diagram showing the comparison of document image correction effects before and after, as provided in an embodiment of the present invention. Detailed Implementation

[0053] The present invention will be further described below with reference to the accompanying drawings and embodiments. It should be understood that the following embodiments are only for explaining the present invention and are not intended to limit the scope of protection of the present invention. Various modifications and substitutions can be made to the present invention by those skilled in the art without departing from the spirit of the present invention, and all such modifications and substitutions should fall within the scope of protection of the present invention.

[0054] Example 1: Overall Method Flow

[0055] like Figure 1 As shown, this embodiment provides a document image correction method, which mainly includes the following steps:

[0056] S1. Acquire the document image to be processed. The document image may originate from a mobile phone, tablet, scanning terminal, industrial camera, or other imaging device. The document image is typically an RGB three-channel image. The document image may contain single-page documents, book pages, invoices, forms, contracts, archives, book pages, or other image objects with text and graphic layout structures.

[0057] In one implementation, the input image can be size-normalized to fit the network's preset input resolution, for example, by adjusting the image to a fixed size before inputting it into the network. Alternatively, the input image can also be standardized.

[0058] S2. Perform convolutional layer and input projection. The input image is fed into the convolutional layer module, which uses 3×3 convolution for shallow feature extraction, obtaining initial embedded features while maintaining the image spatial resolution. Since the convolutional kernels cover overlapping areas, they can preserve the local continuity information between adjacent pixels, which is beneficial for subsequent modeling of document texture, text edges, and crease textures.

[0059] Features following the convolutional layers are further fed into the input layer. The input layer is preferably a receptive field-enhancing convolutional module. This module first generates feature representations with an expanded receptive field through depthwise convolution, then obtains channel recalibration weights through channel attention subunits, generates spatial response weights through spatial weighting subunits, and maps the weighted features to sequential features that can be processed by subsequent base layers. Thus, the network can simultaneously achieve local structure enhancement and key region highlighting capabilities in its shallow stages.

[0060] S3. Multi-level coding feature extraction: The sequence features after input projection are fed into the encoder. The encoder consists of multiple cascaded base layers and downsampling modules located between adjacent base layers. Each base layer consists of one or more transformer blocks, used to model the feature relationships within the window and between adjacent windows at the current scale; the downsampling modules are used to reduce spatial resolution and increase channel dimension, thereby gradually expanding the feature receptive field and extracting stronger semantic information.

[0061] In a preferred embodiment, the encoder includes four encoding stages. The first-level encoding stage processes high-resolution shallow features, focusing more on text edges, local curvature, and contour information; subsequent encoding stages gradually extract mid-to-high-level semantics, modeling the overall page deformation distribution, perspective relationships, and non-linear folding patterns.

[0062] S4. Bottleneck Layer: Deep Feature Modeling. The deepest output of the encoder is fed into the bottleneck layer. The bottleneck layer consists of at least one base layer, with relatively low spatial resolution and high channel dimension. This layer is primarily responsible for deep modeling of the global document structure, long-range dependencies, and complex deformation patterns. Due to the reduced feature resolution at this stage, stronger high-level representation capabilities can be achieved with relatively low computational overhead.

[0063] S5. Multi-level decoding and skip connection recovery: The bottleneck layer output is fed into the decoder. The decoder includes multiple upsampling modules and multiple base layers. Each upsampling module recovers the feature space resolution step by step, and each base layer further models and refines the recovered features. To compensate for the loss of detail caused by multiple downsampling, the decoder uses skip connections to concatenate and fuse the corresponding layer features from the encoder, thereby utilizing both shallow detail information and deep semantic information simultaneously.

[0064] Through this multi-level recovery mechanism, the network can gradually correct problems such as document edge curvature, character arrangement distortion, local folding compression, and page perspective distortion.

[0065] S6. Output projection and residual fusion: The decoder's final output is fed into the output layer. The output layer maps high-dimensional features to output features with the same number of channels as the input image. Preferably, the output projection result is residually fused with shallow embedded features to obtain the final corrected image. This residual reconstruction method helps the network focus on learning the differences between the input and the target, improving convergence speed and reconstruction effect.

[0066] In this embodiment, the mapping relationship between the distorted document image and the target flat document image is directly learned through an encoder-decoder structure. During training, the network learns the implicit geometric correction relationships corresponding to page curvature, perspective distortion, local wrinkles, and text line curvature. The geometric correction of this invention restores the document layout structure, boundary shape, and text line arrangement through end-to-end image reconstruction.

[0067] Figure 4 This is a schematic diagram comparing the document image before and after correction provided in this embodiment. The left side is the document image before correction, and the right side is the document image after correction.

[0068] Example 2: Adaptive Sparse Probabilistic Window Attention

[0069] like Figure 2 As shown, the adaptive sparse probability window attention performs sparse attention modeling on the input features within a local window range.

[0070] Specifically, a linear mapping is first performed on the input window features to obtain a query vector, a key vector, and a value vector. Then, the relevance score between the query vector and the key vector is calculated within a local window, and combined with a relative position bias to form a window relevance matrix. The relative position bias is used to characterize the relative spatial relationship between different positions within the window to enhance the spatial awareness capability of attention computation.

[0071] In this embodiment, the α-entmax sparse probability mapping function is used to normalize the window correlation matrix. Unlike the Softmax function, which generates a dense attention distribution, the α-entmax sparse probability mapping function can compress the attention weights at low response positions to zero or near zero values, while retaining larger weights at high response positions, thus obtaining a sparse attention distribution.

[0072] Furthermore, to enable different attention heads to have different sparse modeling capabilities, this embodiment sets learnable basic parameters for each attention head in each transformer block. For the h-th attention head in the k-th transformer block, let its window correlation matrix be... The sparse attention distribution is as follows:

[0073] = Entmax

[0074] in, Entmax is a sparse probability normalization function, representing the learnable sparsity parameter corresponding to this attention head. To make... Within a stable range, set learnable basic parameters. And obtained through the following constraint mapping:

[0075] = 1+

[0076] in, This is the Sigmoid function. Therefore... It is restricted to a range greater than 1 and less than 2. When When initialized to 0, The initial value is 1.5, corresponding to the 1.5-entmax sparse mapping; during training, The network parameters are updated together, thus enabling different transformer blocks and attention heads to obtain different sparse modeling capabilities.

[0077] When the sparsity parameter of a particular attention head is large, the resulting attention distribution is sparser, highlighting a small number of highly relevant locations within the window. Conversely, when the sparsity parameter of a particular attention head is small, the resulting attention distribution is relatively smooth, preserving more contextual information. Thus, different attention heads can form different attention patterns within the same network layer, improving the diversity and selectivity of feature modeling.

[0078] After obtaining the sparse attention distribution corresponding to each attention head, the value vectors are weighted and aggregated to obtain the output results of each attention head; then the outputs of each attention head are concatenated and linearly projected to form an adaptive sparse probability window attention output, which is then fed into the subsequent residual connection and feedforward branch.

[0079] Through the above mechanism, adaptive sparse probability window attention can directly achieve salient position enhancement and low-correlation position suppression at the probability distribution level. By learning sparsity parameters independently by different attention heads, the functional division of labor among different attention heads is improved, thereby more effectively modeling edges, creases, dense text regions and local nonlinear deformation regions in document images.

[0080] Example 3: Shifted Window and Attention Mask

[0081] In a preferred embodiment, some transformer blocks employ a shift window strategy. Specifically, the input features are cyclically shifted by a preset number of pixels along the height and width directions before window partitioning and local attention calculation are performed. To avoid unreasonable information interaction between different window regions, a shift mask matching the shift method is constructed, and the mask is added to the relevance score during attention calculation. For scenarios requiring restriction of attention relationships in specific regions, an input mask can also be introduced, which, after interpolation and window partitioning, forms the final attention mask together with the shift mask.

[0082] By using the shift window mechanism, the model can not only maintain the computational efficiency of local windows, but also enhance the information exchange capability between adjacent windows, thereby improving the modeling effect of cross-regional deformation.

[0083] Example 4: Parallel Spatial-Channel Attention Module (PSCA)

[0084] like Figure 3 As shown, the parallel spatial-channel attention module includes a two-dimensional normalization subunit, a channel attention subunit, a spatial attention subunit, a first projection convolution, and a second projection convolution.

[0085] The input two-dimensional features are first processed by a two-dimensional normalization subunit to obtain normalized features. Subsequently, the normalized features are fed into the channel attention subunit and the spatial attention subunit in parallel.

[0086] The channel attention subunit performs global average pooling and global max pooling on the normalized features, respectively. Then, it performs feature compression and recovery through a shared mapping branch, and adds the two results to obtain the channel weights. These weights are used to recalibrate the channel dimension of the normalized features.

[0087] The spatial attention subunit calculates the average and maximum responses in the channel dimension, concatenates them, and then obtains the spatial weights through convolution. These spatial weights are used to recalibrate the spatial location of the normalized features.

[0088] The features after channel recalibration are projected through a first 1×1 convolution; the features after spatial recalibration are projected through a second 1×1 convolution; finally, the two projection results are added together to obtain the parallel spatial-channel modulation output.

[0089] This module can maintain a lightweight structure while taking into account both channel-level selection and spatial positioning, making it particularly suitable for the joint enhancement of text edges, crease textures, page boundaries, and high-response deformation areas in document correction tasks.

[0090] Example 5: Parallel Spatial-Channel Attention Feedforward Encapsulation

[0091] In a preferred embodiment, to adapt the parallel spatial-channel attention module to the sequence processing flow of the converter block, it is encapsulated as a feedforward modulation module. Specifically, the sequence features input to the converter block are restored into a two-dimensional feature map according to the spatial dimensions and sent to the parallel spatial-channel attention module for processing; after processing, the two-dimensional features are rearranged into a sequence form and returned to the subsequent converter calculation flow.

[0092] This encapsulation method can introduce two-dimensional attention modulation capability into the feedforward branch without disrupting the overall computational paradigm of the converter, thereby achieving local texture restoration and structural enhancement.

[0093] Example 6: Encoder-Decoder Network Structure

[0094] like Figure 1 As shown, the document image correction network adopts an encoder-decoder architecture.

[0095] The network input first passes through a convolutional layer and an input layer to obtain initial sequence features. The encoder consists of four encoding stages, each consisting of a base layer and a downsampling module. The downsampling module reduces spatial resolution and increases channel dimension through convolution. A bottleneck layer is placed at the end of the encoder to extract deep semantic features. The decoder consists of four decoding stages, each consisting of an upsampling module, feature concatenation, and a base layer. The upsampling module restores spatial resolution through deconvolution, then concatenates the upsampled features with the corresponding encoding layer output features in the channel dimension, and then feeds them into the base layer for fusion and refinement. Finally, the output layer generates the document image correction result.

[0096] In this implementation, the input document image is normalized to a size of 256×256, and the basic number of channels is set to 32. The encoder includes four encoding stages, with the number of feature channels at each stage set to 32, 64, 128, and 256 respectively, and the bottleneck layer having 512 feature channels. The decoder restores the resolution and number of channels step by step in the opposite direction to the encoder. The adaptive sparse probabilistic window attention has a window size of 8×8. When the current feature map resolution is smaller than the preset window size, the window size is adjusted to the smaller of the height and width of the current feature map.

[0097] The number of attention heads in multiple encoding and decoding stages varies with the feature level. The number of attention heads at each level of the encoder is 1, 2, 4, and 8, respectively, while the number of attention heads at the bottleneck layer is 16. The number of attention heads at each level of the decoder is 16, 8, 4, and 2, respectively. Adjacent transformer blocks can alternate between non-shifted windows and shifted windows. The shift amount in a non-shifted window is 0, while the shift amount in a shifted window is half the window size. Invalid associations between different window regions are limited by constructing a shifted window mask.

[0098] Skip connections are implemented using channel-dimensional concatenation, whereby the upsampled features of the current level of the decoder are concatenated with the features of the corresponding level of the encoder along the channel dimension, and then input into the corresponding decoding base layer for fusion and refinement. Residual fusion occurs at the network output, where the output layer maps the decoded high-dimensional features to output features with the same dimension as the shallow embedded features. Subsequently, the output features are added element-wise to the shallow embedded features output by the convolutional layer module to directly reconstruct the corrected document image.

[0099] Example 7: Training and Reasoning

[0100] This invention employs supervised learning to train a document image correction network. The training samples consist of distorted document images and their corresponding target corrected document images, denoted as (…). ),in This indicates an input document image that exhibits curling, wrinkling, perspective distortion, or local nonlinear deformation. This represents the corresponding flat, front-view, or manually labeled / dataset-provided target corrected document image. The training data can come from publicly available document image correction datasets or self-built paired datasets; in practice, one or more of DocUNet, DIR300, or self-built captured document datasets can be used as training and testing data sources. For synthetic data, distorted document images can be generated from flat document images through random perspective transformation, local curvature perturbation, grid deformation, rotation, scale change, brightness perturbation, blurring, noise, and shadow perturbation, and the original flat document image can be used as the target corrected document image.

[0101] In implementation, the dataset is divided into training, validation, and test sets in an 8:1:1 ratio. If a public test set is used, the training and validation sets are used for model training and parameter selection, while the public test set is only used for final performance evaluation. Both input and target images are scaled or cropped to 256×256 pixels and pixel normalized. Data augmentation during training includes one or more of the following: random rotation, random cropping, random perspective transformation, random bending perturbation, brightness perturbation, noise perturbation, and blur perturbation.

[0102] The training epochs are set to 300, and the batch size is set to 2 to 8, with the specific value determined based on the GPU memory capacity. The AdamW optimizer is used, and the initial learning rate is set to... The weight decay is set to The learning rate scheduling strategy employs cosine annealing, and a linear warm-up can be set in the early stages of training. The learning rate and batch size settings are based on the following principles: to ensure stable convergence during training, and to achieve low reconstruction error and high structural similarity on the validation set; the loss weights are selected through ablation experiments on the validation set.

[0103] Distorted document images Input the document image correction network to obtain the predicted corrected image. The training objective is to predict the corrected image. Image of the target correction document Maintain consistency in pixel content, structural morphology, and edge geometry. Preferably, the training loss function includes pixel reconstruction loss, structural similarity loss, edge gradient loss, and perceptual loss, with the following overall form:

[0104]

[0105] in, To predict the pixel-level reconstruction loss between the corrected image and the target corrected document image, L1 loss is used; This is the structural similarity loss, used to constrain the overall page layout structure; Edge gradient loss is used to constrain the geometry of text edges, document boundaries, and crease regions; This is used to constrain the consistency of high-level visual structures, representing a perceptual loss. , , , These correspond to the loss weights. Experimental settings. =0.8, =0.1, =0.05, .

[0106] Furthermore, to enhance the geometric flattening effect, text line structure constraints or boundary straightness constraints can be optionally introduced. The text line structure constraint is used to constrain the text lines in the corrected image to tend to be straight, and the boundary straightness constraint is used to constrain the outer contour boundary of the document to tend to be regular. Through the above supervision signals, the network not only learns image denoising or texture enhancement, but also learns the geometric flattening relationship between the distorted document image and the flattened target document image.

[0107] During training, the parameters in the convolutional layer module, input layer, adaptive sparse probability window attention module, parallel spatial-channel attention module, encoder, bottleneck layer, decoder, and output projection module are updated using the backpropagation algorithm. After training is complete, during the inference phase, the document image to be corrected is input into the trained network to obtain the corrected document image. This method can be deployed on servers, edge devices, mobile terminals, or scanner control terminals for real-time or offline correction processing of document images.

[0108] Example 8: Ablation Experiment

[0109] To verify the effectiveness of the adaptive sparse probability window attention and parallel spatial-channel attention feedforward modulation modules in this invention, ablation experiments were conducted on different network configurations under the same training set, test set, input image size, number of training epochs, and optimizer settings. The ablation experiments included at least the following three model configurations:

[0110] First, the complete model, namely the document image correction network that simultaneously employs adaptive sparse probabilistic window attention and parallel spatial-channel attention feedforward modulation modules.

[0111] Second, the model with the α-entmax sparse probability mapping removed is replaced by the Softmax normalization function in the adaptive sparse probability window attention, while the rest of the structure remains unchanged.

[0112] Third, the model of the parallel spatial-channel attention feedforward modulation module is removed, that is, the parallel spatial-channel attention feedforward modulation module is replaced with a normal feedforward network, while the rest of the structure remains unchanged.

[0113] In one implementation, the model was tested using common evaluation metrics for document image correction, including one or more of structural similarity, local distortion, average distortion, edit distance, and character recognition error rate. Experimental results show that, compared to the model without the α-entmax sparse probability mapping, the complete model more effectively highlights document edges, creases, and local nonlinear deformation regions; compared to the model without the parallel spatial-channel attention feedforward modulation module, the complete model better preserves text regions, document boundaries, and local texture structures. Therefore, both the adaptive sparse probability window attention and the parallel spatial-channel attention feedforward modulation module contribute to improving the document image correction effect.

[0114] The above are preferred embodiments of the present invention. Any changes made to the technical solution of the present invention that do not exceed the scope of the technical solution of the present invention shall fall within the protection scope of the present invention.

Claims

1. A document image correction method based on spatial-channel sparse attention, characterized in that, include: The document image to be corrected is obtained and input into a pre-constructed document image correction network to extract multi-scale features and generate correction results. The document image correction network includes a convolutional layer module, an input layer, an encoder, a bottleneck layer, a decoder, and an output layer connected in sequence. The convolutional layer module is used to extract two-dimensional shallow embedding features from the document image to be corrected, and the input layer is used to enhance the receptive field of the two-dimensional shallow embedding features and rearrange them into sequence features; The encoder includes multiple cascaded base layers and downsampling modules corresponding to the base layers. The decoder includes multiple cascaded base layers and upsampling modules corresponding to the base layers. The decoder also fuses features with the corresponding layers of the encoder through skip connections. The base layer includes a transformer block, which includes an attention branch and a feedforward branch. The attention branch is used to restore the sequence features to two-dimensional features, divide them into multiple local windows according to a preset window size, and perform adaptive sparse probability window attention calculation within each local window. The adaptive sparse probability window attention calculation is based on the query vector, key vector, and value vector to obtain a window relevance matrix, and converts the window relevance matrix into a sparse attention distribution through an α-entmax sparse probability mapping function based on the learnable sparsity parameter of each attention head. The feedforward branch includes a parallel spatial-channel attention feedforward modulation module, which is used to restore the sequence features to two-dimensional features, perform parallel modulation of spatial attention and channel attention, and rearrange the modulation results into sequence features for output. The output layer is used to map the decoded high-dimensional features into output features with the same dimension as the two-dimensional shallow embedding features. The output features and the two-dimensional shallow embedding features are residually fused to directly reconstruct the corrected document image. The document image correction network is a direct image-to-image end-to-end reconstruction network that does not explicitly output deformation field, control point grid, sampling coordinates, perspective matrix, or dense displacement field.

2. The document image correction method based on spatial-channel sparse attention according to claim 1, characterized in that, The input layer includes a receptive field enhancement convolutional unit, which generates extended receptive field features through depthwise convolution, and enhances the input features by combining channel attention and spatial weighting mechanisms, and then maps the enhanced features to sequence features.

3. The document image correction method based on spatial-channel sparse attention according to claim 1, characterized in that, The adaptive sparse probabilistic window attention calculation includes: performing a linear mapping on the input features to obtain a query vector, a key vector, and a value vector; calculating a relevance score based on the query vector and the key vector within a local window, and obtaining a window relevance matrix by combining the relative position bias; and setting a learnable sparsity parameter for each attention head. The window relevance matrix of the corresponding attention head is input into the sparse probability mapping function to generate the sparse attention distribution of the corresponding attention head; wherein, the attention weights of low response positions in the sparse attention distribution are compressed to zero or close to zero, while high response positions retain larger attention weights; the value vectors are weighted and aggregated based on the sparse attention distributions corresponding to each attention head to obtain the adaptive sparse probability window attention output.

4. The document image correction method based on spatial-channel sparse attention according to claim 3, characterized in that, The learnable sparsity parameter is a parameter set independently for each attention head in each transformer block, used to control the sparsity of the attention distribution of the corresponding attention head; for the h-th attention head in the k-th transformer block, a basic parameter is set. The sparsity parameters are obtained through the following constraint mapping. : = 1+ in, Represents the Sigmoid function, such that the learnable sparsity parameter It lies within the interval greater than 1 and less than 2; Window correlation matrix of the h-th attention head in the k-th transformer block Its attention distribution is represented as: = Entmax in, For the sparse attention distribution corresponding to the attention head, Entmax is the sparse probability normalization function.

5. The document image correction method based on spatial-channel sparse attention according to claim 4, characterized in that, The learnable basic parameters Initialize to 0, so that the learnable sparsity parameter The initial value is 1.5; during training, the learnable basic parameters... The network parameters are updated together to enable different transformer blocks and different attention heads to adaptively form attention distributions with different sparsity.

6. The document image correction method based on spatial-channel sparse attention according to claim 1, characterized in that, The parallel spatial-channel attention feedforward modulation module includes a normalization submodule, a channel attention submodule, a spatial attention submodule, a first projection convolution submodule, and a second projection convolution submodule. Its processing includes: The input sequence features are restored into a two-dimensional feature map according to the spatial dimensions; The two-dimensional feature map is normalized to obtain a normalized feature map; The channel attention submodule performs global average pooling and global max pooling on the normalized feature map respectively, and inputs the pooling results into the shared mapping branch to generate channel attention weights. The spatial attention submodule performs average aggregation and maximum aggregation on the normalized feature map in the channel dimension, and then concatenates the aggregation results and inputs them into the convolutional mapping branch to generate spatial attention weights. The channel attention weights are multiplied element-wise with the normalized feature map by the first projection convolution submodule, and then projected by the first 1×1 convolution to obtain the channel modulation features. The spatial attention weights are multiplied element-wise with the normalized feature map by the second projection convolution submodule, and then projected by the second 1×1 convolution to obtain the spatial modulation features. The channel modulation features and spatial modulation features are added and fused to obtain the parallel spatial-channel attention modulation result; The modulation result is rearranged into sequence features and used as the output of the feedforward branch.

7. The document image correction method based on spatial-channel sparse attention according to claim 1, characterized in that, The method further includes a training step, which employs the following loss function: in, To predict the pixel-level reconstruction loss between the corrected image and the target corrected document image, L1 loss is used; This is the structural similarity loss, used to constrain the overall page layout structure; Edge gradient loss is used to constrain the geometry of text edges, document boundaries, and crease regions; This is used to constrain the consistency of high-level visual structures, representing a perceptual loss. , , , The corresponding loss weights.

8. The document image correction method based on spatial-channel sparse attention according to claim 7, characterized in that, =0.8， =0.1， =0.05， 。 9. A document image correction system based on spatial-channel sparse attention that implements the document image correction method as described in any one of claims 1 to 8, characterized in that, include: The image acquisition module is used to acquire images of the document to be corrected. A feature embedding module is used to perform convolutional layers and input projection on the document image to be corrected; The encoding module is used to extract multi-scale features through multiple base layers and a downsampling module; The deep modeling module is used to perform deep modeling of features through adaptive sparse probability window attention and parallel spatial-channel attention modulation. The decoding and reconstruction module is used to gradually restore spatial resolution through multiple upsampling modules, skip connections, and a base layer; The output module is used to output the corrected document image.

10. An electronic device or computer-readable storage medium, characterized in that: The electronic device includes a processor, a memory, and a computer program stored in the memory, or the computer-readable storage medium stores a computer program, which, when executed by the processor, implements the document image correction method according to any one of claims 1 to 8.

Images