A semi-supervised deep learning image restoration enhancement method based on double network cooperation

Abstract

The application provides a semi-supervised deep learning image restoration enhancement method based on double-network cooperation. The method comprises the following steps: constructing a double-task branch image restoration network, training the double-task branch image restoration network by adopting a teacher-student semi-supervised training mode, performing feature extraction and restoration enhancement processing on an image to be restored by a restoration branch in the trained double-task branch image restoration network, outputting the restored image, low-order semantic features and high-order semantic features, fusing the low-order semantic features and the high-order semantic features by an evaluation branch, obtaining a structural similarity index SSIM score matrix by using a self-attention mechanism, and determining the confidence of pseudo-label screening in the teacher-student semi-supervised training mode training process of the double-task branch image restoration network according to the SSIM score matrix. The application fully excavates the value of a large number of low-quality samples and a small number of high-low quality paired samples, generates high-quality pseudo labels, and realizes accurate restoration and enhancement of trackside images.

Description

Technical Field

[0001] This invention relates to the field of image restoration technology, and in particular to a semi-supervised deep learning image restoration and enhancement method based on dual-network collaboration. Background Technology

[0002] Trains (express trains, bullet trains, and high-speed trains) play a crucial role in passenger and freight transport as a commonly used mode of transportation. Therefore, to further improve train transport safety, fundamentally reduce and eliminate train equipment failures, and effectively detect specific faults, it is essential to research an effective train fault detection and diagnosis method. This will lay a solid foundation for improving railway transport safety and play a key protective role. Among these methods, the health of trackside equipment seriously threatens the operational safety of high-speed trains; therefore, timely acquisition and analysis of high-quality trackside images are necessary to ensure safe operation.

[0003] Traditional image restoration techniques are mainly divided into model-based and non-model-based methods. Model-based methods typically utilize hand-designed prior knowledge to estimate unknown parameters in the imaging model, such as transmittance and ambient illumination. In contrast, non-model-based methods rely on adaptive image enhancement techniques, such as CLAHE, Retinex, Fusion, and MMLE. Although these traditional methods have certain advantages in terms of computational speed and data volume dependence, they are generally difficult to handle complex and variable real-world scenarios.

[0004] In early deep learning-based image restoration methods, researchers attempted to combine physical imaging models to achieve more accurate restoration results. These methods primarily utilized neural networks to estimate parameters such as transmittance and ambient light. However, the performance of such methods was limited due to the inaccuracy of parameter estimation. In recent years, the development of deep learning technology has spurred the emergence of new methods that directly learn restoration maps from labeled datasets, freeing them from strong dependence on physical imaging models. For example, one study employed a network that improved restoration results by merging three feature maps processed using traditional non-model-based methods. Another study, by incorporating the best aspects of traditional model-based methods, designed an encoder combining multiple color spaces and a transmittance-guided decoder to improve restoration results.

[0005] Recently, various deep learning-based methods have been proposed to address image restoration problems across a wide range of applications, achieving significant results in this specific field. Due to their data-driven nature, these deep learning-based methods can achieve more competitive restoration results compared to traditional models that rely on manually designed prior knowledge. For example, some image restoration techniques focus on designing complex network architectures, introducing attention mechanisms, and combining multi-view image information to improve restoration quality. Compared to traditional techniques, these deep learning-driven image restoration models typically achieve restoration results that are more consistent with human visual perception.

[0006] While deep learning methods have made significant breakthroughs in image restoration, these methods often require large amounts of paired sample data and rely on fully supervised learning to train image restoration maps. However, in practice, obtaining accurate paired images is often very difficult or even impossible. To address this challenge, research has proposed the CycleGAN framework, which incorporates an attention mechanism, for unsupervised image restoration. To further guide the model in learning restoration mapping patterns specific to underwater images, research has proposed a framework called Semi-UIR, which aims to achieve semi-supervised image restoration by combining a small amount of labeled data with a large amount of unlabeled data. This method utilizes a repository maintained by a non-reference image quality assessment model for semi-supervised guidance. Similarly, for low-light images, research has designed an image quality assessment method based on contrastive self-supervised learning to select the best pseudo-labels from numerous restored images.

[0007] Currently, existing image restoration technologies, especially trackside image restoration, still face the following challenges:

[0008] (1) Lack of real data samples. Unlike simulation experiments, the inherent properties of real trackside scenes make it difficult to obtain high-quality reference images, while methods that rely on artificial augmentation to generate pseudo-labels are costly. In addition, due to domain transfer issues, models trained using synthetic or different domain datasets often lack generalization ability in real industrial scenarios.

[0009] (2) Image quality evaluation is difficult. Although some recent studies have attempted to apply semi-supervised learning strategies to image restoration tasks, existing methods are still in the early stages of operating the image restoration model and the quality assessment model independently. This approach of using a pre-trained non-reference image quality assessment model for quality assessment lacks universality, thus severely limiting the flexibility of exchanging effective information between the two models. Consequently, the quality of the generated pseudo-labels cannot be reliably guaranteed. Summary of the Invention

[0010] This invention provides a semi-supervised deep learning image restoration and enhancement method based on dual-network collaboration, so as to effectively perform image restoration and enhancement processing.

[0011] To achieve the above objectives, the present invention adopts the following technical solution.

[0012] A semi-supervised deep learning image restoration and enhancement method based on dual-network collaboration includes:

[0013] A dual-task branch image restoration network is constructed, which includes a restoration branch and an evaluation branch;

