Lightweight deep learning model-based lithium ore related lithology identification method and system

Abstract

The invention belongs to the technical field of rock image classification of earth science, and discloses a lithium ore related lithology identification method and system based on a lightweight deep learning model. Acquiring lithology data related to pegmatite type lithium ore, and constructing a deep learning data set; coupling a Fmix data enhancement method based on image mixing in the MobileViT model, introducing transfer learning, and constructing a TL-Fmix-MobileViT lithology identification model; performing lithium ore related lithology identification based on a TL-Fmix-MobileViT lithology identification model; and identification result analysis and TL-Fmix-MobileViT lithology identification model performance evaluation are completed. According to the method, information can be captured in a boundary region, features from different images can be fused, the mixed rock features can be better classified, the classification performance is good, the precision reaches 99% or above, and the reasoning speed is high.

Description

Technical Field

[0001] The present invention belongs to the technical field of rock image classification in earth science, and in particular to a lithium ore-related lithology identification method and system based on a lightweight deep learning model. Background Art

[0002] Lithology identification is a fundamental and important issue in the field of geological resource exploration. Accurate and efficient lithology identification has great practical value for mineral resource exploration. In field exploration, geologists will judge the type of rock based on its composition, structure, and color. However, due to inconsistent subjective understanding, there will be certain deviations in the results of lithology judgment. With the development of computer vision technology, existing research on intelligent rock identification has made certain progress, and various rock identification methods and models have been developed.

[0003] Although scholars have proposed many methods for automatic lithology identification, some of them are aimed at the study of rock thin section images, which is inconsistent with the application scenario of automatic rock identification in the field. In addition, deep learning methods consume a lot of computing resources and are difficult to apply in field exploration. In addition, the current lithology classification based on machine learning methods has the disadvantages of insufficient extraction of detailed features and limited classification accuracy. Summary of the invention

[0004] In order to overcome the problems existing in the related art, the embodiments disclosed in the present invention provide a method and system for identifying lithium ore-related lithology based on a lightweight deep learning model, and specifically relate to a method for identifying lithium ore-related lithology based on a TL-Fmix-MobileViT model.

[0005] The technical solution is as follows: A lithium ore-related lithology identification method based on a lightweight deep learning model comprises the following steps:

[0006] S1, obtain image data of rocks related to lithium ore, use image annotation, and build a deep learning dataset;

[0007] S2, coupled the Fmix image mixing data enhancement method in the MobileViT model to improve generalization, introduced transfer learning to accelerate training and improve accuracy, and constructed the TL-Fmix-MobileViT lithology recognition model;

[0008] S3, based on the TL-Fmix-MobileViT lithology recognition model, feature extraction and classification of lithium-related rock images;

[0009] S4, analysis of identification results and performance evaluation of the TL-Fmix-MobileViT lithology identification model.

[0010] In step S1, a deep learning dataset is constructed, including:

[0011] The collected rock samples are rotated, translated, flipped, contrast adjusted, brightness adjusted, chromaticity adjusted, and sharpness adjusted to obtain multiple rock image samples; they are divided into multiple categories according to the collected rock categories, and the multiple rock image samples are divided into training sets and verification sets according to an 8:2 ratio to construct the lithium ore-related lithology classification dataset required for deep learning training.

[0012] In step S2, the Fmix image mixing data enhancement method is coupled to the MobileViT model, including:

[0013] Mix the two images according to the mask. The area with mask 1 comes from the first image, and the area with mask 0 comes from the second image.

[0014] The two input images are x1 and x2, and the corresponding labels are y1 and y2; the mask is M; the calculation formula for the mixed image x and the mixed label y is:

[0015]

[0016] Where λ is the proportion of the first image in the mask M.

[0017] In step S2, transfer learning is introduced, including:

[0018] The original MobileViT network is fully trained on the ImageNet dataset, and the weights are optimized through the stochastic gradient descent SGD optimization algorithm to learn the general feature representation of the image; after the training is completed, the weight parameters of the entire network are saved and a .pth file containing all network layer parameters is generated; after loading the pre-trained weight file, it is applied to the MobileViT model coupled with the Fmix data augmentation module; specifically, it includes:

[0019] Freeze the first few layers of the Fmix-MobileViT model, that is, keep them unchanged to retain the common feature representations learned by the pre-trained model, while only fine-tuning the subsequent layers. This strategy can effectively reduce training time and better adapt to the data distribution of the target task through fine-tuning; through this process, the model is able to achieve higher classification accuracy on the target task.

[0020] In step S2, the TL-Fmix-MobileViT lithology identification model includes the MV2 module and the MobileViT module;

