High-Accuracy Image Segmentation Based on Hybrid Attention Mechanism for Sandstone Analysis

Abstract

Mineral image segmentation based on computer vision is vital to realize automatic mineral analysis. However, current image segmentation methods still cannot effectively solve the problem of sandstone grains that are adjoined and concealed by leaching processes, and the segmentation performance of small and irregular grains still needs to be improved. This investigation explores and designs a Mask R-CNN-based sandstone image segmentation model, including a hybrid attention mechanism, loss function construction, and receptive field enlargement. Simultaneously, we propose a high-quality sandstone dataset with abundant labels named SMISD to facilitate comprehensive training of the model. The experimental results show that the proposed segmentation model has excellent segmentation performance, effectively solving adhesion and overlap between adjacent grains without affecting the classification accuracy. The model has comparable performance to other models on the COCO dataset, and performs better on SMISD than others.

1. Introduction

Most mineral identification and textural characterization in coarse-grained clastic sediments (such as sandstones) necessitate the use of microscopy assessment, which involves interpretation and the gathering of counting statistics using a light microscope [1,2,3,4,5,6]. The analysis of sandstone images is a vital part of geological exploration research, as rocks are rich in hydrocarbon resources [7]. Sedimentary petrology is a specialized field that demands significant expertise to interpret and classify porosity, minerals, matrix materials, and mineral cements, along with their associated subclasses. A range of objective lenses, optical techniques, and light sources are employed to identify and quantify components. Certain components in sandstone can be small, dark, or opaque, posing challenges for interpretation even to an experienced petrologist. In such cases, methods beyond light microscopy, such as scanning electron microscopy or bulk measurement techniques like X-ray diffraction, are often utilized to assist in mineral identification [8]. Traditionally, the analysis of sandstone thin sections is primarily conducted by professionals, which is labor-intensive, expensive, and subjective.

At present, classic mineral image segmentation technology primarily relies on the low-level visual information of image pixels and is mainly categorized into three types: (1) the threshold-based mineral image segmentation algorithm, which compares each pixel in the input images with a preset threshold value to segment the target areas [9,10,11,12]; (2) the region-based mineral image segmentation algorithm that divides the original image into different pixel regions, separating the target areas from the background [13,14,15]; and (3) the specific theory-based mineral image segmentation algorithm that employs more targeted computational methods such as cluster analysis to separate mineral grain images [16]. Although these classic mineral image segmentation methods have improved the efficiency of sandstone analysis, they cannot effectively address the issue of adhesion and overlap between adjacent grains, and their segmentation performance for small and irregular grains is relatively poor. Additionally, the classic image segmentation algorithm requires manual tuning for different types of mineral grains, which lowers the efficiency of mineral image segmentation and increases the time consumption.

Owing to the booming development of deep learning and the exceptional feature extraction capabilities of CNNs (convolutional neural networks), deep learning-based approaches have been increasingly employed in mineral image segmentation. These methods, diverging from traditional ones, adhere to an end-to-end paradigm and significantly surpass conventional methods in accuracy and efficiency. Furthermore, they are data-driven methods (i.e., the larger the dataset or the more refined the data analysis, the higher the achievable accuracy), allowing for enhanced performance with continuous data expansion. Some studies have explored and applied it in mineral image segmentation tasks. An RDU (R: residual connection; DU: DUNet) ore image segmentation model was proposed to estimate the grain size of ore fragments in conveyor belts, which can adjust the receptive field adaptively according to the size and shape of different ore fragments and achieve accurate segmentation [17]. Deep learning-based methods for mineral image segmentation also excel in segmenting adherent, overlapping, and multi-scale mineral grains, effectively addressing the typical challenges encountered in previous approaches [18]. However, existing image segmentation methods mainly focus on semantic segmentation, which cannot meet the requirements to compute specified morphological data in the mineral analysis, such as circularity, particle size, and contact relationship. Furthermore, these models are mostly designed for scenarios with clean image backgrounds, whereas in sandstone analysis scenarios, the image backgrounds are complex and filled with indistinguishable fillers, which makes it a challenge to segment grains out [19].

