import numpy as np

from bokeh.plotting import figure
from bokeh.models import ColumnDataSource


class Plot2D:

    def __init__(self, *args, **kwargs):
        data = {
            "image": [np.zeros((1, 1))]
        }

        self.source = source = ColumnDataSource(data=data)

        self.fig = fig = figure()
        fig.image(source=source, x=0, y=0, dw=1, dh=1, palette="Spectral11")


    def set(self, image):
        self.source.data.update(image=[image])