[0021] The MV2 module is an inverted residual block based on MobileNetV2. It first expands the channels of the input image through 1×1 convolution, and applies batch normalization and ReLU6 activation function to increase the feature dimension. Then, 3×3 convolution is performed independently on each channel, and then the expanded number of channels is compressed back to the number of original input channels through 1×1 convolution. When the number of channels of input and output is the same, the input is directly added to the output using residual connection.

[0022] The MobileViT module adopts a structure that combines CNN and Transformer. For the input data of size H×W×C, a 3×3 convolution is first performed to extract local features. Before entering the Transformer module, the feature map is adjusted through a 1×1 convolution layer to adjust the number of channels, and then the feature map is mapped to a d-dimensional space. After two convolution operations, the input feature map X is transformed into a local feature map X of equal size. L ; The feature map X L Split into N image blocks of size H×W, and then expand into a set of feature sequences X of size P×N×d U , where P = W × H, N = (H × W) / P, where W and H are the preset width and height of the image block, and d is the feature dimension; after the expansion, the feature sequence X U The Transformer encoder is applied on the image block. The Transformer encoder uses the self-attention mechanism to capture the relationship between different feature blocks and realize the modeling of the global context. In the feature fusion part, the encoded image block is restored to the dimension C before encoding through a 1×1 convolution layer to obtain the feature X F , X F The size of is H×W×C, which is superimposed in parallel with the input feature map X to form a new feature map with a dimension of 2C. Finally, a 3×3 convolution kernel is used to fuse the new feature maps after parallel connection, and the new feature map dimension is mapped back to C to obtain the final output feature Y.

[0023] In step S3, feature extraction and classification of lithium ore-related rock images are performed based on the TL-Fmix-MobileViT lithology recognition model, including:

[0024] The first step is to preprocess the data by normalizing, resizing, and performing random horizontal flipping on the image through torchvision.transforms; load the deep learning dataset and use torch.utils.data.DataLoader to create iterators for the training set and validation set;

[0025] The second step is to build the model. Use timm to create the MobileViT-S model and load the pre-trained weights. Use Adam, AdamW and SGD optimizers and three learning rates of 0.002, 0.001 and 0.0002 for ablation experiments. After verification by ablation experiments, the number of iterations is 100, the number of training and verification batches is 16, the optimizer uses SGD, the learning rate is 0.002, the weight decay coefficient is 0.0001, and the momentum parameter is 0.9 to accelerate convergence. The loss function uses cross entropy loss to determine the difference between the actual output and the sample label.

[0026] The third step is to initialize Fmix with data to generate mixed masks for data enhancement; use the custom loss function fmix.loss() provided by Fmix to process mixed samples during training; fmix.loss() is a custom loss function designed specifically for mixed sample training, which uses the weight coefficient (lam) of the mixed mask to perform a weighted summation of the losses of the two sets of target labels. This function supports multiple task requirements, such as cross entropy loss and binary cross entropy loss, and provides flexible training mode selection.

[0027] The fourth step is to train and save the TL-Fmix-MobileViT lithology recognition model. Use the trial interface of torchbearer to encapsulate the model, optimizer, loss function and data generator. Set the training to run on the GPU and monitor the accuracy and loss. After the training is completed, save the trained TL-Fmix-MobileViT lithology recognition model weights to the specified path.

[0028] In step S4, the recognition results are analyzed and the performance of the TL-Fmix-MobileViT lithology recognition model is evaluated, including: using accuracy, precision, recall, F1 value, model inference speed Frames PerSecond and model size as evaluation indicators of the TL-Fmix-MobileViT lithology recognition model; the expression is:

[0029]

[0030] Where TP indicates that the model correctly predicts a positive example as a positive example; FP indicates that the model incorrectly predicts a negative example as a positive example; TN indicates that the model correctly predicts a negative example as a negative example; FN indicates that the model incorrectly predicts a positive example as a negative example;

[0031]

[0032] Where N is the number of frames of samples to be processed.

[0033] Another object of the present invention is to provide a lithium ore-related lithology identification system based on a lightweight deep learning model, the system implements the lithium ore-related lithology identification method based on a lightweight deep learning model, and the system comprises:

[0034] A deep learning dataset construction module is used to obtain image data of rocks related to lithium mines, use image annotation, and construct a deep learning dataset;

[0035] TL-Fmix-MobileViT lithology recognition model building module, which is used to couple the Fmix image mixing data enhancement method in the MobileViT model to improve generalization, introduce transfer learning to accelerate training and improve accuracy, and build the TL-Fmix-MobileViT lithology recognition model;

[0036] The lithology recognition module is used to extract features and classify and identify lithium-related rock images based on the TL-Fmix-MobileViT lithology recognition model;

