Rock Particle Segmentation Using Mask R-CNN

1. Introduction

Computer vision has advanced greatly in recent years with deep learning approaches pushing the boundaries of image analysis and object recognition [1]. These approaches have created opportunities for the segmentation of complex and irregular objects such as rock particles in geological images. Accurate identification and segmentation of rock particles is vital and holds high significance in geology [2], and other industries such as engineering construction [3,4,5], mining [6,7], and the petroleum industry [8], just to mention a few. In geology for example, geological analysis has shown to be essential to understanding the Earth’s history, processes, and composition, thus proving vital in lithology, mineral prospecting, environmental studies, and hazard assessment [9]. Before carrying out this analysis, Geologists carefully study rock formations where they can identify geological markers, from microscopic grains to large slabs of rock, that can provide essential information to understanding ancient ecosystems, geological transitions, and current geological dynamics [10]. By examining the shapes, sizes, and textures of these rock particles, Geologists can utilize the information to reveal key characteristics of the Earth’s geology, minerals, sedimentary processes, and climate, thereby providing a stronger foundation for explaining the Earth’s past and predicting its future [11]. However, carrying out any form of statistical analysis to extract useful information from these rock particles is only possible after the accurate segmentation of the particles.

Traditional methods of rock segmentation have required manually identifying and classifying the rocks either physically or from photographs. This human-subjective particle recognition process is laborious, time-consuming, and prone to standardization errors that may emerge from the growing data volumes, thus rendering the process inefficient and unreliable. This brought forth the incorporation of image segmentation techniques like thresholding, watershed, region and edge-based segmentation, and clustering [12,13], to address the limitations of the traditional method. However, these segmentation techniques have been found to be ineffective when working on a batch of images, most specifically when these images have occluded objects or irregular shapes of different sizes. This is so, because, they expect the objects in the images to exhibit specific characteristic features based on their background, lighting, size, or shape. They also require manual tuning of their parameters and fail to produce accurate segmentations when the objects don’t have clear boundaries, as they are more useful with high-contrast images.

Matthew [14] implemented a multi-Otsu threshold algorithm [15] to segment a synthetic rock image dataset generated from sandstone and shale rock. For the greyscale rock objects with high noise levels, the algorithm was ineffective, while also failing completely at segmentations for non-greyscale (textural contrast) images. Qinpeng et al [5] proposed an adaptive watershed algorithm for segmenting blasted rock particles based on the rock contour solidity threshold. The algorithm was able to reduce errors from blurry edges and overlapping particles. Still, it wasn’t robust enough as it required different solidity thresholds for different rock types which had to be manually set up. Haibo and Jialing [16] combined characteristics of the Fuzzy C-means clustering (FCM) algorithm with a semi-supervised Support Vector Machine (SVM) for segmenting rock images by taking into account the spatial information of the image pixels. Amankwah and Aldrich [17] came up with a novel approach to the watershed algorithm for segmenting rocks, solving the watershed over-segmentation problem by leveraging the adaptive threshold technique to generate object markers. This ensured the watershed segmentation lines correctly matched the object boundaries. Now, although exponential development has occurred with the introduction of these techniques in comparison with prior manual human segmentation methods, they are not robust enough for rock particles of varying features - sizes, and shapes.

Consequently, classical machine learning algorithms like the supervised learning algorithms: K-nearest neighbors (KNN), ensemble learning - Adaboost, Random Forest, SVM, and the unsupervised learning: K-Means Clustering, have shown prospects in recognizing and processing rock particles from digital images [18,19,20,21]. These algorithms require a diverse dataset containing tons of images with particles along with their corresponding labels so the model built from the algorithm can learn the intricate features of the data. Junxing and Roman [22] utilized the Adaboost ensemble algorithm with varying aspect ratios and sliding window orientation angles for recognizing sand particles. However, although the authors used a large training dataset of 85,000 images, the proposed method still failed to detect particles with irregular surfaces and elongated sizes. Therefore, there is a need for a more robust and adaptable system that can learn hierarchical representations captured in complex particle features with varied object boundaries and lightning conditions, thus leading to the implementation of Deep Learning systems [4,23,24,25].