In order to solve the above problems and further improve the application prospects of automatic sandstone analysis, we propose a high-accuracy instance segmentation model based on Mask R-CNN [20] by introducing a hybrid attention mechanism to better adapt to the complex shapes of sandstone particles. Secondly, we design a shape-aware training loss function for the improved model and conduct an ablation experiment to confirm its advantages. Additionally, original convolution is replaced with dilated convolution for the purpose of obtaining more global information. Finally, this experiment establishes a dataset of sandstone instance segmentation with 40,122 grain labels.

2. Methodology

Section 2 mainly introduces various methods and ideas involved in the improved model of the sandstone image segmentation system. In the model-building stage, backbone selection, module setting, and loss function design are the main considered aspects to improve the segmentation model performance and solve edge blurring and incomplete segmentation in the segmentation process.

2.1. Mask R-CNN

The deep learning-based Mask R-CNN network segmentation method has superior performance and can achieve good results in the sandstone microscopic image segmentation task [21]. We customize and optimize the Mask R-CNN network to form a task specified model, based on the distribution characteristics of sandstone grains in sandstone microscopic images.

The Mask R-CNN network is divided into three steps: object localization, object category calculation, and segmentation mask prediction. The process of the Mask R-CNN network is as follows: the image passes through the ResNet backbone network, and different levels of feature maps are obtained using the feature pyramid. This facilitates the model in extracting features at different levels. It then enters the region proposal network to generate candidate areas where grains might be present. On the one hand, a classifier determines whether pixels belong to sandstone grains or the background. On the other hand, a box regressor corrects the boundaries of the sandstone grains. The combination of these two steps forms the candidate target areas. Finally, through a fully convolutional neural network, accurate segmentation results are achieved, completely extracting the sandstone grains. The processing procedure is shown in Figure 1.

2.2. SE-Net

Squeeze-and-Excitation Network (SE-Net [22]) introduces the concept of channel attention. By modeling and assigning weights to feature channels, it forms the parameters that can be learned and updated and continuously increases the weights of useful feature channels to optimize the generalization capability of the model. In this paper, SE-Net is added to the backbone recognition network to improve the spatial information processing capability of the model and therefore achieve better results for sandstone grains at different scales. The specific structure of SE-Net is shown in Figure 2.

The first operation is F_tr conversion, which converts the input x with a feature channel number C₁ by a series of general transformations such as convolution to obtain a feature with a feature channel number C₂. The F_tr operation is shown in Equation (1).

$${\mathrm{F}}_{\mathrm{t}\mathrm{r}}:\mathrm{X}\to \mathrm{U},\hspace{1em}\hspace{1em}\mathrm{X}\in {\mathrm{R}}^{\mathrm{H}\times \mathrm{W}\times {\mathrm{c}}_1},\hspace{1em}\hspace{1em}\mathrm{U}\in {\mathrm{R}}^{\mathrm{H}\times \mathrm{W}\times {\mathrm{c}}_2}$$

(1)

The Squeeze operation is shown in Equation (2), where the two-dimensional feature channels are mapped by compression to obtain a real number, which is connected to form a one-dimensional vector to obtain the global distribution and construct the global receptive field of the model. The dimension number of the vector equals the number of input feature channels.

$${\mathrm{z}}_{\mathrm{c}}={\mathrm{F}}_{\mathrm{s}\mathrm{q}}\left({\mathrm{u}}_{\mathrm{c}}\right)=\frac{1}{\mathrm{H}\times \mathrm{W}}\sum _{\mathrm{i}=1}^{\mathrm{H}}\sum _{\mathrm{j}=1}^{\mathrm{W}}{\mathrm{u}}_{\mathrm{c}}(\mathrm{i},\mathrm{j})$$

(2)

Next is the Excitation operation. The parameter w is used to generate the weights for each feature channel. Multiply W₁ (which dimension is C/r × C) and z to reduce the computational complexity by scale operation. The dimension of the feature map of W₁z is still 1 × 1 × C/r; it is then passed through the ReLU layer and multiplied with W₂ (which dimension is C × C/r) to obtain the feature map dimension 1 × 1 × C. Finally, the Sigmoid is derived to generate the weights (s) of the feature maps (at a total of C). The parameter s incorporates the feature map information of each feature channel and is part of the neural network, which can be learned and optimized.

