append_diff_label.py

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os
from pathlib import Path
from matplotlib.font_manager import FontProperties
import pickle
import pdb


if __name__=='__main__':
    data_list = [os.path.join('merged_heatmap2', x) for x in os.listdir('merged_heatmap2')]
    font_path = 'scripts/Times New Roman.ttf'
    times_new_roman_title = FontProperties(fname=font_path, size=30)

    score_diff = pickle.load(open('heatmap_score_diff.pkl', 'rb'))
    score_prev = pickle.load(open('heatmap_score_prev.pkl', 'rb'))
    score_curr = pickle.load(open('heatmap_score_curr.pkl', 'rb'))

    for idx, path in enumerate(data_list):
        fig, ax = plt.subplots(figsize=(15, 15))
        img = mpimg.imread(path)
        ax.imshow(img)
        ax.axis("off")

        title_text = f'Previous Score: {score_prev[idx][0]:.2f} \u2192 Current Score: {score_curr[idx][0]:.2f} (Score Diff: {score_diff[idx]:.2f})'
        fig.suptitle(title_text, fontproperties=times_new_roman_title, y=0.77)

        plt.tight_layout(pad=0)
        fig.savefig(path, dpi=300, bbox_inches="tight")
        plt.close('all')
        print(idx)