In recent years, Deep Learning computer vision systems have paved the way in capturing global features of objects thus pushing the boundaries of image analysis and object recognition [1], and creating opportunities for the segmentation of complex and irregular objects such as rock particles in geological images. These systems offer robust capabilities in recognizing complex object features by passing in the training data with pixel-wise labels, so the model can learn the unique mappings of the rock objects [26]. Karolina and Sebastian [20] proposed a multi-textural approach in segmenting rock grains using the SVM, KNN, and a fully connected Artificial Neural Network (ANN) with the Levenberg-Marquardt back propagation method. Although the ANN classifier performed the best at a 79% accuracy score, the authors identified the need for post-processing techniques to improve the border thinning and automatic closing of delineated contours of the rock particles. With remarkable achievements being made from the emergence of Convolutional Neural Networks (CNN) [27], the efficiency and accuracy of segmentation techniques enhanced through Semantic segmentation architectures like the Fully Convolutional Network (FCN) and U-Net [7,24,28,29,30,31,32], providing a means of classifying every image pixel into a class label. Huang et al [28] utilized an FCN for the segmentation of cracks and leakage defects in tunnels, resulting in an improved performance in comparison to traditional techniques like Adaptive thresholding and region growing algorithm. In the mining industry, Yang et al [7] leveraged depth-separable convolutions with feature-depth concatenation to improve the rock particle segmentation accuracy, using an enhanced U-Net model architecture. Zhou et al [24] proposed a dual U-Net with multi-scale inputs and side-output for the automatic segmentation of muck images. Even though the U-Net algorithm was improved upon to produce more accurate results, it encountered limitations when segmenting overlapping rock objects which led to under-segmentation and over-segmentation. Additionally, the proposed algorithm did not bear in mind the regional features of the rocks such as the edges, and the rock surface. Liang et al [29] developed a deep convolutional neural network, by leveraging the U-Net architecture to extract particle projects from raw images. However, the authors highlighted the lack of generalization with the models’ predictions when attempting to segment rock types that were different from what was used in the training dataset.

Due to its design, semantic segmentation encounters limitations when the objects are densely packed with overlapping particles, leading to under-segmentation or over-segmentation [24]. Beyond semantic segmentation where segmentation identifies and labels every image pixel in the same class under the same category, instance segmentation has been shown to accurately distinguish overlapping instances by delineating object boundaries, making it easier to carry out further independent actions on the individual rock particles. For example, size calculation, shape identification, or any form of counting/tracking application would be more straightforward. These instance segmentation algorithms - Mask RCNN [33], in the Detectron2 [34] framework, are better suited to identify each individual object instance, making it more adaptable and efficient for rock recognition. Fan et al [4] utilized the Mask R-CNN model to segment rockfill particles by leveraging the ResNeXt101 backbone with the squeeze and excitation block, enhancing the feature extraction capability. Trong et al [35] developed an automatic means of measuring blast fragments in open-pit mines using the Mask R-CNN algorithm, although the authors observed a decline in segmentation accuracy due to poor spatial resolution of the image. As the application of instance segmentation methods to rock recognition is fairly novel, this research study provides a stepping stone for further research.

In this study, Mask R-CNN and Detectron2 models are implemented for recognizing rock particles by leveraging the transfer learning technique, which saves time and delivers excellent results with limited training data [36]. To start off, the quality of the training data is improved through a Gaussian filter and an Illumination Adaptive Transformer (IAT) [37], to denoise the image and to adjust the lighting exposure respectively. Then, the models are trained using the annotated preprocessed data with different hyperparameters, after which the performance of the proposed algorithms is evaluated using the standard Average Precision (AP) evaluation metric on a test dataset with diverse environments, by comparing the identification ability with similar research works. The rest of the paper is organized as follows: Section 2 covers the proposed methodology, introducing the dataset with the segmentation algorithms. Section 3 describes the training process implemented on the preprocessed data as well as the steps involved in evaluating the developed models. In section 4, the results from the training experiments are discussed alongside identified drawbacks. Finally, Section 5 concludes the paper with recommendations for future work.

