A ViT-based electron microscope image registration method

Abstract

The present invention relates to the field of image processing technology, and discloses a method for electron microscope image registration based on ViT, wherein a fixed image and a deformed image are input into a trained image registration model to obtain an optical flow field for mapping the deformed image to the fixed image; the training process of the image registration model comprises the following steps: collecting image pairs consisting of a fixed image and deformed images with different degrees of deformation to obtain a training data set; the image registration model adopts a dual feature encoder, extracts a feature pyramid by layering, and fuses features by cost volume blocks and GRU units, calculates the residual optical flow of each layer, and performs recursive optical flow field prediction from coarse to fine; the training process starts with a data set with a smaller deformation, introduces a data set with a larger deformation layer by layer, and gradually expands the deformation adaptation range of the image registration model, thereby enhancing its robustness. The present invention can analyze visual information of each scale in detail and greatly improve the accuracy of image registration.

Description

Technical Field

[0001] The present invention relates to the technical field of image processing, and in particular to an electron microscope image registration method based on ViT. Background Art

[0002] Electron microscopes are devices that use electron beams to perform high-resolution imaging and have important applications in medical and biological research.

[0003] Electron microscope image registration refers to the precise adjustment and alignment of images through algorithms to compensate for physical deformations caused by sample preparation and imaging. Traditional image registration methods mainly rely on manual features, such as Scale-Invariant Feature Transform (SIFT) and ORB (Oriented FAST and Rotated BRIEF). These methods achieve registration by identifying and matching key points in the image, and combine the Random Sample Consensus (RANSAC) algorithm to handle outliers to improve the robustness of the registration. However, these methods require complex parameter adjustments and manual intervention when dealing with large-scale deformations and complex structures, which limits their flexibility and efficiency.

[0004] With the rapid development of deep learning technology, the field of image registration has undergone significant changes. Optical flow estimation algorithms based on convolutional neural networks (CNNs), such as the DeepFlow algorithm and the FlowNet model, have greatly improved the accuracy and efficiency of registration by automatically learning the deep features of image content. When the optical flow estimation algorithm registers images, two electron microscope images to be registered are given, namely a fixed image and a deformed image. The goal of electron microscope image registration is to calculate an optical flow field, map the deformed image to the fixed image, and strive to make the pixels on the deformed image correspond one-to-one with the pixels on the fixed image.

[0005] However, CNN still has limitations when dealing with large-scale deformed images. To overcome these challenges, researchers have proposed a series of improved methods based on deep learning. For example, Balakrishnan et al. proposed an unsupervised learning strategy that overcomes the traditional method's reliance on registration truth labels by directly learning the deformation field from unlabeled image pairs. The ssEMnet model specifically optimizes ssEM images for dense and changing structures by combining spatial transformers and convolutional autoencoders. However, these methods usually only perform well in small deformations and still face difficulties in dealing with large deformations.

[0006] The Visual Self-Attention Model (Vision Transformer, ViT) is a model that applies the Transformer architecture in the field of computer vision. Summary of the invention

[0007] In order to solve the difficulty of image alignment caused by physical deformation of sample preparation and imaging process in electron microscope image registration, the present invention provides an electron microscope image registration method based on ViT. In order to cope with the challenges of complex structures and large-scale deformation in electron microscope image registration, the present invention extracts feature pyramids in layers through a dual feature encoder (TwinAxialFormer), calculates the residual optical flow of each layer, and performs recursive optical flow field prediction from coarse to fine; in order to overcome the training difficulties of large deformation data sets, a progressive training method is adopted to make the image registration model more stable when processing large deformation images, improve the accuracy and efficiency of image registration, and thus achieve more accurate three-dimensional structure reconstruction in medical and biological research.

[0008] In order to solve the above technical problems, the present invention adopts the following technical solutions:

[0009] A ViT-based electron microscope image registration method includes inputting a fixed image and a deformed image into a trained image registration model to obtain an optical flow field for mapping the deformed image to the fixed image; the training process of the image registration model includes the following steps:

[0010] Step 1: Collect image pairs consisting of fixed images and deformed images with different deformation degrees to obtain a training data set; the fixed image is recorded as the initial fixed coding feature , the deformed image is recorded as the initial deformation encoding feature ; The image registration model includes an N-layer structure and a refinement module. The n-th layer structure of the image registration model includes a dual feature encoder , a cross-spatial attention module and a registration pipeline , ;

[0011] Step 2: For the nth layer structure of the image registration model: convert the fixed coding features obtained in the n-1th layer structure and deformation encoding features Input to the dual feature encoder The convolution blocks are downsampled to obtain multiple fixed image blocks and multiple deformed image patches ;Will and They are input into the axial attention module respectively, and feature enhancement is performed through row attention and column attention to obtain fixed attention features. and deformation attention features ;Will and Jointly input to the criss-cross spatial attention module , get the fixed coding features of the n-th layer structure and deformation encoding features ;

[0012] calculate and The cost volume CV between , and CV input to the registration pipeline The GRU units are fused to obtain the fusion feature ; Then the fusion features And the optical flow map estimated by the n+1th layer structure of the image registration model Input to the registration pipeline In the optical flow estimation block, a preliminary estimated optical flow map is obtained. ;Will Input into the upsampling optical flow block to get the optical flow map ;

[0013] Note that when n=N, that is, in the lowest level of the image registration model, the input of the optical flow estimation block of the registration pipeline is only the fusion feature , because this is already the lowest level structure of the image registration model, and there is no optical flow map estimated by a deeper level structure .

[0014] Step 3: Output the optical flow map of the registration pipeline of the first layer of the image registration model Input to the refinement module to obtain the final optical flow map;

[0015] Step 4, progressive training: perform hierarchical training on the multi-layer structure of the image registration model, that is, after completing the training of the n-th layer structure of the image registration model, freeze the n-th layer structure, and then train the n+1-th layer structure of the image registration model; the degree of deformation of the image pairs used when training the n+1-th layer structure is greater than the degree of deformation of the image pairs used when training the n-th layer structure; after completing the training of all layers of the image registration model structure, unfreeze all layers of the structure and perform overall training on the image registration model.

[0016] Further, in step 2, the and Input the axial attention module respectively, perform feature enhancement through row attention and column attention, and obtain fixed attention features and deformation attention features , specifically including:

[0017] Fixed image block and deformed image patches Image blocks ;

[0018] Row attention calculates the attention weights between image blocks in each row, and enhances the horizontal integration of features by calculating the attention scores of other image blocks in each row;

[0019] Column attention calculates the attention weights between image blocks in each column, and enhances the vertical integration of features by calculating the attention scores of other image blocks in each column.

[0020] Compared with the prior art, the present invention has significant advantages in processing large-scale deformation and complex structure image registration. The specific advantages are as follows:

[0021] 1. High-precision registration: By combining multi-scale feature extraction, the present invention can analyze the visual information of each scale in detail and greatly improve the accuracy of image registration.

[0022] 2. Efficient training: The progressive training method effectively overcomes the training difficulties of large deformation data sets, making the model training process more efficient and stable.

[0023] 3. Strong robustness: By introducing larger deformed data sets layer by layer, the deformation adaptation range of the model is gradually expanded, thus enhancing the robustness of the model.

[0024] 4. High degree of automation: The present invention does not require complicated parameter adjustment and manual intervention, and is suitable for automatic processing of large-scale data sets.

[0025] Experimental results show that the registration accuracy and efficiency of the present invention are significantly better than those of the prior art when processing electron microscope images with large-scale deformation and complex structures. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 Schematic diagram of the structure of the image registration model used in the present invention. DETAILED DESCRIPTION

[0027] A preferred embodiment of the present invention is described in detail below with reference to the accompanying drawings.