[0037] Performance evaluation module, used for identification result analysis and TL-Fmix-MobileViT lithology identification model performance evaluation.

[0038] Furthermore, the lithium ore-related lithology identification system based on a lightweight deep learning model is carried on a computer-readable storage medium, and the computer-readable storage medium stores a computer program. When the computer program is executed by a processor, it can realize the functions of the above-mentioned lithium ore-related lithology identification system based on a lightweight deep learning model.

[0039] Furthermore, the lithium ore-related lithology identification system based on the lightweight deep learning model is installed on an information data processing terminal, and the information data processing terminal is used to provide a user input interface to implement the functions of the lithium ore-related lithology identification system based on the lightweight deep learning model when executed on an electronic device.

[0040] In combination with all the above technical solutions, the beneficial effects of the present invention are as follows:

[0041] Based on the lightweight network MobileViT, the present invention introduces the FMIX data enhancement method and combines the idea of ​​deep transfer learning to propose a new rock image recognition method. The present invention aims to use the improved model to identify and classify the main lithologies related to lithium ore, and compares it with the traditional model to verify its superiority in the field lithology recognition task. Taking a certain area and a certain lithium ore lithology as sample data, the present invention proposes a lightweight deep learning model based on TL-Fmix-MobileViT for efficient identification of the lithology of six types of lithium ore-related rocks. The model combines the inverted residual block of MobileNetv2 and Vision Transformer, which can simultaneously extract local and global features of rock images. By introducing transfer learning and Fmix data enhancement modules, the model can learn richer features. The TL-Fmix-MobileViT lithology recognition model can capture and fuse features from different images in the boundary area, can better classify mixed rock features, has good classification performance, an accuracy of more than 99% and a fast inference speed, which is better than the traditional deep learning model. It demonstrates its superiority in complex rock classification tasks and provides a reference for the classification research of complex rock images. The lightweight network model proposed in this paper provides an efficient and accurate method for lithology identification related to lithium ore.

[0042] The lithium ore-related lithology identification method based on a lightweight deep learning model proposed in the present invention has significant commercial application value. Through accurate and efficient lithology identification technology, the method can improve the accuracy and efficiency of lithium ore exploration, reduce human intervention and misidentification, and save time and cost. The present invention solves the key technical problems in lithology identification, especially in the classification of complex rock images and the extraction of mixed rock features. It also significantly accelerates the reasoning and calculation speed, meeting the industry's demand for efficient and accurate identification. Through Fmix and transfer learning technology, the present invention overcomes the limitations of traditional methods in the classification of complex rock features. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] The accompanying drawings herein are incorporated in and constitute a part of the specification, illustrate embodiments consistent with the present disclosure, and together with the description, serve to explain the principles of the present disclosure;

[0044] Figure 1 It is a flow chart of a lithium ore-related lithology identification method based on a lightweight deep learning model provided by an embodiment of the present invention;

[0045] Figure 2 It is a flow chart of the Fmix method and transfer learning provided by an embodiment of the present invention;

[0046] Figure 3It is a TL-Fmix-MobileViT network structure diagram provided by an embodiment of the present invention;

[0047] Figure 4 Schematic diagram of a lithium ore-related lithology identification system based on a lightweight deep learning model provided in an embodiment of the present invention;

[0048] Figure 5 is a graph showing the accuracy variation of different models provided by an embodiment of the present invention;

[0049] Figure 6 It is a heat map of 8 deep learning models generated by using Grad-CAM provided by an embodiment of the present invention;

[0050] In the figure: 1. Deep learning dataset construction module; 2. TL-Fmix-MobileViT lithology identification model construction module; 3. Lithology identification module; 4. Performance evaluation module. DETAILED DESCRIPTION

[0051] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the specific embodiments of the present invention are described in detail below in conjunction with the accompanying drawings. In the following description, many specific details are set forth to facilitate a full understanding of the present invention. However, the present invention can be implemented in many other ways different from those described herein, and those skilled in the art can make similar improvements without violating the connotation of the present invention, so the present invention is not limited by the specific implementation disclosed below.

[0052] The innovation of the present invention is that: the present invention collects the relevant lithology data of pegmatite-type lithium ore and constructs a deep learning data set; introduces the Fmix data enhancement method in the MobileViT model, and constructs the TL-Fmix-MobileViT lithology recognition model in combination with transfer learning; uses the model to identify the relevant lithology of lithium ore, and analyzes the recognition results and evaluates the model performance. The present invention can capture the mixed characteristics of the boundary area, efficiently identify complex rock features, with a classification accuracy of more than 99%, fast inference speed, and is suitable for lithology recognition tasks in complex geological environments.