$$\mathrm{s}={\mathrm{F}}_{\mathrm{e}\mathrm{x}}\left(\mathrm{z},\mathrm{W}\right)=\mathsf{\sigma }\left(\mathrm{g}\left(\mathrm{z},\mathrm{W}\right)\right)=\mathsf{\sigma }\left({\mathrm{W}}_2\mathsf{\delta }\left({\mathrm{W}}_1\mathrm{z}\right)\right)$$

(3)

Finally, the weights calculated by the model are weighted onto the original feature map by the previous channels through the scale operation, as shown:

$$\widetilde{{\mathrm{x}}_{\mathrm{c}}}={\mathrm{F}}_{\mathrm{s}\mathrm{c}\mathrm{a}\mathrm{l}\mathrm{e}}\left({\mathrm{u}}_{\mathrm{c}},{\mathrm{s}}_{\mathrm{c}}\right)={\mathrm{s}}_{\mathrm{c}}{\mathrm{u}}_{\mathrm{c}}$$

(4)

2.3. Coordinate Attention and Spatial Attention

The channel attention mechanism models global information through channels, allowing the model to effectively extract sandstone grains. When extracting, the location information of the sandstone is directly related to the accuracy of the boundary fit. In this paper, a CA + SP hybrid attention mechanism is proposed to optimize the model by combining the coordinate attention mechanism [23] and the spatial attention mechanism [24].

Coordinate attention takes a similar operation to channel attention in both horizontal and vertical directions to obtain relatively independent feature maps in both directions, effectively preserving one-dimensional location information and establishing spatial long-range dependence in one dimension. This mechanism is very sensitive to coordinate information and can effectively pinpoint spatial coordinate information. Similar to the channel attention mechanism, the coordinate attention mechanism first performs coordinate position encoding and then generates coordinate attention. A coordinate attention module can be seen as a computational unit used to augment feature representation capabilities. It can take any intermediate tensor $\mathrm{X}=[{\mathrm{x}}_1,{\mathrm{x}}_2,\dots ,{\mathrm{x}}_{\mathrm{C}}]\in {\mathrm{R}}^{\mathrm{C}\times \mathrm{H}\times \mathrm{W}}$ as input and produces an output $\mathrm{Y}=[{\mathrm{y}}_1,{\mathrm{y}}_2,\dots ,{\mathrm{y}}_{\mathrm{C}}]$ of the same size with enhanced representational power. The horizontal and vertical directions are treated separately and computed using pooling kernels of dimensions (H, 1) and (1, W), giving the following outputs for the vertical and horizontal channels.

$${\mathrm{z}}_{\mathrm{c}}^{\mathrm{h}}\left(\mathrm{h}\right)=\frac{1}{\mathrm{W}}\sum _{0\le \mathrm{i}<\mathrm{W}}{\mathrm{x}}_{\mathrm{c}}(\mathrm{h},\mathrm{i}){,\mathrm{z}}_{\mathrm{c}}^{\mathrm{w}}\left(\mathrm{w}\right)=\frac{1}{\mathrm{W}}\sum _{0\le \mathrm{j}<\mathrm{W}}{\mathrm{x}}_{\mathrm{c}}(\mathrm{j},\mathrm{w})$$

(5)

The two feature maps generated by the previous module are cascaded and then transformed using a shared 1 × 1 convolution to perform the transformation F₁, expressed as in Equation (5); the generated $\mathrm{f}\in {\mathrm{R}}^{\mathrm{C}/\mathrm{r}\times (\mathrm{H}+\mathrm{W})}$ is the intermediate feature map during the computation, where r denotes the down-sampled ratio, which is used to control the size of the module, like the SE module.

$$\mathrm{f}=\mathsf{\delta }\left({\mathrm{F}}_1\right(\left[{\mathrm{z}}^{\mathrm{h}},{\mathrm{z}}^{\mathrm{w}}\right]\left)\right)$$

(6)