[0028] Figure 1This paper presents an image registration model that combines multi-scale feature extraction and a visual self-attention model (VisionTransformer, ViT) proposed by the present invention, which is used to accurately capture the complex structure and deformation of the image. The image registration model uses a twin feature encoder (TwinAxialFormer) to extract feature pyramids in layers, calculate the residual optical flow of each layer, and perform recursive optical flow field prediction from coarse to fine. Specifically, the twin feature encoder (TwinAxialFormer) downsamples the input image through a convolution block (convpatch) to generate an image block. Then, the axial attention block (AxialAttention Block) enhances the feature integration of the image block through the row attention mechanism and the column attention mechanism respectively, and the cross-spatial attention block (Crossed Spatial Attention Block) further fuses feature map information of different scales. The image registration model effectively solves the registration problem of large-scale deformation and complex structure images through multi-scale feature extraction and progressive training methods, and realizes high-precision alignment of electron microscope images.

[0029] The present invention is described in detail below.

[0030] 1. TwinAxialFormer:

[0031] The TwinAxialFormer is one of the core components of the present invention, which is improved on the basis of the visual self-attention model (ViT). It uses convolution blocks to downsample the input image, generate image blocks, and process these image blocks through a multi-head attention structure. The multi-head attention structure is an important part of the network, which can process multiple information representations at the same time and significantly improves the network's ability to integrate features in different spatial dimensions. The present invention develops an axial attention module specifically for 2D image registration. The axial attention module is nested in the multi-head attention structure and is divided into two main steps: row attention and column attention to optimize feature extraction.

[0032] 1.1, line attention:

[0033] Row attention calculates the attention weights between image blocks in each row, and enhances the horizontal integration of features by calculating the attention scores of other image blocks in each row. The specific formula is as follows:

[0034] ;

[0035] in, represents the row attention score of the i-th row, , and Represents the image block The query, key and value of the row, is the dimension of the key, Represents the softmax function.

[0036] 1.2, column attention:

[0037] Column attention calculates the attention weights between image patches in each column, and enhances the vertical integration of features by calculating the attention scores for other image patches in each column. Combining row attention and column attention, this hierarchical application of axial attention achieves global connectivity between image patches while significantly reducing computational resource requirements.

[0038] After calculating the axial attention, the TwinAxialFormer further processes the features through a multi-layer perceptron (MLP) module. The multi-layer perceptron enhances and refines the transformed features through nonlinear transformation and feature fusion, thereby improving the model's expressiveness and performance. Similar to the classic Transformer architecture, the TwinAxialFormer also combines layer normalization, residual connections, and dropout strategies to ensure the generalization ability of the model and the refinement of feature learning trajectories.

[0039] 2. Alignment Pipeline:

[0040] After extracting features using the TwinAxialFormer and further feature fusion using the Crossed Spatial Attention Block, the present invention estimates the optical flow map through the registration pipeline. , To begin, first set the feature , Input to the Cost Volume Block to calculate the cost volume between the image blocks. The cost volume constructs a three-dimensional data structure that contains the similarity between pixels or image blocks in two images. Next, the cost volume and feature , Input into the GRU unit (GRU Unit) for fusion. Then, the optical flow estimation block (Flow Head Block) processes the estimated optical flow map of the next layer and the fused features of the GRU unit together to obtain a preliminary estimated optical flow map of the current layer; then the preliminary estimated optical flow map is input into the upsampling optical flow block (UPFlow Block) to obtain the estimated optical flow map of the current layer. In this embodiment, the image registration model has a three-layer structure, that is, N=3. Through three recursive optical flow estimates, the optical flow map estimate of the original resolution is finally obtained. Subsequently, a refinement module (Refiner) is used for refinement processing to further improve the accuracy and continuity of the optical flow prediction and obtain the final optical flow map. The refinement module includes a series of convolutional layers, each with a ReLU activation function.

[0041] 3. Progressive training method:

