121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""
|
||
Generate all paper figures from measurement CSVs.
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import numpy as np
|
||
import pandas as pd
|
||
import matplotlib.pyplot as plt
|
||
|
||
|
||
def plot_lpips_bar(df, output_dir):
|
||
order = df.groupby("config")["lpips_mean"].first().sort_values().index
|
||
means = df.set_index("config").reindex(order)
|
||
|
||
fig, ax = plt.subplots(figsize=(12, 5))
|
||
colors = plt.cm.RdYlGn_r(np.linspace(0.2, 0.9, len(order)))
|
||
ax.bar(range(len(order)), means["lpips_mean"], yerr=means["lpips_std"],
|
||
capsize=4, color=colors)
|
||
ax.set_xticks(range(len(order)))
|
||
ax.set_xticklabels(order, rotation=45, ha="right", fontsize=9)
|
||
ax.set_ylabel("LPIPS (vs Q8_0)")
|
||
ax.set_title("Perceptual Distance from Q8_0 Baseline")
|
||
ax.grid(axis="y", alpha=0.3)
|
||
plt.tight_layout()
|
||
plt.savefig(os.path.join(output_dir, "lpips_bar.png"), dpi=200)
|
||
plt.close()
|
||
|
||
|
||
def plot_lpips_heatmap(df, output_dir):
|
||
cats = ["A", "B", "C", "D", "E"]
|
||
cat_cols = [c for c in df.columns if c.startswith("lpips_") and len(c) == 7]
|
||
if not cat_cols:
|
||
return
|
||
|
||
order = df.sort_values("lpips_mean")["config"]
|
||
pivot = df.set_index("config").reindex(order)[cat_cols]
|
||
pivot.columns = [c.replace("lpips_", "Cat ") for c in cat_cols]
|
||
|
||
fig, ax = plt.subplots(figsize=(10, max(4, len(order)*0.5)))
|
||
im = ax.imshow(pivot.values, cmap="YlOrRd", aspect="auto")
|
||
ax.set_xticks(range(len(pivot.columns)))
|
||
ax.set_xticklabels(pivot.columns)
|
||
ax.set_yticks(range(len(order)))
|
||
ax.set_yticklabels(order, fontsize=9)
|
||
for i in range(pivot.shape[0]):
|
||
for j in range(pivot.shape[1]):
|
||
v = pivot.values[i, j]
|
||
if not np.isnan(v):
|
||
ax.text(j, i, f"{v:.3f}", ha="center", va="center", fontsize=8)
|
||
plt.colorbar(im, ax=ax, label="LPIPS")
|
||
ax.set_title("LPIPS by Category × Config")
|
||
plt.tight_layout()
|
||
plt.savefig(os.path.join(output_dir, "lpips_heatmap.png"), dpi=200)
|
||
plt.close()
|
||
|
||
|
||
def plot_quality_vs_lpips(lpips_df, qual_df, output_dir):
|
||
merged = lpips_df.merge(qual_df, on="config", suffixes=("_lp", "_q"))
|
||
|
||
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
|
||
|
||
for ax, metric, label in [
|
||
(axes[0], "quality_mean", "Quality Score"),
|
||
(axes[1], "clip_mean_q" if "clip_mean_q" in merged.columns else "clip_mean", "CLIP Score"),
|
||
]:
|
||
if metric not in merged.columns:
|
||
continue
|
||
uniform = merged[merged["config"].str.startswith("uniform")]
|
||
asym = merged[merged["config"].str.startswith("asym")]
|
||
|
||
ax.scatter(uniform["lpips_mean"], uniform[metric], c="steelblue",
|
||
s=80, label="Uniform", zorder=3)
|
||
ax.scatter(asym["lpips_mean"], asym[metric], c="darkorange",
|
||
marker="D", s=80, label="Asymmetric", zorder=3)
|
||
|
||
for _, row in merged.iterrows():
|
||
ax.annotate(row["config"].replace("uniform_", "").replace("asym_", ""),
|
||
(row["lpips_mean"], row[metric]),
|
||
fontsize=7, alpha=0.7, textcoords="offset points",
|
||
xytext=(5, 5))
|
||
ax.set_xlabel("LPIPS vs Q8_0")
|
||
ax.set_ylabel(label)
|
||
ax.legend()
|
||
ax.grid(alpha=0.3)
|
||
|
||
plt.suptitle("Reference-Based vs Reference-Free Metrics", fontsize=13)
|
||
plt.tight_layout()
|
||
plt.savefig(os.path.join(output_dir, "scatter_quality_vs_lpips.png"), dpi=200)
|
||
plt.close()
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--metrics-dir", required=True)
|
||
parser.add_argument("--output-dir", required=True)
|
||
args = parser.parse_args()
|
||
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
||
# Load summaries
|
||
lp_path = os.path.join(args.metrics_dir, "summary_lpips_clip.csv")
|
||
q_path = os.path.join(args.metrics_dir, "summary_quality_reffree.csv")
|
||
|
||
if os.path.exists(lp_path):
|
||
lp_df = pd.read_csv(lp_path)
|
||
plot_lpips_bar(lp_df, args.output_dir)
|
||
plot_lpips_heatmap(lp_df, args.output_dir)
|
||
print("Generated: lpips_bar.png, lpips_heatmap.png")
|
||
|
||
if os.path.exists(q_path):
|
||
q_df = pd.read_csv(q_path)
|
||
plot_quality_vs_lpips(lp_df, q_df, args.output_dir)
|
||
print("Generated: scatter_quality_vs_lpips.png")
|
||
|
||
print("Done!")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|