[0014] The dual-task branch image restoration network was trained using labeled and unlabeled data in a teacher-student semi-supervised training mode to obtain the trained dual-task branch image restoration network.

[0015] The image to be restored is input into the trained dual-task branch image restoration network. The restoration branch performs feature extraction and restoration enhancement processing on the image to be restored, and outputs the restored image, as well as the low-order semantic features and high-order semantic features of the image to be restored.

[0016] The low-order and high-order semantic features are input into the evaluation branch, which fuses the low-order and high-order semantic features and obtains the Structural Similarity Index (SSIM) score matrix using a self-attention mechanism. Based on the SSIM score matrix, the confidence level of pseudo-label selection during the teacher-student semi-supervised training mode of the dual-task branch image restoration network is determined.

[0017] Preferably, the construction of the dual-task branch image restoration network includes a restoration branch and an evaluation branch, comprising:

[0018] A dual-task branch image restoration network is constructed, comprising a restoration branch and an evaluation branch. The restoration branch includes a multi-scale feature module, an attention feature fusion module, and a shared feature module. The multi-scale feature extraction module uses multi-scale residual blocks as the backbone to obtain parallel flow features at three different scales of the input image. The attention feature fusion module uses channel attention to fuse the flow features with lower-resolution flow features to obtain multi-scale fused features as low-order semantic features. The shared feature module uses feature transition and integration to obtain shared features. Combined with the repeated multi-scale feature extraction module and the attention feature fusion module, higher-order semantic features are further obtained, and the restored image of the input image is obtained.

[0019] The evaluation branch includes an attention feature fusion module and a Transformer module. The attention feature fusion module receives low-order and high-order semantic features obtained from the restoration branch, obtains fused features using an attention mechanism, and inputs the fused features into the Transformer module. The Transformer module uses a self-attention mechanism to mine the complex relationships between features of different blocks and obtains the SSIM score matrix of the restored image.

[0020] Preferably, the step of training the dual-task branch image restoration network using labeled and unlabeled data in a teacher-student semi-supervised training mode to obtain the trained dual-task branch image restoration network includes:

[0021] The training data consists of labeled and unlabeled data. The labeled data includes raw, unprocessed images to be restored, and the unlabeled data also includes raw, unprocessed images to be restored. The dual-task branch image restoration network is trained using a teacher-student semi-supervised training mode. For labeled data, the student network is used for quality restoration and image quality evaluation to obtain the predicted restored image and SSIM score. The loss is calculated by comparing the predicted image with the ground truth obtained from manual image retouching, which serves as the supervised loss. For unlabeled data, a perturbation is first added to the input image to be restored. The strongly perturbated version is input into the student network to obtain the predicted restored image and evaluation score matrix, while the weakly perturbated version is input into the teacher network to obtain the prediction result as a pseudo-label. Then, the loss between the pseudo-label and the student network prediction is calculated as the unsupervised loss. The sum of the losses is used to update the student network parameters, and the teacher network parameters are updated through the EMA of the student network parameters.

[0022] The SSIM score matrix has values ​​between 0 and 1. It is used as the confidence level for pseudo-label selection during the teacher-student semi-supervised training process of the dual-task branch image restoration network. The training process of the dual-task branch image restoration network is comprehensively evaluated and optimized through the total supervised loss function, resulting in a well-trained dual-task branch image restoration network. The total supervised loss function includes a supervised loss function for the restoration branch and a supervised loss function for the image quality fitting task. During semi-supervised training, the evaluation branch generates a "pseudo" image quality evaluation matrix for unlabeled samples, and maintains the current best restoration pseudo-label for each unlabeled data point. The pseudo-labels and corresponding evaluation score matrices are stored in a pseudo-label repository.

[0023] Preferably, the step of inputting the image to be restored into a trained dual-task branch image restoration network, wherein the restoration branch performs feature extraction and restoration enhancement processing on the image to be restored, and outputs the restored image, as well as the low-order and high-order semantic features of the image to be restored, including:

[0024] The image to be restored is input into the trained dual-task branch image restoration network. The restoration branch first uses the multi-scale residual block, multiple dilated convolution block and residual context block in the multi-scale feature extraction module to extract three parallel flow features of different scales from the input image. The small-scale parallel flow features focus on the local details of the image, while the large-scale parallel flow features focus on the overall structure of the image. Then, the attention feature fusion module fuses the flow features with the lower-resolution flow features to obtain low-order semantic features. Then, the shared feature module performs feature transition and integration to obtain shared features. Finally, the multi-scale feature extraction module and the attention feature fusion module further obtain high-order semantic features and the final restored image.

[0025] The restoration branch outputs the obtained low-order and high-order semantic features to the evaluation branch.

[0026] Preferably, the step of inputting the image to be restored into a trained dual-task branch image restoration network, wherein the restoration branch performs feature extraction and restoration enhancement processing on the image to be restored, and outputs the restored image, as well as the low-order and high-order semantic features of the image to be restored, including:

[0027] The evaluation branch receives low-order and high-order semantic features from the restoration branch. It then fuses these features using an attention feature fusion module, highlighting key features and suppressing unimportant ones to obtain fused features. These fused features are input into the Transformer module. The Transformer's block-segmentation strategy module divides the image into 16×16 blocks. Features are extracted from each block using a self-attention mechanism, and corresponding image quality evaluation values ​​are learned. This yields an evaluation score matrix containing location information, known as the SSIM score matrix. Each element in the SSIM score matrix corresponds to the image quality evaluation score of each block in the image, including location information. The values ​​of the SSIM score matrix are between 0 and 1. The SSIM score matrix is ​​used as the confidence level for pseudo-label selection during the teacher-student semi-supervised training process of the dual-task branch image restoration network.

