|
|
|
|
|
""" |
|
|
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() |
|
|
|
|
|
parser.add_argument( |
|
|
"input_file", |
|
|
help="Input image, directory, or npy." |
|
|
) |
|
|
parser.add_argument( |
|
|
"output_file", |
|
|
help="Output npy filename." |
|
|
) |
|
|
|
|
|
parser.add_argument( |
|
|
"--model_def", |
|
|
default=os.path.join(pycaffe_dir, |
|
|
"../models/bvlc_reference_caffenet/deploy.prototxt"), |
|
|
help="Model definition file." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--pretrained_model", |
|
|
default=os.path.join(pycaffe_dir, |
|
|
"../models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel"), |
|
|
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 [Channels x Height x Width] dimensions " + |
|
|
"(numpy array). Set to '' for no mean subtraction." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--input_scale", |
|
|
type=float, |
|
|
help="Multiply input features by this scale to finish preprocessing." |
|
|
) |
|
|
parser.add_argument( |
|
|
"--raw_scale", |
|
|
type=float, |
|
|
default=255.0, |
|
|
help="Multiply raw input by this scale before preprocessing." |
|
|
) |
|
|
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(',')] |
|
|
|
|
|
mean, channel_swap = None, None |
|
|
if args.mean_file: |
|
|
mean = np.load(args.mean_file) |
|
|
if args.channel_swap: |
|
|
channel_swap = [int(s) for s in args.channel_swap.split(',')] |
|
|
|
|
|
if args.gpu: |
|
|
caffe.set_mode_gpu() |
|
|
print("GPU mode") |
|
|
else: |
|
|
caffe.set_mode_cpu() |
|
|
print("CPU mode") |
|
|
|
|
|
|
|
|
classifier = caffe.Classifier(args.model_def, args.pretrained_model, |
|
|
image_dims=image_dims, mean=mean, |
|
|
input_scale=args.input_scale, raw_scale=args.raw_scale, |
|
|
channel_swap=channel_swap) |
|
|
|
|
|
|
|
|
args.input_file = os.path.expanduser(args.input_file) |
|
|
if args.input_file.endswith('npy'): |
|
|
print("Loading file: %s" % args.input_file) |
|
|
inputs = np.load(args.input_file) |
|
|
elif os.path.isdir(args.input_file): |
|
|
print("Loading folder: %s" % args.input_file) |
|
|
inputs =[caffe.io.load_image(im_f) |
|
|
for im_f in glob.glob(args.input_file + '/*.' + args.ext)] |
|
|
else: |
|
|
print("Loading file: %s" % args.input_file) |
|
|
inputs = [caffe.io.load_image(args.input_file)] |
|
|
|
|
|
print("Classifying %d inputs." % len(inputs)) |
|
|
|
|
|
|
|
|
start = time.time() |
|
|
predictions = classifier.predict(inputs, not args.center_only) |
|
|
print("Done in %.2f s." % (time.time() - start)) |
|
|
|
|
|
|
|
|
print("Saving results into %s" % args.output_file) |
|
|
np.save(args.output_file, predictions) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
main(sys.argv) |
|
|
|
