Container damage detection method based on improved U-Net and ViT

Abstract

The invention belongs to the technical field of deep learning image processing, and particularly relates to a container damage detection method based on improved U-Net and ViT, and the method comprises the steps: obtaining the sample image data of a container, and forming an original data set; a CBAMVggUNet model is constructed, and the CBAMVggUNet The CBAMVgUNet model is trained, and a container surface segmentation network model is obtained; obtaining image data only containing the container surface of the container to form a container surface data set; the ViT classification model is trained, and a container damage judgment classification model based on ViT is obtained; performing real-time image acquisition on the container to form a real-time data set; and acquiring real-time image data only containing the container surface, inputting the real-time image data into the container damage judgment and classification model, setting a judgment threshold value, performing damage judgment, and completing container damage detection. According to the invention, the characteristics of the container surface can be better captured, so that the damage detection is more reliable and accurate.

Description

Technical Field

[0001] The present invention belongs to the technical field of deep learning image processing, and particularly relates to a container damage detection method based on improved U-Net and ViT. Background Art

[0002] As an indispensable means of transportation in maritime trade, the integrity of containers is crucial to ensuring cargo safety and reducing logistics costs. With the continuous growth of global trade, higher requirements are placed on container damage detection. However, the existing container damage detection technology has some obvious limitations and shortcomings, as follows: Reliance on manual judgment: Current container damage detection mainly relies on manual judgment, which is time-consuming and costly. Manual judgment is easily affected by subjective factors and fatigue, which reduces detection accuracy and increases the risk of misjudgment and missed judgment.

[0003] Low accuracy of target detection algorithms: Some traditional target detection algorithms have low accuracy in complex backgrounds, especially in large-scale surveillance images. It is difficult to accurately identify defects on the surface of containers, and it is easy to generate false positives or false negatives. This makes the current technology perform poorly in complex scenes.

[0004] Difficulty in detection in complex scenes: Port containers are in complex scenes and may be blocked by other objects, affected by lighting, or have cluttered backgrounds. Traditional algorithms face greater difficulties in accurately detecting container defects, so more robust and adaptable solutions are needed.

[0005] Dependence on data annotation: Deep learning methods have a high dependence on large-scale annotated datasets, which increases the difficulty of obtaining sufficient quantity and quality of annotated data. This poses a challenge to training deep learning models, especially in the absence of sufficient annotated data. Summary of the invention

[0006] In view of the above deficiencies in the prior art, the object of the present invention is to provide a container damage detection method based on improved U-Net and ViT, which can better capture the characteristics of the container surface and make damage detection more reliable and accurate.

[0007] To achieve the above objectives, the present invention provides a container damage detection method based on improved U-Net and ViT, comprising the following steps: S1. Obtain sample image data of the container to form an original data set, and divide the original data set into a training set and a test set; S2. Introduce the attention mechanism CBAM into the U-Net model of the VGG network structure and build the CBAM_Vgg_UNet model; S3. Use the original data set to train the CBAM_Vgg_UNet model to obtain the container surface segmentation network model; S4, inputting the original data set into the container surface segmentation network model to obtain image data containing only the container surface to form a container surface data set, and dividing the container surface data set into a training set and a test set; S5. Using the container surface data set, train the ViT classification model to obtain a container damage judgment classification model based on ViT; S6, collecting real-time images of the container to form a real-time data set; S7, inputting the real-time data set into the container surface segmentation network model to obtain real-time image data containing only the container surface; S8. Input the real-time image data containing only the container surface into the container damage judgment classification model, set the judgment threshold, perform damage judgment, and complete the container damage detection.

[0008] As a preferred solution of the present invention, in S2, the U-Net model of the VGG network structure includes: Encoder: It includes several encoding blocks, each of which includes an input layer, an initial convolution layer, repeated convolution blocks, and a maximum pooling layer; the initial convolution layer uses a 3×3 convolution kernel with a step size of 1; the convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 1, an activation function, and a batch normalization layer; the maximum pooling layer uses a 2×2 maximum pooling layer with a step size of 2; the encoder gradually reduces the size of the feature map and retains the key features; Middle layer: It includes two convolution blocks. The convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 1, an activation function, and a batch normalization layer. The middle layer is used to extract global features. Decoder: Contains the same number of decoding blocks as encoding blocks. Each decoding block includes an upsampling layer, a convolution block and an output layer. The convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 2, an activation function and a batch normalization layer. The output layer is a convolution layer using a Sigmoid activation function, and the number of output channels is equal to the number of categories. The decoder gradually increases the size of the feature map through the upsampling layer, and then performs skip connections to connect the feature maps of the corresponding encoding blocks and decoding blocks together.

