# Automatic Medical Image Segmentation with Vision Transformer
## 3 Method
### Infection Segmentation Network
#### 3.2.2 Edge Attention Module  
As shown in prior work, edge information provides useful constraints to guide feature extraction for segmentation tasks [46; 47; 48]. Since low-level feature representations (e.g., \(f_{l}\) in our model) contain several edge features at a moderate resolution, we fed these low-level feature representations \(f_{l}\) to a newly introduced edge attention (EA) module. This module learns an explicit edge-aware feature representation. Meanwhile, \(f_{l}\) passes through a convolutional layer with a single filter to produce an edge map prediction.  
Figure 1: The framework overview. (**a**) The schematic of transformer layer; (**b**) structure of our proposed framework.  
The edge attention module improves medical image segmentation in the following way. It first identifies edges and boundaries between structures in the input image using gradient information. These detected edges provide vital cues for delineating infection regions from surrounding healthy tissue. The predicted edge maps are then used to refine the convolutional feature representations extracted from the image. By element-wise multiplying the edge maps with the convolutional filters, the features are enhanced at semantically meaningful edges, incorporating critical localization cues. The refined edge-aware features help guide the network's attention to focus on the most relevant regions for accurate segmentation. The edge-guided spatial attention provides fine-grained localization signals to help distinguish infection areas from other tissues. In essence, the edge attention module incorporates explicit edge knowledge to enhance the convolutional features. By refining features and steering attention based on edges, the module enables the more precise delineation of infection regions, thus improving medical image segmentation performance. The edge-guided spatial attention is crucial for accurate localization of infection regions.  
We subsequently assessed the disparity between the edge map prediction from the EA module and the ground-truth (GT) edge map \(\mathrm{G_{e}}\) derived from the annotated segmentation tasks. This was performed by computing the binary cross entropy (BCE) loss function:  
\[\mathcal{L}_{edge}=-\sum_{x=1}^{w}\sum_{y=1}^{h}[G_{e}\log(S_{e})+(1-G_{e})\log (1-S_{e})] \tag{4}\]  
Here, \((x,y)\) indexes the spatial coordinates within the predicted edge map \(S_{e}\) and the ground-truth edge map \(G_{e}\). The ground-truth edge map \(\mathrm{G_{e}}\) is derived by computing the gradient of the annotated segmentation mask. The \(w\) and \(h\) are the corresponding width and height of the edge maps, respectively. The binary cross-entropy loss compares the predicted and ground-truth edge maps at each pixel location, providing localized supervision for predicting salient boundaries. By explicitly guiding the model to focus on edges, it learns to delineate challenging vague transitions between infection regions and healthy tissues.  
We exploited low-level edge information by passing the convolutional features through the edge attention module. By optimizing the edge map predictions against the ground-truth edge maps using the BCE loss, it guides the network to obtain edge-aware feature representations that are beneficial for the segmentation task. The BCE loss compares the predicted and ground-truth edge maps pixel-by-pixel, guiding the network to learn edge-aware features that can effectively capture and represent the boundary information present in the input image.