utils.py

import matplotlib.pyplot as plt
import PIL
import numpy as np


def visualize_results(orig_img:PIL.Image.Image, answer:str, question:str = None):
    """
    Helper function for results visualization

    Parameters:
       orig_img (PIL.Image.Image): original image
       answer (str): model answer in text format.
       question (str, *optional*, None): input question, if not provided answer will be used as caption
    Returns:
       fig (matplotlib.pyplot.Figure): matplotlib generated figure contains drawing result
    """
    fig = plt.figure()
    fig.patch.set_facecolor('white')
    ax = fig.add_subplot(111)
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
    ax.grid(False)
    ax.imshow(np.array(orig_img))
    qa_text = "question: {}\nanswer: {}"
    cap_text = "caption: {}"
    ax.set_title(qa_text.format(question, answer) if question is not None else cap_text.format(answer),
                 y=-0.01, pad=-30 if question is not None else -15)
    return fig