Next, f is cut into two direction-independent tensors, ${\mathrm{f}}^{\mathrm{h}}\in {\mathrm{R}}^{\frac{\mathrm{C}}{\mathrm{r}}\times \mathrm{H}},{\mathrm{f}}^{\mathrm{w}}\in {\mathrm{R}}^{\frac{\mathrm{C}}{\mathrm{r}}\times \mathrm{W}}$, and then using the two 1 × 1 convolutions F_h and F_w, we transform two tensors to the same number of channels as the input and output, as in the following equation:

$${\mathrm{g}}^{\mathrm{h}}=\mathsf{\sigma }\left({\mathrm{F}}_{\mathrm{h}}\left({\mathrm{f}}^{\mathrm{h}}\right)\right),{\mathrm{g}}^{\mathrm{w}}=\mathsf{\sigma }\left({\mathrm{F}}_{\mathrm{h}}\left({\mathrm{f}}^{\mathrm{h}}\right)\right)$$

(7)

The two are then expanded as attention weights, and final output of the CA module is

$$\widetilde{{\mathrm{y}}_{\mathrm{c}}}(\mathrm{i},\mathrm{j})={\mathrm{x}}_{\mathrm{c}}\left(\mathrm{i},\mathrm{j}\right)\times {\mathrm{g}}_{\mathrm{c}}^{\mathrm{h}}\left(\mathrm{i}\right)\times {\mathrm{g}}_{\mathrm{c}}^{\mathrm{w}}\left(\mathrm{j}\right)$$

(8)

The spatial attention mechanism can simulate the function of the human eye and extract the parts of interest to the model, which generates a mask for the space, and draws out one way, which undergoes similar operations as described above to form the spatial attention mechanism, effectively extracting its relative spatial information.

2.4. Dilated Convolution

The up-sampling process of bilinear interpolation has large errors, leading to problems such as distortion when generating sandstone grain contours, and this paper achieves refinement of sandstone contours by introducing dilated convolution. The dilated convolution [25] has the following main functions:

• Expanding the Mask R-CNN network receptive fields more efficiently while taking into account image resolution;
• Changing the size of the convolution kernel and the perceptual field of the model by adjusting the expansion rate (r) to obtain multi-scale global semantic information.

The dilated convolution is shown in Figure 3. Only non-zero elements play a role in the calculation, and the dilated convolution fills the ordinary convolution with zeros to increase the size of the receptive field, as shown in Equation (9).

$$\mathrm{K}=\mathrm{k}+(\mathrm{k}-1)(\mathrm{r}-1)$$

(9)

K is the size of the expanded convolution kernel, k is the size of the original convolution kernel, and r represents the expansion rate. As shown in the figure, when k = 5 and r = 2, compared with the ordinary 5 × 5 convolution on the left, the dilated convolution expands the convolution kernel to 11 × 11, and the range of the sensory field is greatly improved.

3. Methods and Materials

3.1. Experimental Methods

3.1.1. Hybrid Attention Mechanism

In this paper, the SE module is embedded in the backbone network ResNet and its processing flow is shown in Figure 4. In contrast to the SE-Net structure mentioned earlier, Squeeze takes global average pooling and Excitation takes two full connection layers for computation. To maintain the ResNet structure, it is also necessary to ensure that the input and output dimensions are the same: the feature dimension is first reduced to 1/4 of the input, and then the dimension is recovered after ReLU activation. This operation not only complies more with the specification of the Excitation operation but also reduces the computational costs by reducing the dimensionality. The neural network architecture with two fully connected layers allows for better modeling of correlations between channels. Finally, analogous to the SE-Net model, the scale operation is set to weight the weights by channel onto the corresponding original feature maps to form the SE-ResNet model.

Additionally, a CA + SP-ResNet backbone network is proposed to optimize the model by combining the coordinate attention mechanism and the spatial attention mechanism in a re-weighted way, which is shown in Figure 5.

After the spatial attention mechanism, the output section of the module should also multiply the coordinate attention output above by the effect of spatial attention and modify the final output to Equation (10).

$$\mathrm{y}\left(\mathrm{i},\mathrm{j}\right)=\tilde{\mathrm{y}}\left(\mathrm{i},\mathrm{j}\right)\times {\mathrm{g}}_{\mathrm{c}}^{\mathrm{s}}(\mathrm{i},\mathrm{j})$$

(10)

