Center of gravity

#!/usr/bin/env python

from matplotlib import pyplot as plt
from scipy import ndimage
import sys
import os
import numpy as np

def centerOfGravity(src, show=False, invert=False, edgeDetect=True):

    """
    Gets the center of gravity for an image

    Arguments:
    src     --  a 3d numpy array or a string corresponding to an image file

    Keyword arguments:
    show    --  indicates whether the results should be plotted (default=False)
    invert  --  indicates whether the colors should be inverted, which makes
                sense if the background is white (default=True)

    Returns:
    An (x, y) tuple of floats that gives the center of gravity for the image
    """

    # Read if the image is a string and not a numpy array
    if isinstance(src, basestring):
        src = plt.imread(src)

    # Invert the 0-1 values
    if invert:
        src = 1-src

    # Edge detection using Sobel operator, see
    # <http://scipy-lectures.github.com/advanced/image_processing/index.html#edge-detection>
    if edgeDetect:
        sx = ndimage.sobel(src, axis=0, mode='constant')
        sy = ndimage.sobel(src, axis=1, mode='constant')
        src = np.hypot(sx, sy)

    # Get center of mass
    y, x, col = ndimage.measurements.center_of_mass(src)

    # Optionally plot
    if show:
        plt.imshow(src)
        plt.axvline(x)
        plt.axhline(y)
        plt.annotate('%.1f,%.1f' % (x,y), (x,y))
        plt.show()

    return x, y

if __name__ == '__main__':

    path = 'images'
    for src in os.listdir(path):
        x, y = centerOfGravity(os.path.join(path, src), show=True,
            invert='--invert' in sys.argv,
            edgeDetect='--edgeDetect' in sys.argv)
        print 'The center of gravity is %.1f, %.1f' % (x,y)