About some failed compress/decompress results

[JayceJiangWHU]

Discussed in #319

^{Originally posted by JayceJiangWHU November 28, 2024}
I suppose this should be a bug or something I've neglected. When testing with pretrained models(e.g. mbt2018-mean) on Kodak. I firstly use model.compress() to encode the image into bitestream. And then use model.decompress() to reconstruct the image. This works perfectly fine with some Kodak images. However the other results seem to be quite distorted including some blurry and black blocks. And if I run again the locations of distortion are different everytime. However when testing with compressai.utils.eval_model.__main__.py everything goes on well. What could possibly be wrong?
Here are two examples from Kodak.


Here is a snippet of my codes

import torch
import torch.nn.functional as F
from torchvision.transforms import ToTensor, ToPILImage
import compressai
from compressai.zoo import image_models
from compressai.ops import compute_padding
from PIL import Image
from metrics import bit_per_pixel, compute_psnr, compute_ms_ssim
from tools import load_args
import matplotlib.pyplot as plt
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
compressai.set_entropy_coder('ans')


def main():
    filepath = r'saved/mbt2018_mean_0.01_batch64'
    device = torch.device('cuda')

    model = image_models["mbt2018-mean"](quality=2, metric='mse', pretrained=True, progress=False).to(device)
    model.eval()

    test_path = r'../Kodak'
    view_type = 'reconstruction'
    # view_type = 'inference'
    assert view_type in ['reconstruction', 'inference']

    test_res_path = os.path.join(filepath, f'test_{os.path.basename(test_path)}')
    if not os.path.exists(test_res_path):
        os.mkdir(test_res_path)

    for img_name in os.listdir(test_path):
        img_path = os.path.join(test_path, img_name)
        img = Image.open(img_path).convert('RGB')
        with torch.no_grad():
            x = ToTensor()(img).unsqueeze(0).to(device)
            N, _, H, W = x.shape
            num_pixels = N * H * W
            pad, unpad = compute_padding(H, W, min_div=2 ** 6)  # pad to allow 6 strides of 2
            x_padded = F.pad(x, pad, mode="constant", value=0)

            outputs = model(x_padded)
            rate_estimate = bit_per_pixel(outputs)

            out_enc = model.compress(x)

            rate_true = sum(len(s[0]) for s in out_enc["strings"]) * 8.0 / num_pixels

            print(f'Image: {img_name}----estimated rate: {rate_estimate:.4f}\ttrue rate:{rate_true:.4f}')
            out_dec = model.decompress(out_enc['strings'], out_enc['shape'])
            out_dec["x_hat"] = F.pad(out_dec["x_hat"], unpad)
            psnr = compute_psnr(x, out_dec['x_hat'])
            ms_ssim = compute_ms_ssim(x, out_dec['x_hat'], data_range=1.0)

            img_recon = ToPILImage()(out_dec['x_hat'].squeeze())
            img_recon.save(os.path.join(test_res_path, img_name.split('.')[0]+'_Recon.png'))
```</div>