[0028] As can be seen from the technical solutions provided by the embodiments of the present invention above, the present invention fully explores the value of a large number of low-quality samples and a small number of high-low quality paired samples to generate high-quality pseudo-labels, thereby training an efficient image restoration and enhancement model that meets the needs of actual industrial scenarios, achieving accurate restoration and enhancement of trackside images, and providing strong support for actual industrial production and inspection work.

[0029] Additional aspects and advantages of the invention will be set forth in part in the description which follows, and will become apparent from the description or may be learned by practice of the invention. Attached Figure Description

[0030] To more clearly illustrate the technical solutions of the embodiments of the present invention, the drawings used in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0031] Figure 1 This is a flowchart illustrating a semi-supervised deep learning image restoration and enhancement method based on dual-network collaboration proposed in an embodiment of the present invention.

[0032] Figure 2 This is an architecture diagram of a dual-task branch image restoration and enhancement network provided in an embodiment of the present invention;

[0033] Figure 3 This is an architecture diagram of a semi-supervised trackside image restoration network provided in an embodiment of the present invention;

[0034] Figure 4 This is a schematic diagram showing the enhancement effects of different algorithms on a trackside image dataset, as provided in an embodiment of the present invention. Detailed Implementation

[0035] Embodiments of the present invention are described in detail below, examples of which are shown in the accompanying drawings, wherein the same or similar reference numerals denote the same or similar elements or elements having the same or similar functions throughout. The embodiments described below with reference to the accompanying drawings are exemplary and are only used to explain the present invention, and should not be construed as limiting the present invention.

[0036] Those skilled in the art will understand that, unless specifically stated otherwise, the singular forms “a,” “an,” “the,” and “the” used herein may also include the plural forms. It should be further understood that the term “comprising” as used in this specification means the presence of the stated features, integers, steps, operations, elements, and / or components, but does not exclude the presence or addition of one or more other features, integers, steps, operations, elements, components, and / or groups thereof. It should be understood that when we say an element is “connected” or “coupled” to another element, it can be directly connected or coupled to the other element, or there may be intermediate elements. Furthermore, “connected” or “coupled” as used herein can include wireless connections or couplings. The term “and / or” as used herein includes any and all combinations of one or more of the associated listed items.

[0037] It will be understood by those skilled in the art that, unless otherwise defined, all terms used herein (including technical and scientific terms) have the same meaning as commonly understood by one of ordinary skill in the art to which this invention pertains. It should also be understood that terms such as those defined in general dictionaries should be understood to have the same meaning as in the context of the prior art, and should not be interpreted in an idealized or overly formal sense unless defined as herein.

[0038] To facilitate understanding of the embodiments of the present invention, the following will provide further explanation and description with reference to the accompanying drawings and several specific embodiments. These embodiments do not constitute a limitation on the embodiments of the present invention.

[0039] This invention addresses the difficulty of generating quality evaluations for pseudo-labels by constructing a dual-branch network that simultaneously outputs the restored image and its corresponding image quality evaluation value.

[0040] The processing flow of a semi-supervised deep learning image restoration and enhancement method based on dual-network collaboration proposed in this invention is as follows: Figure 1 As shown, the processing steps include the following:

[0041] Step S10: Construct a dual-task branch image restoration network.

[0042] An architecture diagram of a dual-task branch image restoration network provided in this embodiment of the invention is shown below. Figure 2 As shown. Figure 2 As shown, the architecture of this dual-task branch image restoration network includes a dual-branch structure: the first restoration branch (shown in the green area in the figure) is responsible for image restoration and enhancement; the second evaluation branch (marked in the yellow area) focuses on evaluating the image quality. The image mentioned above can be a trackside image.

[0043] The restoration branch includes a multi-scale feature module for low-order and high-order semantics, an attention feature fusion module, and a shared feature module. The multi-scale feature extraction module primarily uses multi-scale residual blocks as its backbone, containing three parallel flows at different scales. First, each flow passes through multiple dilated convolutional blocks, which expand the receptive field by combining four dilated convolutions; then it passes through a residual context block, which uses an attention mechanism to model and transform inter-channel dependencies, extracting useful contextual features. The attention feature fusion module mainly uses channel attention to fuse the features of its current flow with features from a lower-resolution flow, obtaining multi-scale fused features as low-order semantic features. This module fuses global and local information, helping to highlight small local targets. Then, it passes through the shared feature module, which uses the general structure of the restoration model to obtain shared features, playing a role in feature transition and integration. Finally, the multi-scale feature extraction module and the attention feature fusion module are repeated to further obtain high-order semantic features and the final restored image.

[0044] Function: The main function of the restoration branch is to obtain a higher-quality, clearer restored image by performing multi-scale feature extraction and fusion on the original degraded image. It also provides low-order and high-order semantic feature services. An evaluation branch is provided to output image quality evaluation.

[0045] Input / Output: The input to the restoration branch is the original image to be restored, i.e., the image of poor quality. The output is the restored image, i.e., the image of higher quality. Both are the same size.

