LFC-UNet: learned lossless medical image fast compression with U-Net

Traditional image codecs

In the field of image compression, traditional lossless image encoders are commonly categorized into transform-based methods and prediction-based methods. These two categories possess unique characteristics, offering flexible solutions for various application scenarios. In this section, we will introduce several representative traditional lossless image encoders: JPEG-XL (Alakuijala et al., 2019), PNG (Boutell, 1997), WebP (Version, WebEngines Blazer Platform, 2000), and BPG (Fabrice Bellard contributors, 2023). These encoders exhibit significant features in terms of image quality, compression efficiency, and application domains.

JPEG-XL is a transform-based lossless image encoder. Its optimization strategy includes the use of variable-sized DCT (Khayam, 2003) blocks, allowing better adaptation to different frequency components of the image. Local adaptive quantization further enhances image quality, while additional DC coefficient prediction modes improve prediction accuracy. The adoption of asymmetric numeral systems (ANSE) enhances coding efficiency, achieving better compression ratios.

PNG, on the other hand, employs prediction-based lossless image encoding. It represents each pixel by the difference between the current pixel and its surrounding neighboring pixels through differential prediction. This differential data is encoded using various prediction filters, achieving lossless compression. PNG maintains high image quality while offering a high compression ratio, making it suitable for scenarios requiring high-quality image transmission, such as digital art and medical imaging fields.

WebP is a lossless image compression format developed by Google, combining prediction-based methods with techniques such as context modeling. It strikes a balance between small file sizes and high-quality images and finds widespread usage in web images and mobile applications. WebP’s features make it an ideal choice for web image transmission, providing users with faster loading speeds and a better user experience.

BPG is an emerging lossless image compression format that combines prediction-based methods with advanced compression techniques, such as prediction filtering and motion compensation. BPG achieves relatively small file sizes while preserving high-quality images, making it popular in scenarios where maintaining image quality while reducing file size is crucial. These encoders offer diverse and efficient choices for lossless image compression in different domains.

In summary, traditional image codecs commonly employ two main techniques for lossless compression: transformation and prediction. These techniques aim to reduce redundancy in image information and enhance coding efficiency. Transformation, exemplified by methods such as discrete cosine transform (DCT) used in applications like JPEG-XL, WebP, and BPG, alters the representation of image data in the spatial domain to the frequency domain. This organized presentation of frequency components makes image data more compact, forming the basis for subsequent prediction steps, all while ensuring reversibility to avoid information loss. Prediction, on the other hand, leverages regularities or patterns in data to forecast future data. It achieves more efficient compression by encoding residuals that eliminate redundant information. For instance, PNG utilizes predictive methods with filtering based on the idea of differential encoding, reducing redundancy by computing the difference between the current pixel value and adjacent pixel values. Additionally, JPEG-XL employs DC coefficient prediction, analyzing information from neighboring blocks or rows to predict the DC coefficients of the current block. This strategy reduces the dynamic range of DC coefficients, thereby improving the coding efficiency for low-frequency information in the image.

Learning-based image compression

In early methods, image encoding was performed on a pixel-by-pixel basis, where the compression of each pixel relied on previously encoded pixels. For example, PixelRNN and PixelCNN modeled pixels as conditional products, expressed as p(x) = ∏p(x_i∣x₁, …, x_i−1), where x_i represent a single pixel. Subsequently, PixelCNN++ introduced several important optimizations on the original PixelCNN model, including improvements in generation efficiency and image quality. PixelCNN++ employed a segmental logistic mixture model for layer-wise modeling, introduced residual connections, and utilized techniques like multiscale generation to enhance both generation speed and image quality. However, it is worth noting that PixelCNN++ still involved network computations for each pixel, resulting in impractical inference times.

To enhance practicality, subsequent research proposed several improved methods. For instance, L3C introduced a practical compression framework utilizing hierarchical probability models. subimages were implicitly modeled through neural networks and aided by a feature extractor for prediction tasks. In contrast to the pixel-wise predictions of PixelRNN and PixelCNN, L3C predicted probabilities for all pixels, achieving more than two orders of magnitude acceleration.