[0053] Embodiment 1, as Figure 1 As shown, the lithium ore-related lithology identification method based on a lightweight deep learning model provided by an embodiment of the present invention includes:

[0054] S1, obtain image data of rocks related to lithium ore, use image annotation, and build a deep learning dataset;

[0055] Pegmatite samples and their surrounding rocks were collected in three areas of three lithium mines in a certain region and photographed as images. In order to maintain the diversity of the data set, the European pegmatite database developed by the GREENPEG project was introduced, and the pegmatite samples were selected as a supplement. The collected rock images were uniformly adjusted to a size of 224×224, and a series of data enhancement processes such as rotation, translation, flipping, contrast adjustment, brightness adjustment, chromaticity adjustment, and sharpness adjustment were performed on the rock samples, and finally 3138 rock image samples were obtained. According to the collected rock categories, they are divided into 6 categories, and the specific classification is shown in Table 1. It is divided into a training set and a validation set according to 8:2 to construct a lithium ore-related lithology classification data set required for deep learning training.

[0056] Table 1 Lithium ore related lithology classification dataset

[0057] Rock Type serial number Sample size Training set Validation set Dahongliutan spodumene-bearing pegmatite Rock1 512 409 103 Dahongliutan tourmaline-bearing pegmatite Rock2 528 422 106 Dahongliutan Granite Rock3 540 432 108 Dahongliutan Muscovite Pegmatite Rock4 540 432 108 Dahongliutan biotite quartz schist Rock5 560 448 112 Lepidolite-bearing pegmatites in Spain and Portugal Rock6 458 366 92

[0058] S2, coupled the Fmix image mixing data enhancement method in the MobileViT model to improve generalization, introduced transfer learning to accelerate training and improve accuracy, and constructed the TL-Fmix-MobileViT lithology recognition model; specifically, it includes:

[0059] S21, Fmix method and transfer learning technology.

[0060] Fmix is ​​a data augmentation method based on image mixing. The Fmix method mixes two images according to the mask. The area with mask 1 comes from the first image, and the area with mask 0 comes from the second image. The mixed image retains some features of the two original images and also contains new structures and information. The label of the new image is the weighted average of the labels of the two original images, and the weight is determined by the ratio of 1 to 0 in the mask. This method of mixing labels ensures that the model learns the features of two categories at the same time during training. The implementation process is as follows: Figure 2 . Let the two input images be x1 and x2, and the corresponding labels be y1 and y2. The mask is M (Mask). The calculation formula for the mixed image x and the mixed label y is:

[0061]

[0062] Where λ is the proportion of the first image in the mask M.

[0063] The ideal machine learning scenario is to have a large number of labeled training samples that have the same distribution as the test data. However, in many cases, collecting enough training data is time-consuming and labor-intensive, or even unrealistic. Transfer learning can achieve cross-domain knowledge transfer and is an effective way to solve the above problems. The convolutional neural network pre-trained on ImageNet can achieve good performance after fine-tuning on the new task. When training the TL-Fmix-MobileViT lithology recognition model, the initial Fmix-MobileViT network was pre-trained on the large-scale image dataset ImageNet and saved as a weight model, so that the model can learn rich feature representations that can be transferred to the rock classification task. Subsequently, the first few layers of the network (feature extraction layers) are frozen when the weight model is loaded, because these layers have learned common features. Using the method of loading the pre-trained weight model for training on the rock classification dataset can make the TL-Fmix-MobileViT lithology recognition model more generalizable.

[0064] S22, MobileViT model and model improvement.

[0065] MobileViT is a model based on a hybrid architecture of convolutional neural networks and Transformer. Compared with the shortcomings of the ViT model, which has many parameters, high computational requirements, and difficult training, this new network architecture significantly reduces the number of model parameters and computational costs, while maintaining high accuracy and adapting to the computational constraints of mobile devices. To meet the needs of different application scenarios, MobileViT provides three versions with different weights: MobileViT-XXS, MobileViT-XS, and MobileViT-S. You can flexibly select the most suitable model according to the hardware environment and application requirements. Its structure can be seen in Figure 3 , TL-Fmix-MobileViT network structure diagram, mainly including convolutional layer MV2 module and MobileViT module.

[0066] Exemplarily, the MV2 module includes: The MV2 module is an inverted residual block based on MobileNetV2, and its structure is as follows Figure 3As shown in the figure, the channels of the input image are first expanded through 1×1 convolution, and batch normalization and ReLU6 activation function are applied to increase the feature dimension. The subsequent depth-separable convolution can extract features in a higher dimension. Then, 3×3 convolution is performed independently on each channel, which can greatly reduce the amount of parameter calculation. Then, the expanded number of channels is compressed back to the number of original input channels through 1×1 convolution, reducing the output dimension of the model. When the number of input and output channels is the same, the module will use residual connection to add the input directly to the output. This inverted residual structure can alleviate the gradient vanishing problem and help the training of deep networks.