[0046] The evaluation branch includes an attention feature fusion module and a Transformer module. Specifically, the low-order and high-order semantic features obtained from the restoration branch are input into the attention feature fusion module, which uses an attention mechanism to obtain fused features. Then, the fused features are input into the Transformer module, which uses a self-attention mechanism to mine the complex relationships between features from different blocks, obtaining the evaluation score matrix, i.e., the SSIM (Structural Similarity) score matrix.

[0047] The main function of the evaluation branch is to evaluate the image quality of the restored image by region based on the low-order and high-order features of the restoration branch. The evaluation metric is a quantified SSIM score to measure the quality of the restored image. This evaluation result can be used not only as a supervisory signal to guide the supervised learning process, but also as a confidence level to guide pseudo-label selection in semi-supervised training strategies. The input to the evaluation branch is the low-order and high-order semantic features extracted by the restoration branch, and the output is a block evaluation score matrix containing location information. Each block corresponds to an SSIM score (between 0 and 1) to reflect the image quality differences between the blocks.

[0048] Training data: The training data mainly consists of a small amount of labeled data and a large amount of unlabeled data. Labeled data includes unprocessed, raw, degraded images (i.e., images to be restored), which have quality issues such as uneven lighting and blurriness, as well as restored images obtained through manual retouching. These restored images are of higher quality and clearer, and can be used as ground truth (labels) for training. Unlabeled data consists only of unprocessed, raw, degraded images; there are no corresponding ground truth (labels) for training.

[0049] Training Process: Since the training data contains a small amount of labeled data and a large amount of unlabeled data, the problem to be solved is transformed into a semi-supervised learning problem. This invention employs a classic teacher-student semi-supervised training model to directly train a dual-task branch image restoration network. For labeled data, the student network is used for quality restoration and image quality evaluation, obtaining predicted restored images and SSIM scores. These are compared with the ground truth values ​​obtained through manual image retouching to calculate the supervised loss. For unlabeled data, perturbations are first added to the input degraded image. A strongly perturbated version is input into the student network to obtain the predicted restored image and evaluation score matrix, while a weakly perturbated version is input into the teacher network to obtain the prediction results as pseudo-labels. Then, the loss between the pseudo-labels and the student network predictions is calculated as the unsupervised loss. Finally, the sum of the losses is used to update the student network parameters, and the teacher network parameters are updated using the EMA of the student network parameters. This invention employs a semi-supervised learning strategy to train a dual-task branch image restoration method. First, it overcomes the challenge that restoration tasks, unlike classification tasks, do not have a reliable 0-1 confidence level for pseudo-label selection. Therefore, this invention creatively proposes a dual-branch structure: while the restoration branch learns the restoration mapping, an evaluation branch objectively evaluates the restoration's "confidence" and quantifies it as an objective evaluation metric, the SSIM score. The SSIM score is a widely recognized and reliable evaluation metric in image restoration, with values ​​between 0 and 1, suitable for fitting tasks. Furthermore, this invention is the first to use a semi-supervised image restoration method for trackside image restoration and enhancement, while also innovatively utilizing a dual-branch architecture.

[0050] The semi-supervised deep learning approach of this invention is mainly reflected in the training data and training strategy. Firstly, regarding the training data, machine learning methods are generally divided into three types: fully supervised learning (where each input corresponds to a ground truth output for supervised learning); unsupervised learning (where there is only input and no ground truth output, such as clustering algorithms); and semi-supervised learning (where a small number of inputs contain ground truth outputs, while a large number of inputs do not, and optimization strategies are used to achieve a fully supervised learning level). In this invention, because the collected trackside equipment images are often not clear or balanced, manual image retouching to obtain the restored image (i.e., ground truth) is time-consuming and laborious. Therefore, we often label only a very small portion of the data, resulting in a large amount of data being unlabeled. The aim is to train a model that learns the mapping between input and output, thereby achieving automated image restoration. Secondly, regarding the training strategy, models trained on very little labeled data are prone to overfitting. Therefore, this invention designs semi-supervised learning methods such as pseudo-labels and consistency regularization strategies to achieve trackside image restoration and enhancement, enabling the training of a restoration and enhancement model with only a small number of labeled (retouched) samples. Specifically, the pseudo-labeling strategy involves using a teacher-student network model and a stable teacher network output to filter reliable pseudo-labels for unlabeled data using a repository strategy. The consistency regularization strategy forces the teacher-student network predictions corresponding to the weakly perturbated and strongly perturbated versions of the original input to be consistent. Therefore, this invention employs a semi-supervised strategy to achieve trackside image restoration and enhancement, enabling the training of a restoration and enhancement model with only a small number of labeled (edited) samples.

[0051] Step S20: Input the image to be restored into the trained dual-task branch image restoration network. The restoration branch performs feature extraction and restoration enhancement processing on the image to be restored.

[0052] The restoration branch first employs a multi-scale feature extraction module (multi-scale residual blocks, multiple dilated convolutional blocks, and residual context blocks) to extract three parallel flows at different scales from the input image. These three feature flows at different scales capture detailed information at different levels of the image; the small-scale feature extraction unit focuses on local details, while the large-scale feature extraction unit focuses on the overall structure of the image. Then, the attention feature fusion module fuses the features of the current flow with the features of the lower-resolution flow to obtain multi-scale fused features as low-order semantic features. These are then further processed and integrated by a shared feature module to obtain shared features. Finally, the multi-scale feature extraction module and the attention feature fusion module further refine the high-order semantic features and the final restored image.

