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.
278 lines
11 KiB
278 lines
11 KiB
# Copyright (c) 2016 Baidu, Inc. All Rights Reserved
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import sys
|
|
import cPickle
|
|
import logging
|
|
from PIL import Image
|
|
import numpy as np
|
|
from optparse import OptionParser
|
|
|
|
import paddle.utils.image_util as image_util
|
|
|
|
from py_paddle import swig_paddle, DataProviderConverter
|
|
from paddle.trainer.PyDataProvider2 import dense_vector
|
|
from paddle.trainer.config_parser import parse_config
|
|
|
|
logging.basicConfig(format='[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s')
|
|
logging.getLogger().setLevel(logging.INFO)
|
|
|
|
class ImageClassifier():
|
|
def __init__(self, train_conf, model_dir=None,
|
|
resize_dim=256, crop_dim=224,
|
|
use_gpu=True,
|
|
mean_file=None,
|
|
output_layer=None,
|
|
oversample=False, is_color=True):
|
|
"""
|
|
train_conf: network configure.
|
|
model_dir: string, directory of model.
|
|
resize_dim: int, resized image size.
|
|
crop_dim: int, crop size.
|
|
mean_file: string, image mean file.
|
|
oversample: bool, oversample means multiple crops, namely five
|
|
patches (the four corner patches and the center
|
|
patch) as well as their horizontal reflections,
|
|
ten crops in all.
|
|
"""
|
|
self.train_conf = train_conf
|
|
self.model_dir = model_dir
|
|
if model_dir is None:
|
|
self.model_dir = os.path.dirname(train_conf)
|
|
|
|
self.resize_dim = resize_dim
|
|
self.crop_dims = [crop_dim, crop_dim]
|
|
self.oversample = oversample
|
|
self.is_color = is_color
|
|
|
|
self.output_layer = output_layer
|
|
if self.output_layer:
|
|
assert isinstance(self.output_layer, basestring)
|
|
self.output_layer = self.output_layer.split(",")
|
|
|
|
self.transformer = image_util.ImageTransformer(is_color = is_color)
|
|
self.transformer.set_transpose((2,0,1))
|
|
self.transformer.set_channel_swap((2,1,0))
|
|
|
|
self.mean_file = mean_file
|
|
if self.mean_file is not None:
|
|
mean = np.load(self.mean_file)['data_mean']
|
|
mean = mean.reshape(3, self.crop_dims[0], self.crop_dims[1])
|
|
self.transformer.set_mean(mean) # mean pixel
|
|
else:
|
|
# if you use three mean value, set like:
|
|
# this three mean value is calculated from ImageNet.
|
|
self.transformer.set_mean(np.array([103.939,116.779,123.68]))
|
|
|
|
conf_args = "is_test=1,use_gpu=%d,is_predict=1" % (int(use_gpu))
|
|
conf = parse_config(train_conf, conf_args)
|
|
swig_paddle.initPaddle("--use_gpu=%d" % (int(use_gpu)))
|
|
self.network = swig_paddle.GradientMachine.createFromConfigProto(conf.model_config)
|
|
assert isinstance(self.network, swig_paddle.GradientMachine)
|
|
self.network.loadParameters(self.model_dir)
|
|
|
|
data_size = 3 * self.crop_dims[0] * self.crop_dims[1]
|
|
slots = [dense_vector(data_size)]
|
|
self.converter = DataProviderConverter(slots)
|
|
|
|
def get_data(self, img_path):
|
|
"""
|
|
1. load image from img_path.
|
|
2. resize or oversampling.
|
|
3. transformer data: transpose, channel swap, sub mean.
|
|
return K x H x W ndarray.
|
|
|
|
img_path: image path.
|
|
"""
|
|
image = image_util.load_image(img_path, self.is_color)
|
|
# Another way to extract oversampled features is that
|
|
# cropping and averaging from large feature map which is
|
|
# calculated by large size of image.
|
|
# This way reduces the computation.
|
|
if self.oversample:
|
|
# image_util.resize_image: short side is self.resize_dim
|
|
image = image_util.resize_image(image, self.resize_dim)
|
|
image = np.array(image)
|
|
input = np.zeros((1, image.shape[0], image.shape[1], 3),
|
|
dtype=np.float32)
|
|
input[0] = image.astype(np.float32)
|
|
input = image_util.oversample(input, self.crop_dims)
|
|
else:
|
|
image = image.resize(self.crop_dims, Image.ANTIALIAS)
|
|
input = np.zeros((1, self.crop_dims[0], self.crop_dims[1], 3),
|
|
dtype=np.float32)
|
|
input[0] = np.array(image).astype(np.float32)
|
|
|
|
data_in = []
|
|
for img in input:
|
|
img = self.transformer.transformer(img).flatten()
|
|
data_in.append([img.tolist()])
|
|
# paddle input: [[[]],[[]],...], [[]] is one sample.
|
|
return data_in
|
|
|
|
def forward(self, input_data):
|
|
"""
|
|
return output arguments which are the Outputs() in network configure.
|
|
|
|
input_data: py_paddle input data.
|
|
call forward.
|
|
"""
|
|
in_arg = self.converter(input_data)
|
|
return self.network.forwardTest(in_arg)
|
|
|
|
def forward(self, data, output_layer):
|
|
"""
|
|
return output arguments which are the Outputs() in network configure.
|
|
|
|
input_data: py_paddle input data.
|
|
call forward.
|
|
"""
|
|
input = self.converter(data)
|
|
self.network.forwardTest(input)
|
|
output = self.network.getLayerOutputs(output_layer)
|
|
res = {}
|
|
if isinstance(output_layer, basestring):
|
|
output_layer = [output_layer]
|
|
for name in output_layer:
|
|
# For oversampling, average predictions across crops.
|
|
# If not, the shape of output[name]: (1, class_number),
|
|
# the mean is also applicable.
|
|
res[name] = output[name].mean(0)
|
|
|
|
return res
|
|
|
|
def predict(self, data_file):
|
|
"""
|
|
call forward and predicting.
|
|
|
|
data_file: input image list.
|
|
"""
|
|
image_files = open(data_file, 'rb').readlines()
|
|
results = {}
|
|
if self.output_layer is None:
|
|
self.output_layer = ["output"]
|
|
for line in image_files:
|
|
image = line.split()[0]
|
|
data = self.get_data(image)
|
|
prob = self.forward(data, self.output_layer)
|
|
lab = np.argsort(-prob[self.output_layer[0]])
|
|
results[image] = lab[0]
|
|
logging.info("Label of %s is: %d", image, lab[0])
|
|
return results
|
|
|
|
def extract(self, data_file, output_dir, batch_size = 10000):
|
|
"""
|
|
extract and save features of output layers, which are
|
|
specify in Outputs() in network configure.
|
|
|
|
data_file: file name of input data.
|
|
output_dir: saved directory of extracted features.
|
|
batch_size: sample number of one batch file.
|
|
"""
|
|
if not os.path.exists(output_dir):
|
|
os.mkdir(output_dir)
|
|
|
|
sample_num = 0
|
|
batch_num = 0
|
|
image_feature = {}
|
|
image_files = open(data_file, 'rb').readlines()
|
|
for idx, line in enumerate(image_files):
|
|
image = line.split()[0]
|
|
data = self.get_data(image)
|
|
feature = self.forward(data, self.output_layer)
|
|
# save extracted features
|
|
file_name = image.split("/")[-1]
|
|
image_feature[file_name] = feature
|
|
sample_num += 1
|
|
if sample_num == batch_size:
|
|
batch_name = os.path.join(output_dir, 'batch_%d' %(batch_num))
|
|
self.save_file(image_feature, batch_name)
|
|
logging.info('Finish batch %d', batch_num)
|
|
batch_num += 1
|
|
sample_num = 0
|
|
image_feature = {}
|
|
if idx % 1000 == 0:
|
|
logging.info('%d/%d, %s', idx, len(image_files), file_name)
|
|
if sample_num > 0:
|
|
batch_name = os.path.join(output_dir, 'batch_%d' %(batch_num))
|
|
self.save_file(image_feature, batch_name)
|
|
logging.info('Finish batch %d', batch_num)
|
|
logging.info('Done: make image feature batch')
|
|
|
|
def save_file(self, data, file):
|
|
of = open(file, 'wb')
|
|
cPickle.dump(data, of, protocol=cPickle.HIGHEST_PROTOCOL)
|
|
|
|
def option_parser():
|
|
"""
|
|
Main entry for predciting
|
|
"""
|
|
usage = "%prog -c config -i data_list -w model_dir [options]"
|
|
parser = OptionParser(usage="usage: %s" % usage)
|
|
parser.add_option("-j", "--job",
|
|
action="store", dest="job_type",
|
|
help="job type: predict, extract\
|
|
predict: predicting,\
|
|
extract: extract features")
|
|
parser.add_option("-c", "--conf",
|
|
action="store", dest="train_conf",
|
|
help="network config")
|
|
parser.add_option("-i", "--data",
|
|
action="store", dest="data_file",
|
|
help="image list")
|
|
parser.add_option("-w", "--model",
|
|
action="store", dest="model_path",
|
|
default=None, help="model path")
|
|
parser.add_option("-g", "--use_gpu", action="store",
|
|
dest="use_gpu", default=True,
|
|
help="Whether to use gpu mode.")
|
|
parser.add_option("-o", "--output_dir",
|
|
action="store", dest="output_dir",
|
|
default="output", help="output path")
|
|
parser.add_option("-m", "--mean", action="store",
|
|
dest="mean", default=None,
|
|
help="mean file.")
|
|
parser.add_option("-p", "--multi_crop", action="store_true",
|
|
dest="multi_crop", default=False,
|
|
help="Wether to use multiple crops on image.")
|
|
parser.add_option("-l", "--output_layer", action="store",
|
|
dest="output_layer", default=None,
|
|
help="--job=extract, specify layers to extract "\
|
|
"features, --job=predict, specify layer of "
|
|
"classification probability, output in resnet.py.")
|
|
return parser.parse_args()
|
|
|
|
def main():
|
|
"""
|
|
1. parse input arguments.
|
|
2. predicting or extract features according job type.
|
|
"""
|
|
options, args = option_parser()
|
|
obj = ImageClassifier(options.train_conf,
|
|
options.model_path,
|
|
use_gpu=options.use_gpu,
|
|
mean_file=options.mean,
|
|
output_layer=options.output_layer,
|
|
oversample=options.multi_crop)
|
|
if options.job_type == "predict":
|
|
obj.predict(options.data_file)
|
|
|
|
elif options.job_type == "extract":
|
|
obj.extract(options.data_file,
|
|
options.output_dir)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|