зеркало из https://github.com/microsoft/caffe.git
131 строка
3.8 KiB
Python
Executable File
131 строка
3.8 KiB
Python
Executable File
#!/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,
|
|
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(',')]
|
|
|
|
# Make classifier.
|
|
classifier = caffe.Classifier(args.model_def, args.pretrained_model,
|
|
image_dims=image_dims, gpu=args.gpu, mean=mean,
|
|
input_scale=args.input_scale, raw_scale=args.raw_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)
|