Transmission line foreign object segmentation based on RB-UNet algorithm

The classical U-Net network model

Segmentation of foreign objects in transmission lines is a challenging task that requires an advanced semantic segmentation model. The U-Net network is selected as the foundational network for this segmentation task. The U-Net network consists of an encoder and a decoder. The encoder is responsible for extracting semantic features from the input images, while the decoder gradually upscales the low-resolution feature maps from the encoder back to the original image size. The encoder and decoder are connected through skip connections, allowing the decoder to access feature representations from the encoder and incorporate additional semantic information. The architecture of the U-Net network is illustrated in Fig. 1, showcasing the classical U-Net network architecture at the center, the EncodeBlock architecture on the left, and the DecodeBlock architecture on the right.

ResNet50 network

In order to enhance the feature extraction capability of the U-Net network, it was necessary to replace the encoder section of the U-Net network with the ResNet50 network. ResNet50 is a specific type of residual network consisting of 50 convolutional neural network layers. It is characterized by residual blocks, global average pooling layers, and fully connected layers. The ResNet50 network is divided into five stages. The first stage, conv1, comprises a convolutional layer, a BN layer, and a ReLu layer. The convolutional layer uses a 7 × 7 kernel with a stride of 2. The subsequent four stages, conv2-x, conv3-x, conv4-x, and conv5-x, all contain residual blocks, with varying numbers of blocks in each stage. Specifically, there are 3, 4, 6, and 3 residual blocks in conv2-x, conv3-x, conv4-x, and conv5-x, respectively, totaling 16 residual blocks. The ResNet50 network is commonly employed for image classification tasks due to its exceptional feature extraction capability stemming from its residual structure.

BN layer

In order to address the issue of slow convergence in the U-Net network, the properties of the BN layer are utilized. The BN layer is a technique that can expedite network convergence and enhance generalization capabilities. Training becomes more challenging with an excessively deep neural network, leading to significantly longer convergence times. Therefore, the BN layer is typically placed between the convolutional and activation layers to mitigate this issue. It standardizes each small batch of input data to have a mean close to 0 and a variance close to 1. The BN layer enhances both model stability and generalization capacity.

RB-UNet network

The U-Net network was originally designed for the recognition and segmentation of medical images. However, it faced challenges when dealing with images with complex backgrounds due to limitations in feature extraction, category imbalance, and model overfitting. To address these issues and achieve better performance in recognizing foreign objects on power transmission lines with intricate backgrounds, a more sophisticated network is necessary. The proposed semantic segmentation model, RB-UNet, builds upon improvements to the original U-Net. Figure 2 illustrates the framework of RB-UNet, highlighting the enhanced and replaced sections in blue. In contrast to the original U-Net, the RB-UNet network utilizes ResNet50 as the feature extraction network for the encoder section, incorporating the residual structure to create the RB-UNet’s encoding block, Res-EncodeBlock. To enhance model stability, a BN layer is inserted between each convolutional layer and the activation layer in the DecodeBlock of the decoder section, resulting in BN-DecodeBlock. Additionally, to address the class imbalance problem, a hybrid loss function combining Dice loss and Focal loss is introduced.

Res-EncodeBlock

ResNet50 is introduced as the U-Net network’s encoder to improve the extraction of semantic features from images, lessen reliance on a large amount of labeled data, and enhance model performance and generalization ability. This addresses the issue of the classic U-Net network’s weak feature extraction capability. The ResNet50 network is composed of five stages: conv1, conv2-x, conv3-x, conv4-x, and conv5-x. As illustrated in Fig. 2, the first two convolutional layers of the original U-Net network are first replaced by the convolutional layers in conv1. The U-Net network’s EncodeBlock is then replaced with each of the final four stages of the ResNet50 network, creating the encoding block Res-EncodeBlock. Each Res-EncodeBlock contains a different number of ResNet50 residual blocks; the number of residual blocks are 3, 4, 6, and 3 in that order, to be used for better extraction of image feature information. Figure 3 displays the ResNet50 residual blocks’ structure. The residual block on the left is made up of three convolutional layers: a 1 × 1 convolutional layer for dimensionality reduction, a 3 × 3 convolutional layer for feature extraction, and a 1 × 1 convolutional layer for dimensionality increase, as well as the corresponding BN layer, ReLu layer, and skip connection. An extra 1x1 convolutional layer and BN layer are included in the residual block on the right, above the skip connection. They are used to modify the feature map dimensions in cases where the input and output dimensions do not match.

BN-DecodeBlock

To address the limitations of the classic U-Net network in terms of stability and generalization capacity, this paper introduces the BN layer to the original U-Net model. As illustrated in Fig. 4, incorporating a BN layer after the convolutional layers in the DecodeBlock normalizes the outputs of the convolutional layers. This transformation results in a standard normal distribution with a mean of 0 and variance of 1. This approach can help the model to increase stability and converge quickly.