[0009] As a preferred solution of the present invention, in the U-Net model of the VGG network structure, the encoder includes four encoding blocks, and the decoder is correspondingly provided with four decoding blocks. The specific process of obtaining the output of the U-Net model of the VGG network structure is as follows: Encoder: In the first encoding block, the input image is the original image of [512×512×3], where [512×512] represents length×width, and 3 represents the number of channels; after two [3×3] 64-channel convolutions, the length and width of the image remain unchanged after each [3×3] convolution, and the number of channels corresponds to the number of channels of the convolution kernel, that is, 64, to obtain a [512×512×64] feature layer; then [2×2] maximum pooling is performed, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [256×256×64] feature layer is obtained; The input of the second encoding block is the output of the previous encoding block, that is, the [256×256×64] feature layer; it then undergoes two [3×3] 128-channel convolutions to obtain a [256×256×128] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [128×128×128] feature layer is obtained; The input of the third encoding block is the output of the previous encoding block, that is, the [128×128×128] feature layer; it then undergoes two [3×3] 256-channel convolutions to obtain a [128×128×256] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [64×64×256] feature layer is obtained; The input of the fourth encoding block is the output of the previous encoding block, that is, the [64×64×256] feature layer; it then undergoes two [3×3] 512-channel convolutions to obtain a [64×64×512] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [32×32×512] feature layer is obtained; Middle layer: The output of the fourth-layer encoding block is used as the input of the middle layer, and two [3×3] 512-channel convolutions are performed to obtain a [32×32×512] feature layer; Decoder: In the first decoding block, the input is the [32×32×512] feature layer obtained from the intermediate layer; then it is upsampled by [2×2]. After each upsampling by [2×2], the number of channels remains unchanged, and the length and width become twice as large as before, obtaining a [64×64×512] feature layer; after a jump connection, this feature layer is concatenated with the [64×64×512] feature layer output from the fourth encoding block to obtain a [64×64×1024] feature layer; after two [3×3] 512-channel convolutions, a [64×64×512] feature layer is obtained; then after [2×2] upsampling, a [128×128×512] feature layer is obtained; The input of the second decoding block is the [128×128×512] feature layer output by the first decoding block, which is then skipped and concatenated with the [128×128×256] feature layer output by the third encoding block to obtain a [128×128×768] feature layer. After two [3×3] 256-channel convolutions, a [128×128×256] feature layer is obtained. After [2×2] upsampling, a [256×256×256] feature layer is obtained. The input of the third decoding block is the [256×256×256] feature layer output by the second decoding block, which is then skipped and concatenated with the [256×256×128] feature layer output by the second encoding block to obtain a [256×256×384] feature layer. After two [3×3] 128-channel convolutions, a [256×256×128] feature layer is obtained. After [2×2] upsampling, a [512×512×128] feature layer is obtained. The input of the fourth decoding block is the feature layer of size [512×512×128] output by the third decoding block, which is then skipped and concatenated with the feature layer of [512×512×64] output by the second encoding block to obtain a feature layer of [512×512×192]. After two [3×3] 64-channel convolutions, a feature layer of [512×512×64] is obtained. Finally, after a [3×3] 2-channel convolution, the output is a mask image of size [512×512×2].

[0010] As a preferred solution of the present invention, in the U-Net model of the VGG network structure, an attention mechanism CBAM is introduced during the jump connection, and the CBAM includes: Channel attention module: Through wide-based global average pooling and global maximum pooling operations, the average and maximum values ​​in the channel dimension are obtained respectively, and two 1×1×C feature maps are generated. These two feature maps are sent to a two-layer shared multi-layer perceptron MLP, where the number of neurons in the first layer is C / r and the activation function is Relu; the number of neurons in the second layer is C; the features output by the two MLPs are element-wise added, and the channel attention feature Mc is generated after the Sigmoid activation function. Mc and the input feature map F are element-wise multiplied to generate the input feature F' required by the spatial attention module; where C is the number of channels in the feature map, and r is the reduction rate; Spatial attention module: For the input feature F', perform channel-based global maximum pooling and global average pooling operations to obtain two H×W×1 feature maps; concatenate these two feature maps in the channel dimension to obtain a H×W×2 feature map; through a 7×7 convolution operation, reduce the dimension to one channel, i.e. H×W×1; after the Sigmoid activation function, generate the spatial attention feature M_s; perform element-wise multiplication operation on M_s and the input feature map F' to generate the final feature map; where H and W are the height and width of the feature map, respectively.

[0011] As a preferred solution of the present invention, in S3, the original data set is sample images of containers obtained based on ImageNet. After the CBAM_Vgg_UNet model is trained, the pre-trained weights are saved, and then the CBAM_Vgg_UNet model is trained again using container images obtained based on historical collection, and the model parameters are fine-tuned to finally obtain a container surface segmentation network model.

