show_image.py

import cv2
import numpy
import matplotlib
from matplotlib import pyplot as plt


def plot_image(img, cmap=None):
    """Plot an image to the screen.
    :param img: numpy.ndarray
    :param cmap: str
    :return:
    """

    image_type = type(img)

    dimensions = numpy.shape(img)

    if image_type == numpy.ndarray:
        # If the image is color then OpenCV stores it as BGR, we plot it as RGB
        if len(dimensions) == 3:
            plt.figure()
            plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            plt.show()

        elif cmap is None and len(dimensions) == 2:
            plt.figure()
            plt.imshow(img, cmap="gray")
            plt.show()

        elif cmap is not None and len(dimensions) == 2:
            plt.figure()
            plt.imshow(img, cmap=cmap)
            plt.show()

    elif image_type == matplotlib.figure.Figure:
        print("Error, matplotlib Figure not supported. Instead try running without plot_image.")

    # Plot if the image is a plotnine ggplot image
    elif str(image_type) == "<class 'plotnine.ggplot.ggplot'>":
        print(img)