You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
120 lines
3.5 KiB
120 lines
3.5 KiB
#!/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)
|
|
|