[0042] In order to cope with the training difficulties of large deformed data sets, the present invention adopts a progressive training method with a phased strategy, which is particularly suitable for models with hierarchical structures, such as pyramid network architectures. The training process starts with a data set with smaller deformations, and mainly learns the precise registration of subtle features. As the training progresses, data sets with larger deformations are introduced layer by layer, gradually expanding the deformation adaptation range of the image registration model, thereby enhancing its robustness; after completing the training of each layer, its parameters are frozen to lock the features learned by the layer. Finally, after the parameter training of all layers is completed, all layers are unfrozen and the global fine-tuning stage is entered to optimize the relationship between layers to achieve accurate registration performance within a specific deformation range.

[0043] Through this progressive training method, the model can gradually adapt to large deformation datasets and improve the registration ability of complex structures and large-scale deformation images. This method not only improves the accuracy and efficiency of electron microscope image registration, but also enhances the robustness and adaptability of the model.

[0044] 4. The present invention specifically comprises the following steps:

[0045] 1) Image preprocessing: Standardize and normalize the input electron microscope image to eliminate noise and inhomogeneity and improve image quality.

[0046] 2) Feature extraction: The twin-feature encoder (TwinAxialFormer) is used to extract multi-scale features from the preprocessed image. The row attention and column attention respectively calculate the attention weights within the row and column to enhance the horizontal and vertical integration of features. Specifically, the preprocessed electron microscope image is firstly subjected to feature extraction. In the twin-feature encoder, the input image is first downsampled by a convolutional block (convpatch) to generate image blocks. Then, the row attention of the axial attention block (Axial Attention Block) enhances the horizontal integration of features by calculating the attention weights between image blocks within each row; similarly, the column attention of the axial attention block (Axial Attention Block) enhances the vertical integration of features by calculating the attention weights between image blocks within each column. Next, the cross-spatial attention block (CrossedSpatial Attention Block) performs information fusion between feature maps of different scales to further improve the effect of feature extraction.

[0047] The structure of the cross-spatial attention module is specifically referenced in "Electron Microscopy Image Registration using Correlation Volume (ISBI 2023)".

[0048] 3) Optical flow field estimation: The residual optical flow is calculated by fusing features through the cost volume block and the GRU unit, and recursive optical flow field prediction is performed from coarse to fine, and finally the full-resolution optical flow field estimation is obtained. Specifically, the present invention performs iterative estimation of the optical flow field through the alignment pipeline. First, the cost volume is constructed through the extracted feature pyramid. Then, the cost volume and the feature map are input into the GRU unit for feature fusion. The optical flow estimation block (Flow Head Block) combines the estimated optical flow field and the fused features of the next layer to perform optical flow estimation. The upsampling optical flow block (UPFlow Block) recursively upsamples the optical flow field, passes the optical flow field estimation result to the previous layer, and refines the optical flow field estimation layer by layer. This process is recursively performed under a multi-layer structure, and finally a full-resolution optical flow field estimation is obtained.

[0049] 4) Refinement: Use the refinement model (Refiner) to refine the optical flow field estimation to further improve the accuracy and continuity of the optical flow field estimation. The refinement module includes a series of convolutional layers, each followed by a ReLU activation function, which expands the receptive field layer by layer to enhance the accuracy of the optical flow estimation.

[0050] 5) Progressive training: Using a progressive training method, we introduce larger deformed datasets layer by layer, gradually expand the deformation adaptation range of the model, and optimize the inter-layer relationship through the global fine-tuning stage to achieve accurate alignment.

[0051] 5. The electron microscope image registration system corresponding to the present invention includes the following modules:

[0052] Image preprocessing module: includes standardization and normalization processing functions to improve image quality.

[0053] Dual-feature encoder: It consists of convolution blocks, axial attention modules, and multi-layer perceptrons for multi-scale feature extraction.

[0054] Cross-spatial attention module: Information fusion between feature maps of different scales to further improve the effect of feature extraction.