[0053] The restoration branch outputs the obtained low-order and high-order semantic features to the evaluation branch. The restoration branch also outputs feature maps, namely low-order semantic features (closer to the input) and high-order semantic features (closer to the output). These feature maps record intermediate state information of the image during multi-scale feature extraction and attention feature fusion, helping the evaluation branch analyze the quality of the restored image from different perspectives. Low-order semantic features (closer to the input) mainly contain information from the original image after preliminary multi-scale feature extraction. At this stage, the feature maps retain some basic image features, such as preliminary representations of edges and textures. Feature maps at different scales reflect the presentation of these basic features at different resolutions. High-order semantic features (closer to the output) are deep features obtained after one multi-scale feature extraction, attention feature fusion, and shared feature module processing, followed by another multi-scale feature extraction and attention feature fusion. These feature maps integrate important information from different levels of the image and highlight features that play a key role in image restoration and enhancement. They are a high-level summary and generalization of the image processing by the restoration branch, providing important reference value for the evaluation branch to accurately assess image quality. In this invention, we use the feature maps corresponding to the attention feature fusion module for fusion.

[0054] Step S30: The evaluation branch's task is to assist semi-supervised learning in reliable pseudo-label selection, and its output is the evaluation score matrix, i.e., the SSIM score matrix. By using an objective and quantitative SSIM evaluation index (with values ​​between 0 and 1) as the confidence level for pseudo-label selection (confidence levels exist in classification tasks, but there is no such 0-1 index for pseudo-label selection in restoration and enhancement tasks), the reliability of pseudo-labels can be effectively improved.

[0055] The evaluation branch receives low-order semantic features (closer to the input) and high-order semantic features (closer to the output) from the restoration branch. Then, the two sets of features are fused using an attention feature fusion module to obtain a fused feature. This module filters and fuses input features based on their importance, highlighting key features and suppressing unimportant ones, making the features more representative for subsequent processing. Finally, the fused feature is input into the Transformer module, which uses a block-based strategy to obtain an evaluation score matrix containing positional information, i.e., the SSIM score matrix.

[0056] Whether to output an evaluation score matrix containing location information: The evaluation branch will output an evaluation score matrix containing location information. The Transformer's block-segmentation strategy module divides the image into 16×16 blocks, extracts features from each block through a self-attention mechanism, and learns the corresponding image quality evaluation value, ultimately forming this evaluation score matrix. Each element in the matrix corresponds to the image quality evaluation score of each block in the image, including location information. When the evaluation value of a block in the image increases, the corresponding evaluation value is updated, and the corresponding block in the target image is replaced on a block-by-block basis.

[0057] The role of the evaluation score matrix: Due to the uncertainty of lighting conditions and dynamic / static distance in the trackside environment, the image quality of a single image may vary significantly in different areas. The evaluation score matrix uses the recognized objective image quality evaluation index SSIM (between 0 and 1) to help the model more meticulously evaluate the quality of different blocks in the restored image. Therefore, for labeled data, it can serve as a supervised learning signal to guide the learning quality of the restored image. Simultaneously, the evaluation score matrix can serve as a confidence level to guide the selection of pseudo-labels for unlabeled data, obtaining more reliable pseudo-labels for unsupervised learning.

[0058] The role of the overall supervision loss function: The overall supervision loss function is used to comprehensively evaluate and optimize the training process of a dual-task branch image restoration network. It consists of several parts, such as the supervision loss function for the restoration branch, which ensures the consistency between the restored image and the target image at the pixel level through reconstruction loss; perceptual loss evaluates the similarity between the restored image and the target image at the high-level feature level; and high-frequency gradient loss preserves the texture details of the trackside restored image. Through the weighted combination of these sub-losses, the restoration branch is guided to generate a high-quality image that is closer to the target. For the evaluation branch, by constructing a supervision loss function oriented towards image quality fitting, the local SSIM score is accurately fitted, improving the evaluation branch's ability to evaluate image structural similarity, thereby more accurately revealing image quality. The overall supervision loss function can effectively integrate the information of the dual-branch network, coordinate the training of the two branches, and make the entire network update in a better direction.

[0059] The Transformer's block-based strategy module divides the image into 16×16 blocks (a common hyperparameter setting used by Transformers), and then uses a self-attention mechanism to extract and analyze features from each block. This approach takes into account the local features of different regions of the image, fully exploring the relationships between different parts of the image. Due to the uncertainty of lighting conditions and dynamic / static distances in a trackside environment, the image quality in a single image may exhibit significant differences in different regions. Block processing can learn corresponding image quality evaluation values ​​for each region, avoiding the problem of ignoring local detail differences during overall processing. This results in more detailed and accurate evaluation results, thereby improving the model's accuracy in assessing image quality. Compared to outputting a single evaluation value and updating the evaluation of the entire image, this block-based strategy is more accurate and better suited for situations with uneven image lighting and quality.

[0060] The role of the pseudo-label repository B: The pseudo-label repository stores relevant information about unlabeled data. During semi-supervised training, the evaluation branch generates a "pseudo" image quality evaluation matrix for unlabeled samples and maintains the current best restored pseudo-label for each unlabeled data point. The pseudo-labels and evaluation score matrices in the repository participate in pseudo-label semi-supervised training, providing the student network with more training information, helping the model better utilize unlabeled data for learning, thereby improving the model's generalization ability and performance. Simultaneously, the repository is continuously updated based on the output of the evaluation branch, ensuring that the stored pseudo-labels and evaluation score matrices are always optimal, allowing the model to continuously optimize during training.

