image classification in python

python/caffe/__init__.py

@@ -1,2 +1,3 @@
 from .pycaffe import Net, SGDSolver
+from .classifier import Classifier
 import io

python/caffe/classifier.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+"""
+Classifier is an image classifier specialization of Net.
+"""
+
+import numpy as np
+
+import caffe
+
+
+class Classifier(caffe.Net):
+    """
+    Classifier extends Net for image class prediction
+    by scaling, center cropping, or oversampling.
+    """
+    def __init__(self, model_file, pretrained_file, image_dims=None,
+                 gpu=False, mean_file=None, input_scale=None, channel_swap=None):
+        """
+        Take
+        image_dims: dimensions to scale input for cropping/sampling.
+                    Default is to scale to net input size for whole-image crop.
+        gpu, mean_file, input_scale, channel_swap: convenience params for
+            setting mode, mean, input scale, and channel order.
+        """
+        caffe.Net.__init__(self, model_file, pretrained_file)
+        self.set_phase_test()
+
+        if gpu:
+            self.set_mode_gpu()
+        else:
+            self.set_mode_cpu()
+
+        if mean_file:
+            self.set_mean(self.inputs[0], mean_file)
+        if input_scale:
+            self.set_input_scale(self.inputs[0], input_scale)
+        if channel_swap:
+            self.set_channel_swap(self.inputs[0], channel_swap)
+
+        self.crop_dims = np.array(self.blobs[self.inputs[0]].data.shape[2:])
+        if not image_dims:
+            image_dims = self.crop_dims
+        self.image_dims = image_dims
+
+
+    def predict(self, inputs, oversample=True):
+        """
+        Predict classification probabilities of inputs.
+
+        Take
+        inputs: iterable of (H x W x K) input ndarrays.
+        oversample: average predictions across center, corners, and mirrors
+                    when True (default). Center-only prediction when False.
+
+        Give
+        predictions: (N x C) ndarray of class probabilities
+                     for N images and C classes.
+        """
+        # Scale to standardize input dimensions.
+        inputs = np.asarray([caffe.io.resize_image(im, self.image_dims)
+                             for im in inputs])
+
+        if oversample:
+            # Generate center, corner, and mirrored crops.
+            inputs = caffe.io.oversample(inputs, self.crop_dims)
+        else:
+            # Take center crop.
+            center = np.array(self.image_dims) / 2.0
+            crop = np.tile(center, (1, 2))[0] + np.concatenate([
+                -self.crop_dims / 2.0,
+                self.crop_dims / 2.0
+            ])
+            inputs = inputs[:, crop[0]:crop[2], crop[1]:crop[3], :]
+
+        # Classify
+        caffe_in = self.preprocess(self.inputs[0], inputs)
+        out = self.forward_all(**{self.inputs[0]: caffe_in})
+        predictions = out[self.outputs[0]].squeeze(axis=(2,3))
+
+        # For oversampling, average predictions across crops.
+        if oversample:
+            predictions = predictions.reshape((len(predictions) / 10, 10, -1))
+            predictions = predictions.mean(1)
+
+        return predictions

python/classify.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python
+"""
+classify.py is an out-of-the-box image classifer callable from the command line.
+
+By default it configures and runs the Caffe reference ImageNet model.
+"""
+import numpy as np
+import os
+import sys
+import argparse
+import glob
+import time
+
+import caffe
+
+
+def main(argv):
+    pycaffe_dir = os.path.dirname(__file__)
+
+    parser = argparse.ArgumentParser()
+    # Required arguments: input and output files.
+    parser.add_argument(
+        "input_file",
+        help="Input image, directory, or npy."
+    )
+    parser.add_argument(
+        "output_file",
+        help="Output npy filename."
+    )
+    # Optional arguments.
+    parser.add_argument(
+        "--model_def",
+        default=os.path.join(pycaffe_dir,
+                "../examples/imagenet/imagenet_deploy.prototxt"),
+        help="Model definition file."
+    )
+    parser.add_argument(
+        "--pretrained_model",
+        default=os.path.join(pycaffe_dir,
+                "../examples/imagenet/caffe_reference_imagenet_model"),
+        help="Trained model weights file."
+    )
+    parser.add_argument(
+        "--gpu",
+        action='store_true',
+        help="Switch for gpu computation."
+    )
+    parser.add_argument(
+        "--center_only",
+        action='store_true',
+        help="Switch for prediction from center crop alone instead of " +
+             "averaging predictions across crops (default)."
+    )
+    parser.add_argument(
+        "--images_dim",
+        default='256,256',
+        help="Canonical 'height,width' dimensions of input images."
+    )
+    parser.add_argument(
+        "--mean_file",
+        default=os.path.join(pycaffe_dir,
+                             'caffe/imagenet/ilsvrc_2012_mean.npy'),
+        help="Data set image mean of H x W x K dimensions (numpy array). " +
+             "Set to '' for no mean subtraction."
+    )
+    parser.add_argument(
+        "--input_scale",
+        type=float,
+        default=255,
+        help="Multiply input features by this scale before input to net"
+    )
+    parser.add_argument(
+        "--channel_swap",
+        default='2,1,0',
+        help="Order to permute input channels. The default converts " +
+             "RGB -> BGR since BGR is the Caffe default by way of OpenCV."
+
+    )
+    parser.add_argument(
+        "--ext",
+        default='jpg',
+        help="Image file extension to take as input when a directory " +
+             "is given as the input file."
+    )
+    args = parser.parse_args()
+
+    image_dims = [int(s) for s in args.images_dim.split(',')]
+    channel_swap = [int(s) for s in args.channel_swap.split(',')]
+
+    # Make classifier.
+    classifier = caffe.Classifier(args.model_def, args.pretrained_model,
+            image_dims=image_dims, gpu=args.gpu, mean_file=args.mean_file,
+            input_scale=args.input_scale, channel_swap=channel_swap)
+
+    if args.gpu:
+        print 'GPU mode'
+
+    # Load numpy array (.npy), directory glob (*.jpg), or image file.
+    args.input_file = os.path.expanduser(args.input_file)
+    if args.input_file.endswith('npy'):
+        inputs = np.load(args.input_file)
+    elif os.path.isdir(args.input_file):
+        inputs =[caffe.io.load_image(im_f)
+                 for im_f in glob.glob(args.input_file + '/*.' + args.ext)]
+    else:
+        inputs = [caffe.io.load_image(args.input_file)]
+
+    print "Classifying %d inputs." % len(inputs)
+
+    # Classify.
+    start = time.time()
+    predictions = classifier.predict(inputs, not args.center_only)
+    print "Done in %.2f s." % (time.time() - start)
+
+    # Save
+    np.save(args.output_file, predictions)
+
+
+if __name__ == '__main__':
+    main(sys.argv)