Multitask image coding method and device based on multi-branch aggregation

Abstract

The invention provides a multi-task image coding method and device based on multi-branch aggregation. The method comprises the following steps: step 1, training a basic image coding model, a multi-branch aggregation module and an importance predictor; 2, the features are sent into corresponding importance predictors, and the predictors calculate and output binary masks according to the importance of different features to different tasks; 3, splitting the features into two groups of sub-features, namely main path features and bypass features; 4, selecting a corresponding bypass according to a specified task, and sending the main path characteristics and the bypass characteristics to the main path and the selected bypass respectively; and step 5, recombining the two groups of sub-features after main path and bypass calculation according to the positions and sizes of the original features, and completing decoding to output a complete image. According to the method, an independent specific task model does not need to be optimized, and the optimization difficulty and the training cost can be greatly reduced.

Description

Technical Field

[0001] The present invention relates to the field of image coding, and in particular to a multi-task image coding method and device based on multi-branch aggregation. Background Art

[0002] In recent years, research on multi-task coding that meets human perception and machine vision has made some progress. In order to achieve the best performance on various tasks, an intuitive and simple solution is to optimize different coding models for different tasks and adopt multiple bitstreams (independent or scalable). However, this approach will lead to significant parameter overhead and redundant bitstream storage. Therefore, some studies have tried to enhance the efficiency of multi-task compression by extracting a unified compact representation from the image. Generally, these methods rely on a general encoder for feature extraction and deploy different decoding models to support the corresponding visual tasks. However, these methods still require the use of multiple task-specific decoders, which has the problem of low parameter efficiency. The latest research has begun to explore the design of a unified compression model for multiple tasks, in which the encoder generates a unified representation of the image, while the decoder focuses on task-oriented reconstruction. Common techniques include conditional generation and residual prediction, which achieve user-controllable adjustments to the decoded image through task-driven guidance. Although these methods are close to the multi-decoder model method in performance, they only focus on reconstruction based on human vision requirements, such as perceptual quality-oriented and fidelity-oriented. More importantly, in the context of multi-task learning, unified compression models for multiple tasks also face the challenge of multi-objective optimization, and their performance is difficult to achieve the results of independent optimization for each task. The key to this problem is that existing research treats unified representations between different tasks indiscriminately, which requires finding an accurate Pareto front between tasks. This usually requires the use of many complex techniques for optimization. Once the optimal solution is not found, it may lead to a significant decline in performance, and as the number of tasks increases, the difficulty of optimization search also increases sharply. So far, few methods consider the correlation of features between tasks to ensure efficient multi-task collaboration, while optimizing features considering the specificity of different tasks to achieve efficient multi-task expansion. Summary of the invention

[0003] Purpose of the invention: The technical problem to be solved by the present invention is to provide a multi-task image coding method based on multi-branch aggregation in view of the deficiencies of the prior art. In the backbone network of the existing MetaFormer structure, the original single-channel MLP (Multi-Layer Perceptron) module is replaced with a multi-branch aggregation module, which is composed of multiple MLPs. Each MLP branch is optimized for a target task, and features are allocated to each MLP by designing an additional predictor to achieve migration between multiple targets. Experiments have shown that the scheme can achieve the most advanced performance in multi-task coding scenarios such as high fidelity, high authenticity, classification, and semantic segmentation, and can support smooth migration between various reconstruction tasks by adjusting the feature aggregation ratio.

[0004] The method of the present invention comprises the following steps:

[0005] Step 1, training phase: for the perceptual loss and visual task loss, the correlation of semantic features between machine vision tasks is used to train the basic image coding model, multi-branch aggregation module and importance predictor;

[0006] Step 2, testing phase: according to the specified task and feature aggregation ratio, the features are fed into the corresponding importance predictor, which calculates the output binary mask based on the importance of different features for different tasks;

[0007] Step 3: The multi-branch aggregation module splits the features into two groups of sub-features, namely main path features and bypass features, according to the binary mask;

[0008] Step 4: According to the specified task, select the corresponding bypass, and send the main road characteristics and bypass characteristics to the main road and the selected bypass respectively;

[0009] Step 5: The two groups of sub-features calculated by the main path and the bypass path are recombined according to the position and size of the original features, and subsequent decoding is completed to output the complete image.

[0010] Step 1 includes: considering the correlation of semantic features between various machine vision tasks, optimizing the basic image coding model and its main path by using an image coding training method optimized for perceptual quality loss, so that the basic image coding model can extract rich semantic information from the original image during the encoding process, realize feature reuse between multiple tasks, and reconstruct images with excellent perceptual quality during decoding; after obtaining the perceptual quality optimized basic image coding model, adding a bypass for a specific task side by side with the main path to form a multi-branch aggregation module, and fine-tuning some parameters for the loss function of the specific task to obtain a multi-task image coding model that supports multi-task target decoding.

