first

.gitattributes

@@ -1 +1,2 @@
 *.png filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

README_demo.md

@@ -0,0 +1,5 @@
+# Light Field Image Compression using VAE
+
+The more details about our training and model will be made available after acceptance of our paper.
+
+[Project Page (Will be active after paper acceptance)](https://takhtardeshirsoheib.github.io/DUALF_D/index.html)

app.py

@@ -0,0 +1,338 @@
+# app.py
+# 1. IMPORTS
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import os
+import math
+import warnings
+from compressai.layers import GDN, conv3x3, subpel_conv3x3
+from compressai.entropy_models import EntropyBottleneck, GaussianConditional
+from skimage.metrics import structural_similarity as ssim
+from skimage.metrics import peak_signal_noise_ratio as psnr
+
+'''
+01 - Best for Low Bit Rates ModelvLowBit
+005 - Mid Level for Low Bit Rates ModelvMidBit
+001 - Mid Level for High Bit Rates ModelvHighBit
+0001 - Best for High Bit Rates ModelvBestHighBit
+'''
+
+'''
+QP - Smaller value is worst quality but best for storage
+'''
+
+warnings.filterwarnings("ignore", "Inputs have mismatched dtype", UserWarning)
+filt_n = 128
+latent_channels = 128
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+save_path = "./checkpoint/"
+
+# 3. MODEL DEFINITIONS (from model.py)
+def get_scale_table(min_val=0.11, max_val=256, levels=64):
+    """Get the scale table as a list of floats."""
+    return [float(f) for f in torch.exp(torch.linspace(math.log(min_val), math.log(max_val), levels))]
+
+class SpatialEncoder(nn.Module):
+    def __init__(self):
+        super(SpatialEncoder, self).__init__()
+        self.conv_layers_S1 = nn.Sequential(nn.Conv2d(3, filt_n, kernel_size=5, stride=1, padding=1, dilation=3), GDN(filt_n))
+        self.conv_layers_S2 = nn.Sequential(nn.Conv2d(filt_n, filt_n, kernel_size=5, stride=2, padding=1), GDN(filt_n))
+        self.conv_layers_S3 = nn.Sequential(nn.Conv2d(filt_n, filt_n, kernel_size=5, stride=2, padding=1), GDN(filt_n))
+        self.conv_layers_S4 = nn.Sequential(nn.Conv2d(filt_n, filt_n, kernel_size=5, stride=2, padding=1), GDN(filt_n))
+        self.conv_layers_S5 = nn.Sequential(nn.Conv2d(filt_n, 64, kernel_size=5, stride=3, padding=1), GDN(64))
+    def forward(self, x):
+        x = self.conv_layers_S1(x)
+        x = self.conv_layers_S2(x)
+        x = self.conv_layers_S3(x)
+        x = self.conv_layers_S4(x)
+        x = self.conv_layers_S5(x)
+        return x
+
+class AngularEncoder(nn.Module):
+    def __init__(self):
+        super(AngularEncoder, self).__init__()
+        self.conv_layers_A1 = nn.Sequential(nn.Conv2d(3, filt_n, kernel_size=3, stride=3, padding=1), GDN(filt_n))
+        self.conv_layers_A2 = nn.Sequential(nn.Conv2d(filt_n, filt_n, kernel_size=5, stride=2, padding=1), GDN(filt_n))
+        self.conv_layers_A3 = nn.Sequential(nn.Conv2d(filt_n, filt_n, kernel_size=5, stride=2, padding=1), GDN(filt_n))
+        self.conv_layers_A4 = nn.Sequential(nn.Conv2d(filt_n, 64, kernel_size=5, stride=2, padding=1), GDN(64))
+    def forward(self, x):
+        x = self.conv_layers_A1(x)
+        x = self.conv_layers_A2(x)
+        x = self.conv_layers_A3(x)
+        x = self.conv_layers_A4(x)
+        return x
+
+class HyperpriorNetwork(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.entropy_bottleneck = EntropyBottleneck(channels)
+        self.h_a = nn.Sequential(
+            conv3x3(channels, channels), nn.LeakyReLU(inplace=True),
+            conv3x3(channels, channels, stride=2), nn.LeakyReLU(inplace=True),
+            conv3x3(channels, channels, stride=2),
+        )
+        self.h_s = nn.Sequential(
+            conv3x3(channels, channels), nn.LeakyReLU(inplace=True),
+            subpel_conv3x3(channels, channels, 2), nn.LeakyReLU(inplace=True),
+            subpel_conv3x3(channels, channels, 2), nn.LeakyReLU(inplace=True),
+            conv3x3(channels, channels),
+        )
+    def forward(self, x):
+        z = self.h_a(x)
+        z_hat, z_likelihoods = self.entropy_bottleneck(z)
+        scales = torch.exp(self.h_s(z_hat))
+        return scales, z_likelihoods
+
+class Encoder(nn.Module):
+    def __init__(self, latent_channels):
+        super().__init__()
+        self.spatial_encoder = SpatialEncoder()
+        self.angular_encoder = AngularEncoder()
+        self.spatial_hyperprior = HyperpriorNetwork(64)
+        self.angular_hyperprior = HyperpriorNetwork(64)
+        self.entropy_model_s = GaussianConditional(get_scale_table())
+        self.entropy_model_a = GaussianConditional(get_scale_table())
+    def forward(self, x):
+        y_s = self.spatial_encoder(x)
+        y_a = self.angular_encoder(x)
+        scales_s, z_likelihood_s = self.spatial_hyperprior(y_s)
+        scales_a, z_likelihood_a = self.angular_hyperprior(y_a)
+        z_s, likelihood_s = self.entropy_model_s(y_s, scales_s)
+        z_a, likelihood_a = self.entropy_model_a(y_a, scales_a)
+        concatenated = torch.cat((z_s, z_a), dim=1)
+        return {
+            "y_hat": concatenated,
+            "latents": {"y_s": y_s, "y_a": y_a},
+            "likelihoods": {"y_s": likelihood_s, "y_a": likelihood_a, "z_s": z_likelihood_s, "z_a": z_likelihood_a}
+        }
+
+class Decoder(nn.Module):
+    def __init__(self, latent_channels):
+        super().__init__()
+        self.initial_layer = nn.Sequential(nn.ConvTranspose2d(latent_channels, filt_n, kernel_size=5, stride=3, padding=0), GDN(filt_n, inverse=True))
+        self.conv_layers_D1 = nn.Sequential(nn.ConvTranspose2d(filt_n, filt_n, kernel_size=4, stride=2, padding=0), GDN(filt_n, inverse=True))
+        self.conv_layers_D2 = nn.Sequential(nn.ConvTranspose2d(filt_n, filt_n, kernel_size=4, stride=2, padding=1), GDN(filt_n, inverse=True))
+        self.conv_layers_D3 = nn.Sequential(nn.ConvTranspose2d(filt_n, 3, kernel_size=4, stride=2, padding=1), nn.Sigmoid())
+    def forward(self, z):
+        x = self.initial_layer(z)
+        x = self.conv_layers_D1(x)
+        x = self.conv_layers_D2(x)
+        x = self.conv_layers_D3(x)
+        return x
+
+class VAE(nn.Module):
+    def __init__(self, latent_channels):
+        super().__init__()
+        self.encoder = Encoder(latent_channels)
+        self.decoder = Decoder(latent_channels)
+    def forward(self, x):
+        enc_out = self.encoder(x)
+        dec_out = self.decoder(enc_out["y_hat"])
+        return {"x_hat": dec_out, "likelihoods": enc_out["likelihoods"], "latents": enc_out["latents"]}
+
+def extract_patches(image, patch_size=(216, 312), step_size=(180, 260)):
+    patches = []
+    img_width, img_height = image.size
+    for y in range(0, img_height - patch_size[0] + 1, step_size[0]):
+        for x in range(0, img_width - patch_size[1] + 1, step_size[1]):
+            box = (x, y, x + patch_size[1], y + patch_size[0])
+            patch = image.crop(box)
+            patches.append(patch)
+            if len(patches) == 49:
+                return patches
+    return patches
+
+def reassemble_image(patches, original_size, patch_size, step_size):
+    original_width, original_height = original_size
+    reconstructed = torch.zeros((3, original_height, original_width), device='cpu')
+    counts = torch.zeros_like(reconstructed)
+    patch_idx = 0
+    for y in range(0, original_height - patch_size[0] + 1, step_size[0]):
+        for x in range(0, original_width - patch_size[1] + 1, step_size[1]):
+            if patch_idx >= len(patches):
+                break
+            patch = patches[patch_idx]
+            reconstructed[:, y:y + patch_size[0], x:x + patch_size[1]] += patch
+            counts[:, y:y + patch_size[0], x:x + patch_size[1]] += 1
+            patch_idx += 1
+    reconstructed /= counts.clamp(min=1)
+    return reconstructed
+
+def rgb_to_ycbcr(rgb_image):
+    if isinstance(rgb_image, torch.Tensor):
+        rgb_image = rgb_image.cpu().numpy()
+    if rgb_image.shape[0] == 3:
+        rgb_image = np.transpose(rgb_image, (1, 2, 0))
+    R, G, B = rgb_image[:, :, 0], rgb_image[:, :, 1], rgb_image[:, :, 2]
+    Y = 0.299 * R + 0.587 * G + 0.114 * B
+    return Y
+
+def calculate_metrics(original, reconstructed):
+    original_np = original.cpu().numpy()
+    reconstructed_np = reconstructed.cpu().numpy()
+    if original_np.shape[0] == 3:
+        original_np_hwc = np.transpose(original_np, (1, 2, 0))
+        reconstructed_np_hwc = np.transpose(reconstructed_np, (1, 2, 0))
+    else:
+        original_np_hwc = original_np
+        reconstructed_np_hwc = reconstructed_np
+
+    psnr_rgb = psnr(original_np_hwc, reconstructed_np_hwc, data_range=1.0)
+    ssim_rgb = ssim(original_np_hwc, reconstructed_np_hwc, channel_axis=2, data_range=1.0, win_size=11)
+    
+    y_original = rgb_to_ycbcr(original_np)
+    y_reconstructed = rgb_to_ycbcr(reconstructed_np)
+    
+    psnr_y = psnr(y_original, y_reconstructed, data_range=1.0)
+    ssim_y = ssim(y_original, y_reconstructed, data_range=1.0, win_size=11)
+    
+    return {'PSNR_RGB': psnr_rgb, 'SSIM_RGB': ssim_rgb, 'PSNR_Y': psnr_y, 'SSIM_Y': ssim_y}
+
+def calculate_entropy(tensor):
+    symbols = tensor.flatten()
+    _, counts = torch.unique(symbols, return_counts=True)
+    probs = counts.float() / symbols.numel()
+    entropy = -torch.sum(probs * torch.log2(probs + 1e-10))
+    return entropy * symbols.numel()
+
+MODEL_LIST = ['DUALF_D_v_Best_High_Bit_Rate.pth', 'DUALF_D_v_Low_Bit_Rate.pth', 'DUALF_D_v_High_Bit_Rate.pth', 'DUALF_D_v_Mid_Bit_Rate.pth']
+QP_LIST = [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45,  0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9, 3.0]
+model_cache = {}
+
+def load_model_for_gradio(model_filename):
+    if model_filename in model_cache:
+        return model_cache[model_filename]
+
+    model = VAE(latent_channels).to(device)
+    model_path = os.path.join(save_path, model_filename)
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file not found at {model_path}. Please place models in the '{save_path}' directory.")
+
+    state_dict = torch.load(model_path, map_location=device)
+    
+    try:
+        spatial_cdf_size = state_dict['encoder.spatial_hyperprior.entropy_bottleneck._quantized_cdf'].size(1)
+        angular_cdf_size = state_dict['encoder.angular_hyperprior.entropy_bottleneck._quantized_cdf'].size(1)
+        
+        model.encoder.spatial_hyperprior.entropy_bottleneck._offset = torch.zeros(64, device=device)
+        model.encoder.spatial_hyperprior.entropy_bottleneck._quantized_cdf = torch.zeros(64, spatial_cdf_size, device=device)
+        model.encoder.spatial_hyperprior.entropy_bottleneck._cdf_length = torch.zeros(64, dtype=torch.int32, device=device)
+        
+        model.encoder.angular_hyperprior.entropy_bottleneck._offset = torch.zeros(64, device=device)
+        model.encoder.angular_hyperprior.entropy_bottleneck._quantized_cdf = torch.zeros(64, angular_cdf_size, device=device)
+        model.encoder.angular_hyperprior.entropy_bottleneck._cdf_length = torch.zeros(64, dtype=torch.int32, device=device)
+    except KeyError as e:
+        print(f"Warning: Could not find key {e} in state_dict. This may happen with older models. Trying to load without it.")
+
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    model_cache[model_filename] = model
+    return model
+
+def compress_and_display(image_pil, model_filename, qp_value):
+    print(f"Processing with model: {model_filename} and QP: {qp_value}")
+    
+    model = load_model_for_gradio(model_filename)
+
+    original_tensor = transforms.ToTensor()(image_pil)
+    patch_size_config = (216, 312)
+    step_size_config = (180, 260)
+    patches = extract_patches(image_pil, patch_size=patch_size_config, step_size=step_size_config)
+    patches_tensor = [transforms.ToTensor()(p) for p in patches]
+
+    total_bits = 0
+    reconstructed_patches = []
+    with torch.no_grad():
+        for patch in patches_tensor:
+            patch = patch.unsqueeze(0).to(device)
+            enc_out = model.encoder(patch)
+            y_s = enc_out["latents"]["y_s"]
+            y_a = enc_out["latents"]["y_a"]
+
+            step_size = 1.0 / qp_value
+            y_s_quantized = torch.round(y_s / step_size)
+            y_a_quantized = torch.round(y_a / step_size)
+
+            y_s_dequantized = y_s_quantized * step_size
+            y_a_dequantized = y_a_quantized * step_size
+            latents_dequantized = torch.cat((y_s_dequantized, y_a_dequantized), dim=1)
+            
+            reconstructed = model.decoder(latents_dequantized)
+            reconstructed_patches.append(reconstructed.squeeze(0).cpu())
+
+            bits_spatial = calculate_entropy(y_s_quantized)
+            bits_angular = calculate_entropy(y_a_quantized)
+            total_bits += bits_spatial.item() + bits_angular.item()
+
+    reconstructed_tensor = reassemble_image(reconstructed_patches, image_pil.size, patch_size_config, step_size_config)
+    reconstructed_tensor = reconstructed_tensor.clamp(0, 1)
+
+    total_pixels = image_pil.width * image_pil.height
+    bpp = total_bits / total_pixels
+    metrics_dict = calculate_metrics(original_tensor, reconstructed_tensor)
+    metrics_dict['BPP'] = bpp
+
+    original_np = (original_tensor.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
+    reconstructed_np = (reconstructed_tensor.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
+
+    comparison_image = np.hstack((original_np, reconstructed_np))
+    metrics_str = (
+        f"Bits Per Pixel (BPP): {metrics_dict['BPP']:.4f}\n\n"
+        f"--- RGB Metrics ---\n"
+        f"PSNR (RGB): {metrics_dict['PSNR_RGB']:.2f} dB\n"
+        f"SSIM (RGB): {metrics_dict['SSIM_RGB']:.4f}\n\n"
+        f"--- Luma (Y) Metrics ---\n"
+        f"PSNR (Y): {metrics_dict['PSNR_Y']:.2f} dB\n"
+        f"SSIM (Y): {metrics_dict['SSIM_Y']:.4f}"
+    )
+    return comparison_image, metrics_str
+
+title = "Light Field Image Compression with DUALF_D VAE"
+description = """
+Upload a macropixel image (e.g., a 3x3 view light field image taken with Lytro Illum 2.0) to compress and decompress it using a VAE-based neural network.
+
+* You can select different pre-trained model checkpoints and adjust the Quantization Parameter (QP) to see its effect on quality and bitrate.
+* A lower QP generally results in lower quality and a lower storage requirement, while a higher QP means better quality but requires more storage for image.
+* Please refer to our [GitHub Page](https://takhtardeshirsoheib.github.io/DUALF_D/index.html) for more details (it will be public after acceptance of our paper)
+"""
+with gr.Blocks() as demo:
+    gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="pil", label="Upload Macropixel Image")
+            model_selector = gr.Dropdown(choices=MODEL_LIST, value=MODEL_LIST[3], label="Selected Model")
+            qp_selector = gr.Dropdown(choices=QP_LIST, value=1.0, label="Selected Quantization Parameter (QP)")
+            submit_button = gr.Button("Compress and Analyze")
+            
+        with gr.Column(scale=2):
+            output_comparison = gr.Image(label="Original vs. Compressed")
+            output_metrics = gr.Textbox(label="Performance Metrics")
+
+    submit_button.click(
+        fn=compress_and_display,
+        inputs=[input_image, model_selector, qp_selector],
+        outputs=[output_comparison, output_metrics]
+    )
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                ["./samples/macropixel_059.png", MODEL_LIST[3], 0.5],
+                ["./samples/macropixel_033.png", MODEL_LIST[2], 0.5],
+                ["./samples/macropixel_028.png", MODEL_LIST[3], 2.0],
+                ["./samples/macropixel_026.png", MODEL_LIST[3], 2.5],
+                ["./samples/macropixel_019.png", MODEL_LIST[3], 2.6],
+                ["./samples/macropixel_203.png", MODEL_LIST[3], 2.8],
+                ["./samples/macropixel_923.png", MODEL_LIST[3], 3.0]
+            ],
+            inputs=[input_image, model_selector, qp_selector]
+        )
+
+if __name__ == "__main__":
+    demo.launch()

checkpoint/DUALF_D_v_Best_High_Bit_Rate.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:663df825eaf0e0359cf18ff02b65cf1ee827873113862f2cd430d28ec655575b
+size 18339498

checkpoint/DUALF_D_v_High_Bit_Rate.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a44171f24a1aaaf69b8437c46ee0558e8a04da0bb1cabb6600bd034ce464c79
+size 18344362

checkpoint/DUALF_D_v_Low_Bit_Rate.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8f3cc3aa976755fe9cde13316dfadb9cfd0e946999cc68c6065eba26178754a
+size 18332842

checkpoint/DUALF_D_v_Mid_Bit_Rate.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:945df8ca6caf60f9ef83d4dc9dd104c456d0cc550d2178f9dee62ffe0cdab2cc
+size 18339242

requirements.txt

@@ -0,0 +1,7 @@
+torch>=1.9.0
+torchvision>=0.10.0
+gradio>=3.50.0
+Pillow>=8.0.0
+numpy>=1.21.0
+scikit-image>=0.18.0
+compressai>=1.7.0