<jupyter_start><jupyter_text>Aggregate Statistics<jupyter_code>psnr_base_dir = "/scratch/gf332/CWOQ/balle/compression/results/image_compression/kodak/PSNR_sRGB_RGB/"
ms_ssim_base_dir = "/scratch/gf332/CWOQ/balle/compression/results/image_compression/kodak/MS-SSIM_sRGB_RGB/"
psnr_method_names = {
    "jpeg420.txt": ("JPEG (4:2:0)", "MSE", 8),
    "bpg444.txt": ("BPG (4:4:4)", "MSE", 8),
    "theis-2017-iclr.txt": ("Theis (2017)", "MSE", 9),
    "balle-2018-iclr-opt-mse.txt": ("Ballé (2018) (opt. for MSE)", "MSE", 9),
    #"balle-2018-iclr-opt-msssim.txt": ("Ballé (2018) (opt. for MS-SSIM)", "MS-SSIM", 9),
    #"minnen-2018-neurips.txt": ("Minnen (2018)", "MSE", 9),
    "minnen-2020-icip.txt": ("Minnen (2020)", "MSE", 9),
}

ms_ssim_method_names = {
    "jpeg420.txt": ("JPEG (4:2:0)", "MSE", 10),
    "bpg444.txt": ("BPG (4:4:4)", "MSE", 10),
    "theis-2017-iclr.txt": ("Theis (2017)", "MSE", 11),
    #"balle-2018-iclr-opt-mse.txt": ("Ballé (2018) (opt. for MSE)", "MSE", 11),
    "balle-2018-iclr-opt-msssim.txt": ("Ballé (2018) (opt. for MS-SSIM)", "MS-SSIM",  11),
    #"minnen-2018-neurips.txt": ("Minnen (2018)", "MSE", 11),
}
def prepare_method_df(metric, method_name, method_legend, skiprows, opt_objective):
    
    base_dir = psnr_base_dir if metric == "psnr" else ms_ssim_base_dir
    
    df = pd.read_csv(f"{base_dir}/{method_name}", 
                                   skiprows=skiprows, 
                                   header=0, 
                                   names=["bpp", metric])
    
    df["objective"] = opt_objective
    df["method"] = method_legend
    
    df = df.astype({metric: float,
                    "bpp": float})
    
    if metric == "ms_ssim":
        df["ms_ssim"] = to_db(df["ms_ssim"])
    
    df = df[df["bpp"] < 2]
    
    return df

to_db = lambda x: -10. * np.log10(1. - x)
def get_rec_psnr_average_df(rec_dfs, kind):
    aggregate_df = pd.concat(rec_dfs, axis=1)
    aggregate_df = aggregate_df.T.astype({
        "comp_lossy_BPP": float,
        "lossy_BPP": float,
        "ideal_PSNR": float,
        "PSNR": float,
        "ideal_MS_SSIM": float,
        "MS_SSIM": float,
    })

    psnr_avg_df = aggregate_df.groupby("target_bpp").mean().rename({"PSNR": "psnr",
                                                                 "comp_lossy_BPP": "bpp",
                                                                 "MS_SSIM": "ms_ssim"}, axis=1)
    psnr_avg_df["model"] = f"iREC {kind}"
    psnr_avg_df["type"] = f"actual"
    

    ideal_psnr_avg_df = aggregate_df.groupby("target_bpp").mean().rename({"ideal_PSNR": "psnr",
                                                                 "lossy_BPP": "bpp",
                                                                 "ideal_MS_SSIM": "ms_ssim"}, axis=1)
    ideal_psnr_avg_df["model"] = f"iREC {kind}"
    ideal_psnr_avg_df["type"] = f"ideal"
    
    return pd.concat([psnr_avg_df, ideal_psnr_avg_df])
rec_l1_gaussian_psnr_avg_df = get_rec_psnr_average_df(level_1_gaussian_rec_dfs, "(Ours) (opt. for MSE)")
rec_gaussian_psnr_avg_df  = get_rec_psnr_average_df(gaussian_rec_dfs, "(Ours) (opt. for MSE)")
rec_ms_ssim_psnr_avg_df = get_rec_psnr_average_df(ms_ssim_rec_dfs, "(Ours) (opt. for MS-SSIM)")
psnr_results = pd.concat([prepare_method_df("psnr", method_name, method_legend, skiprows, opt_objective)
                for method_name, (method_legend, opt_objective, skiprows) in psnr_method_names.items()])