[0011] In step 1, the multi-task image coding model includes 1 basic image coding model and 6 multi-branch aggregation modules, wherein the basic image coding model is an existing learning image coding model such as TinyLIC, including 1 encoder, 1 decoder and 1 entropy model; the function of the encoder is to encode the image into a binary code stream, the function of the entropy model is to predict the entropy of the binary code stream, and the function of the decoder is to decode the binary code stream into an image; the multi-task image coding model is constructed based on the basic image coding model, the encoder includes 3 multi-branch aggregation modules, the decoder includes 3 multi-branch aggregation modules, and the function of the multi-branch aggregation module is to enable the codec to support multi-task encoding and decoding;

[0012] Each multi-branch aggregation module contains 1 main path, 3 side paths and 3 importance predictors. The main path supports high-fidelity image reconstruction, and the 3 side paths support low-distortion image reconstruction, classification image reconstruction, and semantic segmentation image reconstruction respectively. The importance predictor is used to assign semantic features to the main path and side paths.

[0013] The main path and the branch path each contain 2 fully connected layers and 1 activation layer, and the importance predictor contains 3 fully connected layers and 2 activation layers; the main path, the branch path, and the importance predictor are denoted as φ main ,φ side ,φ pred .

[0014] In step 1, the training phase includes:

[0015] Step 1-1, Phase 1: Using the Variational Auto-Encoder (VAE) and Generative Adversarial Network (GAN), combined with perceptual quality loss, the basic image coding model and its main path are rate-distortion optimized to obtain the basic image coding model for multi-task optimization;

[0016] Step 1-2, Phase 2: Add bypasses according to the preset tasks to form a multi-branch aggregation module, and use the loss function of each task to optimize the corresponding bypasses and importance predictors separately.

[0017] Step 1-1 includes: in the calculation process of each training batch Epoch, optimizing the basic image coding model according to the following loss function:

[0018]

[0019] Among them, the proportional loss function Used to optimize the importance predictor, generator loss function Used to optimize the generator G(·) and discriminator loss functions of the generative adversarial network GAN Used to optimize the discriminator D(·) of the generative adversarial network GAN, image coding loss function The encoder E(·), generator G(·) and entropy model P(·) used to optimize the basic image coding model; x, y、 They are the input image, the output reconstructed image, the VAE latent space features before quantization, and the VAE latent space features after quantization; It means that when x follows the probability distribution p x expectations at the time; Indicates when Obey the probability distribution p y ; S represents the number of stages using multi-branch aggregation in the basic image coding model, M (s) represents the mask M,H generated by the importance predictor of the s-th stage of the multi-task image coding model (s) and W (s) They represent the spatial height and width of the s-th stage of the multi-task image coding model, and ρ is the feature aggregation ratio; Represents the entropy model pair in the image coding model The predicted value of the bit rate size, represents x and The mean square error between ; is the learned patch image perceptual similarity (LPIPS); the bit rate loss weight when the encoding quality is q Generator loss weight λ G , perceptual quality loss weight λ perc , proportional loss weight λ ratio is a hyperparameter for adjusting the loss weight.

[0020] Step 1-2 includes: in the calculation process of each training batch Epoch, optimizing the multi-branch aggregation module according to the following optimization objectives:

[0021]

[0022] Among them, the bypass parameter set and the importance predictor parameter set Bypass φ side Optimal parameters and importance predictor φ pred The optimal parameters of the task loss function is the loss function of the machine vision task selected in the current training batch, and the task loss weight λ task is a hyperparameter for adjusting the loss weight.

[0023] In step 1, the loss function for a specific task is fine-tuned for some parameters, specifically including: in step 1-1, all model parameters are updated during the training process, and in step 1-2, only φ side and φ pred The parameters of are updated during the training process, and the rest of the parameters remain unchanged.

[0024] Step 2 includes: first, the user selects the corresponding importance predictor according to the task that the decoded image needs to support, and sets the feature aggregation ratio ρ according to the accuracy of the task to be supported. Then, the optimized importance predictor can generate the optimal binary mask M suitable for the current task and accuracy setting according to the importance of different features for different tasks; the importance predictor uses a non-parametric partial averaging layer To fuse multi-scale information; given the importance predictor intermediate feature tensor u with spatial size H×W×C′ in the real domain, the partial average layer The second half of the intermediate feature is globally averaged and then re-expanded to the spatial size of the feature as global information, while the first half remains unchanged as local information. The process is formulated as:

[0025]

[0026] Among them, u(h,w,c) and u(i,j,c) represent the real value corresponding to the spatial coordinates (h,w,c) and (i,j,c) in the tensor u respectively;