Finally, the output of SE-ResNet and CA + SP-ResNet also uses the similar re-weighted method to obtain the final result.

3.1.2. Loss Function

Sandstone grains are often dense and complex in shape. To reduce segmentation error, an accurate sandstone grain boundary is of great importance. In fact, it is necessary to compare the real grain area with the predictions to ensure that their boundaries and contours fit perfectly. Thus, we use DiceLoss [26] to calculate the segmentation error, which is used to calculate the Dice (the graphical similarity between the true and predicted segmentation results of the sandstone grains). Let the true result of segmentation be |T| and the predicted result be |P|, and then |T∩P| be the dot product of the binary plot of the predicted segmentation result and the true segmentation result. The elements in |T| and |P| directly take the summation. The similarity coefficient Dice and segmentation error DiceLoss are as follows:

$$\mathrm{D}\mathrm{i}\mathrm{c}\mathrm{e}\left(\mathrm{P},\mathrm{T}\right)=\frac{|\mathrm{P}\cap \mathrm{T}|}{\left(\left|\mathrm{P}+\mathrm{T}\right|\right)/2}=\frac{2\ast \mathrm{g}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}\cap \mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}}{\mathrm{g}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}+\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}}$$

(11)

$$\mathrm{D}\mathrm{i}\mathrm{c}\mathrm{e}\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}(\mathrm{P},\mathrm{T})=1-2\frac{|\mathrm{P}\cap \mathrm{T}|}{\left|\mathrm{P}\right|+\left|\mathrm{T}\right|}$$

(12)

To avoid the denominator of the equation being zero and to reduce over-fitting, smoothing is added to the numerator and denominator at the same time:

$$\mathrm{D}\mathrm{i}\mathrm{c}\mathrm{e}\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}(\mathrm{P},\mathrm{T})=1-2\frac{\left|\mathrm{P}\cap \mathrm{T}\right|+\mathrm{s}\mathrm{m}\mathrm{o}\mathrm{o}\mathrm{t}\mathrm{h}}{\left|\mathrm{P}\right|+\left|\mathrm{T}\right|+\mathrm{s}\mathrm{m}\mathrm{o}\mathrm{o}\mathrm{t}\mathrm{h}}$$

(13)

Since cross-entropy is only a form of proxy, we directly adopt DiceLoss as the loss function. The real target of the segmentation is to maximize the overlap between the predicted results and the true segmentation results, i.e., the Dice coefficient, thus minimizing DiceLoss. Also, the prediction of segmentation network accuracy performance often needs to be tested with IoU, which is very similar to DiceLoss; this operation allows the performance metric to be trained and optimized directly as a component of the loss function. Due to the non-convexity of the DiceLoss function, there can be problems such as gradient explosion, which typically leads to problems such as more unstable training and difficulty in converging the training. Through several experiments, this paper uses Log-Cosh DiceLoss, based on the Log-Cosh function, to smooth the DiceLoss. The Cosh function as well as the Log-Cosh DiceLoss equation, i.e., L_lc-dce, are shown in Equations (14) and (15).

$$\mathrm{C}\mathrm{o}\mathrm{s}\mathrm{h} \mathrm{x}=\frac{{\mathrm{e}}^{\mathrm{x}}+{\mathrm{e}}^{-\mathrm{x}}}{2}$$

(14)

$${\mathrm{L}}_{\mathrm{l}\mathrm{c}-\mathrm{d}\mathrm{c}\mathrm{e}}=\log \left(\cosh \left(\mathrm{D}\mathrm{i}\mathrm{c}\mathrm{e}\mathrm{L}\mathrm{o}\mathrm{s}\mathrm{s}\right)\right)$$

(15)

By changing the segmentation loss from L_mask to L_lc-dce, we obtain a more refined simulation of sandstone microscopic image boundaries, effectively improving the segmentation effect and contour refinement for irregular sandstone grains.

3.2. Material Preparation

Physical thin sections were acquired in a standard manner, in which blue epoxy (variable hue) was impregnated into the sandstone (a strictly coarse-grained clastic sediment), an approximately 30 μm thick thin section on a glass slide was prepared, and a mechanical polish was applied to the surface. Immediately prior to imaging, mineral oil and a cover slip were applied to the thin section. Generally, few label sets for petrology disciplines exist and their generation is extremely time-consuming. The images, used for labeling, were acquired on several different optical microscopes using objective lenses with 5× or 10× magnification.