[0061] The parameters of the two models are trained together and influence each other during training. The evaluation results of the evaluation branch play a decisive role in the selection of pseudo-labels, and pseudo-labels are an important part of training the semi-supervised recovery branch, affecting the calculation of semi-supervised loss, and thus affecting the parameter updates of both branches.

[0062] To ensure the accuracy of the evaluation, this invention employs a region-based SSIM (Structural Similarity Index) as the evaluation metric, which also serves as a supervisory signal to guide the learning process. Given the uncertainties in lighting conditions and dynamic / static distances in the trackside environment, the image quality in a single image may exhibit significant differences across different regions, reflecting the regional variations in the reconstructed image. To effectively utilize this characteristic, this invention introduces a Transformer block-based strategy (generally dividing the image into 16×16 blocks), extracts features from each block using a self-attention mechanism, learns the corresponding image quality evaluation value, and finally outputs an evaluation score matrix containing location information to reflect the image quality differences between each block.

[0063] When constructing the evaluation branch, this invention extracts relevant feature maps near the input and output ends of the restoration branch and uses them as the input of the evaluation branch. This invention also employs an attention feature fusion module to integrate these two sets of features. To simplify the subsequent description, the image quality scoring Transformer module is denoted as T(·), and the overall network of the restoration branch is represented as R(·).

[0064] When processing the restoration branch, its supervised loss function will consist of three parts to comprehensively evaluate the restoration effect: reconstruction loss L str To ensure pixel-level consistency between the restored image and the target image; perceptual loss L per The similarity between the restored image and the target image at a high-level feature level is evaluated; high-frequency gradient loss L gra This preserves the texture details of the trackside restored image. They are represented as follows:

[0065]

[0066] in This represents the low-quality input image and its corresponding high-quality output label in image restoration. B represents the batch size. G(·) represents the features extracted by the j-th layer of the pre-trained base model, and G(·) denotes gradient extraction on the image. i represents the index of the i-th sample in the current batch, and J is the total number of effective feature layers in the pre-trained base model.

[0067] In summary, the supervised loss function L for the recovery branch sup_r The weighted sum of the three sub-losses in formula (1) is expressed as:

[0068] L sup_r =L str +L per +L gra (2)

[0069] For unlabeled data, this invention constructs a semi-supervised loss L based on pseudo-labels. unsup_r The loss function is consistent with the form described in formula (1), but the difference lies in that it is applied to unlabeled data. The model cannot directly obtain its label values; instead, it needs to use an estimation method to obtain its pseudo-labels.

[0070] In the design of the evaluation branch, this invention first constructs a supervised loss function L oriented towards the image quality fitting task. aseAs shown in Equation (3), the local SSIM score is accurately fitted. This method aims to improve the evaluation branch's ability to effectively assess the similarity of image structures, thereby revealing image quality more accurately. Furthermore, the detailed evaluation of local image quality is particularly suitable for the complexity and variability of trackside environments, providing a more detailed image quality assessment.

[0071]

[0072] In formula (3), Indicates the restored image Find the tag corresponding to the k-th block. The structural similarity index is given by ε(·), which is the binary cross-entropy loss function. Therefore, the main function of the evaluation branch is to evaluate the quality of the restored image on a block-by-block basis.

[0073] Regarding the unsupervised loss component of the evaluation branch, this invention considers the fact that lower-quality images are often more difficult to restore, and the restoration effect is usually not as good as that of images with higher original quality. Based on this observation, this invention proposes a novel soft-ranking loss function L. sr As shown in Equation (4), it does not rely on the labels of the original data, but instead uses relative ranking for loss constraints. The purpose of this loss function is to enhance the reliability of the evaluation branch output, ensuring that it can provide a more reasonable and stable assessment of the image restoration quality.

[0074]

[0075] in This indicates an unlabeled, low-quality input image. For images subjected to simple perturbations, such as color jitter or blurring According to formula (4), this setting is designed to encourage the output of image quality scores for original low-quality images. Not lower than the image quality score output of the corresponding sample after perturbation

[0076] Furthermore, to further stabilize the model's output, this invention also incorporates the regularization term commonly used in image restoration tasks into the unsupervised loss constraint, namely the contrastive loss L. tra The semantic features used to constrain the output are closer to At the same time, move away from the original low-quality image.

[0077] In summary, the total supervision loss function in this invention is expressed as:

[0078] L sup =L sup_r +ω3L ase (5)

[0079] The total unsupervised loss function is expressed as:

[0080] L unsup =L tra +ω4L unsup_r +ω5L sr (6)

[0081] The total loss is expressed as:

[0082] L total =L sup +T(t)L unsup (7)

[0083] Where T(t) represents the weight adjustment factor that changes with the number of training epochs t. In order to effectively utilize the information of the labeled data and at the same time use the potential information of the unlabeled data to update the network, T(t) is a linear increment function from 0 to 1. This ensures the stability of the initial training and prevents the target from stagnating due to inaccurate SSIM predictions, while also ensuring that the network can learn the restoration mapping better from the unlabeled data in the later stages.