[0012] As a preferred solution of the present invention, when training the CBAM_Vgg_UNet model, sample images of containers are preprocessed and labeled using labelme, and the labeled images are matched one-to-one with the preprocessed images as input for training; The preprocessing process includes scaling the input image to a uniform size, normalizing the image to ensure that the image data has a consistent range and distribution, and applying data enhancement techniques including random rotation, flipping, scaling, and translation to increase the diversity of training data.

[0013] As a preferred embodiment of the present invention, in S5, the classification process of the ViT classification model is: S51. The input image is divided into P×P blocks. Each block is converted into a d-dimensional vector through a linear mapping, and a position embedding vector is added to each block to obtain a set of vector sequences Z representing the image: ; In the formula, represents the tile after adding the position embedding vector, and N is the number of tiles; S52, introduce multiple Transformer modules, process the vector sequence Z through the self-attention mechanism and feedforward network, and gradually extract and integrate image features; S53, through residual connection and layer normalization, the vector sequence Z is connected to the output of the feedforward network Add together and get the result Pass it directly to the next layer. The calculation is expressed as: ; In the formula, Representation layer normalization operation; S54, through the fully connected layer Mapped to the category space, Softmax is used for probability normalization to obtain the final classification result, which is expressed as: ; Where W O is the output mapping weight; z is the intermediate representation mapped to the category space; ; In the formula, SoftMax is the activation function, O represents the final output, and represents the probability of each category.

[0014] As a preferred embodiment of the present invention, in S52, the specific steps of extracting and integrating image features are: S521. Self-attention mechanism It is expressed as: ; Where Q, K, and V represent the linear mapping of query, key, and value, respectively. k is the mapping dimension; Use a multi-head self-attention mechanism, that is, calculate multiple groups of Q, K, V, and then concatenate the results, expressed as: ; In the formula, MultiHeadSelfAttention represents multi-head self-attention, It means concatenating the outputs of all heads in the feature dimension, and the number of heads is h; S522. The output of the self-attention mechanism is further processed by a feed-forward network. The feed-forward network includes two fully connected layers, and the middle of the two fully connected layers is a ReLU activation function, which is expressed as: ; Where Y represents the output of the self-attention mechanism, represents the function of the feedforward network processing Y, ReLU is the activation function, M is the transposed weight matrix after the first fully connected layer, W F1 and b F1 are the weight matrix and bias vector of the first fully connected layer, W F2 and b F2 are the weight matrix and bias vector of the second fully connected layer respectively.

[0015] As a preferred solution of the present invention, in said S5, when training the ViT classification model, the images in the box surface data set are preprocessed respectively, and the broken and unbroken images are classified manually, and the classified images are matched one by one with the preprocessed images as input for training; In S6, photos are taken at the gates passing through the port when entering or leaving the port, and container image data is collected in real time.

[0016] As a preferred embodiment of the present invention, in S8, the output of the container damage judgment classification model is the classification result and the confidence level associated with the classification result. If the classification result is greater than or equal to the judgment threshold, it is judged to be damaged; if the classification result is less than the judgment threshold, it is judged to be undamaged. When it is determined to be damaged, a confidence level is determined. If the confidence level is greater than or equal to the confidence level threshold, an alarm is issued. If the confidence level is less than the confidence level threshold, a manual inspection is prompted. When it is determined that there is no damage, a confidence judgment is performed. If the confidence is greater than or equal to the confidence threshold, the process returns to S6 to continue collecting real-time images. If the confidence is less than the confidence threshold, a manual inspection is prompted.

[0017] The model and algorithm involved in the present invention can be executed by an electronic device, which includes a memory, a processor, and a computer program stored in the memory and executable on the processor. The processor executes the software to implement the above-mentioned model and algorithm calculation.

[0018] The beneficial effects of the present invention are: The present invention introduces a container surface segmentation network model, which can separate the container surface from the background and improve the accuracy of damage detection; the ViT classification model is more likely to obtain accurate results when only processing the container surface and is not affected by environmental factors; the introduction of jump connections with an attention mechanism makes the segmentation network more sensitive and precise, and can better capture the characteristics of the container surface, which optimizes the network performance, improves the accuracy of segmentation, and provides more reliable data for subsequent classification.

[0019] The present invention introduces a threshold setting step and adjusts the sensitivity of classification according to actual conditions, thereby improving the robustness of detection, reducing the false alarm rate, and making damage detection more reliable and accurate; the introduction of automated classification threshold setting reduces dependence on operators, making the detection process more intelligent, improving autonomy, and reducing the need for manual intervention, thereby effectively saving human resources. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 It is a schematic diagram of the process of the present invention; Figure 2 It is a schematic diagram of the principle of the CBAM_Vgg_UNet model of the present invention. DETAILED DESCRIPTION