Recent studies like LC-FDNet took into account the helpfulness of texture information in improving probability predictions. It adopted a method of predicting image splitting, dividing the image into high-frequency and low-frequency parts. It encoded the low-frequency region first, then utilized prior information from the low-frequency region to predict the high-frequency region, thereby improving the effectiveness of probability prediction.

In the field of medicine, Wang et al. (2023) and colleagues have proposed a new perspective. They divided the image into two types of subimages: the most significant bytes (MSB) subimage, representing the high-order bits, and the least significant bytes (LSB) subimage, representing the low-order bits. Due to the significantly lower bit rates required for compressed MSB subimages compared to compressed LSB subimages, they chose conventional methods to encode MSB. Simultaneously, they used MSB as input for a neural network to predict LSB. Subsequently, they encoded LSB.

In learning-based image compression methods, the primary research goal is to find an approach that achieves a balance between probability prediction accuracy and prediction time. From the initial autoregressive methods, predicting pixel by pixel, to the layered context (L3C) employing implicit modeling for hierarchical prediction, and further to segmented prediction methods proposed by researchers such as LC-FDNet and Kai Wang, it is evident that researchers have consistently pursued this balance.

Image decomposer

In this stage, we divide the image into subimages following the pattern illustrated in Fig. 1, where pixels in odd rows and odd columns are categorized as subimage A, pixels in odd rows and even columns are categorized as subimage B, and so on, resulting in four subimages. Subsequently, we perform lossless compression on subimage A pixels. We chose to use JPEG-XL to compress the subimage A images. This choice was made after a comprehensive comparison between traditional lossless codecs and learning-based ones, where we found that JPEG-XL yielded the best results. In the realm of traditional lossless codecs, JPEG-XL demonstrated superior encoding efficiency. Regarding learning-based lossless codecs, except for Iwave (Ma et al., 2020), other methods either performed worse than JPEG-XL or did not show a significant advantage over it. However, the runtime of Iwave was deemed unacceptable. For example, on an image with dimensions 768 × 768, in an environment with GeForce GTX 3060 and Intel Core i7- 12700, JPEG-XL took approximately 1.1 s, while Iwave’s processing time extended to a lengthy 160.4 s.

Learning-based compressor

In this section, we introduce the architecture of LFC-UNet as depicted in Fig. 2, which provides a detailed overview of the details outlined in Fig. 1. The objective of LFC-UNet is to compress the N-th subimage Y given the input x_in, which includes the N-1 previously compressed subimages Y. As described in Fig. 1, we employ three independent LFC-UNet networks, each tailored to handle a different number of input subimages. Consequently, these networks do not share parameters, as each of them focuses on a specific subimage.

Preprocessing: During the preprocessing stage, we first employ a neural network to extract features from the input subimages $x_{in}\in R^{\frac{W}{2}\times \frac{H}{2}\times N}$, where N represents the number of previously compressed subimages (N ∈  {1, 2, 3}). Subsequently, the extracted features are transformed into an intermediate matrix in the latent space, referred to as the “Middle Matrix.” This matrix serves as the input for subsequent probability predictions and image predictions.

Subimage prediction: In Fig. 2, “Predict Image” represents the prediction of the “Real Image.” After obtaining this prediction, we calculate the residual image residual = Real Image - Predict Image, which signifies the difference between the real subimage and the prediction. As residuals might include negative and floating-point numbers, we add 255 to all residuals and then quantize them (round to the nearest integer), ultimately constraining the range of residuals within integers from 0 to 511. This quantized residual, denoted as Q, is then passed on to the entropy coder for further processing.

Probability distribution prediction: In Fig. 2, “Probability” represents the probability prediction of the quantized residual image Q. In this step, we use the U-Net architecture shown in Fig. 3 for prediction. The input Middle Matrix undergoes four down process and four up process operations, and finally, through the Softmax function, the predicted probability distribution ensures that the sum of probabilities at each position equals 1. This results in the probability prediction of the quantized residual Q.

Loss function

In this section, we introduce the loss functions used in the training of LFC-UNet. It consists of two components: subimage prediction loss and probability prediction loss.