[0084] This invention further introduces a semi-supervised training mode to make full use of labeled and unlabeled data. Figure 3 This is an architecture diagram of a semi-supervised trackside image restoration network provided in an embodiment of the present invention. Figure 3 As shown, this invention employs a classic teacher-student semi-supervised learning model, denoted as a "two-branch network (T)" and a "two-branch network (S)," with the aim of stabilizing pseudo-label output. In this architecture, the parameters of the teacher network (T) are not obtained through independent training, but are updated by applying an exponential moving average (EMA) to the parameters of the student network (S). Simultaneously, the parameters of the student network (S) are iteratively updated using a standard backpropagation method. This mechanism not only ensures that the teacher network can smoothly learn knowledge from the student network but also maintains the consistency of system updates. Furthermore, Figure 2 The model also includes a key component—a repository. The basic idea is to generate a "pseudo" image quality evaluation matrix for unlabeled samples through evaluation branches, and maintain the current best restored pseudo-label for each unlabeled data in the repository, thereby participating in pseudo-label semi-supervised training.

[0085] Specifically, for semi-supervised models, the training process for a batch is as follows:

[0086]

[0087] As shown in the steps above, the core idea of ​​this invention lies in using a dual-branch network architecture, where the evaluation branch is responsible for outputting image quality scores to maintain a pseudo-label repository B. This repository stores the optimal pseudo-label outputs for all unlabeled samples. and the corresponding evaluation score matrix By leveraging a teacher-student semi-supervised training model and an end-to-end dual-task collaborative training mechanism, this invention achieves an effective and efficient semi-supervised image restoration process.

[0088] In summary, this invention addresses the inherent challenges of capturing paired data and the complexity and variability of trackside image acquisition equipment by proposing a semi-supervised image restoration model based on a dual-task branch network. Compared to traditional image restoration and enhancement models, this invention achieves effective training without requiring the annotation of a large number of real paired samples. Furthermore, compared to existing semi-supervised image restoration and enhancement models, this invention can simultaneously perform collaborative training of the image restoration model and the quality assessment model. This allows the algorithm involved in this invention to more effectively cope with data incompleteness and the variability of the trackside environment, achieving high-quality trackside image enhancement and restoration.

[0089] The proposed semi-supervised deep learning-based trackside image restoration and enhancement method based on dual-network collaboration is compared with existing image restoration and enhancement methods on three datasets. The results are shown in Table 1. The datasets include two publicly available multi-exposure image datasets, MSEC and LCDP, as well as a real trackside image dataset collected during the testing of this invention. The comparison metrics include PSNR and SSIM.

[0090] Table 1. Quantitative comparison results of each method on the MSEC, LCDP, and trackside image datasets.

[0091]

[0092] As shown in Table 1, the method proposed in this invention consistently outperforms existing classical methods in both PSNR and SSIM metrics on the MSEC and LCDP datasets, achieving optimal overall performance. Particularly noteworthy is that the proposed method significantly outperforms existing classical methods in PSNR metrics on a self-built trackside image dataset, achieving an improvement of over 10%.

[0093] To further illustrate the effectiveness of each module in the proposed semi-supervised deep learning-based trackside image restoration and enhancement method based on dual-network collaboration, ablation experiments were conducted, and the results are shown in Table 2. Comparison metrics include PSNR and SSIM.

[0094] Table 2 Ablation experiments of the proposed method on a trackside image dataset.

[0095]

[0096] As shown in Table 2, in trackside image processing technology, the introduction of the image quality evaluation function has greatly improved the model restoration effect. Although it achieves a baseline effect of 27.17 PSNR and 0.89 SSIM when used alone, there is still significant room for improvement. By adopting an MLP structure to replace the traditional Transformer module, the model effectively reduces the computational complexity while maintaining the continuity of spatial features, resulting in a 6.15% jump in PSNR and an increase in SSIM to 0.91, verifying the breakthrough in the adaptability of the Transformer architecture to track scenes. Finally, the introduced dual-branch architecture further breaks through the 30 mark to reach 31.57 PSNR and increases SSIM to 0.96. The organic combination of the three forms a hierarchical and progressive technological innovation system.

[0097] Figure 4 This diagram illustrates the comparison of enhancement effects of different algorithms on a trackside image dataset, as provided in an embodiment of the present invention. It can be clearly seen that the method of the present invention has a better restoration effect when dealing with single-channel trackside images with local overexposure. It can effectively improve the texture details of local overexposure areas of the image without losing image resolution and texture details, and the overall color tone of the image is harmonious.

[0098] Those skilled in the art will understand that the accompanying drawings are merely schematic diagrams of one embodiment, and the modules or processes shown in the drawings are not necessarily essential for implementing the present invention.

[0099] As can be seen from the above description of the embodiments, those skilled in the art can clearly understand that the present invention can be implemented by means of software plus necessary general-purpose hardware platforms. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, can be embodied in the form of a software product. This computer software product can be stored in a storage medium, such as ROM / RAM, magnetic disk, optical disk, etc., and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute the methods described in various embodiments or some parts of the embodiments of the present invention.

[0100] The various embodiments in this specification are described in a progressive manner. Similar or identical parts between embodiments can be referred to mutually. Each embodiment focuses on describing the differences from other embodiments. In particular, for apparatus or system embodiments, since they are basically similar to method embodiments, the description is relatively simple; relevant parts can be referred to the descriptions in the method embodiments. The apparatus and system embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Those skilled in the art can understand and implement this without creative effort.

[0101] The above description is merely a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the technical scope disclosed in the present invention should be included within the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be determined by the scope of the claims.

Claims

