image_provider.py

# 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 io
import random

import paddle.utils.image_util as image_util
from paddle.trainer.PyDataProvider2 import *


#
# {'img_size': 32,
# 'settings': <paddle.trainer.PyDataProviderWrapper.Cls instance at 0x7fea27cb6050>,
# 'color': True,
# 'mean_img_size': 32,
# 'meta': './data/cifar-out/batches/batches.meta',
# 'num_classes': 10,
# 'file_list': ('./data/cifar-out/batches/train_batch_000',),
# 'use_jpeg': True}
def hook(settings, img_size, mean_img_size, num_classes, color, meta, use_jpeg,
         is_train, **kwargs):
    settings.mean_img_size = mean_img_size
    settings.img_size = img_size
    settings.num_classes = num_classes
    settings.color = color
    settings.is_train = is_train

    if settings.color:
        settings.img_raw_size = settings.img_size * settings.img_size * 3
    else:
        settings.img_raw_size = settings.img_size * settings.img_size

    settings.meta_path = meta
    settings.use_jpeg = use_jpeg

    settings.img_mean = image_util.load_meta(settings.meta_path,
                                             settings.mean_img_size,
                                             settings.img_size,
                                             settings.color)

    settings.logger.info('Image size: %s', settings.img_size)
    settings.logger.info('Meta path: %s', settings.meta_path)
    settings.input_types = [
        dense_vector(settings.img_raw_size),  # image feature
        integer_value(settings.num_classes)]  # labels

    settings.logger.info('DataProvider Initialization finished')


@provider(init_hook=hook)
def processData(settings, file_name):
    """
    The main function for loading data.
    Load the batch, iterate all the images and labels in this batch.
    file_name: the batch file name.
    """
    data = cPickle.load(io.open(file_name, 'rb'))
    indexes = list(range(len(data['images'])))
    if settings.is_train:
        random.shuffle(indexes)
    for i in indexes:
        if settings.use_jpeg == 1:
            img = image_util.decode_jpeg(data['images'][i])
        else:
            img = data['images'][i]
        img_feat = image_util.preprocess_img(img, settings.img_mean,
                                             settings.img_size, settings.is_train,
                                             settings.color)
        label = data['labels'][i]
        yield img_feat.tolist(), int(label)