Subimage prediction loss: We define the subimage prediction loss as the absolute difference between the predicted pixel values and the original pixel values. Mathematically, it is defined as follows: ${\displaystyle L_{SPL}\left(y,\hat{y}\right)=\frac{1}{N}\sum _{i=1}^N\left|y_i-{\hat{y}}_i\right|}$ Where N represents the total number of pixels in the image, y represents the original pixel values, and $\hat{y}$ represents the predicted pixel values by the neural network.

Probability prediction loss: In the probability loss function, we employ weighted cross-entropy, defined mathematically as follows: ${\displaystyle L_{PPL}\left(p,\hat{p}\right)=-\frac{1}{N}\sum _{i=1}^N\sum _{j=1}^C{w}_j\cdot p_{i,j}\cdot log\left({\hat{p}}_{i,j}\right)}$ where N represents the total number of pixels in the image, C represents the number of classes, denoting the magnitude of residual values in the article, with a range from 0 to 510, p_i,j represents the probability that the i-th pixel actually belongs to class j. More precisely, p_i,j is equal to 1 if the i-th pixel belongs to class j, and 0 otherwise. ${\hat{p}}_{i,j}$ is the predicted probability of the i-th pixel belonging to class j by the model, and w_j is the weight for class j. This is a variant of cross-entropy loss, commonly used to address class imbalance issues. In our experiments, the neural network predicts pixel values, where 90% of the predictions differ from the actual pixel values by no more than 50. This results in a large concentration of label values around 255. Therefore, conventional cross-entropy loss performs sub-optimally in such situations. Altogether, we train our network with the loss: ${\displaystyle L={\lambda }_{SPL}\cdot L_{SPL}+{\lambda }_{PPL}\cdot L_{PPL}}$ where λ_SPL and λ_PPL are the balancing hyperparmeters. In our experiments, we set both λ_SPL and λ_PPL as 1.We opt for the normalization of the hyperparameters of the loss functions for each module to avoid unnecessary model optimization, thereby facilitating a more straightforward training and evaluation of the model. This approach represents a balance between parameter tuning and the pursuit of optimal performance.

Experimental setup

Compression result

Table 1 presents the comparative results obtained from the evaluation dataset as described. It is evident that our method excels in terms of performance, whether dealing with medical images or natural images in grayscale format, outperforming existing traditional codecs as well as learning-based codecs.

Notably, in datasets such as Lung CT, Head CT, and Lumbar Spine CT, there is a significant performance gap between traditional encoders and learning-based encoders. This disparity can be attributed to the broader grayscale range typically found in CT images, showcasing density variations within tissues, while natural images usually operate within the dynamic range of a camera, exhibiting a narrower grayscale spectrum. Traditional encoders often struggle with processing images featuring such extensive grayscale ranges. Our learning-based approach outperforms other algorithms, indicating its high practicality in the medical field.

Furthermore, our method achieves state-of-the-art performance in diverse grayscale natural images, demonstrating the model’s generalizability and robustness.

Specifically, when compared to the best learning-based methods: In the Chest X-ray dataset, our method surpasses LC-FDNet by 1.2%. In the Lung CT dataset, our method outperforms LC-FDNet by 1.6%. In the Head CT dataset, our method exceeds LC-FDNet by 3.8%. In the Lumbar Spine CT dataset, our method surpasses LC-FDNet by 2.1%. In the COCO dataset, our method outperforms LC-FDNet by 1.1%.

Additionally, when compared to the best traditional methods: In the Chest X-ray dataset, our method outperforms JPEG-XL by 2.4%. In the Lung CT dataset, our method exceeds FLIF by 6.0%. In the Head CT dataset, our method surpasses JPEG-XL by 11.9%. In the Lumbar Spine CT dataset, our method outperforms JPEG-XL by 8.9%. In the COCO dataset, our method excels by 3.7% compared to JPEG-XL.

Time result