[0067] The MobileViT module includes: The MobileViT module is the core part of the MobileViT model, which adopts the structure of CNN and Transformer. For the input data of size H×W×C, a 3×3 convolution is first performed to extract local features. Before entering the Transformer module, the feature map will pass through a 1×1 convolution layer to adjust the number of channels, and then map the feature map to d-dimensional space (d>C). After two convolution operations, the input feature map X is transformed into a local feature map X of equal size. L . The feature map X L Split into N image blocks of size H×W, and then expand them into a set of feature sequences X of size P×N×d U , where P = W × H, N = (H × W) / P (W and H are the preset width and height of the image block, and d is the feature dimension). U The Transformer encoder is applied on the image block. The Transformer uses the self-attention mechanism to capture the relationship between different feature blocks, thereby realizing the modeling of the global context. In the feature fusion part, in order to avoid the loss of information between image blocks, it is also necessary to restore the encoded image block to the dimension C before encoding through a 1×1 convolution layer to obtain the feature X F , at this time X F The size of is H×W×C, which can be superimposed in parallel with the input feature map X to form a new feature map with a dimension of 2C. Finally, a 3×3 convolution kernel is used to fuse the new feature maps after parallel connection, and the new feature map dimension is mapped back to C to obtain the final output feature Y.

[0068] Exemplarily, based on the MobileViT model, the present invention introduces the Fmix image mixing-based data enhancement method (data enhancement module) and transfer learning technology, improves the MobileViT model, and constructs a rock classification model based on TL-Fmix-MobileViT. Fmix randomly mixes two different categories of rock images, and the model can find key features related to the category in complex backgrounds and foregrounds. When the image background is complex or the image is not clear, this helps the model to more accurately identify the rock category in the real environment. Transfer learning significantly reduces the dependence on a large amount of labeled data in the rock classification task. Higher accuracy can be achieved through fine-tuning with a small amount of labeled data, while also reducing the demand for computing resources.

[0069] S3, based on the TL-Fmix-MobileViT lithology recognition model, feature extraction and classification of lithium-related rock images;

[0070] The code was built using Python language to implement a PyTorch-based image classification task training process, using the Fmix image mixing-based data enhancement method and the MobileViT model.

[0071] Specifically, the first step is to preprocess the data, normalize, resize, and perform random horizontal flipping on the image through torchvision.transforms. Load a custom dataset (containing 6 types of rock images), and use torch.utils.data.DataLoader to create iterators for the training set and validation set.

[0072] The second step is to build the model, use timm to create the MobileViT-S model, and load the pre-trained weights. The Adam, AdamW and SGD optimizers and three learning rates of 0.002, 0.001, and 0.0002 were used for ablation experiments. After verification by the ablation experiment, the present invention sets the number of iterations to 100, the number of batches for training and verification to 16, the optimizer uses SGD, the learning rate is 0.002, the weight decay coefficient is 0.0001, and the momentum parameter is 0.9, which is used to accelerate convergence. The loss function uses cross entropy loss to determine the difference between the actual output and the sample label. When loading data, 8 subprocesses are used to speed up data loading.

[0073] In the third step, data is initialized with Fmix to generate mixed masks for data augmentation. The custom loss function fmix.loss() provided by Fmix is ​​used to enable the model to handle mixed samples during training. The last step is to train and save the model, using the torchbearer Trial interface to encapsulate the model, optimizer, loss function, and data generator. Set the training to run on the GPU, monitor the accuracy and loss, and train for 100 epochs. After training is complete, save the trained model weights to the specified path.

[0074] S4, analysis of identification results and performance evaluation of the TL-Fmix-MobileViT lithology identification model.

[0075] Accuracy, precision, recall and F1 score are used as model evaluation indicators. In addition, model reasoning speed Frames Per Second and model size are also important indicators for evaluating the complexity of rock classification models.

[0076]

[0077]

[0078]

[0079]

[0080] In the formula, TP (True Positive) indicates that the model correctly predicts a positive example as a positive example; FP (False Positive) indicates that the model incorrectly predicts a negative example as a positive example; TN (True Negative) indicates that the model correctly predicts a negative example as a negative example; FN (False Negative) indicates that the model incorrectly predicts a positive example as a negative example.

[0081]

[0082] Where N refers to the number of frames of samples to be processed.

[0083] Embodiment 2, as Figure 4 As shown, an embodiment of the present invention provides a lithium ore-related lithology identification system based on a lightweight deep learning model, including:

