img_utils.py

"""Utility functions for process and visualize images"""
import numpy as np
import matplotlib.pyplot as plt


def show_img(img):
    print img.shape, img.dtype
    plt.imshow(img[:, :, 0])
    plt.ion()
    plt.show()
    raw_input()


def heatmap2d(hm_mat, title='', block=True, fig_num=1, text=True):
    """
    Display heatmap
    input:
      hm_mat:   mxn 2d np array
    """
    print 'map shape: {}, data type: {}'.format(hm_mat.shape, hm_mat.dtype)

    if block:
        plt.figure(fig_num)
        plt.clf()

    # plt.imshow(hm_mat, cmap='hot', interpolation='nearest')
    plt.imshow(hm_mat, interpolation='nearest')
    plt.title(title)
    plt.colorbar()

    if text:
        for y in range(hm_mat.shape[0]):
            for x in range(hm_mat.shape[1]):
                plt.text(x, y, '%.1f' % hm_mat[y, x],
                         horizontalalignment='center',
                         verticalalignment='center',
                         )

    if block:
        plt.ion()
        print 'press enter to continue'
        plt.show()
        raw_input()


def heatmap3d(hm_mat, title=''):
    from mpl_toolkits.mplot3d import Axes3D
    import matplotlib.pyplot as plt
    import numpy as np
    #
    # Assuming you have "2D" dataset like the following that you need
    # to plot.
    #
    data_2d = hm_mat
    #
    # Convert it into an numpy array.
    #
    data_array = np.array(data_2d)
    #
    # Create a figure for plotting the data as a 3D histogram.
    #
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    plt.title(title)
    # _, ax = fig.add_subplot(111)
    #
    # Create an X-Y mesh of the same dimension as the 2D data. You can
    # think of this as the floor of the plot.
    #
    x_data, y_data = np.meshgrid(np.arange(data_array.shape[1]),
                                 np.arange(data_array.shape[0]))
    #
    # Flatten out the arrays so that they may be passed to "ax.bar3d".
    # Basically, ax.bar3d expects three one-dimensional arrays:
    # x_data, y_data, z_data. The following call boils down to picking
    # one entry from each array and plotting a bar to from
    # (x_data[i], y_data[i], 0) to (x_data[i], y_data[i], z_data[i]).
    #
    x_data = x_data.flatten()
    y_data = y_data.flatten()
    z_data = data_array.flatten()
    ax.bar3d(x_data,
             y_data,
             np.zeros(len(z_data)),
             1, 1, z_data)
    #
    # Finally, display the plot.
    #
    plt.show()
    raw_input()