Aligned Intra Prediction and Hyper Scale Decoder Under Multistage Context Model for JPEG AI

Learning-based image compression has raised increasing interests in the last few years. Currently, Joint Photographic Experts Group (JPEG) is working on the standardization of learning-based image compression as JPEG AI. It adopts a deep neural network based encoder-decoder architecture with hyperprior based probability formulation for entropy coding. JPEG AI currently contains two coding profiles: Base Operating Point (BaseOP) and High Operating Point (HighOP). Among the various techniques developed in JPEG AI, Multistage Context Model (MCM) was adopted as the context model to perform intra prediction in HighOP. It transforms the spatially progressive context prediction into sub-image feature prediction among channels via feature down-shuffling. However, in this prediction process, sub-image features are not spatially aligned to each other, and directly using the neighboring sub-image features cannot provide accurate prediction. Moreover, the distributions of residual features generated by MCM are also not consistent with that of the hyper scale decoder, which is used to construct the probability model in the entropy coding of residual features, leading to suboptimal residual coding. To address the above problems, we propose an Aligned Intra Prediction (AIP) and Aligned Hyper Scale Decoder (AHSD) under MCM for JPEG AI coding. AIP aligns the reference sub-image features to the to-be-predicted feature in MCM. AHSD further generates hyper scale features with matched distributions to the residual features. Experimental results demonstrate that the proposed method improves the coding performance by 1.3% in terms of BD-rate saving over the JPEG AI reference software.