ColorModel.py

# SimpleCV Color Model Library
#load required libraries
from SimpleCV.base import *
from SimpleCV.ImageClass import *


class ColorModel:
    """
    **SUMMARY**

    The color model is used to model the color of foreground and background objects
    by using a a training set of images.

    You can create the color model with any number of "training" images, or
    add images to the model with add() and remove().  Then for your data images,
    you can useThresholdImage() to return a segmented picture.

    """
    #TODO: Discretize the colorspace into smaller intervals,eg r=[0-7][8-15] etc
    #TODO: Work in HSV space
    mIsBackground = True
    mData = {}
    mBits = 1

    def __init__(self, data = None, isBackground=True):
        self.mIsBackground = isBackground
        self.mData = {}
        self.mBits = 1

        if data:
            try:
                [ self.add(d) for d in data ]
            except TypeError:
                self.add(data)


    def _makeCanonical(self, data):
        """
        Turn input types in a common form used by the rest of the class -- a
        4-bit shifted list of unique colors
        """
        ret = ''

        #first cast everything to a numpy array
        if(data.__class__.__name__ == 'Image'):
            ret =  data.getNumpy().reshape(-1, 3)
        elif(data.__class__.__name__ == 'cvmat'):
            ret = np.array(data).reshape(-1, 3)
        elif(data.__class__.__name__ == 'list'  ):
            temp = []
            for d in data: #do the bgr conversion
                t = (d[2],d[1],d[0])
                temp.append(t)
            ret = np.array(temp,dtype='uint8')
        elif (data.__class__.__name__=='tuple'):
            ret = np.array((data[2],data[1],data[0]),'uint8')
        elif(data.__class__.__name__=='np.array'):
            ret = data
        else:
            logger.warning("ColorModel: color is not in an accepted format!")
            return None

        rs = np.right_shift(ret, self.mBits)  #right shift 4 bits

        if( len(rs.shape) > 1 ):
            uniques = np.unique(rs.view([('',rs.dtype)]*rs.shape[1])).view(rs.dtype).reshape(-1, 3)
        else:
            uniques = [rs]
        #create a unique set of colors.  I had to look this one up

        #create a dict of encoded strings
        return dict.fromkeys(list(map(np.ndarray.tostring, uniques)), 1)

    def reset(self):
        """
        **SUMMARY**
        Resets the color model. I.e. clears it out the stored values.


        **RETURNS**

        Nothing.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.add(Image("lenna))
        >>> cm.clear()

        """
        self.mData = {}

    def add(self, data):
        """
        **SUMMARY**

        Add an image, array, or tuple to the color model.

        **PARAMETERS**

        * *data* - An image, array, or tupple of values to the color model.

        **RETURNS**

        Nothings.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.add(Image("lenna))
        >>> cm.clear()

        """
        self.mData.update(self._makeCanonical(data))

    def remove(self, data):
        """
        **SUMMARY**

        Remove an image, array, or tuple from the model.

        **PARAMETERS**

        * *data* - An image, array, or tupple of value.

        **RETURNS**

        Nothings.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.add(Image("lenna))
        >>> cm.remove(Color.BLACK)

        """
        self.mData = dict.fromkeys(set(self.mData) ^ set(self._makeCanonical(data)), 1)

    def threshold(self, img):
        """
        **SUMMARY**

        Perform a threshold operation on the given image. This involves iterating
        over the image and comparing each pixel to the model. If the pixel is in the
        model it is set to be either the foreground (white) or background (black) based
        on the setting of mIsBackground.

        **PARAMETERS**

        * *img* - the image to perform the threshold on.

        **RETURNS**

        The thresholded image.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.add(color.RED)
        >>> cm.add(color.BLUE)
        >>> result = cm.threshold(Image("lenna")
        >>> result.show()

        """
        a = 0
        b = 255
        if( self.mIsBackground == False ):
            a = 255
            b = 0

        rs = np.right_shift(img.getNumpy(), self.mBits).reshape(-1, 3) #bitshift down and reshape to Nx3
        mapped = np.array(list(map(self.mData.has_key, list(map(np.ndarray.tostring, rs))))) #map to True/False based on the model
        thresh = np.where(mapped, a, b) #replace True and False with fg and bg
        return Image(thresh.reshape(img.width, img.height))

    def contains(self, c):
        """
        **SUMMARY**

        Return true if a particular color is in our color model.

        **PARAMETERS**

        * *c* - A three value color tupple.

        **RETURNS**

        Returns True if the color is in the model, False otherwise.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.add(Color.RED)
        >>> cm.add(Color.BLUE)
        >>> if( cm.contains(Color.RED) )
        >>>   print "Yo - we gots red y'all."


       """
        #reverse the color, cast to uint8, right shift, convert to string, check dict
        return np.right_shift(np.cast['uint8'](c[::-1]), self.mBits).tostring() in self.mData

    def setIsForeground(self):
        """
        **SUMMARY**

        Set our model as being foreground imagery. I.e. things in the model are the foreground
        and will be marked as white during the threhsold operation.

        **RETURNS**

        Nothing.

        """
        mIsBackground = False

    def setIsBackground(self):
        """
        **SUMMARY**

        Set our model as being background imagery. I.e. things in the model are the background
        and will be marked as black during the threhsold operation.


        **RETURNS**

        Nothing.

        """
        mIsBackground = True

    def load(self, filename):
        """
        **SUMMARY**

        Load the color model from the specified file.

        **TO DO**

        This should be converted to pickle.

        **PARAMETERS**

        * *filename* - The file name and path to load the data from.

        **RETURNS**

        Nothing.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.load("myColors.txt")
        >>> cm.add(Color.RED)
        >>> cm.add(Color.BLUE)
        >>> cm.save("mymodel)

        """
        self.mData =  load(open(filename))

    def save(self, filename):
        """
        **SUMMARY**

        Save a color model file.

        **PARAMETERS**

        * *filename* - The file name and path to save the data to.

        **RETURNS**

        Nothing.

        **EXAMPLE**

        >>> cm = ColorModel()
        >>> cm.add(Color.RED)
        >>> cm.add(Color.BLUE)
        >>> cm.save("mymodel.txt")

        **TO DO**

        This should be converted to pickle.

        """
        dump(self.mData, open(filename, "wb"))