[0021] The embodiments of the present invention are further described below in conjunction with the accompanying drawings: like Figure 1 As shown in the figure, the container damage detection method based on improved U-Net and ViT includes the following steps: S1. Obtain sample image data of the container to form an original data set, and divide the original data set into a training set and a test set; S2. Introduce the attention mechanism CBAM into the U-Net model of the VGG network structure and build the CBAM_Vgg_UNet model; S3. Use the original data set to train the CBAM_Vgg_UNet model to obtain the container surface segmentation network model; S4, inputting the original data set into the container surface segmentation network model to obtain image data containing only the container surface to form a container surface data set, and dividing the container surface data set into a training set and a test set; S5. Using the container surface data set, train the ViT classification model to obtain a container damage judgment classification model based on ViT; S6, collecting real-time images of the container to form a real-time data set; S7, inputting the real-time data set into the container surface segmentation network model to obtain real-time image data containing only the container surface; S8. Input the real-time image data containing only the container surface into the container damage judgment classification model, set the judgment threshold, perform damage judgment, and complete the container damage detection.

[0022] In S2, the U-Net model of the VGG network structure includes: Encoder: It includes several encoding blocks, each of which includes an input layer, an initial convolution layer, repeated convolution blocks, and a maximum pooling layer; the initial convolution layer uses a 3×3 convolution kernel with a step size of 1; the convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 1, an activation function, and a batch normalization layer; the maximum pooling layer uses a 2×2 maximum pooling layer with a step size of 2; the encoder gradually reduces the size of the feature map and retains the key features; Middle layer: It includes two convolution blocks. The convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 1, an activation function, and a batch normalization layer. The middle layer is used to extract global features. Decoder: Contains the same number of decoding blocks as encoding blocks. Each decoding block includes an upsampling layer, a convolution block and an output layer. The convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 2, an activation function and a batch normalization layer. The output layer is a convolution layer using a Sigmoid activation function, and the number of output channels is equal to the number of categories. The decoder gradually increases the size of the feature map through the upsampling layer, and then performs skip connections to connect the feature maps of the corresponding encoding blocks and decoding blocks together.

[0023] In the U-Net model of the VGG network structure, the encoder includes four encoding blocks, and the decoder has four decoding blocks. The specific process of obtaining the output of the U-Net model of the VGG network structure is as follows: Encoder: In the first encoding block, the input image is the original image of [512×512×3], where [512×512] represents length × width, and 3 represents the number of channels; after two [3×3] 64-channel convolutions ([3×3] represents the length and width of the convolution kernel, and 64 is the number of channels of the convolution kernel, which can be understood as 64 3×3 convolution kernels. The meaning here is that the size of the convolution kernel in the initial convolution layer is [3×3] 64 channels, and the repeated convolution block is two such convolutions), the length and width of the image remain unchanged after each [3×3] convolution, and the number of channels corresponds to the number of channels of the convolution kernel, that is, 64, to obtain a [512×512×64] feature layer; then [2×2] maximum pooling is performed, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [256×256×64] feature layer is obtained; The input of the second encoding block is the output of the previous encoding block, that is, the [256×256×64] feature layer; it then undergoes two [3×3] 128-channel convolutions to obtain a [256×256×128] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [128×128×128] feature layer is obtained; The input of the third encoding block is the output of the previous encoding block, that is, the [128×128×128] feature layer; it then undergoes two [3×3] 256-channel convolutions to obtain a [128×128×256] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [64×64×256] feature layer is obtained; The input of the fourth encoding block is the output of the previous encoding block, that is, the [64×64×256] feature layer; it then undergoes two [3×3] 512-channel convolutions to obtain a [64×64×512] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [32×32×512] feature layer is obtained; Middle layer: The output of the fourth-layer encoding block is used as the input of the middle layer, and two [3×3] 512-channel convolutions are performed to obtain a [32×32×512] feature layer; Decoder: In the first decoding block, the input is the [32×32×512] feature layer obtained from the intermediate layer; then it is upsampled by [2×2]. After each upsampling by [2×2], the number of channels remains unchanged, and the length and width become twice as large as before, obtaining a [64×64×512] feature layer; after a jump connection, this feature layer is concatenated with the [64×64×512] feature layer output from the fourth encoding block to obtain a [64×64×1024] feature layer; after two [3×3] 512-channel convolutions, a [64×64×512] feature layer is obtained; then after [2×2] upsampling, a [128×128×512] feature layer is obtained; The input of the second decoding block is the [128×128×512] feature layer output by the first decoding block, which is then skipped and concatenated with the [128×128×256] feature layer output by the third encoding block to obtain a [128×128×768] feature layer. After two [3×3] 256-channel convolutions, a [128×128×256] feature layer is obtained. After [2×2] upsampling, a [256×256×256] feature layer is obtained. The input of the third decoding block is the [256×256×256] feature layer output by the second decoding block, which is then skipped and concatenated with the [256×256×128] feature layer output by the second encoding block to obtain a [256×256×384] feature layer. After two [3×3] 128-channel convolutions, a [256×256×128] feature layer is obtained. After [2×2] upsampling, a [512×512×128] feature layer is obtained. The input of the fourth decoding block is the feature layer of size [512×512×128] output by the third decoding block, which is then skipped and concatenated with the feature layer of [512×512×64] output by the second encoding block to obtain a feature layer of [512×512×192]. After two [3×3] 64-channel convolutions, a feature layer of [512×512×64] is obtained. Finally, after a [3×3] 2-channel convolution, the output is a mask image of size [512×512×2].

