About some failed decompress result

[JayceJiangWHU]

When testing with pretrained models(e.g. mbt2018-mean) on Kodak. I ran into some random distorted compress/decompress results like these. And if I ran these codes again I got different distortion everytime. However when testing with compressai.utils.eval_model.__main__.py everything went on well. What could possibly be wrong


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'))