The dataset SMISD has 288 typical labeled, single-polarized sandstone grain images with a cumulative total of 40,122 sandstone grains. An approach was taken to expand the dataset by rotating the images counterclockwise by different degrees, as shown in Figure 6, to a total of 1121 images with a total of 153,466 grains.

In order to obtain better generalization ability, in this paper, 85% of them are divided into a training set and the remaining 15% are divided into a test set.

4. Results

4.1. Experiments and Analysis

To test the performance of the improved Mask R-CNN model based on the above theory for sandstone microscopic image segmentation, several experiments were conducted based on the traditional Mask R-CNN network and the improved Mask R-CNN network, and the experimental environment is shown in

Table 1. Table of experimental environment.

Software/Hardware	Configuration
Operating system	Ubuntu 20.04
Memory	32 GB
CPU	Intel(R) Core(TM) i9-10920X CPU @ 3.50 GHz
GPU	NVIDIA GeForce RTX 3090
Related software	Python3.8/Torch1.8.0/cuda11.1

.

The basic parameters of the Mask R-CNN network were adjusted to obtain the best experimental results, and the hyper-parameters were set as follows: the model threshold was set to 0.5, the MINI_MASK was set to (28,28) (to facilitate the detection of fine grains), ResNet50 was used for the skeleton network (to increase the model processing speed), and the maximum number of grains present in a single image was set to 1000 (taking into account the dense distribution of grains in some sandstone micro-graphs).

4.2. Performance Metric

For the sandstone microscopic image segmentation task, the IoU is fixed and the average precision rate and average recall rate are taken as evaluation metrics.

4.2.1. Accuracy and Recall Rate, AP

To introduce the performance metric of sandstone microscopic image segmentation models (i.e., the generalization ability), the following four concepts are introduced: TP and FP, respectively, refer to the number of correct and incorrect determinations for all samples judged to be sandstone grains. TN and FN, respectively, refer to the number of correct and incorrect determinations for samples judged to be non-sandstone grains. Accuracy is the probability that all samples predicted to be sandstone grains are indeed sandstone grains, and it addresses the accuracy of the prediction results for a specific sample and sets the accuracy rate as P, given by the following equation:

$$\mathrm{P}=\frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{P}}$$

(16)

The recall rate is the probability that the model correctly predicts a sample that is indeed sandstone grains and it focuses on the prediction of the specific sample. If we set the recall as R, the equation is as follows:

$$\mathrm{R}=\frac{\mathrm{T}\mathrm{P}}{\mathrm{T}\mathrm{P}+\mathrm{F}\mathrm{N}}$$

(17)

Accuracy and recall rate are interactive, with both being high when the model generalization is strong. The two are usually negatively correlated. If the X-axis is the recall rate and the Y-axis is the precision rate, the PR curve is obtained by tracing the points. The average accuracy (AP) is the PR curve integrated from 0 to 1. The calculation formula is as follows:

$$\mathrm{A}\mathrm{P}={\int }_0^1\mathrm{P}\left(\mathrm{r}\right)\mathrm{d}\mathrm{r}$$

(18)

The calculation is simplified by smoothing, which is calculated as follows:

$$\mathrm{A}\mathrm{P}=\frac{1}{11}\sum _{\mathrm{i}=\mathrm{0,0.1},\dots ,1.0}\mathrm{s}\mathrm{m}\mathrm{o}\mathrm{o}\mathrm{t}\mathrm{h}\left(\mathrm{i}\right)$$

(19)

4.2.2. IoU

IoU [27] represents the overlap ratio between the generated and real image segmentation regions. The calculation of IoU is shown in Equation (20), where C is the real image segmentation area and G is the segmentation area generated by the model. A diagram of IoU is shown in Figure 7. The left side of the image is the intersection of the two and the right side is the merging of the two.