[0027] During the training process, since it is not differentiable to directly use binarization to obtain a binary mask from the importance score according to the feature aggregation ratio ρ, a set of learnable bias parameters b are first used to bias the output u′ of the last fully connected layer in the importance predictor to obtain an intermediate variable corresponding to each discrete ρ, and then Gumbel noise is introduced in the binarization process with a threshold τ = 0.5, that is, the Gumbel-Sigmoid technique is used to soften the sampling process, as shown in the following formula:

[0028] M = Gumbel-Sigmoid(u′+b,τ)

[0029] After training, the importance predictor is able to binarize the score map according to the feature aggregation ratio ρ, and the bias layer will be discarded.

[0030] In step 3, the multi-branch aggregation module divides the features sent to the multi-branch aggregation module into two parts according to the optimal binary mask M generated in step 2, wherein the features corresponding to 1 in the optimal binary mask M are main features, and the features corresponding to 0 are bypass features, and there is no intersection between the main features and the bypass features;

[0031] In step 4, a bypass φ optimized for the task is selected in the multi-branch aggregation module according to the task that the decoded image needs to support. side , and send the main road features and bypass features separated in step 3 to the main road φ main and bypass φ side ; Main road φ main and each bypass φ side The complexity of can be designed differently according to different task requirements, so that two groups of sub-features can be calculated independently and in parallel by multi-layer perceptrons in different branches;

[0032] In step 5, after the main path features and bypass features are calculated in the multi-branch aggregation module, they will be re-aggregated into new features of the same size and corresponding positions as the original features according to the spatial size and feature positions before splitting; the new features output by the multi-branch aggregation module will be sent to the subsequent multi-branch aggregation module, and steps 2 to 5 will be repeated until the decoding process is completed and the complete image is output.

[0033] The present invention also provides a multi-task image encoding device based on multi-branch aggregation, including a multi-branch aggregation module, an importance predictor and a corresponding training strategy;

[0034] The multi-branch aggregation module training strategy is used to perform the following steps:

[0035] Step 1, training phase: targeting the perceptual quality loss and visual task loss, the correlation of semantic features between machine vision tasks such as image reconstruction, image classification, and image segmentation is used to train the basic image coding model, multi-branch aggregation module, and importance predictor to form a multi-task image coding model;

[0036] The importance predictor is used to perform the following steps:

[0037] Step 2, testing phase: according to the specified task and feature aggregation ratio, the features are fed into the corresponding importance predictor, which calculates the output binary mask based on the importance of different features for different tasks;

[0038] The multi-branch aggregation module is used to perform the following steps:

[0039] Step 3: The multi-branch aggregation module splits the features into two groups of sub-features, namely main path features and bypass features, according to the binary mask;

[0040] Step 4: According to the specified task, select the corresponding bypass, and send the main road characteristics and bypass characteristics to the main road and the selected bypass respectively;

[0041] Step 5: The two groups of sub-features calculated by the main path and the bypass path are recombined according to the position and size of the original features, and subsequent decoding is completed to output the complete image.

[0042] Beneficial effects: Considering that the features in the image coding model have different importance for different tasks, the present invention can process the features non-uniformly. By using the importance predictor proposed in the present invention, the model can utilize the importance of different features in different tasks, reuse the common semantic features of the tasks, and allocate features in the main path and bypass according to the required tasks and precision, so that the model can support seamless conversion of multiple tasks in an integrated manner. At the same time, by utilizing the feature correlation between tasks, the present invention proposes a two-stage optimization strategy with partial parameter fine-tuning function to overcome the multi-objective optimization problem in multi-task coding, so that the multi-branch aggregation module can support new tasks by adding bypasses and corresponding importance predictors and using partial parameter optimization, without optimizing separate specific task models, greatly reducing the optimization difficulty and training cost. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] The present invention will be further described in detail below in conjunction with the accompanying drawings and specific embodiments, and the above and / or other advantages of the present invention will become more clear.

[0044] Figure 1 It is a flow chart of the overall method of the present invention.

[0045] Figure 2 This is a diagram of the TinyLIC system using a multi-branch aggregation module.

[0046] Figure 3a It is a schematic diagram of a single branch block.

[0047] Figure 3b is a schematic diagram of a multi-branch block.

[0048] Figure 3c is a schematic diagram of a multi-branch aggregation module.

[0049] Figure 3d is a schematic diagram of the importance prediction module.

[0050] Figure 4a , Figure 4b , Figure 4c , Figure 4d , Figure 4e , Figure 4f These are the quantitative test results of the present invention on various data sets.

[0051] Figure 5 This is the visualization effect of the present invention on the multi-realistic image reconstruction task.

[0052] Figure 6This is the visualization effect of the present invention on the image classification task.