[0084] Deep learning data set construction module 1 is used to obtain image data of rocks related to lithium mines, use image annotation, and construct a deep learning data set;

[0085] TL-Fmix-MobileViT lithology recognition model construction module 2 is used to couple the Fmix image mixing data enhancement method in the MobileViT model to improve generalization, introduce transfer learning to accelerate training and improve accuracy, and build the TL-Fmix-MobileViT lithology recognition model;

[0086] Lithology recognition module 3 is used to extract features and classify and recognize lithium-related rock images based on the TL-Fmix-MobileViT lithology recognition model;

[0087] Performance evaluation module 4 is used for analysis of identification results and performance evaluation of the TL-Fmix-MobileViT lithology identification model.

[0088] In order to verify the performance advantages of the lightweight lithium-containing mineral-related lithology identification method, the present invention trained 7 comparative deep learning models. Including classic convolutional neural networks VGG-16, GoogLe-NeT, ResNet-50, as well as lightweight networks MobileNeT, as well as Vision Transformer and Swin Transformer networks based on Transformer architecture and the original MobileViT model. All deep learning models used in the comparative experiments were pre-trained on the ImageNet dataset and then trained using fine-tuning, so they can all achieve high classification accuracy. The model proposed in the present invention couples the FMIX data enhancement method to the head of the MobileViT model, so the pre-trained model of MobileViT can also be used. During the training process, the data reading method of the pre-trained model head can be modified. The present invention selects the model corresponding to the round with the highest accuracy in the experimental process as the final model. The accuracy change curves of all models in the comparative experiment are shown in FIG. Figure 5 Accuracy change curves of different models.

[0089] from Figure 5It can be seen that with the increase in the number of training iterations, the model accuracy continues to improve. Among them, the model coupled with the Fmix method converges faster and has a significantly higher classification accuracy than other models. This shows that the model performs well on the rock classification dataset, has strong learning ability and fast convergence speed. Models such as Vision Transformer, SwinTransformer, MobileViT and Resnet50 also perform well, tend to be stable after about the 40th epoch, and maintain a high accuracy. The traditional convolutional neural network GoogleNet, VGG16 and some lightweight networks MobileNeT perform relatively poorly, and did not converge after 100 epochs, with large fluctuations and low accuracy. These results show that the hybrid model combining transformer and CNN may have better generalization ability and accuracy in image classification tasks, and the Fmix method is helpful to improve the classification accuracy.

[0090] The accuracy, precision, recall and F1 value can be calculated through the confusion matrix. In addition, we also calculated the weight and inference speed of the model. The calculation results are shown in Table 2. From Table 2, it can be concluded that the FMIX-MobileViT model has the highest accuracy, and the model has a faster inference speed and a smaller weight, which is suitable for lightweight applications; the MobileNeT model has the fastest inference speed, but its classification accuracy is the lowest, only 95.4%; the weight of the VGG16 model reaches 512MB, which requires a lot of memory; although the classification accuracy of the Vision Transformer model is second only to the FMIX-MobileViT model, its inference speed is only 2.92FPS, and the weight model is large, which requires high computing resources.

[0091] As can be seen from Table 2, the FMIX-MobileViT model proposed in this study has good accuracy, and its accuracy is 0.6 percentage points higher than that of the initial MobileViT model, reaching 99.8%. The precision, recall and F1 values ​​are 99.9%, 99.8% and 99.8% respectively, which shows that the model not only performs well overall, but also performs particularly well when processing positive samples. The model can accurately predict positive samples and identify as many positive samples as possible, while the proportion of wrong predictions is also very low. The weight of the model is 19.0MB, and the inference speed is 12.16FPS. Without increasing the weight, the inference speed is only 0.02FPS lower than the original MobileViT model, and the classification accuracy has been greatly improved.

[0092] The TL-Fmix-MobileViT lithology recognition model combines CNN and Transformer. It not only has the advantage of being lightweight, but also can capture the dependencies between distant pixels on rock images, thereby understanding the overall structure and pattern of the image, and surpasses the other seven classic models in overall performance indicators.

[0093] Table 2 Classification performance of different models

[0094]

[0095]