ms_ssim_results = pd.concat([prepare_method_df("ms_ssim", method_name, method_legend, skiprows, opt_objective)
                for method_name, (method_legend, opt_objective, skiprows) in ms_ssim_method_names.items()])
sns.set_palette("bright")

fig, ax = plt.subplots(figsize=(7.4, 6.5))

sns.lineplot(x="bpp", 
             y="psnr", 
             data=psnr_results,
             markers=False,
             dashes=False,
             hue='method',
             ax=ax)

rec_psnr_df = rec_gaussian_psnr_avg_df[rec_gaussian_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
ax.plot(rec_psnr_df["bpp"], rec_psnr_df["psnr"], 
        label=rec_psnr_df["method"].values[0], marker="o", linewidth=3, c='k')

# rec_ssim_df = rec_ms_ssim_psnr_avg_df[rec_ms_ssim_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
# ax.plot(rec_ssim_df["bpp"], rec_ssim_df["psnr"], label=rec_ssim_df["method"].values[0], marker="o", linewidth=3, c='r')

handles, labels = ax.get_legend_handles_labels()
ax.legend(handles=handles[1:][::-1], 
          labels=labels[1:][::-1], 
          fontsize=18,
          loc="upper left",
          title=None,
          fancybox=False,
          frameon=False)

ax.set_xlim([0, 1.5])
ax.set_ylim([25, 42])
ax.set_ylabel(None)
ax.set_xlabel("Bits per Pixel", fontsize=18)

ax.tick_params(axis='both', which='major', labelsize=18)

x_major_ticks = np.arange(0., 1.5, 0.2)
x_minor_ticks = np.arange(0., 1.5, 0.05)

ax.set_xticks(x_major_ticks)
ax.set_xticks(x_minor_ticks, minor=True)

y_major_ticks = np.arange(25., 42, 5)
y_minor_ticks = np.arange(25., 42, 1)

ax.set_yticks(y_major_ticks)
ax.set_yticks(y_minor_ticks, minor=True)

ax.grid(which='major', alpha=1)
ax.grid(which='minor', c="#eeeeee", alpha=1)

ax.text(.45, .02, 'PSNR (RGB) on Kodak',
        horizontalalignment='left',
        transform=ax.transAxes,
        fontsize=24)

fig.tight_layout()

fig.savefig(f"{save_dir}/method_psnr_comparison.png")
sns.set_palette("bright")

fig, ax = plt.subplots(figsize=(7.4, 6.5))


sns.lineplot(x="bpp", 
             y="ms_ssim", 
             data=ms_ssim_results,
             markers=False,
             dashes=False,
             hue='method',
             ax=ax)

# rec_psnr_df = rec_gaussian_psnr_avg_df[rec_gaussian_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
# ax.plot(rec_psnr_df["bpp"], to_db(rec_psnr_df["ms_ssim"]), label=rec_psnr_df["method"].values[0], marker="o", linewidth=3, c='b')

rec_ssim_df = rec_ms_ssim_psnr_avg_df[rec_ms_ssim_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, 
                                                                                          axis=1)
ax.plot(rec_ssim_df["bpp"], to_db(rec_ssim_df["ms_ssim"]), 
        label=rec_ssim_df["method"].values[0], marker="o", linewidth=3, c='k')

handles, labels = ax.get_legend_handles_labels()
ax.legend(handles=handles[1:][::-1], 
          labels=labels[1:][::-1], 
          fontsize=18,
          loc="upper left",
          title=None,
          fancybox=False,
          frameon=False)

ax.set_xlim([0, 1.5])
ax.set_ylim([8, 27])
ax.set_xlabel("Bits per Pixel", fontsize=18)
ax.set_ylabel(None)

ax.tick_params(axis='both', which='major', labelsize=18)

x_major_ticks = np.arange(0., 1.5, 0.2)
x_minor_ticks = np.arange(0., 1.5, 0.05)

ax.set_xticks(x_major_ticks)
ax.set_xticks(x_minor_ticks, minor=True)

y_major_ticks = np.arange(8., 27, 2)
y_minor_ticks = np.arange(8., 27, 0.5)

ax.set_yticks(y_major_ticks)
ax.set_yticks(y_minor_ticks, minor=True)

ax.grid(which='major', alpha=1)
ax.grid(which='minor', c="#eeeeee", alpha=1)

ax.text(.25, .02, 'MS-SSIM (dB) (RGB) on Kodak',
        horizontalalignment='left',
        transform=ax.transAxes,
        fontsize=24)

fig.tight_layout()

fig.savefig(f"{save_dir}/method_ms_ssim_comparison.eps")
fig.savefig(f"{save_dir}/method_ms_ssim_comparison.png")
fig, ax = plt.subplots(figsize=(10, 13))

balle_psnr_results = psnr_results[psnr_results["method"].str.contains("Ballé")]

sns.lineplot(x="bpp", 
             y="psnr", 
             data=balle_psnr_results,
             markers=False,
             dashes=False,
             hue='method',
             ax=ax)

rec_psnr_df = rec_gaussian_psnr_avg_df[rec_gaussian_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
ax.plot(rec_psnr_df["bpp"], 
        rec_psnr_df["psnr"], 
        label="iREC (Ours) (opt. for MSE)", 
        linewidth=2, 
        c='g')

ideal_rec_psnr_df = rec_gaussian_psnr_avg_df[rec_gaussian_psnr_avg_df["type"] == "ideal"].rename({"model": "method"}, axis=1)
ax.plot(ideal_rec_psnr_df["bpp"], 
        ideal_rec_psnr_df["psnr"], 
        label="Ideal iREC (Ours) (opt. for MSE)",
        marker="", 
        linewidth=2, 
        c='g', 
        linestyle="--")


rec_psnr_df = rec_ms_ssim_psnr_avg_df[rec_ms_ssim_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
ax.plot(rec_psnr_df["bpp"], 
        rec_psnr_df["psnr"], 
        label="iREC (Ours) (opt. for MS-SSIM)", 
        linewidth=2, 
        c='r')

ideal_rec_psnr_df = rec_ms_ssim_psnr_avg_df[rec_ms_ssim_psnr_avg_df["type"] == "ideal"].rename({"model": "method"}, axis=1)
ax.plot(ideal_rec_psnr_df["bpp"], 
        ideal_rec_psnr_df["psnr"], 
        label="Ideal iREC (Ours) (opt. for MS-SSIM)",
        marker="", 
        linewidth=2, 
        c='r', 
        linestyle="--")

handles, labels = ax.get_legend_handles_labels()
ax.legend(handles=handles[1:], 
          labels=labels[1:], 
          fontsize=24,
          loc="lower right",
          title=None,
          fancybox=False,
          frameon=False)

ax.set_xlim([0, 1.5])
ax.set_ylim([24, 40])
ax.set_xlabel("Bits per Pixel", fontsize=24)
ax.set_ylabel(None)

ax.tick_params(axis='both', which='major', labelsize=24)

x_major_ticks = np.arange(0., 1.5, 0.2)
x_minor_ticks = np.arange(0., 1.5, 0.05)

ax.set_xticks(x_major_ticks)
ax.set_xticks(x_minor_ticks, minor=True)

y_major_ticks = np.arange(25., 40, 5)
y_minor_ticks = np.arange(24., 40, 1)

ax.set_yticks(y_major_ticks)
ax.set_yticks(y_minor_ticks, minor=True)

ax.grid(which='major', alpha=1)
ax.grid(which='minor', c="#eeeeee", alpha=1)

ax.text(.03, .96, 'PSNR (RGB) on Kodak',
        horizontalalignment='left',
        transform=ax.transAxes,
        fontsize=28)

fig.tight_layout()
fig.savefig(f"{save_dir}/actual_vs_ideal_psnr_comparison.eps")
fig, ax = plt.subplots(figsize=(10, 13))

balle_ms_ssim_results = ms_ssim_results[ms_ssim_results["method"].str.contains("Ballé")]


sns.lineplot(x="bpp", 
             y="ms_ssim", 
             data=balle_ms_ssim_results,
             markers=False,
             dashes=False,
             hue='method',
             ax=ax)

rec_psnr_df = rec_gaussian_psnr_avg_df[rec_gaussian_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
ax.plot(rec_psnr_df["bpp"], 
        to_db(rec_psnr_df["ms_ssim"]), 
        label="iREC (Ours) (opt. for MSE)", 
        linewidth=2, 
        c='g')

ideal_rec_psnr_df = rec_gaussian_psnr_avg_df[rec_gaussian_psnr_avg_df["type"] == "ideal"].rename({"model": "method"}, axis=1)
ax.plot(ideal_rec_psnr_df["bpp"], 
        to_db(ideal_rec_psnr_df["ms_ssim"]), 
        label="Ideal iREC (Ours) (opt. for MSE)",
        marker="", 
        linewidth=2, 
        c='g', 
        linestyle="--")


rec_psnr_df = rec_ms_ssim_psnr_avg_df[rec_ms_ssim_psnr_avg_df["type"] == "actual"].rename({"model": "method"}, axis=1)
ax.plot(rec_psnr_df["bpp"], 
        to_db(rec_psnr_df["ms_ssim"]), 
        label="iREC (Ours) (opt. for MS-SSIM)", 
        linewidth=2, 
        c='r')

ideal_rec_psnr_df = rec_ms_ssim_psnr_avg_df[rec_ms_ssim_psnr_avg_df["type"] == "ideal"].rename({"model": "method"}, axis=1)
ax.plot(ideal_rec_psnr_df["bpp"], 
        to_db(ideal_rec_psnr_df["ms_ssim"]), 
        label="iREC (Ours) (opt. for MS-SSIM)",
        marker="", 
        linewidth=2, 
        c='r', 
        linestyle="--")

handles, labels = ax.get_legend_handles_labels()
ax.legend(handles=handles[1:], 
          labels=labels[1:], 
          fontsize=24,
          loc="lower right",
          title=None,
          fancybox=False,
          frameon=False)

ax.set_xlim([0, 1.5])
ax.set_ylim([9, 25])
ax.set_xlabel("Bits / Pixel", fontsize=24)
ax.set_ylabel(None)

ax.tick_params(axis='both', which='major', labelsize=24)

x_major_ticks = np.arange(0., 1.5, 0.2)
x_minor_ticks = np.arange(0., 1.5, 0.05)

ax.set_xticks(x_major_ticks)
ax.set_xticks(x_minor_ticks, minor=True)

y_major_ticks = np.arange(10., 25, 2)
y_minor_ticks = np.arange(10., 25, 0.5)

ax.set_yticks(y_major_ticks)
ax.set_yticks(y_minor_ticks, minor=True)

ax.grid(which='major', alpha=1)
ax.grid(which='minor', c="#eeeeee", alpha=1)

ax.text(.35, .96, 'MS-SSIM (dB) (RGB) on Kodak',
        horizontalalignment='left',
        transform=ax.transAxes,
        fontsize=28)

fig.tight_layout()
fig.savefig(f"{save_dir}/actual_vs_ideal_ms_ssim_comparison.eps")<jupyter_output><empty_output><jupyter_text>Individual Images<jupyter_code>jpeg_stats_file = "../experimental_data/jpeg/stats.json"
balle_clic_stats_file = "../experimental_data/kodak_balle/stats.json"

with open(jpeg_stats_file, 'r') as f:
    jpeg_df = pd.read_json(f)
    
with open(balle_clic_stats_file, 'r') as f:
    balle_df = pd.read_json(f)
    balle_df = balle_df.rename(lambda x: x[1:], axis=1)
def get_image_dfs(img_name):
    jpeg_img_df = pd.read_json(jpeg_df[img_name].to_json())
    balle_img_df = pd.read_json(balle_df[img_name].to_json())

    jpeg_img_df.loc["model"] = "JPEG (4:2:0)"
    balle_img_df.loc["model"] = "Ballé (2018) (opt. for MSE)"

    # Put the columns in order
    jpeg_img_df = jpeg_img_df.reindex(sorted(jpeg_img_df.columns, key=lambda c: int(c.split('_')[1])), axis=1)
    balle_img_df = balle_img_df.reindex(sorted(balle_img_df.columns, 
                                               key=lambda c: float(f"0.{c.split('_')[1]}")), axis=1)

    actual_rec_df = pd.concat(gaussian_rec_dfs, axis=1)[img_name].T.sort_values(by=["target_bpp"]).reset_index(drop=True)
    actual_rec_df = actual_rec_df.rename({"PSNR": "psnr",
                                         "comp_lossy_BPP": "bpp",
                                         "MS_SSIM": "ms_ssim"}, axis=1)[["psnr", "bpp", "ms_ssim"]]
    actual_rec_df = actual_rec_df.astype({"psnr":float,
                                        "bpp":float,
                                        "ms_ssim":float})
    actual_rec_df["model"] = "actual_rec"

    joint_df = pd.concat([jpeg_img_df, 
                          balle_img_df], axis=1).T.astype({"bpp": float,
                                                            "psnr": float,
                                                            "ms_ssim": float})

    # Convert MS-SSIM statistic to decibels
    joint_df['ms_ssim'] = to_db(joint_df['ms_ssim'])
    
    return actual_rec_df, joint_df
def plot_kodak_ms_ssim_image_stats(img_name, actual_rec_df, joint_df):
    fig, ax = plt.subplots(figsize=(10, 13))

    sns.lineplot(x="bpp", 
                 y="ms_ssim", 
                 data=joint_df,
                 markers=False,
                 dashes=False,
                 hue='model',
                 ax=ax)

    ax.plot(actual_rec_df["bpp"], 
            to_db(actual_rec_df["ms_ssim"]), 
            label="iREC (Ours) (opt. for MSE)", 
            linewidth=3,
            marker='o',
            c='k')


    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles=handles[1:], 
              labels=labels[1:], 
              fontsize=24,
              loc="lower right",
              title=None,
              fancybox=False,
              frameon=False)

    ax.set_xlim([0, 1.5])
    ax.set_ylim([9, 25])
    ax.set_xlabel("Bits / Pixel", fontsize=24)
    ax.set_ylabel(None)

    ax.tick_params(axis='both', which='major', labelsize=24)

    x_major_ticks = np.arange(0., 1.5, 0.2)
    x_minor_ticks = np.arange(0., 1.5, 0.05)

    ax.set_xticks(x_major_ticks)
    ax.set_xticks(x_minor_ticks, minor=True)

    y_major_ticks = np.arange(8., 25, 2)
    y_minor_ticks = np.arange(8., 25, 0.5)

    ax.set_yticks(y_major_ticks)
    ax.set_yticks(y_minor_ticks, minor=True)

    ax.grid(which='major', alpha=1)
    ax.grid(which='minor', c="#eeeeee", alpha=1)

    ax.text(.35, .96, 'MS-SSIM (dB) (RGB) on Kodak',
            horizontalalignment='left',
            transform=ax.transAxes,
            fontsize=28)

    fig.tight_layout()
    fig.savefig(f"{save_dir}/kodak_comparisons/{img_name}_ms_ssim.eps")
    plt.close()
def plot_kodak_psnr_image_stats(img_name, actual_rec_df, joint_df):
    fig, ax = plt.subplots(figsize=(10, 13))

    sns.lineplot(x="bpp", 
                 y="psnr", 
                 data=joint_df,
                 markers=False,
                 dashes=False,
                 hue='model',
                 ax=ax)

    ax.plot(actual_rec_df["bpp"], 
            actual_rec_df["psnr"], 
            label="iREC (Ours) (opt. for MSE)", 
            linewidth=3,
            marker='o',
            c='k')


    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles=handles[1:], 
              labels=labels[1:], 
              fontsize=24,
              loc="lower right",
              title=None,
              fancybox=False,
              frameon=False)

    ax.set_xlim([0, 1.5])
    ax.set_xlabel("Bits per Pixel", fontsize=24)
    ax.set_ylabel(None)

    ax.tick_params(axis='both', which='major', labelsize=24)

    x_major_ticks = np.arange(0., 2., 0.2)
    x_minor_ticks = np.arange(0., 2., 0.05)

    ax.set_xticks(x_major_ticks)
    ax.set_xticks(x_minor_ticks, minor=True)

    y_major_ticks = np.arange(20., 45, 5)
    y_minor_ticks = np.arange(20., 45, 1)

    ax.set_yticks(y_major_ticks)
    ax.set_yticks(y_minor_ticks, minor=True)

    ax.grid(which='major', alpha=1)
    ax.grid(which='minor', c="#eeeeee", alpha=1)

    ax.text(.03, .96, 'PSNR (RGB) on Kodak',
        horizontalalignment='left',
        transform=ax.transAxes,
        fontsize=28)

    fig.tight_layout()
    fig.savefig(f"{save_dir}/kodak_comparisons/{img_name}_psnr.eps")
    plt.close()
for i in range(1, 25):

    img_name = f"kodim{i:02}.png"
    
    image_dfs = get_image_dfs(img_name)

    plot_kodak_ms_ssim_image_stats(img_name, *image_dfs)
    plot_kodak_psnr_image_stats(img_name, *image_dfs)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16317499439838673, 'fraction_numerical': 0.026159533945776383, 'mean_word_length': 2.871828236824984, 'pattern_counts': {'":': 48, '<': 10, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 27, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '19240', 'license': 'MIT', 'n_tokens_mistral': 7482, 'n_tokens_neox': 6983, 'n_words': 962, 'path': 'notebooks/Lossy Compression Performance Visualization.ipynb', 'repo_name': 'gergely-flamich/relative-entropy-coding', 'size': '17852'}