[0053] Figure 7 This is the visualization effect of the present invention on the semantic segmentation task. DETAILED DESCRIPTION

[0054] like Figure 1 As shown, an embodiment of the present invention provides a multi-task image encoding method based on multi-branch aggregation, comprising the following steps:

[0055] Step 1, training phase: targeting the perceptual quality loss and visual task loss, the correlation of semantic features between machine vision tasks such as image reconstruction, image classification, and image segmentation is used to train the basic image coding model, multi-branch aggregation module, and importance predictor to form a multi-task image coding model;

[0056] Step 2, testing phase: according to the specified task and feature aggregation ratio, the features are fed into the corresponding importance predictor, which calculates the output binary mask based on the importance of different features for different tasks;

[0057] Step 3: The multi-branch aggregation module splits the features into two groups of sub-features, namely main path features and bypass features, according to the binary mask;

[0058] Step 4: According to the specified task, select the corresponding bypass, and send the main road characteristics and bypass characteristics to the main road and the selected bypass respectively;

[0059] Step 5: The two groups of sub-features calculated by the main path and the bypass path are recombined according to the position and size of the original features, and subsequent decoding is completed to output the complete image.

[0060] Step 1 includes: considering the correlation between semantic features of various machine vision tasks, optimizing the basic image coding model and its main path by using an image coding training method optimized for perceptual quality loss, so that the basic image coding model can extract rich semantic information from the original image during the encoding process, realize feature reuse among multiple tasks, and reconstruct images with excellent perceptual quality during decoding; after obtaining the perceptual quality optimized basic image coding model, adding a bypass for a specific task in parallel with the main path to form a multi-branch aggregation module, and fine-tuning some parameters for the loss function of the specific task to obtain a multi-task image coding model that supports multi-task target decoding.

[0061] The multi-task image coding model includes a basic image coding model and six multi-branch aggregation modules, wherein the basic image coding model is an existing learning image coding model such as TinyLIC, including an encoder, a decoder and an entropy model; Figure 2As shown in Figure 1, there are 3 multi-branch aggregation modules in the encoder and 3 multi-branch aggregation modules in the decoder. The function of the encoder is to encode the image into a binary code stream, the function of the entropy model is to predict the entropy of the binary code stream, the function of the decoder is to decode the binary code stream into an image, and the role of the multi-branch aggregation module is to enable the codec to support multi-task encoding and decoding. Each multi-branch aggregation module contains 1 main path, 3 bypass paths and 3 importance predictors, as shown in Figure 1. Figure 3c As shown in the figure, the main path supports high-fidelity image reconstruction, and the three side paths support low-distortion image reconstruction, classification image reconstruction, and semantic segmentation image reconstruction. The importance predictor is used to assign semantic features to the main path and the side paths. Both the main path and the side paths contain 2 fully connected layers and 1 activation layer, and the importance predictor consists of 3 fully connected layers and 2 activation layers. The main path, the side path, and the importance predictor are denoted as φ main ,φ side ,φ pred .

[0062] In the training phase of step 1, in order to stabilize the training process, the present invention is divided into two training phases. First, the basic image coding model and its main path are optimized for perceptual quality loss. Then, a bypass is added in parallel with the main path for different tasks to form a multi-branch aggregation module, and the corresponding bypass and importance predictor are optimized for different tasks, specifically including:

[0063] Step 1-1, using a variational auto-encoder (VAE) and a generative adversarial network (GAN), combined with perceptual quality loss, to perform rate-distortion optimization on the basic image coding model and its main path, and obtain a basic image coding model for multi-task optimization;

[0064] Step 1-2, add bypasses according to the preset tasks to form a multi-branch aggregation module, and use the loss function of each task to optimize the corresponding bypasses and importance predictors respectively.

[0065] For step 1-1, during the calculation of each training batch Epoch, the basic image encoding model is optimized according to the following loss function:

[0066]

[0067] Among them, the proportional loss function Used to optimize the importance predictor, generator loss function Used to optimize the generator G(·) and discriminator loss functions of the generative adversarial network GAN Used to optimize the discriminator D(·) of the generative adversarial network GAN, image coding loss function The encoder E(·), generator G(·) and entropy model P(·) used to optimize the basic image coding model; x, y、 They are the input image, the output reconstructed image, the VAE latent space features before quantization, and the VAE latent space features after quantization; and It means that when x and Subject to probability distribution p x and p y ; S represents the number of stages using multi-branch aggregation in the basic image coding model, M (s) represents the mask M,H generated by the importance predictor of the s-th stage of the multi-task image coding model (s) and W (s) They represent the spatial height and width of the s-th stage of the multi-task image coding model, and ρ is the feature aggregation ratio; Represents the entropy model pair in the image coding model The predicted value of the bit rate size, represents x and The mean square error between ; is the learned patch image perceptual similarity (LPIPS); the bit rate loss weight when the encoding quality is q Generator loss weight λ G , perceptual quality loss weight λ perc , proportional loss weight λ ratio is a hyperparameter for adjusting the loss weight.