[0096] Grad-CAM model visualization enables the present invention to intuitively see why the model makes certain predictions, thereby improving the interpretability of the model. Figure 6 The heatmaps of eight deep learning models generated using Grad-CAM are shown, illustrating the image regions that the eight different models focus on when making classification decisions for six different rock types. Figure 6 It can be concluded that the MobileViT model has a good balance in capturing global and local features, and can better identify the main features of rocks; the MobileNeT model's attention area is relatively scattered and there is insufficient local coverage, and its ability to extract complex features may be slightly insufficient; the ResNet50 model's attention area can cover the key feature parts of the rock, and its performance is relatively stable; the GoogLeNet model's attention area is concentrated on the main part of the rock, with a wider coverage; the VGG16 model has a problem of excessive concentration on local areas, which leads to the neglect of global information; the Vision Transformer model has a strong focus on both the global and the details, and can capture most of the important features of the rock; the Swin Transformer model's attention area has comprehensive coverage, and can effectively maintain the capture of global information.

[0097] Compared with the focus areas of other models, the FMIX-MobileViT model shows obvious differences. The images generated using the FMIX data augmentation technology are mixed by two rock images according to random masks. This mixture usually causes the final image to contain multiple regional information from different images, so the model does not only focus on the central area during training. The thermal area of ​​the FMIX-MobileViT model is concentrated on the boundary of the FMIX enhanced mask image, so that the model can capture and fuse features from different images in the boundary area. This learning mechanism helps improve the performance of the model in complex tasks, enabling it to better understand and classify mixed features, thereby enhancing the generalization and robustness of the model.

[0098] The above description is only a preferred specific implementation manner of the present invention, but the protection scope of the present invention is not limited thereto. Any modifications, equivalent substitutions and improvements made by any technician familiar with the technical field within the technical scope disclosed by the present invention and within the spirit and principles of the present invention should be covered within the protection scope of the present invention.

Claims

1. A lithium ore-related lithology identification method based on a lightweight deep learning model, characterized in that: The method comprises the following steps: S1, obtain image data of rocks related to lithium ore, use image annotation, and build a deep learning dataset; S2, coupled the Fmix image mixing data enhancement method in the MobileViT model to improve generalization, introduced transfer learning to accelerate training and improve accuracy, and constructed the TL-Fmix-MobileViT lithology recognition model; S3, based on the TL-Fmix-MobileViT lithology recognition model, feature extraction and classification of lithium-related rock images; S4, analysis of identification results and performance evaluation of the TL-Fmix-MobileViT lithology identification model.

2. The lithium ore-related lithology identification method based on a lightweight deep learning model according to claim 1 is characterized in that: In step S1, a deep learning dataset is constructed, including: The collected rock samples are rotated, translated, flipped, contrast adjusted, brightness adjusted, chromaticity adjusted, and sharpness adjusted to obtain multiple rock image samples; they are divided into multiple categories according to the collected rock categories, and the multiple rock image samples are divided into training sets and verification sets according to an 8:2 ratio to construct the lithium ore-related lithology classification dataset required for deep learning training.

3. The lithium ore-related lithology identification method based on a lightweight deep learning model according to claim 1 is characterized in that: In step S2, the Fmix image mixing data enhancement method is coupled to the MobileViT model, including: Mix the two images according to the mask. The area with mask 1 comes from the first image, and the area with mask 0 comes from the second image. The two input images are x1 and x2, and the corresponding labels are y1 and y2; the mask is M; the calculation formula for the mixed image x and the mixed label y is: Where λ is the proportion of the first image in the mask M.

4. The lithium ore-related lithology identification method based on a lightweight deep learning model according to claim 1 is characterized in that: In step S2, transfer learning is introduced, including: The original MobileViT network is fully trained on the ImageNet dataset, and the weights are optimized through the stochastic gradient descent (SGD) optimization algorithm to learn the general feature representation of the image. After the training is completed, the weight parameters of the entire network are saved to generate a .pth file containing all network layer parameters. After loading the pre-trained weight file, it is applied to the MobileViT model coupled with the Fmix data augmentation module.

5. The lithium ore-related lithology identification method based on a lightweight deep learning model according to claim 1 is characterized in that: In step S2, the TL-Fmix-MobileViT lithology identification model includes the MV2 module and the MobileViT module; The MV2 module is an inverted residual block based on MobileNetV2. It expands the channels of the input image through 1×1 convolution, and applies batch normalization and ReLU6 activation function to increase the feature dimension. It performs 3×3 convolution independently on each channel, and then compresses the expanded number of channels back to the number of original input channels through 1×1 convolution. When the number of channels of input and output is the same, the input is directly added to the output using residual connection. The MobileViT module adopts a structure that combines CNN and Transformer. A 3×3 convolution is performed on the input data of size H×W×C to extract local features. Before entering the Transformer module, the feature map is adjusted through a 1×1 convolution layer to adjust the number of channels. The feature map is mapped to a d-dimensional space. After two convolution operations, the input feature map X is transformed into a local feature map X of equal size. L ; The feature map X L Split into N image blocks of size H×W, and then expand into a set of feature sequences X of size P×N×d, expressed as: P=W×H N=(H×W) / P Where W and H are the preset width and height of the image block, and d is the feature dimension; After the expansion, the feature sequence X U The Transformer encoder is applied on the image block. The Transformer encoder uses the self-attention mechanism to capture the relationship between different feature blocks and realize the modeling of the global context. In the feature fusion part, the encoded image block is restored to the dimension C before encoding through a 1×1 convolution layer to obtain the feature X F , X F The size of is H×W×C, which is superimposed in parallel with the input feature map X to form a new feature map with a dimension of 2C. A 3×3 convolution kernel is used to fuse the new feature maps after parallel connection, and the new feature map dimension is mapped back to C to obtain the final output feature Y.