[0024] In the U-Net model of the VGG network structure, the attention mechanism CBAM is introduced during the jump connection. CBAM includes: Channel attention module: Through wide-based global average pooling and global maximum pooling operations, the average and maximum values ​​in the channel dimension are obtained respectively, and two 1×1×C feature maps are generated. These two feature maps are sent to a two-layer shared multi-layer perceptron MLP, where the number of neurons in the first layer is C / r and the activation function is Relu; the number of neurons in the second layer is C; the features output by the two MLPs are element-wise added, and the channel attention feature Mc is generated after the Sigmoid activation function. Mc and the input feature map F are element-wise multiplied to generate the input feature F' required by the spatial attention module; where C is the number of channels in the feature map, and r is the reduction rate; Spatial attention module: For the input feature F', perform channel-based global maximum pooling and global average pooling operations to obtain two H×W×1 feature maps; concatenate these two feature maps in the channel dimension to obtain a H×W×2 feature map; through a 7×7 convolution operation, reduce the dimension to one channel, i.e. H×W×1; after the Sigmoid activation function, generate the spatial attention feature M_s; perform element-wise multiplication operation on M_s and the input feature map F' to generate the final feature map; where H and W are the height and width of the feature map, respectively.

[0025] Combined with CBAM, the principle of the complete CBAM_Vgg_UNet model is as follows Figure 2 As shown, the arrows represent the process of image processing.

[0026] In S3, the original data set is sample images of containers obtained based on ImageNet. After training the CBAM_Vgg_UNet model, the pre-training weights are saved, and then the CBAM_Vgg_UNet model is trained again using container images obtained based on historical collection, and the model parameters are fine-tuned to finally obtain the container surface segmentation network model. This process makes full use of the rich image data to improve the generalization ability of the model.

[0027] When training the CBAM_Vgg_UNet model, 100 epochs (rounds) of training are performed. Sample images of containers are preprocessed and labeled using labelme. The labeled images are matched one-to-one with the preprocessed images as input for training. The preprocessing process includes scaling the input image to a uniform size, normalizing the image to ensure that the image data has a consistent range and distribution, and applying data enhancement techniques including random rotation, flipping, scaling, and translation to increase the diversity of training data.

[0028] In S5, the classification process of the ViT classification model is: S51. The input image is divided into P×P blocks. Each block is converted into a d-dimensional vector through a linear mapping, and a position embedding vector is added to each block to obtain a set of vector sequences Z representing the image: ; In the formula, represents the tile after adding the position embedding vector, and N is the number of tiles; S52, introduce multiple Transformer modules, process the vector sequence Z through the self-attention mechanism and feedforward network, and gradually extract and integrate image features; S53, through residual connection and layer normalization, the vector sequence Z is connected to the output of the feedforward network Add together and get the result Pass it directly to the next layer. The calculation is expressed as: ; In the formula, Representation layer normalization operation; S54, through the fully connected layer Mapped to the category space, Softmax is used for probability normalization to obtain the final classification result, which is expressed as: ; Where WO is the output mapping weight; z is the intermediate representation mapped to the category space; ; In the formula, SoftMax is the activation function, O represents the final output, and represents the probability of each category.

[0029] In S52, the specific steps of extracting and integrating image features are: S521. Self-attention mechanism It is expressed as: ; Where Q, K, and V represent the linear mapping of query, key, and value, respectively. k is the mapping dimension; Use a multi-head self-attention mechanism, that is, calculate multiple groups of Q, K, V, and then concatenate the results, expressed as: ; In the formula, MultiHeadSelfAttention represents multi-head self-attention, It means concatenating the outputs of all heads in the feature dimension, and the number of heads is h; S522. The output of the self-attention mechanism is further processed by a feed-forward network. The feed-forward network includes two fully connected layers, and the middle of the two fully connected layers is a ReLU activation function, which is expressed as: ; Where Y represents the output of the self-attention mechanism, represents the function of the feedforward network processing Y, ReLU is the activation function, M is the transposed weight matrix after the first fully connected layer, W F1 and b F1 are the weight matrix and bias vector of the first fully connected layer, W F2 and b F2 are the weight matrix and bias vector of the second fully connected layer respectively.