[0055] Registration pipeline: It consists of cost volume block, GRU unit, optical flow estimation module and refinement model, which is used for optical flow estimation and refinement processing.

[0056] Progressive training module: adopts a phased training strategy to gradually expand the deformation adaptation range of the model and optimizes the inter-layer relationship through a global fine-tuning stage.

[0057] It is obvious to those skilled in the art that the present invention is not limited to the details of the exemplary embodiments described above, and that the present invention can be implemented in other specific forms without departing from the spirit or essential features of the present invention. Therefore, from any point of view, the embodiments should be regarded as exemplary and non-limiting, and the scope of the present invention is defined by the appended claims rather than the above description, and it is intended that all changes falling within the meaning and scope of the equivalent elements of the claims are included in the present invention, and any reference numerals in the claims should not be regarded as limiting the claims involved.

[0058] In addition, it should be understood that although the present specification is described in terms of implementation modes, not every implementation mode contains only one independent technical solution. This narrative method of the specification is only for the sake of clarity. Those skilled in the art should regard the specification as a whole. The technical solutions in each embodiment may also be appropriately combined to form other implementation modes that can be understood by those skilled in the art.

Claims

1. A ViT-based electron microscope image registration method, characterized in that: The fixed image and the deformed image are input into the trained image registration model to obtain an optical flow field for mapping the deformed image to the fixed image. The training process of the image registration model includes the following steps: Step 1: Collect image pairs consisting of fixed images and deformed images with different deformation degrees to obtain a training data set; the fixed image is recorded as the initial fixed coding feature The deformed image is recorded as the initial deformation encoding feature The image registration model includes an N-layer structure and a refinement module. The n-th layer structure of the image registration model includes a dual feature encoder TA n , a cross-spatial attention module C n and a registration pipeline P n , 1≤n≤N; the refinement module includes a series of convolutional layers, each followed by a ReLU activation function, which enhances the accuracy of optical flow estimation by expanding the receptive field layer by layer; Step 2: For the nth layer structure of the image registration model: convert the fixed coding features obtained in the n-1th layer structure and deformation encoding features Input to the dual feature encoder TA respectively n The convolution blocks are downsampled to obtain multiple fixed image blocks and multiple deformed image patches Will and They are input into the axial attention module respectively, and feature enhancement is performed through row attention and column attention to obtain fixed attention features. and deformation attention features Will and Commonly input to the cross spatial attention module C n , get the fixed coding features of the n-th layer structure and deformation encoding features calculate and The cost volume CV between and CV input to the registration pipeline P n The GRU units are fused to obtain the fusion feature Then the fusion features The optical flow map DFV estimated by the n+1th layer structure of the image registration model n+1 Input to the registration pipeline P n In the optical flow estimation block, a preliminary estimated optical flow map is obtained Will Input into the upsampling optical flow block to get the optical flow map DFV n ; Step 3, the optical flow map DFV1 output by the registration pipeline of the first layer structure of the image registration model is input into the refinement module to obtain the final optical flow map; Step 4, progressive training: perform hierarchical training on the multi-layer structure of the image registration model, that is, after completing the training of the n-th layer structure of the image registration model, freeze the n-th layer structure, and then train the n+1-th layer structure of the image registration model; the degree of deformation of the image pair used when training the n+1-th layer structure is greater than the degree of deformation of the image pair used when training the n-th layer structure; after completing the training of all layers of the image registration model structure, unfreeze all layers of the structure and perform overall training on the image registration model; In step 2, the and Input the axial attention module respectively, perform feature enhancement through row attention and column attention, and obtain fixed attention features and deformation attention features Specifically include: Fixed image block and deformed image patches Image blocks Row attention calculates the attention weights between image blocks in each row, and enhances the horizontal integration of features by calculating the attention scores of other image blocks in each row; Column attention calculates the attention weights between image blocks in each column, and enhances the vertical integration of features by calculating the attention scores of other image blocks in each column.

Images