2. Methodology

2.1. Data Collection

The dataset comprises of 69 images containing multiple rock particles, each obtained from a drilling company. These images, as seen in Figure 1 exhibit diverse environmental conditions, spanning underwater and outdoor settings, with variations in lighting, humidity, and surface wetness. Encompassing a wide spectrum of particle sizes, shapes (from gravel to cobbles), and colors, the dataset reflects real-world scenarios, thus making it suitable for robust analysis.

However, it should be noted that the raw images were beset with certain challenges that required rectification for reliable analysis:

• Blurriness: the resolution of several images was unclear, hampering the image clarity.
• Indistinct Boundaries: Distinguishing particle boundaries or edges in the images proved to be difficult.
• Limited data size: The dataset was small, comprising only 69 images to work with.

2.2. Data Quality Enhancement and Preprocessing

To mitigate the data issues highlighted, 2 key techniques were carried out:

2.2.1. Image Denoising

Due to the amount of noise and blur in the dataset, a Gaussian filter was utilized to denoise and smoothen the images. The filter uses the weighted mean of the pixels surrounding the central pixel to replace the central pixel - the pixels further from the central pixel will contribute less weight to the average while those closer would have a more weighted effect. This helps us better preserve the edges of the rock objects in the images. Equation 1 shows the equation for a Gaussian function in two directions (x and y axis), where x and y are the respective distances to the horizontal and vertical center of the kernel, and $\sigma$ is the standard deviation of the Gaussian kernel.

$$G(x,y)=\frac{1}{2\pi {\sigma }^2}{e}^{-\frac{x^2+y^2}{2{\sigma }^2}}$$

(1)

Afterward, the image was sharpened by blending the blurry and denoised images through a weighted sum of the weights from the blurry image with the weights of the denoised image.

2.2.2. Exposure Correction

To adjust the lighting exposure on some of the images that had either too bright or too low a contrast, a transformer network called IAT, which uses attention queries to control the image signal processing (ISP) parameters like color transform matrix, and gamma value, was employed thereby improving the overall image quality. The resulting preprocessed data is shown in Figure 2.

2.3. Data Annotation

For the final preparation of the dataset, the data was annotated using the makesense.ai labeling tool to obtain the ground truth mask labels and bounding box coordinates which the model will use to learn, as seen in Figure 3. A dedicated script was developed to parse and load the exported JSON files from the labeling tool. This script was instrumental in overlaying the masks and bounding boxes on the images, ensuring the accuracy and correctness of the labeling process.

2.4. MASK R-CNN

Mask R-CNN is a multi-stage instance segmentation technique that detects objects in an image and generates a mask to segment each object instance. It extends the Faster R-CNN algorithm with an additional branch to output a binary mask for each Region of Interest (ROI) in parallel with the object class and bounding box prediction as shown in Figure 4. ROIs are boxes drawn to tell us where an object might be present in the image. They are extracted through the Region Proposal Network (RPN) which was initially introduced in the Faster R-CNN framework as a neural network with attention. It takes in feature maps of raw images extracted from a Feature Pyramid Network (FPN) with the ResNet backbone while leveraging anchor boxes to allow the ROI to be detected at different scales and aspect ratios. The final stage involves using the Fully Connected (FC) layers to predict the bounding box coordinates of the object with the category of the object class, and a mask branch which uses an FCN leveraged from semantic segmentation, to extract and maintain the spatial layout of the mask unlike the FC layers from Faster R-CNN that collapses the spatial structure into a vector representation. With the FCN, the mask is generated using fewer parameters and more accurately.