[0030] In S5, when training the ViT classification model, 50 epochs of training were performed. The images in the box surface dataset were preprocessed and manually classified into broken and unbroken images. The classified images were matched one by one with the preprocessed images as input for training; In S6, photos are taken at the gates passing through the port when entering or leaving the port, and container image data is collected in real time.

[0031] In S8, the output of the container damage judgment classification model is the classification result and the confidence level associated with the classification result. If the classification result is greater than or equal to the judgment threshold, it is judged to be damaged; if the classification result is less than the judgment threshold, it is judged to be undamaged; When it is determined to be damaged, a confidence level is determined. If the confidence level is greater than or equal to the confidence level threshold, an alarm is issued. If the confidence level is less than the confidence level threshold, a manual inspection is prompted. When it is determined that there is no damage, a confidence judgment is performed. If the confidence is greater than or equal to the confidence threshold, the process returns to S6 to continue collecting real-time images. If the confidence is less than the confidence threshold, a manual inspection is prompted.

[0032] In this embodiment, the judgment threshold and the confidence threshold are both set to 0.75.

[0033] The model setting of the CBAM_Vgg_UNet+ViT of the present invention is compared and verified with existing models and their combinations. The comparison models include target detection models YOLO v5 and YOLO v7, classification models Vggnet, Resnet and ViT, and segmentation + classification model combinations Unet+vit, CBAM_Vgg_UNet+Vggnet, CBAM_Vgg_UNet+Resnet, and CBAM_Vgg_UNet+ViT.

[0034] The confusion matrix is ​​a table used in machine learning and statistics to evaluate the performance of classification models. It shows the performance of a classifier on different categories, especially in binary classification problems. The four key terms of the confusion matrix are true positive (TP), true negative (TN), false positive (FP), and false negative (FN): True positive: indicates that the model correctly classifies positive examples (positive categories) as positive examples; true negative: indicates that the model correctly classifies negative examples (negative categories) as negative examples; false positive: indicates that the model incorrectly classifies negative examples as positive examples; false negative: indicates that the model incorrectly classifies positive examples as negative examples.

[0035] Through the confusion matrix, a series of evaluation indicators can be calculated, including accuracy, precision, recall and F1 score.

[0036] Accuracy refers to the ratio of the number of samples correctly classified by the model to the total number of samples: ; Precision refers to the proportion of true positives among the samples predicted by the model as positive: ; Recall refers to the proportion of all actual positive examples that the model successfully predicts as positive examples: ; The F1 score is the harmonic mean of precision and recall and is useful for balancing precision and recall: ; Table 1 Evaluation index results of each model

[0037] As shown in Table 1, the model setting of CBAM_Vgg_UNet+ViT of the present invention achieved the best evaluation value.

Claims

1. Container damage detection method based on improved U-Net and Vit, Features The following steps are involved: S1. Obtain sample image data of the container to form an original data set, and divide the original data set into a training set and a test set; S2. Introduce the attention mechanism CBAM into the U-Net model of the VGG network structure and build the CBAM_Vgg_UNet model; S3. Use the original data set to train the CBAM_Vgg_UNet model to obtain the container surface segmentation network model; S4, inputting the original data set into the container surface segmentation network model to obtain image data containing only the container surface to form a container surface data set, and dividing the container surface data set into a training set and a test set; S5. Using the container surface data set, train the ViT classification model to obtain a container damage judgment classification model based on ViT; S6, collecting real-time images of the container to form a real-time data set; S7, inputting the real-time data set into the container surface segmentation network model to obtain real-time image data containing only the container surface; S8. Input the real-time image data containing only the container surface into the container damage judgment classification model, set the judgment threshold, perform damage judgment, and complete the container damage detection.