6. The lithium ore-related lithology identification method based on a lightweight deep learning model according to claim 1 is characterized in that: In step S3, feature extraction and classification of lithium ore-related rock images are performed based on the TL-Fmix-MobileViT lithology recognition model, including: The first step is to preprocess the data by normalizing, resizing, and performing random horizontal flipping on the image through torchvision.transforms; load the deep learning dataset and use torch.utils.data.DataLoader to create iterators for the training set and validation set; The second step is to build the model. Use timm to create the MobileViT-S model and load the pre-trained weights. Use Adam, AdamW and SGD optimizers and three learning rates of 0.002, 0.001 and 0.0002 for ablation experiments. After verification by ablation experiments, the number of iterations is 100, the number of training and verification batches is 16, the optimizer uses SGD, the learning rate is 0.002, the weight decay coefficient is 0.0001, and the momentum parameter is 0.9 to accelerate convergence. The loss function uses cross entropy loss to determine the difference between the actual output and the sample label. The third step is to initialize Fmix with data to generate mixed masks for data augmentation. Use the custom loss function fmix.loss() provided by Fmix to process mixed samples during training. fmix.loss() is a custom loss function designed specifically for mixed sample training. It uses the weight coefficient of the mixed mask to perform a weighted summation of the losses of the two sets of target labels. The fourth step is to train and save the TL-Fmix-MobileViT lithology recognition model. Use the trial interface of torchbearer to encapsulate the model, optimizer, loss function and data generator. Set the training to run on the GPU and monitor the accuracy and loss. After the training is completed, save the trained TL-Fmix-MobileViT lithology recognition model weights to the specified path.

7. The lithium ore-related lithology identification method based on a lightweight deep learning model according to claim 1 is characterized in that: In step S4, the recognition results are analyzed and the performance of the TL-Fmix-MobileViT lithology recognition model is evaluated, including: using accuracy, precision, recall, F1 value, model inference speed Frames PerSecond and model size as evaluation indicators of the TL-Fmix-MobileViT lithology recognition model; the expression is: Where TP indicates that the model correctly predicts a positive example as a positive example; FP indicates that the model incorrectly predicts a negative example as a positive example; TN indicates that the model correctly predicts a negative example as a negative example; FN indicates that the model incorrectly predicts a positive example as a negative example; Where N is the number of frames of samples to be processed.

8. A lithium ore-related lithology identification system based on a lightweight deep learning model, characterized in that: The system implements the method for identifying lithium ore-related lithology based on a lightweight deep learning model as described in any one of claims 1 to 7, and the system comprises: A deep learning dataset construction module (1) is used to obtain image data of rocks related to lithium ore, use image annotation, and construct a deep learning dataset; TL-Fmix-MobileViT lithology recognition model construction module (2), which is used to couple the Fmix image mixing data enhancement method in the MobileViT model to improve generalization, introduce transfer learning to accelerate training and improve accuracy, and construct the TL-Fmix-MobileViT lithology recognition model; A lithology recognition module (3) is used to extract features and classify and recognize lithium-related rock images based on the TL-Fmix-MobileViT lithology recognition model; The performance evaluation module (4) is used for analysis of identification results and performance evaluation of the TL-Fmix-MobileViT lithology identification model.

9. The lithium ore-related lithology identification system based on a lightweight deep learning model according to claim 8, characterized in that: The lithium ore-related lithology identification system based on a lightweight deep learning model is carried on a computer-readable storage medium, and the computer-readable storage medium stores a computer program. When the computer program is executed by a processor, it can realize the functions of the above-mentioned lithium ore-related lithology identification system based on a lightweight deep learning model.

10. The lithium ore-related lithology identification system based on a lightweight deep learning model according to claim 8, characterized in that: The lithium ore-related lithology identification system based on a lightweight deep learning model is installed on an information data processing terminal. The information data processing terminal is used to provide a user input interface to implement the functions of the lithium ore-related lithology identification system based on a lightweight deep learning model when executed on an electronic device.

Images