When Mask R-CNN was introduced, it was developed to solve some of the drawbacks of the Faster R-CNN - the ROIPool layer which introduced misalignments between the ROI and the extracted features that affected the accuracy of the predicted pixel masks. These misalignments were a result of the quantization of the ROI boundaries/bins while extracting the features, such that a continuous coordinate x becomes [x/16] due to quantization. To resolve this, the ROIAlign layer was introduced to preserve the spatial dimensions for pixel-to-pixel alignment of the features. As the Mask R-CNN has 3 branches, it has a multi-task loss function (2), for each of the branches - the classifier loss, bounding box loss, and the mask loss.

$$L=L_{cls}+L_{box}+L_{mask}$$

(2)

For every ROI with ground truth, k, the mask loss is defined on only the $K_{th}$ mask (loss is independent of other mask outputs). Masks are generated for every object without competition using the class predicted to select the mask rather than how prior systems were built where the classification depends on the mask predictions, or using a per-pixel sigmoid and binary loss rather than a per-pixel softmax and multinomial cross-entropy loss from previous systems.

2.5. Detectron2

Detectron2 framework was built and leveraged from the Mask R-CNN architecture. It was developed by Facebook AI Research for its application in object detection and instance segmentation tasks. With similarity to the Mask R-CNN, it has support for more feature extraction backbones from the FPN, ResNet, and ResNeXt than Mask R-CNN; in addition to its modular and extensible design. This makes Detectron2 invaluable for this study as it can be utilized in extracting essential morphological features of individual rock particles, thereby producing instance segmentation masks. The utilization of the Detectron2 framework permitted us the opportunity to ensure that our system was robust since we could evaluate and compare the different algorithm’s performance.

3. Implementation Details

Before the model training began, the dataset was split into training, validation, and test sets to make it easy to evaluate the mask r-cnn and detectron2’s performance not just on the data it has learned from (training and validation data), but on data yet to be seen (test data). This approach is a standard method recognized as it facilitates model generalization on the patterns present in the dataset, preventing potential cases of false predictions.

3.1. Evaluation Metrics

To evaluate the models, the Average Precision (AP) standard COCO evaluation metric [33] was used at different thresholds, where this threshold is determined by the Intersection over Union (IoU) value - how close the ground truth mask/bounding box values match with the predicted masks/bounding boxes generated by the model. In this study, the AP is evaluated at the default IoU threshold of 0.5 - AP50, 0.75 IoU threshold - AP75, and across all 10 IoU thresholds (0.5 to 0.95 with increments of 0.05) - AP. The closer the value is to 1 (100%), the more accurate the segmentation.

3.2. Transfer Learning

In order to adapt to the limited data size that the drilling company provided us with, the transfer learning technique was leveraged in the model training of both the Mask R-CNN and Detectron2 models. This technique prevents training the models from scratch which would be a daunting task especially when the dataset size is low, as it can lead to poor performance since the network expects tons of images to learn and generalize from. Instead, the knowledge garnered from feature representations of models that have been trained on large volumes of data for solving different problems was leveraged. By doing this, the knowledge can be transferred to the specific use case with a limited data size, while getting good results. In this paper, the pre-trained mask r-cnn model built to recognize 80 different classes from the COCO dataset such as cars, dogs, traffic lights, etc. [33,39], is leveraged in recognizing the rock particles.

3.3. Training the Mask R-CNN Model

Several hyperparameters were experimented with by altering the backbones, learning rate, detection confidence threshold, and more, to finalize at this model with the following features.