$$\mathrm{I}\mathrm{o}\mathrm{U}=\frac{\mathrm{A}\mathrm{r}\mathrm{e}\mathrm{a}\left(\mathrm{C}\right)\cap \mathrm{A}\mathrm{r}\mathrm{e}\mathrm{a}\left(\mathrm{G}\right)}{\mathrm{A}\mathrm{r}\mathrm{e}\mathrm{a}\left(\mathrm{c}\right)\cup \mathrm{A}\mathrm{r}\mathrm{e}\mathrm{a}\left(\mathrm{G}\right)}$$

(20)

The following six metrics for characterizing the performance of image segmentation can be obtained based on the above concepts, separately shown in

Table 2. Image segmentation performance indicators.

Average Precision (AP)
$\mathrm{A}\mathrm{P}$	$\mathrm{A}\mathrm{P}$ When IoU = 0.50:0.05:0.95
${\mathrm{A}\mathrm{P}}_{50}$	$\mathrm{A}\mathrm{P}$ When IoU = 0.50
${\mathrm{A}\mathrm{P}}_{75}$	AP When IoU = 0.75
AP Across Scales
$\mathrm{A}\mathrm{P}\mathrm{s}$	$\mathrm{A}\mathrm{P}$ When the grain is small: pixel area < 322
${\mathrm{A}\mathrm{P}}_{\mathrm{M}}$	$\mathrm{A}\mathrm{P}$ When the grain is medium: 322 < pixel area < 962
${\mathrm{A}\mathrm{P}}_{\mathrm{L}}$	$\mathrm{A}\mathrm{P}$ When the grain is large: pixel area > 962

.

${\mathrm{A}\mathrm{P}}_{50}$ means that the model is judged to be correct when the IoU of the predicted and real grain area is greater than 0.5; then, the classification AP of the model is calculated. ${\mathrm{A}\mathrm{P}}_{75}$ needs a higher IoU threshold; thus, the segmentation needs to be more refined.

AP, i.e., AP:IoU = 0.50:0.05:0.95:, means, respectively, setting the correct determination IoU to 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, or 0.95 and then averaging it after calculating the AP value to obtain the final AP.

At the same time, the test results are classified according to the area occupied by the pixels. AP Across Scales means that the detection targets are classified according to area into three categories—small, medium, and large—and then the predictive ability of the model is judged.

The accuracy rate of the model was calculated separately by several indicators such as IoU accuracy and target size to obtain the final image segmentation performance metrics, as shown in Table 2.

Compared to the conventional model, the improved model identified significantly more grains, segmented grain contours were significantly finer, and AP values were improved by about 5%.

The effectiveness of the hybrid attention mechanism, its shape awareness, and its dilated convolution are shown in

Table 3. The effectiveness of the attention mechanism.

Backbone	SMISD (%)
	${\mathbf{A}\mathbf{P}}_{\mathbf{S}}$	${\mathbf{A}\mathbf{P}}_{\mathbf{M}}$	${\mathbf{A}\mathbf{P}}_{\mathbf{L}}$
ResNet	15.4	37.4	41.2
+SE	18.9	38.1	42.3
+X attention	19.3	38.9	41.9
+Y attention	19.2	38.7	42.0
+CA	20.3	39.5	43.1
+CA + SP	20.8	39.7	43.2

,

Table 4. The effectiveness of the shape-aware loss function.

Loss Function	SMISD (%)
	$\mathbf{A}\mathbf{P}$	${\mathbf{A}\mathbf{P}}_{50}$	${\mathbf{A}\mathbf{P}}_{75}$
Lmask	33.2	38.6	29.1
Ldice	36.3	41.3	37.3

and

Table 5. The effectiveness of dilated convolution.

Type	SMISD (%)
	$\mathbf{A}\mathbf{P}$	${\mathbf{A}\mathbf{P}}_{50}$	${\mathbf{A}\mathbf{P}}_{75}$
Original convolution	33.2	38.6	29.1
Dilated convolution	37.9	43.2	34.2

, respectively.

4.2.3. Results and Visualization

The sandstone micrographic images used to visualize the effectiveness our model are shown in Figure 8. The results of segmentation when directly using the original Mask R-CNN network are shown in Figure 9; the contour fitting was poor for sandstone grains of varying sizes and shapes. The network can only fit about 10% when grains are small.

