iirwi/plot.py

import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import matplotlib.pyplot as plt
import matplotlib.patches as patches

LINEWIDTH = 2
EDGECOLOR = 'r'

class MemePlotter:
    def __init__(self):
        pass

    def plot(self, input_img, similar_img, input_img_crop_coords, show=False, save_path=None):
        similar_img_crop_coords = similar_img.bbox
        similar_img_parent = similar_img.load_parent()
        fig, (ax1, ax2) = plt.subplots(1, 2)
        fig.subplots_adjust(top=0.8, bottom=0.1)
        ax1.imshow(similar_img_parent)
        ax2.imshow(input_img)
        radius1 = (similar_img_crop_coords[2] - similar_img_crop_coords[0]) / 2
        rect1 = patches.Circle(((similar_img_crop_coords[0] + similar_img_crop_coords[2]) / 2, 
                            (similar_img_crop_coords[1] + similar_img_crop_coords[3]) / 2), 
                            radius1, linewidth=LINEWIDTH, edgecolor=EDGECOLOR, facecolor='none')
        ax1.add_patch(rect1)
        
        radius2 = (input_img_crop_coords[2] - input_img_crop_coords[0]) / 2
        rect2 = patches.Circle(((input_img_crop_coords[0] + input_img_crop_coords[2]) / 2, 
                            (input_img_crop_coords[1] + input_img_crop_coords[3]) / 2),
                            radius2, linewidth=LINEWIDTH, edgecolor=EDGECOLOR, facecolor='none')
        ax2.add_patch(rect2)

        ax1.axis('off')
        ax2.axis('off')
        fig.suptitle('Is It Really Worth It?', fontsize=20, weight='bold')
        if save_path:
            plt.savefig(save_path)
            plt.close()
        if show:
            plt.show()
        canvas = FigureCanvas(fig)
        canvas.draw()
        img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(canvas.get_width_height()[::-1] + (3,))
        return img