[0068] For steps 1-2, during the calculation of each training batch Epoch, the multi-branch aggregation module is optimized according to the following optimization objectives:

[0069]

[0070] Among them, the bypass parameter set and the importance predictor parameter set Bypass φ side and the importance predictor φ pred The optimal parameters of the task loss function is the loss function of the machine vision task selected in the current training batch, and the task loss weight λ task is a hyperparameter for adjusting the loss weight.

[0071] In step 1-1, all model parameters are updated during the training process, while in step 1-2, only φ side and φpred The parameters of are updated during the training process, and the rest of the parameters remain unchanged.

[0072] Step 2 includes: first, manually select the corresponding importance predictor according to the task that the decoded image needs to support, and manually set the feature aggregation ratio ρ according to the accuracy of the task that needs to be supported. Then, the optimized importance predictor can generate the optimal binary mask M suitable for the current task and accuracy setting according to the importance of different features for different tasks. Among them, the importance predictor is composed of multiple cascaded fully connected layers and activation layers, such as Figure 3d As shown, the features in the input image coding model are outputted as the corresponding binary mask M. In addition to the three parameterized fully connected layers, the present invention also uses a non-parametric partial average layer To fuse multi-scale information. Given an intermediate feature tensor u of the importance predictor with a spatial size of H×W×C′ in the real domain, the partial averaging layer The second half of the intermediate feature is globally averaged and then re-expanded to the spatial size of the feature as global information, while the first half remains unchanged as local information. The process is formulated as:

[0073]

[0074] Among them, u(h,w,c) and u(i,j,c) represent the real value corresponding to the spatial coordinates (h,w,c) and (i,j,c) in the tensor u respectively;

[0075] During the training process, since it is not differentiable to directly use binarization to obtain a binary mask from the importance score according to the feature aggregation ratio ρ, a set of learnable bias parameters b are first used to bias the output u′ of the last fully connected layer in the importance predictor to obtain an intermediate variable corresponding to each discrete ρ, and then Gumbel noise is introduced in the binarization process with a threshold τ = 0.5, that is, the Gumbel-Sigmoid technique is used to soften the sampling process, as shown in the following formula:

[0076] M = Gumbel-Sigmoid(u′+b,τ)

[0077] After training, the importance predictor can binarize the score map according to the feature aggregation ratio ρ. Figure 3d The bias layer of the medium importance predictor will be discarded.

[0078] In step 3, the multi-branch aggregation module divides the features sent to the multi-branch aggregation module into two parts according to the optimal binary mask M generated in step 2, wherein the features corresponding to 1 in the optimal binary mask M are main features, and the features corresponding to 0 are bypass features, and there is no intersection between the main features and the bypass features;

[0079] In step 4, according to the task that needs to be supported by the decoded image, a bypass φ optimized for the task is selected in the multi-branch aggregation module. side , and send the main road features and bypass features separated in step 3 to the main road φ main and bypass φ side , main road φ main and each bypass φ side The complexity of can be designed differently according to different task requirements, so that the two groups of sub-features can be calculated independently and in parallel by multi-layer perceptrons in different branches.

[0080] In step 5, after the two groups of sub-features are calculated in the multi-branch aggregation module, they will be re-aggregated into new features of the same size and corresponding positions as the original features according to the spatial size and feature position before splitting. The new features output by the multi-branch aggregation module will be sent to the subsequent multi-branch aggregation module, and steps 2 to 5 will be repeated until the decoding process is completed and the complete image is output.

[0081] The present invention also provides a multi-task image encoding device based on multi-branch aggregation, including a multi-branch aggregation module, an importance predictor and a corresponding training strategy;

[0082] The multi-branch aggregation module training strategy is used to perform the following steps:

[0083] Step 1, training phase: targeting the perceptual quality loss and visual task loss, the correlation of semantic features between machine vision tasks such as image reconstruction, image classification, and image segmentation is used to train the basic image coding model, multi-branch aggregation module, and importance predictor to form a multi-task image coding model;

[0084] The importance predictor is used to perform the following steps:

[0085] Step 2, testing phase: according to the specified task and feature aggregation ratio, the features are fed into the corresponding importance predictor, which calculates the output binary mask based on the importance of different features for different tasks;

[0086] The multi-branch aggregation module is used to perform the following steps:

[0087] Step 3: The multi-branch aggregation module splits the features into two groups of sub-features, namely main path features and bypass features, according to the binary mask;