After the above modification, new segmented images are shown in Figure 10. Obviously, the model can identify about 80% of the grains regardless of their size, and the contours fit more closely. The mis-segmentation and under-segmentation are greatly reduced, and the generalization ability of the model is greatly improved.

Comparison of the Results of the Improved Mask R-CNN Network in Fitting Irregular Sandstone Grain Images

For the irregular, narrow sandstone grains, the new model is far superior to the original model in segmentation integrity as well as fit, which is shown in Figure 11.

Experiments on the Performance of the Improved Segmentation Network

The performance of the proposed image segmentation network is predicted by testing 238 sandstone images in the test set. Experiments were conducted using Mask R-CNN, HTC, PointRend, the latest Refine Mask, and the improved image segmentation algorithm (labeled Hybrid + Loss in the table), based on the publicly available COCO image segmentation dataset [28] and the sandstone microscopic dataset in this paper. For the best results, the recognition backbone networks all adopted ResNet101.

As shown in

Table 6. Performance comparison of segmentation algorithms.

	Dataset	COCO (%)				SMISD (%)
Algorithm		$\mathbf{A}\mathbf{P}$	${\mathbf{A}\mathbf{P}}_{\mathbf{S}}$	${\mathbf{A}\mathbf{P}}_{\mathbf{M}}$	${\mathbf{A}\mathbf{P}}_{\mathbf{L}}$	$\mathbf{A}\mathbf{P}$	${\mathbf{A}\mathbf{P}}_{\mathbf{S}}$	${\mathbf{A}\mathbf{P}}_{\mathbf{M}}$	${\mathbf{A}\mathbf{P}}_{\mathbf{L}}$
Mask R-CNN		39.6	27.2	49.0	57.7	32.3	15.4	37.4	41.2
HTC		41.2	27.2	51.9	61.5	33.9	15.6	38.9	44.0
PointRend		41.1	27.8	52.0	62.0	35.1	16.3	39.9	45.7
RefineMask		41.8	28.6	53.1	62.8	36.7	18.0	41.1	47.3
Hybrid + Loss		41.7	28.9	52.7	62.5	37.9	19.7	42.6	48.1

, our algorithm achieved the best results out of all the algorithms in the sandstone microscopic image segmentation task, based on SMISD. Compared to the conventional Mask R-CNN network, the improvement is 4.3%, 5.2%, and 6.9% for each of the three grain types (large, medium, and small), which is slightly better than the state-of-the-art Refine Mask network. Hybrid + Loss gained significant improvement in AP values for all sizes of grains, which indicates that the image segmentation algorithm tailored to the sandstone microscopic image segmentation task can effectively identify and extract sandstone grains, and can better fit the contours to achieve fine segmentation of sandstone grains.

5. Discussion

Channel attention can enhance the network’s ability to extract image information. Coordinate attention can improve the model’s ability to locate boundaries. Spatial attention can enhance the model’s receptive capability and optimize the model’s generalization ability, that is, its performance on images that have not been used for training. By introducing the hybrid attention mechanism, the model proposed in this paper surpasses other models on SMISD. With the goal of the task—to make the boundaries of grain segmentation results more precise—as the performance evaluation index, a shape-aware loss function can improve the model’s segmentation effect on grain contours, especially irregular grains. Therefore, the model adopts Log-Cosh DiceLoss as the model’s loss function. Finally, the dilated convolution can expand the receptive field and segment by combining more regional information surrounding the grains, which makes model perform better. Additionally, the dataset with rich labeled grains helps the model fully learn the segment rules, which is vital for data-driven methods.

6. Conclusions

Intelligent analysis of images of thin sandstone sections is of high research and application value, but the search for automatic segmentation and recognition is challenging because of the variety and complexity of sandstone images. This paper designed a high-accuracy sandstone segmentation model by mixing several attention mechanisms and applying appropriate loss functions and dilated convolution. The segmented particle images can be classified using other methods to obtain their mineralogical categories. The research on the automatic analysis of thin sandstone sections has made some progress in this paper, but it is difficult to identify filler, especially heterogeneous groups and cementation, and there is still much room for improvement in recognition accuracy.