python/classify.py

   1 #!/usr/bin/env python
   2 """
   3 classify.py is an out-of-the-box image classifer callable from the command line.
   4
   5 By default it configures and runs the Caffe reference ImageNet model.
   6 """
   7 import numpy as np
   8 import os
   9 import sys
  10 import argparse
  11 import glob
  12 import time
  13
  14 import caffe
  15
  16
  17 def main(argv):
  18     pycaffe_dir = os.path.dirname(__file__)
  19
  20     parser = argparse.ArgumentParser()
  21     # Required arguments: input and output files.
  22     parser.add_argument(
  23         "input_file",
  24         help="Input image, directory, or npy."
  25     )
  26     parser.add_argument(
  27         "output_file",
  28         help="Output npy filename."
  29     )
  30     # Optional arguments.
  31     parser.add_argument(
  32         "--model_def",
  33         default=os.path.join(pycaffe_dir,
  34                 "../models/bvlc_reference_caffenet/deploy.prototxt"),
  35         help="Model definition file."
  36     )
  37     parser.add_argument(
  38         "--pretrained_model",
  39         default=os.path.join(pycaffe_dir,
  40                 "../models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel"),
  41         help="Trained model weights file."
  42     )
  43     parser.add_argument(
  44         "--gpu",
  45         action='store_true',
  46         help="Switch for gpu computation."
  47     )
  48     parser.add_argument(
  49         "--center_only",
  50         action='store_true',
  51         help="Switch for prediction from center crop alone instead of " +
  52              "averaging predictions across crops (default)."
  53     )
  54     parser.add_argument(
  55         "--images_dim",
  56         default='256,256',
  57         help="Canonical 'height,width' dimensions of input images."
  58     )
  59     parser.add_argument(
  60         "--mean_file",
  61         default=os.path.join(pycaffe_dir,
  62                              'caffe/imagenet/ilsvrc_2012_mean.npy'),
  63         help="Data set image mean of [Channels x Height x Width] dimensions " +
  64              "(numpy array). Set to '' for no mean subtraction."
  65     )
  66     parser.add_argument(
  67         "--input_scale",
  68         type=float,
  69         help="Multiply input features by this scale to finish preprocessing."
  70     )
  71     parser.add_argument(
  72         "--raw_scale",
  73         type=float,
  74         default=255.0,
  75         help="Multiply raw input by this scale before preprocessing."
  76     )
  77     parser.add_argument(
  78         "--channel_swap",
  79         default='2,1,0',
  80         help="Order to permute input channels. The default converts " +
  81              "RGB -> BGR since BGR is the Caffe default by way of OpenCV."
  82     )
  83     parser.add_argument(
  84         "--ext",
  85         default='jpg',
  86         help="Image file extension to take as input when a directory " +
  87              "is given as the input file."
  88     )
  89     args = parser.parse_args()
  90
  91     image_dims = [int(s) for s in args.images_dim.split(',')]
  92
  93     mean, channel_swap = None, None
  94     if args.mean_file:
  95         mean = np.load(args.mean_file)
  96     if args.channel_swap:
  97         channel_swap = [int(s) for s in args.channel_swap.split(',')]
  98
  99     if args.gpu:
 100         caffe.set_mode_gpu()
 101         print("GPU mode")
 102     else:
 103         caffe.set_mode_cpu()
 104         print("CPU mode")
 105
 106     # Make classifier.
 107     classifier = caffe.Classifier(args.model_def, args.pretrained_model,
 108             image_dims=image_dims, mean=mean,
 109             input_scale=args.input_scale, raw_scale=args.raw_scale,
 110             channel_swap=channel_swap)
 111
 112     # Load numpy array (.npy), directory glob (*.jpg), or image file.
 113     args.input_file = os.path.expanduser(args.input_file)
 114     if args.input_file.endswith('npy'):
 115         print("Loading file: %s" % args.input_file)
 116         inputs = np.load(args.input_file)
 117     elif os.path.isdir(args.input_file):
 118         print("Loading folder: %s" % args.input_file)
 119         inputs =[caffe.io.load_image(im_f)
 120                  for im_f in glob.glob(args.input_file + '/*.' + args.ext)]
 121     else:
 122         print("Loading file: %s" % args.input_file)
 123         inputs = [caffe.io.load_image(args.input_file)]
 124
 125     print("Classifying %d inputs." % len(inputs))
 126
 127     # Classify.
 128     start = time.time()
 129     predictions = classifier.predict(inputs, not args.center_only)
 130     print("Done in %.2f s." % (time.time() - start))
 131
 132     # Save
 133     print("Saving results into %s" % args.output_file)
 134     np.save(args.output_file, predictions)
 135
 136
 137 if __name__ == '__main__':
 138     main(sys.argv)