[0088] Step 4: According to the specified task, select the corresponding bypass, and send the main road characteristics and bypass characteristics to the main road and the selected bypass respectively;

[0089] Step 5: The two groups of sub-features calculated by the main path and the bypass path are recombined according to the position and size of the original features, and subsequent decoding is completed to output the complete image.

[0090] In a specific embodiment of the present invention, a multi-task image encoding method based on multi-branch aggregation is provided, comprising:

[0091] (1) Figure 2 , Figure 3a and Figure 3b As shown in the figure, a scaling factor (SF) is added to TinyLIC to support variable bit rate, and a multi-branch aggregation module is added to the 1st, 2nd and 3rd stages of the codec, including a perception optimization main path, a mean square error optimization bypass, a classification optimization bypass and a semantic segmentation optimization bypass, which support high fidelity, low distortion, classification and semantic segmentation tasks respectively. The rest of the components remain consistent with the original TinyLIC.

[0092] During the training process of step 1-1, the perceptual quality loss The image coding model is optimized to obtain the basic image coding model for subsequent multi-task expansion. During the training process of step 1-2, some parameters of the model are fine-tuned, and only φ is updated. side and φ pred The parameters of the training set are changed, and the other parameters remain unchanged. In this embodiment, a comprehensive dataset is used for perceptual quality and mean square error optimization, including Flicker2W, DIV2K and CLIC training sets; ImageNet-1K dataset is used for classification task optimization; ADE20K dataset is used for semantic segmentation task optimization. The hyperparameters are set as follows: λ G =2.56,λ perc =4.26,λ task =1,λ ratio = 10. The code rate point q is randomly selected from integers between 1 and 8, and the aggregation ratio ρ is randomly sampled from the real number interval [0,1]. The number of iterations of step 1-1 is set to 3M, and the first 1.5M does not include The second half contains The number of iterations of steps 1-2 is set to 500K. In each training step, the initial learning rate is set to 10 -4 , decays to 10 in the last 25% of iterations -5 The training images are randomly cropped to 256×256×3 and random horizontal flipping is applied. ConvNeXt-Tiny is used as the classification model when training the classification bypass, and PSPNet50 is used as the segmentation model when training the semantic model.

[0093] (2) When the model trained in step (1) is actually used, it is necessary to specify the task, manually select the corresponding importance predictor, and manually set the aggregation ratio ρ according to the required task accuracy. Then the optimized importance predictor can generate the optimal binary mask M suitable for the current task and accuracy setting according to the importance of different features for different tasks, such as Figure 3d shown.

[0094] (3) The multi-branch aggregation module will divide the features sent to the module into two parts according to the optimal binary mask M generated in step (2), where the features corresponding to 1 in M ​​are main features, and the features corresponding to 0 are side features. There is no intersection between the main features and the side features, such as Figure 3c shown.

[0095] (4) According to the task that the decoded image needs to support, a bypass φ optimized for the task is selected in the multi-branch aggregation module. side , and send the main path features and bypass features separated in step (3) to the main path φ main and bypass φ side , main road φ main and each bypass φ side The complexity of can be designed differently according to different task requirements, so that the two groups of sub-features can be calculated independently and in parallel by multi-layer perceptrons in different branches, such as Figure 3c shown.

[0096] (5) After the two groups of sub-features are calculated in the multi-branch aggregation module, they will be re-aggregated into new features of the same size and corresponding position as the original features according to the spatial size and feature position before splitting, such as Figure 3c The new features output by the multi-branch aggregation module will be sent to the subsequent multi-branch aggregation module, and steps 2 to 5 will be repeated until the decoding process is completed and the complete image is output.

[0097] The performance of this embodiment can be seen Figure 4a , Figure 4b , Figure 4c , Figure 4d , Figure 4e , Figure 4f , Figure 5 , Figure 6 and Figure 7Among them, the indicator Bit-rate is the bit rate, the unit is bpp (bit per pixel), the indicators PSNR, LPIPS, FID, Top-1Acc, and mIoU represent peak signal-to-noise ratio, learnable image block perceptual similarity, Fréchet Inception distance, top-1 accuracy, and average intersection-over-union ratio, respectively. CRDR, DIRAC, MRIC, HiFiC, TransTIC, SFT, and VTM are the most advanced baseline comparison methods. After the present invention is added to the image coding model TinyLIC, the model can support smooth migration between multi-task reconstruction targets by adjusting the feature aggregation ratio. The final test shows that on the CLIC dataset, when the bit rate is 0.0789bpp, the variable range of PSNR (high quality index) is 29.49-30.83dB, and the variable range of FID (low quality index) is 4.52-45.76; on the Kodak dataset, when the bit rate is 0.1121bpp, the variable range of PSNR is 26.61-28.10dB, and the variable range of LPIPS (low quality index) is 0.1084-0.2758; on the ImageNet dataset, when the bit rate is 0.1521bpp, the variable range of classification accuracy (high quality index) is 60.59-76.77%; on the ADE20K dataset, when the bit rate is 0.0948bpp, the variable range of mIoU (high quality index) is 28.34-37.41%.

