repo_description: An absolute beginner's guide to Machine Learning and Image Classification with Neural Networks
repo_topics: caffe image-classification machine-learning neural-network tutorial
code: # python classify-samples.py -c ~/repos/caffe/ -m ../data/googlenet-dolphins-and-seahorses
  
  import numpy as np
  import os
  import sys
  import argparse
  
  class ImageClassifier:
      def __init__(self, model_dir):
          self.model_dir = model_dir
          deploy_file = os.path.join(model_dir, 'deploy.prototxt')
          weights_file = os.path.join(model_dir, 'snapshot_iter_90.caffemodel')
          self.net = caffe.Net(deploy_file, caffe.TEST, weights=weights_file)
  
      def setup(self):
          mean_file = os.path.join(self.model_dir, 'mean.binaryproto')
          labels_file = os.path.join(self.model_dir, 'labels.txt')
  
          self.transformer = caffe.io.Transformer({'data': self.net.blobs['data'].data.shape})
          self.transformer.set_transpose('data', (2, 0, 1))
          self.transformer.set_raw_scale('data', 255)
          self.transformer.set_channel_swap('data', (2, 1, 0))
  
          # set mean pixel
          with open(mean_file, 'rb') as infile:
              blob = caffe_pb2.BlobProto()
              blob.MergeFromString(infile.read())
              if blob.HasField('shape'):
                  blob_dims = blob.shape
                  assert len(blob_dims) == 4, 'Shape should have 4 dimensions - shape is %s' % blob.shape
              elif blob.HasField('num') and blob.HasField('channels') and \
                      blob.HasField('height') and blob.HasField('width'):
                  blob_dims = (blob.num, blob.channels, blob.height, blob.width)
              else:
                  raise ValueError('blob does not provide shape or 4d dimensions')
              pixel = np.reshape(blob.data, blob_dims[1:]).mean(1).mean(1)
              self.transformer.set_mean('data', pixel)
  
          # This is overkill here, since we only have 2 labels, but here's how we might read them.
          # Later we'd grab the label we want based on position (e.g., 0=dolphin, 1=seahorse)
          self.labels = np.loadtxt(labels_file, str, delimiter='\n')
  
      def classify(self, fullpath):
          # Load the image from disk using caffe's built-in I/O module
          image = caffe.io.load_image(fullpath)
          # Preprocess the image into the proper format for feeding into the model
          self.net.blobs['data'].data[...] = self.transformer.preprocess('data', image)
          # Run the image's pixel data through the network
          out = self.net.forward()
          # Extract the probabilities of our two categories from the final layer
          softmax_layer = out['softmax']
          # Here we're converting to Python types from ndarray floats
          dolphin_prob = softmax_layer.item(0)
          seahorse_prob = softmax_layer.item(1)
  
          # Print the results. I'm using labels
          label = self.labels[0] if dolphin_prob > seahorse_prob else self.labels[1]
          filename = os.path.basename(fullpath)
          print '%s is a %s dolphin=%.3f%% seahorse=%.3f%%' % (filename, label, dolphin_prob*100, seahorse_prob*100)
  
  def setup_caffe(caffe_root):
      # Load Caffe's Python interface from the specified path
      sys.path.insert(0, os.path.join(caffe_root, 'python'))
      global caffe
      global caffe_pb2
      import caffe
      from caffe.proto import caffe_pb2
  
      # Set Caffe to use CPU mode so this will work on as many machines as possible.
      caffe.set_mode_cpu()
  
  def main():
      parser = argparse.ArgumentParser(
          description='Classify images of dolphins and seahorses using trained Caffe model'
      )
      parser.add_argument('-c', '--caffe_root', help='CAFFE_ROOT dir, if not defined in env')
      parser.add_argument('-m', '--model_dir', help='Trained model dir, downloaded from DIGITS')
      parser.add_argument('-d', '--images_dir', help='Directory of images to classify')
  
      args = parser.parse_args()
  
      # Prefer $CAFFE_ROOT in the env if it exists, otherwise get from args
      caffe_root = os.getenv('CAFFE_ROOT') or args.caffe_root
      if not caffe_root:
          print 'Error: Missing CAFFE_ROOT dir. Set env variable or pass via --caffe_root'
          parser.print_help()
          sys.exit(1)
      setup_caffe(caffe_root)
      
      model_dir = args.model_dir
      if not model_dir or not os.path.isdir(model_dir):
          print 'Error: Unable to find model files. Pass dir via --model_dir'
          parser.print_help()
          sys.exit(1)
      classifier = ImageClassifier(model_dir)
      classifier.setup()
  
      # Allow passing images dir, or use ../data/untrained-samples by default
      cwd = os.path.dirname(os.path.abspath(__file__))
      untrained_samples = os.path.join(cwd, '..', 'data', 'untrained-samples')
      images_dir = args.images_dir or untrained_samples
      if not os.path.isdir(images_dir):
          print 'Error: Unable to find images for classification. Pass dir via --images_dir'
          parser.print_help()        
          sys.exit(1)
      
      # Classify all images in images_dir using our trained network
      for filename in os.listdir(images_dir):
          if filename.endswith('.jpg'):
              classifier.classify(os.path.join(images_dir, filename))
  
  if __name__ == '__main__':
      main()
  