2. According to claim 1, the container damage detection method based on improved U-Net and Vit, Features: In S2, the U-Net model of the VGG network structure includes: Encoder: It includes several encoding blocks, each of which includes an input layer, an initial convolution layer, repeated convolution blocks, and a maximum pooling layer; the initial convolution layer uses a 3×3 convolution kernel with a step size of 1; the convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 1, an activation function, and a batch normalization layer; the maximum pooling layer uses a 2×2 maximum pooling layer with a step size of 2; the encoder gradually reduces the size of the feature map and retains the key features; Middle layer: It includes two convolution blocks. The convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 1, an activation function, and a batch normalization layer. The middle layer is used to extract global features. Decoder: Contains the same number of decoding blocks as encoding blocks. Each decoding block includes an upsampling layer, a convolution block and an output layer. The convolution block includes a 3×3 convolution kernel, a convolution layer with a step size of 2, an activation function and a batch normalization layer. The output layer is a convolution layer using a Sigmoid activation function, and the number of output channels is equal to the number of categories. The decoder gradually increases the size of the feature map through the upsampling layer, and then performs skip connections to connect the feature maps of the corresponding encoding blocks and decoding blocks together.

3. According to claim 2, the container damage detection method based on improved U-Net and Vit, Features: In the U-Net model of the VGG network structure, the encoder includes four encoding blocks, and the decoder has four decoding blocks. The specific process of obtaining the output of the U-Net model of the VGG network structure is as follows: Encoder: In the first encoding block, the input image is the original image of [512×512×3], where [512×512] represents length×width, and 3 represents the number of channels; after two [3×3] 64-channel convolutions, the length and width of the image remain unchanged after each [3×3] convolution, and the number of channels corresponds to the number of channels of the convolution kernel, that is, 64, to obtain a [512×512×64] feature layer; then [2×2] maximum pooling is performed, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [256×256×64] feature layer is obtained; The input of the second encoding block is the output of the previous encoding block, that is, the [256×256×64] feature layer; it then undergoes two [3×3] 128-channel convolutions to obtain a [256×256×128] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [128×128×128] feature layer is obtained; The input of the third encoding block is the output of the previous encoding block, that is, the [128×128×128] feature layer; it then undergoes two [3×3] 256-channel convolutions to obtain a [128×128×256] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [64×64×256] feature layer is obtained; The input of the fourth encoding block is the output of the previous encoding block, that is, the [64×64×256] feature layer; it then undergoes two [3×3] 512-channel convolutions to obtain a [64×64×512] feature layer; it then undergoes a [2×2] maximum pooling, and after each [2×2] maximum pooling, the number of channels remains unchanged, the length and width are halved, and a [32×32×512] feature layer is obtained; Middle layer: The output of the fourth-layer encoding block is used as the input of the middle layer, and two [3×3] 512-channel convolutions are performed to obtain a [32×32×512] feature layer; Decoder: In the first decoding block, the input is the [32×32×512] feature layer obtained by the intermediate layer; After that, it is upsampled by [2×2]. After each upsampling by [2×2], the number of channels remains unchanged, and the length and width become twice as large as before, obtaining a feature layer of [64×64×512]. After skip connection, the feature layer is concatenated with the [64×64×512] feature layer output in the fourth encoding block to obtain a [64×64×1024] feature layer; after two [3×3] 512-channel convolutions, a [64×64×512] feature layer is obtained; and after [2×2] upsampling, a [128×128×512] feature layer is obtained; The input of the second decoding block is the [128×128×512] feature layer output by the first decoding block, which is then skipped and concatenated with the [128×128×256] feature layer output by the third encoding block to obtain a [128×128×768] feature layer. After two [3×3] 256-channel convolutions, a [128×128×256] feature layer is obtained. After [2×2] upsampling, a [256×256×256] feature layer is obtained. The input of the third decoding block is the [256×256×256] feature layer output by the second decoding block, which is then skipped and concatenated with the [256×256×128] feature layer output by the second encoding block to obtain a [256×256×384] feature layer. After two [3×3] 128-channel convolutions, a [256×256×128] feature layer is obtained. After [2×2] upsampling, a [512×512×128] feature layer is obtained. The input of the fourth decoding block is the feature layer of size [512×512×128] output by the third decoding block, which is then skipped and concatenated with the feature layer of [512×512×64] output by the second encoding block to obtain a feature layer of [512×512×192]. After two [3×3] 64-channel convolutions, a feature layer of [512×512×64] is obtained. Finally, after a [3×3] 2-channel convolution, the output is a mask image of size [512×512×2].

4. According to claim 3, the container damage detection method based on improved U-Net and Vit, Features: In the U-Net model of the VGG network structure, the attention mechanism CBAM is introduced during the jump connection. CBAM includes: Channel attention module: Through wide-based global average pooling and global maximum pooling operations, the average and maximum values ​​in the channel dimension are obtained respectively, and two 1×1×C feature maps are generated. These two feature maps are sent to a two-layer shared multi-layer perceptron MLP, where the number of neurons in the first layer is C / r and the activation function is Relu; the number of neurons in the second layer is C; the features output by the two MLPs are element-wise added, and the channel attention feature Mc is generated after the Sigmoid activation function. Mc and the input feature map F are element-wise multiplied to generate the input feature F' required by the spatial attention module; where C is the number of channels in the feature map, and r is the reduction rate; Spatial attention module: For the input feature F', perform channel-based global maximum pooling and global average pooling operations to obtain two H×W×1 feature maps; concatenate these two feature maps in the channel dimension to obtain a H×W×2 feature map; through a 7×7 convolution operation, reduce the dimension to one channel, i.e. H×W×1; after the Sigmoid activation function, generate the spatial attention feature M_s; perform element-wise multiplication operation on M_s and the input feature map F' to generate the final feature map; where H and W are the height and width of the feature map, respectively.