[0098] An embodiment of the present application also provides a computer program product, including computer instructions, which are stored in a computer-readable storage medium. When a processor of an electronic device reads the computer instructions from the computer-readable storage medium, the processor executes the computer instructions, so that the electronic device performs the steps of the method.

[0099] The present invention provides a multi-task image coding method and device based on multi-branch aggregation. There are many methods and ways to implement the technical solution. The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the principle of the present invention. These improvements and modifications should also be regarded as the protection scope of the present invention. All components not specified in this embodiment can be implemented by existing technologies.

Claims

1. A multi-task image coding method based on multi-branch aggregation, characterized in that: The following steps are involved: Step 1, training phase: for the perceptual loss and visual task loss, the correlation of semantic features between machine vision tasks is used to train the basic image coding model, multi-branch aggregation module and importance predictor; Step 2, testing phase: according to the specified task and feature aggregation ratio, the features are fed into the corresponding importance predictor, which calculates the output binary mask based on the importance of different features for different tasks; Step 3: The multi-branch aggregation module splits the features into two groups of sub-features, namely main path features and bypass features, according to the binary mask; Step 4: According to the specified task, select the corresponding bypass, and send the main road characteristics and bypass characteristics to the main road and the selected bypass respectively; Step 5: The two groups of sub-features calculated by the main path and the bypass path are recombined according to the position and size of the original features, and subsequent decoding is completed to output the complete image.

2. The method according to claim 1, characterized in that Step 1 includes: considering the correlation of semantic features between various machine vision tasks, optimizing the basic image coding model and its main path by using an image coding training method optimized for perceptual quality loss, so that the basic image coding model can extract rich semantic information from the original image during the encoding process, realize feature reuse between multiple tasks, and reconstruct images with excellent perceptual quality during decoding; after obtaining the perceptual quality optimized basic image coding model, adding a bypass for a specific task side by side with the main path to form a multi-branch aggregation module, and fine-tuning some parameters for the loss function of the specific task to obtain a multi-task image coding model that supports multi-task target decoding.

3. The method according to claim 2, characterized in that In step 1, the multi-task image coding model includes 1 basic image coding model and 6 multi-branch aggregation modules, wherein the basic image coding model includes 1 encoder, 1 decoder and 1 entropy model; the function of the encoder is to encode the image into a binary code stream, the function of the entropy model is to predict the entropy of the binary code stream, and the function of the decoder is to decode the binary code stream into an image; the multi-task image coding model is constructed based on the basic image coding model, the encoder includes 3 multi-branch aggregation modules, the decoder includes 3 multi-branch aggregation modules, and the function of the multi-branch aggregation module is to enable the codec to support multi-task encoding and decoding; Each multi-branch aggregation module contains 1 main path, 3 side paths and 3 importance predictors. The main path supports high-fidelity image reconstruction, and the 3 side paths support low-distortion image reconstruction, classification image reconstruction, and semantic segmentation image reconstruction respectively. The importance predictor is used to assign semantic features to the main path and side paths. The main path and the branch path each contain 2 fully connected layers and 1 activation layer, and the importance predictor contains 3 fully connected layers and 2 activation layers; the main path, the branch path, and the importance predictor are denoted as φ main ,φ side ,φ pred .

4. The method according to claim 3, characterized in that In step 1, the training phase includes: Step 1-1, stage 1: using variational autoencoders and generative adversarial networks, combined with perceptual quality loss, rate-distortion optimization is performed on the basic image coding model and its main path to obtain a basic image coding model for multi-task optimization; Step 1-2, Phase 2: Add bypasses according to the preset tasks to form a multi-branch aggregation module, and use the loss function of each task to optimize the corresponding bypasses and importance predictors separately.

5. The method according to claim 4, characterized in that Step 1-1 includes: in the calculation process of each training batch Epoch, optimizing the basic image coding model according to the following loss function: Among them, the proportional loss function Used to optimize the importance predictor, generator loss function Loss function for optimizing the generator G(·) and discriminator of the generative adversarial network Used to optimize the discriminator D(·) of the generative adversarial network, image coding loss function The encoder E(·), generator G(·) and entropy model P(·) used to optimize the basic image coding model; x, y、 They are the input image, the output reconstructed image, the VAE latent space features before quantization, and the VAE latent space features after quantization; It means that when x follows the probability distribution p x expectations at the time; Indicates when Subject to probability distribution p y ; S represents the number of stages using multi-branch aggregation in the basic image coding model, M (s) represents the mask M,H generated by the importance predictor of the s-th stage of the multi-task image coding model (s) and W (s) They represent the spatial height and width of the s-th stage of the multi-task image coding model, and ρ is the feature aggregation ratio; Represents the entropy model pair in the image coding model The predicted value of the bit rate size, represents x and The mean square error between ; is the perceptual similarity of learnable image blocks; the bit rate loss weight when the encoding quality is q Generator loss weight λ G , perceptual quality loss weight λ perc , proportional loss weight λ ratio is a hyperparameter for adjusting the loss weight.