The mixed loss function

One of the loss functions frequently used in semantic segmentation is CE loss. The objective of this loss function is to minimize the difference between the predicted and true labels. (1)${\displaystyle L_{CE\left(p,y\right)}=-\sum _i{y}_{i}log\left(p_i\right)}$ where L_CE represents the cross-entropy loss value, y_i denotes the true label value of class i, and p_i represents the predicted label value of class i.

When used in semantic segmentation tasks, the Dice loss function is frequently employed to evaluate model performance by quantifying the overlap between true and predicted labels. By effectively handling class imbalance in semantic segmentation, the Dice loss function enhances model Precision and mIoU. (2)${\displaystyle L_D=1-\frac{2\times TP}{2\times TP+FP+FN}}$

To enhance accuracy and address class imbalance, the Focal loss function assigns greater weight to harder-to-classify and less frequent categories. This approach assists the model in managing the distribution imbalance between the numerous background categories and the limited target categories in semantic segmentation tasks, thereby enhancing the model’s learning capacity. (3)${\displaystyle L_F\left(p_t\right)=-{\alpha }_t\left(1-p_t\right)\gamma log\left(p_t\right)}$ where p_t represents the probability of predicting this category, α_t denotes the weight of this, and γ represents the exponential parameter.

To tackle the issue of class imbalance, one approach is to address it through the loss function. CE loss primarily emphasizes pixel-level classification accuracy, but its effectiveness diminishes on imbalanced datasets and small targets. In contrast, Dice loss evaluates the level of overlap between the predicted outcomes and the actual labels, resulting in superior performance with balanced datasets. Focal loss excels in handling class imbalance by assigning greater weights to challenging samples, thereby mitigating the issue of low accuracy for scarce samples. Therefore, combining Dice loss and Focal loss into a composite loss function can better solve the category imbalance problem. (4)${\displaystyle L_{mix}=L_D+L_F.}$

Evaluation metrics

Based on the confusion matrix, the metrics mIoU, Recall, F1-score, Precision, and PA are frequently used in semantic segmentation. These five indicators allow for a comprehensive evaluation of the accuracy and performance of the semantic segmentation model. They provide insights into various aspects of the model’s actual performance.

(1) mIoU: Indicates the degree of overlap between the predicted segmentation results and the true segmentation results by calculation. (5)${\displaystyle mIoU=\frac{1}{k+1}\sum _{i=0}^k\frac{TP}{TP+FP+FN}}$ (2) Recall: It shows the percentage of real positive samples that were accurately predicted to be positive out of true positive samples. (6)${\displaystyle \text{Recall}=\frac{TP}{TP+FN}}$ (3) F1-score: It represents the harmonic mean of Recall and Precision. (7)${\displaystyle F1socre=2\times \frac{\text{Recall}\times \text{Precision}}{\text{Recall}+\text{Precision}}}$ (4) Precision: It shows the likelihood that samples that are expected to be positive will, in fact, be positive. (8)${\displaystyle \text{Precision}=\frac{TP}{TP+FP}}$ (5) PA: It represents the percentage of samples that were correctly classified out of all the samples. (9)${\displaystyle PA=\frac{TP+TN}{TP+TN+FP+FN}}$

Dataset

Due to the unavailability of a publicly accessible dataset of foreign objects on transmission lines for training purposes, this study collected images of foreign object interference on transmission lines from the internet to create a dataset. Although the foreign object dataset in this paper is collected from the internet, the images in the dataset are real-world images, so there is no need to worry about the disconnect between the internet and the real world. This dataset includes four categories: bird nests, balloons, kites, and rubbish. These images vary in the distance and size of the target object. A total of 400 photos were chosen, with 100 for each category, as illustrated in Table 1. From these images, 80 were selected for the training set and 20 for the test set. Using labeling tools like Labelme, the dataset was tagged at the pixel level, resulting in 400 pixel-level annotated images—100 for each category. Unlike image classification and object detection tasks, semantic segmentation models focus more on the number of pixels. The number of object pixels supplied to the model varies for each image in the foreign object dataset on transmission lines due to variables such as the size and distance at which the objects were photographed. This implies that the issues of small sample numbers and class imbalance need to be addressed in the semantic segmentation model discussed in this paper. Although the dataset is relatively small and encounters the issue of class imbalance, it still presents various scenarios that can assist in training the model to effectively identify and segment foreign objects in different situations.

Training models and experimental settings

The training hardware utilized an NVIDIA GeForce RTX 4090 GPU for training. The dataset was divided into training and test sets using an 8:2 ratio. Throughout the experiments, the input image resolution was standardized to 512 × 512. The training process employed an Adam optimizer with a batch size of 24 and lasted for 100 epochs. The minimum learning rate was set to be equal to the maximum learning rate multiplied by 0.01, with the maximum learning rate set to 1e−4.