utils.py

import io

from PIL import Image
import matplotlib.pyplot as plt
import omegaconf


def patch_target(config):
    """Update the _target_ of cfg from src_v2 to src"""
    for key, value in config.items():
        if isinstance(value, omegaconf.DictConfig):
            # Recursive call if the value is another DictConfig
            patch_target(value)
        elif isinstance(value, str) and value.startswith("src_v2"):
            # Update the value if it matches the old_value
            config[key] = value.replace("src_v2", "src")


def ax_to_pil(ax):
    fig = ax.figure
    buf = io.BytesIO()
    fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
    buf.seek(0)

    # Load the image data completely before closing the buffer
    pil_image = Image.open(buf)
    pil_image_copy = pil_image.copy()

    # Now we can safely close everything
    pil_image.close()
    buf.close()
    plt.close(fig)

    return pil_image_copy