show_energy.py

import matplotlib.pyplot as plt
import seaborn as sns
from skimage.io import imread
from skimage.filters import scharr
from argparse import ArgumentParser

from cropping.edge_stats import get_edge_statistics
from cropping.shrink import crop
from cropping.normalize import normalize


def main():
    parser = ArgumentParser()
    parser.add_argument('file', metavar='path', help='The file to load.')
    parser.add_argument('-t', '--threshold', metavar='value', dest='threshold', default=None, type=float,
                        help='The threshold to use.')
    args = parser.parse_args()

    image = imread(args.file, as_grey=True)

    height = image.shape[0]
    width = image.shape[1]

    # determine the energy
    energy = scharr(image)
    energy = normalize(energy)
    stats = get_edge_statistics(energy, edge_width=50)

    # crop the image
    cropped = crop(image, energy, args.threshold, stats)

    fig, (ax0, ax1) = plt.subplots(ncols=2, figsize=(14, 7), sharex=True, sharey=True)
    plt.gray()

    with sns.axes_style('white'):
        img = ax0.imshow(image, cmap=plt.cm.gray)
        ent = ax1.imshow(energy, cmap=plt.cm.gray)

    ax0.grid(False)
    ax0.axis([0, width, height, 0])
    ax0.set_title('Image')

    ax1.grid(False)
    ax1.set_title('Local energy')

    fig.colorbar(img, ax=ax0)
    fig.colorbar(ent, ax=ax1)

    fig.tight_layout()
    sns.despine()

    if len(cropped) > 0:
        fig, (ax0) = plt.subplots(ncols=1, figsize=(7, 7))
        plt.gray()

        with sns.axes_style('white'):
            img = ax0.imshow(cropped, cmap=plt.cm.gray)

        ax0.grid(False)
        ax0.axis([0, cropped.shape[1], cropped.shape[0], 0])
        ax0.set_title('Cropped Image')

        fig.colorbar(img, ax=ax0)

        fig.tight_layout()
        sns.despine()
    else:
        print('All values below threshold; image cropped to zero-size.')

    plt.show()


if __name__ == '__main__':
    main()