content: from SimpleCV.base import *
  from SimpleCV.ImageClass import Image, ImageSet
  from SimpleCV.DrawingLayer import *
  from SimpleCV.Features import FeatureExtractorBase
  
  
  """
  This class is encapsulates almost everything needed to train, test, and deploy a
  multiclass decision tree / forest image classifier. Training data should
  be stored in separate directories for each class. This class uses the feature
  extractor base class to  convert images into a feature vector. The basic workflow
  is as follows.
  1. Get data.
  2. Setup Feature Extractors (roll your own or use the ones I have written).
  3. Train the classifier.
  4. Test the classifier.
  5. Tweak parameters as necessary.
  6. Repeat until you reach the desired accuracy.
  7. Save the classifier.
  8. Deploy using the classify method.
  """
  class TreeClassifier:
      """
      This method encapsulates a number of tree-based machine learning approaches
      and associated meta algorithms.
  
      Decision trees:
      http://en.wikipedia.org/wiki/Decision_trees
  
      boosted adpative decision trees
      http://en.wikipedia.org/wiki/Adaboost
  
      random forrests
      http://en.wikipedia.org/wiki/Random_forest
  
      bagging (bootstrap aggregating)
      http://en.wikipedia.org/wiki/Bootstrap_aggregating
      """
      mClassNames = []
      mDataSetRaw = []
      mDataSetOrange = []
      mClassifier = None
      mLearner = None
      mTree = None
      mFeatureExtractors = None
      mOrangeDomain = None
      mFlavorParams = None
  
      mTreeTypeDict = {
          "Tree":0,    # A vanilla classification tree
          "Bagged":1, # Bootstrap aggregating aka bagging - make new data sets and test on them
          "Forest":2, # Lots of little trees
          "Boosted":3 # Highly optimized trees.
      }
  
      mforestFlavorDict = {
          "NTrees":100, #number of trees in our forest
          "NAttributes":None # number of attributes per split sqrt(features) is default
       }
      mBoostedFlavorDict = {
          "NClassifiers":10, #number of learners
      }
      mBaggedFlavorDict = {
          "NClassifiers":10, #numbers of classifiers / tree splits
      }
  
      def __init__(self,featureExtractors=[],flavor='Tree',flavorDict=None):
          """
          dist = distance algorithm
          k = number of nearest neighbors
          """
          if not ORANGE_ENABLED:
              logger.warning("I'm sorry, but you need the orange machine learning library installed to use this")
              return None
  
          self.mClassNames = []
          self.mDataSetRaw = []
          self.mDataSetOrange = []
          self.mClassifier = None
          self.mLearner = None
          self.mTree = None
          self.mFeatureExtractors = None
          self.mOrangeDomain = None
          self.mFlavorParams = None
          self.mFlavor = self.mTreeTypeDict[flavor]
          if(flavorDict is None):
              if(self.mFlavor == self.mTreeTypeDict["Bagged"]):
                  self.mFlavorParams = self.mBaggedFlavorDict
              elif(self.mFlavor == self.mTreeTypeDict["Forest"]):
                  self.mFlavorParams = self.mforestFlavorDict #mmmm tastes like   pinecones and squirrels
              elif(self.mFlavor == self.mTreeTypeDict["Boosted"]):
                  self.mFlavorParams = self.mBoostedFlavorDict
          else:
              self.mFlavorParams = flavorDict
          self.mFeatureExtractors =  featureExtractors
  
      def load(cls, fname):
          """
          Load the classifier from file
          """
          return pickle.load(file(fname))
      load = classmethod(load)
  
  
      def save(self, fname):
          """
          Save the classifier to file
          """
          output = open(fname, 'wb')
          pickle.dump(self,output,2) # use two otherwise it w
          output.close()
  
      def __getstate__(self):
          mydict = self.__dict__.copy()
          self.mDataSetOrange = None
          del mydict['mDataSetOrange']
          self.mOrangeDomain = None
          del mydict['mOrangeDomain']
          self.mLearner = None
          del mydict['mLearner']
          self.mTree = None
          del mydict['mTree']
          return mydict
  
      def __setstate__(self, mydict):
          self.__dict__ = mydict
          colNames = []
          for extractor in self.mFeatureExtractors:
              colNames.extend(extractor.getFieldNames())
          self.mOrangeDomain = orange.Domain(map(orange.FloatVariable,colNames),orange.EnumVariable("type",values=self.mClassNames))
          self.mDataSetOrange = orange.ExampleTable(self.mOrangeDomain,self.mDataSetRaw)
          if(self.mFlavor == 0):
              self.mLearner =  orange.TreeLearner()
              self.mClassifier = self.mLearner(self.mDataSetOrange)
          elif(self.mFlavor == 1): #bagged
              self.mTree =  orange.TreeLearner()
              self.mLearner = orngEnsemble.BaggedLearner(self.mTree,t=self.mFlavorParams["NClassifiers"])
              self.mClassifier = self.mLearner(self.mDataSetOrange)
          elif(self.mFlavor == 2):#forest
              self.mTree =  orange.TreeLearner()
              self.mLearner =  orngEnsemble.RandomForestLearner(trees=self.mFlavorParams["NTrees"], attributes=self.mFlavorParams["NAttributes"])
              self.mClassifier = self.mLearner(self.mDataSetOrange)
          elif(self.mFlavor == 3):#boosted
              self.mTree =  orange.TreeLearner()
              self.mLearner = orngEnsemble.BoostedLearner(self.mTree,t=self.mFlavorParams["NClassifiers"])
              self.mClassifier = self.mLearner(self.mDataSetOrange)
  
  
      def classify(self, image):
          """
          Classify a single image. Takes in an image and returns the string
          of the classification.
  
          Make sure you haved loaded the feauture extractors and the training data.
  
          """
          featureVector = []
          for extractor in self.mFeatureExtractors: #get the features
              feats = extractor.extract(image)
              if( feats is not None ):
                  featureVector.extend(feats)
          featureVector.extend([self.mClassNames[0]])
          test = orange.ExampleTable(self.mOrangeDomain,[featureVector])
          c = self.mClassifier(test[0]) #classify
          return str(c) #return to class name
  
  
      def setFeatureExtractors(self, extractors):
          """
          Add a list of feature extractors to the classifier. These feature extractors
          must match the ones used to train the classifier. If the classifier is already
          trained then this method will require that you retrain the data.
          """
          self.mFeatureExtractors = extractors
          return None
  
      def _trainPath(self,path,className,subset,disp,verbose):
          count = 0
          files = []
          for ext in IMAGE_FORMATS:
              files.extend(glob.glob( os.path.join(path, ext)))
          if(subset > 0):
              nfiles = min(subset,len(files))
          else:
              nfiles = len(files)
          badFeat = False
          for i in range(nfiles):
              infile = files[i]
              if verbose:
                  print "Opening file: " + infile
              img = Image(infile)
              featureVector = []
              for extractor in self.mFeatureExtractors:
                  feats = extractor.extract(img)
                  if( feats is not None ):
                      featureVector.extend(feats)
                  else:
                      badFeat = True
  
              if(badFeat):
                  badFeat = False
                  continue
  
              featureVector.extend([className])
              self.mDataSetRaw.append(featureVector)
              text = 'Training: ' + className
              self._WriteText(disp,img,text,Color.WHITE)
              count = count + 1
              del img
          return count
  
      def _trainImageSet(self,imageset,className,subset,disp,verbose):
          count = 0
          badFeat = False
          if (subset>0):
              imageset = imageset[0:subset]   
          for img in imageset:
              if verbose:
                  print "Opening file: " + img.filename
              featureVector = []
              for extractor in self.mFeatureExtractors:
                  feats = extractor.extract(img)
                  if( feats is not None ):
                      featureVector.extend(feats)
                  else:
                      badFeat = True
                      
              if(badFeat):
                  badFeat = False
                  continue
              
              featureVector.extend([className])
              self.mDataSetRaw.append(featureVector)
              text = 'Training: ' + className
              self._WriteText(disp,img,text,Color.WHITE)
              count = count + 1
              del img
          return count
  
      def train(self,images,classNames,disp=None,subset=-1,savedata=None,verbose=True):
          """
          Train the classifier.
          images paramater can take in a list of paths or a list of imagesets
          images - the order of the paths or imagesets must be in the same order as the class type
  
          - Note all image classes must be in seperate directories
          - The class names must also align to the directories or imagesets
  
          disp - if display is a display we show images and class label,
          otherwise nothing is done.
  
          subset - if subset = -1 we use the whole dataset. If subset = # then we
          use min(#images,subset)
  
          savedata - if save data is None nothing is saved. If savedata is a file
          name we save the data to a tab delimited file.
  
          verbose - print confusion matrix and file names
          returns [%Correct %Incorrect Confusion_Matrix]
          """
          #if( (self.mFlavor == 1 or self.mFlavor == 3) and len(classNames) > 2):
          #    logger.warning("Boosting / Bagging only works for binary classification tasks!!!")
  
          count = 0
          self.mClassNames = classNames
          # for each class, get all of the data in the path and train
          for i in range(len(classNames)):
              if ( isinstance(images[i], str) ):
                  count = count + self._trainPath(images[i],classNames[i],subset,disp,verbose)
              else:
                  count = count + self._trainImageSet(images[i],classNames[i],subset,disp,verbose)
  
          colNames = []
          for extractor in self.mFeatureExtractors:
              colNames.extend(extractor.getFieldNames())
  
          if(count <= 0):
              logger.warning("No features extracted - bailing")
              return None
  
          self.mOrangeDomain = orange.Domain(map(orange.FloatVariable,colNames),orange.EnumVariable("type",values=self.mClassNames))
          self.mDataSetOrange = orange.ExampleTable(self.mOrangeDomain,self.mDataSetRaw)
          if(savedata is not None):
              orange.saveTabDelimited (savedata, self.mDataSetOrange)
  
          if(self.mFlavor == 0):
              self.mLearner =  orange.TreeLearner()
              self.mClassifier = self.mLearner(self.mDataSetOrange)
          elif(self.mFlavor == 1): #bagged
              self.mTree =  orange.TreeLearner()
              self.mLearner = orngEnsemble.BaggedLearner(self.mTree,t=self.mFlavorParams["NClassifiers"])
              self.mClassifier = self.mLearner(self.mDataSetOrange)
          elif(self.mFlavor == 2):#forest
              self.mTree =  orange.TreeLearner()
              self.mLearner =  orngEnsemble.RandomForestLearner(trees=self.mFlavorParams["NTrees"], attributes=self.mFlavorParams["NAttributes"])
              self.mClassifier = self.mLearner(self.mDataSetOrange)
          elif(self.mFlavor == 3):#boosted
              self.mTree =  orange.TreeLearner()
              self.mLearner = orngEnsemble.BoostedLearner(self.mTree,t=self.mFlavorParams["NClassifiers"])
              self.mClassifier = self.mLearner(self.mDataSetOrange)
  
          correct = 0
          incorrect = 0
          for i in range(count):
              c = self.mClassifier(self.mDataSetOrange[i])
              test = self.mDataSetOrange[i].getclass()
              if verbose:
                  print "original", test, "classified as", c
              if(test==c):
                  correct = correct + 1
              else:
                  incorrect = incorrect + 1
  
          good = 100*(float(correct)/float(count))
          bad = 100*(float(incorrect)/float(count))
  
          confusion = 0
          if( len(self.mClassNames) > 2 ):
              crossValidator = orngTest.learnAndTestOnLearnData([self.mLearner],self.mDataSetOrange)
              confusion = orngStat.confusionMatrices(crossValidator)[0]
  
          if verbose:
              print("Correct: "+str(good))
              print("Incorrect: "+str(bad))
              if( confusion != 0 ):
                  classes = self.mDataSetOrange.domain.classVar.values
                  print "\t"+"\t".join(classes)
                  for className, classConfusions in zip(classes, confusion):
                      print ("%s" + ("\t%i" * len(classes))) % ((className, ) + tuple(    classConfusions))
  
          if(self.mFlavor == 0):
              self._PrintTree(self.mClassifier)
  
          return [good, bad, confusion]
  
  
  
  
      def test(self,images,classNames,disp=None,subset=-1,savedata=None,verbose=True):
          """
          Test the classifier.
          images paramater can take in a list of paths or a list of imagesets
          images - the order of the paths or imagesets must be in the same order as the class type
  
          - Note all image classes must be in seperate directories
          - The class names must also align to the directories or imagesets
  
          disp - if display is a display we show images and class label,
          otherwise nothing is done.
  
          subset - if subset = -1 we use the whole dataset. If subset = # then we
          use min(#images,subset)
  
          savedata - if save data is None nothing is saved. If savedata is a file
          name we save the data to a tab delimited file.
  
          verbose - print confusion matrix and file names
          returns [%Correct %Incorrect Confusion_Matrix]
          """
          count = 0
          correct = 0
          self.mClassNames = classNames
          colNames = []
          for extractor in self.mFeatureExtractors:
              colNames.extend(extractor.getFieldNames())
              if(self.mOrangeDomain is None):
                  self.mOrangeDomain = orange.Domain(map(orange.FloatVariable,colNames),orange.EnumVariable("type",values=self.mClassNames))
  
          dataset = []
          for i in range(len(classNames)):
              if ( isinstance(images[i],str) ):
                  [dataset,cnt,crct] =self._testPath(images[i],classNames[i],dataset,subset,disp,verbose)
                  count = count + cnt
                  correct = correct + crct
              else:
                  [dataset,cnt,crct] =self._testImageSet(images[i],classNames[i],dataset,subset,disp,verbose)
                  count = count + cnt
                  correct = correct + crct
  
  
          testData = orange.ExampleTable(self.mOrangeDomain,dataset)
  
          if savedata is not None:
              orange.saveTabDelimited (savedata, testData)
  
          confusion = 0
          if( len(self.mClassNames) > 2 ):
              crossValidator = orngTest.learnAndTestOnTestData([self.mLearner],self.mDataSetOrange,testData)
              confusion = orngStat.confusionMatrices(crossValidator)[0]
  
          good = 100*(float(correct)/float(count))
          bad = 100*(float(count-correct)/float(count))
          if verbose:
              print("Correct: "+str(good))
              print("Incorrect: "+str(bad))
              if( confusion != 0 ):
                  classes = self.mDataSetOrange.domain.classVar.values
                  print "\t"+"\t".join(classes)
                  for className, classConfusions in zip(classes, confusion):
                      print ("%s" + ("\t%i" * len(classes))) % ((className, ) + tuple(    classConfusions))
          return [good, bad, confusion]
  
      def _testPath(self,path,className,dataset,subset,disp,verbose):
          count = 0
          correct = 0
          badFeat = False
          files = []
          for ext in IMAGE_FORMATS:
              files.extend(glob.glob( os.path.join(path, ext)))
          if(subset > 0):
              nfiles = min(subset,len(files))
          else:
              nfiles = len(files)
          for i in range(nfiles):
              infile = files[i]
              if verbose:
                  print "Opening file: " + infile
              img = Image(infile)
              featureVector = []
              for extractor in self.mFeatureExtractors:
                  feats = extractor.extract(img)
                  if( feats is not None ):
                      featureVector.extend(feats)
                  else:
                      badFeat = True
              if( badFeat ):
                  del img
                  badFeat = False
                  continue
              featureVector.extend([className])
              dataset.append(featureVector)
              test = orange.ExampleTable(self.mOrangeDomain,[featureVector])
              c = self.mClassifier(test[0])
              testClass = test[0].getclass()
              if(testClass==c):
                  text =  "Classified as " + str(c)
                  self._WriteText(disp,img,text, Color.GREEN)
                  correct = correct + 1
              else:
                  text =  "Mislassified as " + str(c)
                  self._WriteText(disp,img,text, Color.RED)
              count = count + 1
              del img
  
          return([dataset,count,correct])
  
      def _testImageSet(self,imageset,className,dataset,subset,disp,verbose):
          count = 0
          correct = 0
          badFeat = False
          if(subset > 0):
              imageset = imageset[0:subset]
          for img in imageset:
              if verbose:
                  print "Opening file: " + img.filename
              featureVector = []
              for extractor in self.mFeatureExtractors:
                  feats = extractor.extract(img)
                  if( feats is not None ):
                      featureVector.extend(feats)
                  else:
                      badFeat = True
              if( badFeat ):
                  del img
                  badFeat = False
                  continue 
              featureVector.extend([className])
              dataset.append(featureVector)
              test = orange.ExampleTable(self.mOrangeDomain,[featureVector])
              c = self.mClassifier(test[0])
              testClass = test[0].getclass()
              if(testClass==c):
                  text =  "Classified as " + str(c)
                  self._WriteText(disp,img,text, Color.GREEN)
                  correct = correct + 1
              else:   
                  text =  "Mislassified as " + str(c)
                  self._WriteText(disp,img,text, Color.RED)
              count = count + 1
              del img
              
          return([dataset,count,correct])
  
      def _WriteText(self, disp, img, txt,color):
          if(disp is not None):
              txt = ' ' + txt + ' '
              img = img.adaptiveScale(disp.resolution)
              layer = DrawingLayer((img.width,img.height))
              layer.setFontSize(60)
              layer.ezViewText(txt,(20,20),fgcolor=color)
              img.addDrawingLayer(layer)
              img.applyLayers()
              img.save(disp)
  
      def _PrintTree(self,x):
          #adapted from the orange documentation
          if type(x) == orange.TreeClassifier:
              self._PrintTree0(x.tree, 0)
          elif type(x) == orange.TreeNode:
              self._PrintTree0(x, 0)
          else:
              raise TypeError, "invalid parameter"
  
      def _PrintTree0(self,node,level):
          #adapted from the orange documentation
          if not node:
              print " "*level + "<null node>"
              return
  
          if node.branchSelector:
              nodeDesc = node.branchSelector.classVar.name
              nodeCont = node.distribution
              print "\n" + "   "*level + "%s (%s)" % (nodeDesc, nodeCont),
              for i in range(len(node.branches)):
                  print "\n" + "   "*level + ": %s" % node.branchDescriptions[i],
                  self._PrintTree0(node.branches[i], level+1)
          else:
              nodeCont = node.distribution
              majorClass = node.nodeClassifier.defaultValue
              print "--> %s (%s) " % (majorClass, nodeCont)
  