ResNet101 performed better than ResNet50 [33] so it was selected as the backbone network, with a learning rate of 0.003 after experimenting with values ranging from the default - 0.001 to 0.00025. The confidence threshold for an ROI detected to be accepted was increased from the default - 0.7 to 0.9. Although this helped reduce false positive ROI instances, it increased the false negative occurrence rate. Then, an early stopping strategy was also added by implementing a terminating callback that will stop the training when validation losses plateaued, thereby preventing overfitting and ensuring model generalization. Finally, to actualize transfer learning for convergence and optimal training, a four-phase training approach was explored:

• Phase 1: Training the network head for 40 epochs with an increased learning rate at 0.006 (learning rate * 2)
• Phase 2: Fine-tuning layers 3+ (ResNet stage 3 and up) for 120 epochs
• Phase 3: Fine-tuning layers 4+ (ResNet stage 4 and up) for 160 epochs
• Phase 4: Fine-tuning all layers for 200 epochs with a reduced learning rate at 0.0003 (learning rate / 10)

3.4. Training the Detectron2 Model

Here, five different baseline models were experimented with, all having the ResNet+FPN backbone, and 3-times schedule: ResNet50, ResNet101, ResNeXt101-32x8d trained with caffe2, ResNet101 using Large-Scale Jitter and Longer Training Schedule [40], and Cascade R-CNN - ResNet50 with Cascade ROI heads. All models utilized a learning rate of 0.0025, batch size of 512, and maximum iterations of 10,000. They were all evaluated against the test data using the AP metric.

4. Results

An Nvidia Quadro RTX 8000 GPU Linux server with 48GB memory was used as the training environment for the developed models. The GPU provided us with an environment to experiment with different hyperparameters more effectively and efficiently. Table 1 reports the metric evaluation of the trained Mask R-CNN model with the ResNet101 backbone against the Detectron2 models with 5 different backbones, on the test dataset.

Based on the results across all metrics especially the AP since it accounts for all IoU thresholds, the Mask R-CNN model had the best performance overall while the Detectron2 with the ResNet50 backbone performed the best across the Detectron2 models. The Cascade R-CNN model and the Resnet 101 with Large Scale Jitter, had the worst performances across all the metrics. Looking at the predictions visually in Figure 5, it can be seen that although the model was able to do a decent job recognizing the individual rock particles regardless of the low AP values, there are few cases of over-segmentation and under segmentation which unsurprisingly, was seen less in the mask r-cnn predictions.

It was also observed that the models experienced more difficulty in recognizing the particles for images that had several particles in comparison to the ones with few particles, images with tiny rocks in comparison to the ones with huge rock slabs, and images that weren’t very clear. Owing to the complexities and variations of the rock particle sizes, and the laborious nature of annotating data, it was recognized that the quality and accuracy of the ground truth labels used to train the models might diminish since the mask annotations were labeled by a human manually drawing the polygons on each rock particle. This poses a huge problem as it can affect the performance of any model trained with the data. Furthermore, the model performance could be further improved by expanding the dataset via data augmentation [41,42,43]. By transforming the limited data, the model is able to generalize and adapt better to new rock images, making it more robust to imaging scenarios from different geological contexts, and less prone to overfitting [44].

5. Conclusions

In this paper, a framework for recognizing rock particles from digital images was presented using the pre-trained Mask R-CNN and Detectron2 deep learning algorithms. These algorithms were trained and validated on an annotated preprocessed dataset using the transfer learning technique, providing a more efficient and optimal training approach. Thereafter, the segmented particles were evaluated using the Average Precision standard coco metric across several IoU thresholds. The Mask R-CNN model with the ResNet101 backbone performed the best with an AP value of 55.5 for the 0.5 IoU threshold, on the test dataset. This shows its potential for use in segmentation tasks and applications across several engineering domains. In light of the results obtained, expansion of the dataset either by means of a re-collection from the drilling company or by employing data augmentation methods such as image rotation, translation, flipping, cropping, brightness control, saturation control, etc., can provide more features for the algorithm to learn from, thereby improving its performance. This can be applied in future works alongside experimenting with other training hyperparameter values that can produce a more robust recognition capability.