1. A semi-supervised deep learning image restoration and enhancement method based on dual-network collaboration, characterized in that, include: A dual-task branch image restoration network is constructed, which includes a restoration branch and an evaluation branch; The dual-task branch image restoration network was trained using labeled and unlabeled data in a teacher-student semi-supervised training mode to obtain the trained dual-task branch image restoration network. The image to be restored is input into the trained dual-task branch image restoration network. The restoration branch performs feature extraction and restoration enhancement processing on the image to be restored, and outputs the restored image, as well as the low-order semantic features and high-order semantic features of the image to be restored. The low-order semantic features and high-order semantic features are input into the evaluation branch, which fuses the low-order semantic features and high-order semantic features and obtains the Structural Similarity Index (SSIM) score matrix using a self-attention mechanism. Based on the SSIM score matrix, the confidence level of pseudo-label selection during the teacher-student semi-supervised training mode of the dual-task branch image restoration network is determined. The method of training the dual-task branch image restoration network using labeled and unlabeled data in a teacher-student semi-supervised training mode to obtain the trained dual-task branch image restoration network includes: The training data consists of labeled and unlabeled data. The labeled data includes the processed raw images to be restored, and the unlabeled data includes the unprocessed raw images to be restored. The dual-task branch image restoration network is trained using a teacher-student semi-supervised training mode. For labeled data, the student network is used for quality restoration and image quality evaluation to obtain the predicted restored image and SSIM score. The loss is calculated by comparing the predicted value with the ground truth obtained by manual image retouching, which serves as the supervised loss. For unlabeled data, a perturbation is first added to the input image to be restored. The strongly perturbated version is input into the student network to obtain the predicted restored image and evaluation score matrix, while the weakly perturbated version is input into the teacher network to obtain the prediction result as a pseudo-label. Then, the loss between the pseudo-label and the student network prediction is calculated as the unsupervised loss. The sum of the losses is used to update the student network parameters, and the teacher network parameters are updated through the EMA of the student network parameters. The SSIM score matrix has values ​​between 0 and 1. It is used as the confidence level for pseudo-label selection during the teacher-student semi-supervised training process of the dual-task branch image restoration network. The training process of the dual-task branch image restoration network is comprehensively evaluated and optimized through the total supervised loss function, resulting in a well-trained dual-task branch image restoration network. The total supervised loss function includes a supervised loss function for the restoration branch and a supervised loss function for the image quality fitting task. During semi-supervised training, the evaluation branch generates a "pseudo" image quality evaluation matrix for unlabeled samples, and maintains the current best restoration pseudo-label for each unlabeled data point. The pseudo-labels and corresponding evaluation score matrices are stored in a pseudo-label repository.

2. The method according to claim 1, characterized in that, The constructed dual-task branch image restoration network includes a restoration branch and an evaluation branch, comprising: A dual-task branch image restoration network is constructed, comprising a restoration branch and an evaluation branch. The restoration branch includes a multi-scale feature extraction module, an attention feature fusion module, and a shared feature module. The multi-scale feature extraction module uses multi-scale residual blocks as the backbone to obtain parallel flow features at three different scales of the input image. The attention feature fusion module uses channel attention to fuse the flow features with lower-resolution flow features to obtain multi-scale fused features as low-order semantic features. The shared feature module uses feature transition and integration to obtain shared features. Combined with the repeated multi-scale feature extraction module and attention feature fusion module, higher-order semantic features are further obtained, and the restored image of the input image is obtained. The evaluation branch includes an attention feature fusion module and a Transformer module. The attention feature fusion module receives low-order and high-order semantic features obtained from the restoration branch, obtains fused features using an attention mechanism, and inputs the fused features into the Transformer module. The Transformer module uses a self-attention mechanism to mine the complex relationships between features of different blocks and obtains the SSIM score matrix of the restored image.

3. The method according to claim 1, characterized in that, The process involves inputting the image to be restored into a trained dual-task branch image restoration network. The restoration branch performs feature extraction and restoration enhancement on the image to be restored, outputting the restored image, as well as the low-order and high-order semantic features of the image to be restored, including: The image to be restored is input into the trained dual-task branch image restoration network. The restoration branch first uses the multi-scale residual block, multiple dilated convolution block and residual context block in the multi-scale feature extraction module to extract three parallel flow features of different scales from the input image. The small-scale parallel flow features focus on the local details of the image, while the large-scale parallel flow features focus on the overall structure of the image. Then, the attention feature fusion module fuses the flow features with the lower-resolution flow features to obtain low-order semantic features. Then, the shared feature module performs feature transition and integration to obtain shared features. Finally, the multi-scale feature extraction module and the attention feature fusion module further obtain high-order semantic features and the final restored image. The restoration branch outputs the obtained low-order and high-order semantic features to the evaluation branch.

4. The method according to claim 3, characterized in that, The process involves inputting the image to be restored into a trained dual-task branch image restoration network. The restoration branch performs feature extraction and restoration enhancement on the image to be restored, outputting the restored image, as well as the low-order and high-order semantic features of the image to be restored, including: The evaluation branch receives low-order and high-order semantic features from the restoration branch. It then fuses these features using an attention feature fusion module, highlighting key features and suppressing unimportant ones to obtain fused features. These fused features are input into the Transformer module. The Transformer's block-segmentation strategy module divides the image into 16×16 blocks. Features are extracted from each block using a self-attention mechanism, and corresponding image quality evaluation values ​​are learned. This yields an evaluation score matrix containing location information, known as the SSIM score matrix. Each element in the SSIM score matrix corresponds to the image quality evaluation score of each block in the image, including location information. The values ​​of the SSIM score matrix are between 0 and 1. The SSIM score matrix is ​​used as the confidence level for pseudo-label selection during the teacher-student semi-supervised training process of the dual-task branch image restoration network.

Images