Our setup consists of GeForce GTX 3060, Intel Core i7-12700, and 32GB of RAM. Table 2 presents the results of time comparison conducted on the specified evaluation dataset using CPU, with the unit measured in seconds.It is worth noting that, in the CPU comparison experiments, we exclusively compared our model with other learning-based methods and did not involve comparisons with traditional methods. This decision was made because traditional compression algorithms, such as JPEG XL, have already been optimized for CPU performance. These optimizations include the application of new instruction sets, multi-threading optimization, and adaptation to CPU hardware characteristics. In contrast, specialized instruction sets for neural network computations, such as CUDA, are designed explicitly for GPUs, and neural network designs involve a significant amount of matrix operations, which are not the strengths of CPUs. Therefore, we believe that conducting time comparisons with traditional methods lacks practical significance, and as a result, we chose not to perform such comparisons.

In Table 3, we exhibit the time comparison results using GPU on the same evaluation dataset, also measured in seconds. It is essential to note that L3C cannot operate on NVIDIA 30-series GPUs; therefore, our comparison was exclusively conducted with LC-FDNet.

In terms of the encoding method, we employed a serial encoding approach, encoding images one by one. Specifically, encoding input images of size x ∈ R^B×1×W×H , where B represents the batch size, W represents the image width, and H represents the image length. In this experiment, we set B to 1. In practical applications, the batch size will be adjusted based on the specific GPU memory configuration, potentially leading to even faster processing speeds.

Evidently from the data, our method significantly outperforms L3C and LC. This disparity can be attributed to our model’s remarkable ability to obtain all probability estimates and predicted pixels in a single inference, whereas L3C’s hierarchical design or LC’s high-low frequency decomposition necessitates multiple inferences.

Residual analyze

In Fig. 4, we utilized a 3D surface plot to illustrate the data, where the X and Y axes represent the two-dimensional coordinates of the residual plots, and the Z axis represents the Residual values. This graphical representation intuitively displays the residual distributions of the three subimages. Taking LungCT as an example, it can be observed that the residual values of all subimages are concentrated around 255. This indicates that the neural network has successfully transformed widely distributed pixel values into concentrated residual values. Additionally, it is noteworthy that the residuals of subimage D exhibit significant fluctuations, whereas the residual distributions of subimage C are relatively smooth. This implies that as the encoded subimages provide more information, the network’s accuracy in pixel estimation gradually improves. This observation effectively explains the sequence of improved compression efficiency as d → b → c. For encoded content that is relatively smooth, the neural network can more easily predict its probability estimation. Conversely, for content with significant fluctuations, the network struggles to obtain accurate probability estimates.

In Fig. 5, we presented the data using a histogram. Here, the X-axis represents the residual values, and the Y-axis represents the frequency of occurrence corresponding to each residual value. This histogram precisely illustrates the distribution of residuals and the frequency of different residual values. This further demonstrates that with an increase in the encoded content, the network’s estimation of pixel values becomes more accurate.

Residual probability analyze

We replaced the original images and selected relatively smooth residuals for encoding. Although this aids in the network’s probability estimation, it also leads to an excessive concentration of residual types, as demonstrated effectively in Fig. 6. Consequently, using a conventional cross-entropy loss function would make it challenging for the neural network to learn probability estimates for residual types that occur less frequently. To address this issue, we adopted a weighted cross-entropy loss function. In this experiment, we set the weights as illustrated in Fig. 6. Within the range of 0–99, representing cases where the absolute difference between the predicted and actual values falls between 156–255, these samples are extremely rare, so we set their weight to 8. In the range of 200–300, indicating cases where the absolute difference between predicted and actual values falls between 0–55, these samples are very common, so we set their weight to 1. The weights for other intervals are set in a similar manner, following this logic.

Loss function ablation study

We conducted ablation experiments on the Head CT dataset. In this experiment, Model 1 utilized a conventional cross-entropy loss function, while Model 2 employed a weighted cross-entropy loss function. Simultaneously, the neural network was trained in the same manner on the Flickr2k dataset. We evaluated the compression performance of subimages using these two models, measured in BPP, where smaller BPP values indicate more accurate probability estimation by the network. For detailed experimental results, please refer to Table 4. It is evident that the model using the weighted cross-entropy loss function outperformed the model using the conventional cross-entropy loss function by 3.5%.