5. According to claim 1, the container damage detection method based on improved U-Net and Vit, Features: In the S3, the original data set is sample images of containers obtained based on ImageNet. After the CBAM_Vgg_UNet model is trained, the pre-trained weights are saved, and then the CBAM_Vgg_UNet model is trained again using container images obtained based on historical collection, and the model parameters are fine-tuned to finally obtain the container surface segmentation network model.

6. According to claim 5, the container damage detection method based on improved U-Net and Vit, Features: When training the CBAM_Vgg_UNet model, the sample images of the container are preprocessed and labeled using labelme, and the labeled images are matched one-to-one with the preprocessed images as input for training; The preprocessing process includes scaling the input image to a uniform size, normalizing the image to ensure that the image data has a consistent range and distribution, and applying data enhancement techniques including random rotation, flipping, scaling, and translation to increase the diversity of training data.

7. According to claim 1, the container damage detection method based on improved U-Net and Vit, Features: In S5, the classification process of the ViT classification model is: S51. The input image is divided into P×P blocks. Each block is converted into a d-dimensional vector through a linear mapping, and a position embedding vector is added to each block to obtain a set of vector sequences Z representing the image: ； In the formula, represents the tile after adding the position embedding vector, and N is the number of tiles; S52, introduce multiple Transformer modules, process the vector sequence Z through the self-attention mechanism and feedforward network, and gradually extract and integrate image features; S53, through residual connection and layer normalization, the vector sequence Z is connected to the output of the feedforward network Add together and get the result Pass it directly to the next layer. The calculation is expressed as: ； In the formula, Representation layer normalization operation; S54, through the fully connected layer Mapped to the category space, Softmax is used for probability normalization to obtain the final classification result, which is expressed as: ； Where W O is the output mapping weight; z is the intermediate representation mapped to the category space; ； In the formula, SoftMax is the activation function, O represents the final output, and represents the probability of each category.

8. According to claim 7, the container damage detection method based on improved U-Net and Vit, Features: In the above S52, the specific steps of extracting and integrating image features are: S521. Self-attention mechanism It is expressed as: ； Where Q, K, and V represent the linear mapping of query, key, and value, respectively. k is the mapping dimension; Use a multi-head self-attention mechanism, that is, calculate multiple groups of Q, K, V, and then concatenate the results, expressed as: ； In the formula, MultiHeadSelfAttention represents multi-head self-attention, It means concatenating the outputs of all heads in the feature dimension, and the number of heads is h; S522. The output of the self-attention mechanism is further processed by a feed-forward network. The feed-forward network includes two fully connected layers, and the middle of the two fully connected layers is a ReLU activation function, which is expressed as: ； Where Y represents the output of the self-attention mechanism, represents the function of the feedforward network processing Y, ReLU is the activation function, M is the transposed weight matrix after the first fully connected layer, W F1 and b F1 are the weight matrix and bias vector of the first fully connected layer, W F2 and b F2 are the weight matrix and bias vector of the second fully connected layer respectively.

9. The container damage detection method based on improved U-Net and Vit according to claim 1, Features: In the above S5, when training the ViT classification model, the images in the box surface data set are preprocessed respectively, and the broken and unbroken images are classified manually, and the classified images are matched one by one with the preprocessed images as input for training; In S6, photos are taken at the gates passing through the port when entering or leaving the port, and container image data is collected in real time.

10. The container damage detection method based on improved U-Net and Vit according to claim 1, Features: In the above S8, the output of the container damage judgment classification model is the classification result and the confidence related to the classification result. If the classification result is greater than or equal to the judgment threshold, it is judged to be damaged; if the classification result is less than the judgment threshold, it is judged to be not damaged; When it is determined to be damaged, a confidence level is determined. If the confidence level is greater than or equal to the confidence level threshold, an alarm is issued. If the confidence level is less than the confidence level threshold, a manual inspection is prompted. When it is determined that there is no damage, a confidence judgment is performed. If the confidence is greater than or equal to the confidence threshold, the process returns to S6 to continue collecting real-time images. If the confidence is less than the confidence threshold, a manual inspection is prompted.

Images