evaluate.py

import os
import glob
import tensorflow as tf
import tensorflow_compression as tfc
from tensorflow.keras.optimizers.legacy import Adam
from zyc2022 import ZYC2022Model  
import torch
import lpips
import numpy as np
#from DISTS_tf import DISTS
from DISTS.DISTS_pytorch import DISTS
from torchmetrics.image.fid import FrechetInceptionDistance
from basicsr.metrics.niqe import calculate_niqe
from PIL import Image
D = DISTS()
def lossfid(a: torch.Tensor, b: torch.Tensor) -> float:
    fid = FrechetInceptionDistance()
    a2=torch.clamp(a * 255.0, 0, 255).to(torch.uint8)
    b2=torch.clamp(b * 255.0, 0, 255).to(torch.uint8)
    fid.update(a2, real=True)
    fid.update(b2, real=False)
    return fid.compute().item()
def lossdists(a: torch.Tensor, b: torch.Tensor) -> float:
    return D(a, b).item()
def main(file):
    tf.config.run_functions_eagerly(True)

    #D = DISTS()
    lpips_model = lpips.LPIPS(net='alex')  # or 'vgg'
    lpips_model.eval()

    # Load model
    # model_path = "model_path"
    # model = tf.keras.models.load_model(model_path)
    # print("✔️ Model loaded")
    model = ZYC2022Model(
        patchsize=256,
        lmbda=0.01,
        num_filters=192,
        latent_depth=320,
        hyperprior_depth=192,
        num_slices=10,
        max_support_slices=10,
        num_scales=64,
        scale_min=0.11,
        scale_max=256.0
    )
    model.compile(optimizer=Adam(learning_rate=1e-5))

    # Restore from checkpoint
    ckpt = tf.train.Checkpoint(model=model, optimizer=model.optimizer)
    checkpoint_path = f"./checkpoints/{file}"
    ckpt.restore(tf.train.latest_checkpoint(checkpoint_path)).expect_partial()
    print("✔️ Restored from checkpoint")
    model.em_z = tfc.ContinuousBatchedEntropyModel(
        model.hyperprior, coding_rank=3, compression=True, offset_heuristic=False
    )
    model.em_y = tfc.LocationScaleIndexedEntropyModel(
        tfc.NoisyNormal,
        num_scales=model.num_scales,
        scale_fn=model.scale_fn,
        coding_rank=3,
        compression=True
    )
    # Folder with input images
    # input_folder = "/compression/stfattn/val2017/test"   
    input_folder = "./validation"
    image_paths = sorted(glob.glob(os.path.join(input_folder, "*.png")))

    total_psnr = 0
    total_msssim = 0
    total_bpp = 0
    total_dists = 0
    total_lpips = 0
    total_niqe = 0
    count = 0

    original_images = []
    compressed_images = []

    for image_path in image_paths:
        print(f"Processing: {os.path.basename(image_path)}")

        # Read image
        x = tf.io.read_file(image_path)
        x = tf.image.decode_png(x, channels=3)
        
        h, w = tf.shape(x)[0], tf.shape(x)[1]
        if h % 256 != 0 or w % 256 != 0:
            print("Resizing image to 256x256")
            x = tf.image.resize(x, [256, 256])
            x = tf.cast(x, tf.uint8)

        # Compress
        tensors = model.compress(x)
        packed = tfc.PackedTensors()
        packed.pack(tensors)
        compressed_string = packed.string

        # Decompress
        x_hat = model.decompress(*tensors)

        # Compute metrics
        x_f = tf.cast(x, tf.float32)
        x_hat_f = tf.cast(x_hat, tf.float32)

        # FID - PyTorch implementation
        x_fid = tf.image.resize(x_f, [299, 299])
        x_hat_fid = tf.image.resize(x_hat_f, [299, 299])
        x_fid_torch = torch.from_numpy(x_fid.numpy()).permute(2, 0, 1).unsqueeze(0)
        x_hat_fid_torch = torch.from_numpy(x_hat_fid.numpy()).permute(2, 0, 1).unsqueeze(0)
        original_images.append(x_fid_torch)
        compressed_images.append(x_hat_fid_torch)

        # NIQE 
        niqe_score = calculate_niqe(x_hat_f.numpy(), 0)
        total_niqe += niqe_score

        psnr = tf.image.psnr(x_f, x_hat_f, 255.0)
        msssim = tf.image.ssim_multiscale(x_f, x_hat_f, 255.0)
        num_pixels = tf.reduce_prod(tf.shape(x)[:-1])
        bpp = len(compressed_string) * 8 / float(num_pixels)
        x_f = tf.cast(x, tf.float32) / 255.0
        x_hat_f = tf.cast(x_hat, tf.float32) / 255.0
        dists_value =tf.reduce_mean(D.get_score(x_f, x_hat_f))

        # Convert to NumPy → PyTorch
        x_torch = torch.from_numpy(x_f.numpy()).permute(2, 0, 1).unsqueeze(0) * 2 - 1
        x_hat_torch = torch.from_numpy(x_hat_f.numpy()).permute(2, 0, 1).unsqueeze(0) * 2 - 1
        # Compute LPIPS
        with torch.no_grad():
            lpips_score = lpips_model(x_torch, x_hat_torch).item()

        # print(f"LPIPS: {lpips_score:.4f}")
        total_psnr += psnr.numpy()
        total_msssim += msssim.numpy()
        total_bpp += bpp
        total_dists+= dists_value.numpy()
        total_lpips += lpips_score
        
        count += 1

        # Save decompressed image

        # x_hat_np = x_hat.numpy()
        # x_hat_pil = Image.fromarray(x_hat_np.astype(np.uint8))
        # image_path = os.path.basename(image_path)
        # output_folder = f"technick40_output/{file}"
        # output_path = os.path.join(output_folder, image_path)
        # os.makedirs(os.path.dirname(output_path), exist_ok=True)
        # x_hat_pil.save(output_path)
        # print(f"Saved decompressed image: {output_path}")

        encoded_img = tf.image.encode_png(x_hat)
        tf.io.write_file(f"{file}/{os.path.basename(image_path)}", encoded_img)

        # print("-----------------")
        # print(f"PSNR: {psnr.numpy():.2f}, MS-SSIM: {msssim.numpy():.4f}, BPP: {bpp:.4f}, DISTS: {dists_value.numpy():.4f}, NIQE: {niqe_score:.4f}, LPIPS: {lpips_score:.4f}")

    #FID calculation
    orig_stacked = torch.cat(original_images, dim=0)
    comp_stacked = torch.cat(compressed_images, dim=0)
    fid = lossfid(orig_stacked, comp_stacked)

    print("===========================")
    print(f"Evaluation Summary: {file}")
    print(f"Average PSNR: {total_psnr / count:.2f}")
    print(f"Average MS-SSIM: {total_msssim / count:.4f}")
    print(f"Average BPP: {total_bpp / count:.4f}")
    print(f"Average LPIPS: {total_lpips / count:.4f}")
    print(f"Average DISTS: {total_dists / count:.4f}")
    print(f"Average NIQE: {total_niqe / count:.4f}")
    print(f"FID: {fid:.4f}")

if __name__ == "__main__":
    # file_dir = ["0.01", "0.001", "0.001_attn", "0.1_Attn", "0.01_Attn", "0.01_lpip", "0.003", "0.003_Attn", 
    #             "0.003_Attn_10lpip", "0.0003_Attn_lpips", "0.007", "0.007_Attn", "0.007_Attn_new"]
    file_dir = ["0.007"]

    for file in file_dir:
        main(file)