6. The method according to claim 5, characterized in that Step 1-2 includes: in the calculation process of each training batch Epoch, optimizing the multi-branch aggregation module according to the following optimization objectives: Among them, the bypass parameter set and the importance predictor parameter set Bypass φ side Optimal parameters and importance predictor φ pred The optimal parameters of the task loss function is the loss function of the machine vision task selected in the current training batch, and the task loss weight λ task is a hyperparameter for adjusting the loss weight.

7. The method according to claim 6, characterized in that In step 1, the loss function for a specific task is fine-tuned for some parameters, specifically including: in step 1-1, all model parameters are updated during the training process, and in step 1-2, only φ side and φ pred The parameters of are updated during the training process.

8. The method according to claim 7, characterized in that Step 2 includes: first, the user selects the corresponding importance predictor according to the task that the decoded image needs to support, and sets the feature aggregation ratio ρ according to the accuracy of the task to be supported. Then, the optimized importance predictor can generate the optimal binary mask M suitable for the current task and accuracy setting according to the importance of different features for different tasks; the importance predictor uses a non-parametric partial averaging layer To fuse multi-scale information; given the real domain space size is H×W×C ′ The importance predictor intermediate feature tensor u, partial averaging layer The second half of the intermediate feature is globally averaged and then re-expanded to the spatial size of the feature as global information, while the first half remains unchanged as local information. The process is formulated as: Among them, u(h,w,c) and u(i,j,c) represent the real value corresponding to the spatial coordinates (h,w,c) and (i,j,c) in the tensor u respectively; During training, a set of learnable bias parameters b is first used to bias the output u of the last fully connected layer in the importance predictor ′ , thereby obtaining the intermediate variable corresponding to each discrete ρ, and then introducing Gumbel noise in the binarization process of the threshold τ, that is, using the Gumbel-Sigmoid technique to soften the sampling process, as shown in the following formula: M=Gumbel-Sigmoid(u ′ +b,τ) After training, the importance predictor is able to binarize the score map according to the feature aggregation ratio ρ, and the bias layer will be discarded.

9. The method according to claim 8, characterized in that In step 3, the multi-branch aggregation module divides the features sent to the multi-branch aggregation module into two parts according to the optimal binary mask M generated in step 2, wherein the features corresponding to 1 in the optimal binary mask M are main features, and the features corresponding to 0 are bypass features, and there is no intersection between the main features and the bypass features; In step 4, a bypass φ optimized for the task is selected in the multi-branch aggregation module according to the task that the decoded image needs to support. side , and send the main road features and bypass features separated in step 3 to the main road φ main and bypass v side ; In step 5, after the main path features and bypass features are calculated in the multi-branch aggregation module, they will be re-aggregated into new features of the same size and corresponding positions as the original features according to the spatial size and feature positions before splitting; the new features output by the multi-branch aggregation module will be sent to the subsequent multi-branch aggregation module, and steps 2 to 5 will be repeated until the decoding process is completed and the complete image is output.

10. A multi-task image encoding device based on multi-branch aggregation implemented by the method according to any one of claims 1 to 9, characterized in that: Including multi-branch aggregation module and importance predictor; The multi-branch aggregation module training strategy is used to perform the following steps: Step 1, training phase: targeting the perceptual quality loss and visual task loss, the correlation of semantic features between machine vision tasks such as image reconstruction, image classification, and image segmentation is used to train the basic image coding model, multi-branch aggregation module, and importance predictor to form a multi-task image coding model; The importance predictor is used to perform the following steps: Step 2, testing phase: according to the specified task and feature aggregation ratio, the features are fed into the corresponding importance predictor, which calculates the output binary mask based on the importance of different features for different tasks; The multi-branch aggregation module is used to perform the following steps: Step 3: The multi-branch aggregation module splits the features into two groups of sub-features, namely main path features and bypass features, according to the binary mask; Step 4: According to the specified task, select the corresponding bypass, and send the main road characteristics and bypass characteristics to the main road and the selected bypass respectively; Step 5: The two groups of sub-features calculated by the main path and the bypass path are recombined according to the position and size of the original features, and subsequent decoding is completed to output the complete image.

Images