#copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.

#

#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 math

import random

import functools

import numpy as np

from PIL import Image, ImageEnhance

import paddle

random.seed(0)

np.random.seed(0)

DATA_DIM = 224

THREAD = 8

BUF_SIZE = 2048

DATA_DIR = 'data/ILSVRC2012'

img_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))

img_std = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))

def resize_short(img, target_size):

percent = float(target_size) / min(img.size[0], img.size[1])

resized_width = int(round(img.size[0] * percent))

resized_height = int(round(img.size[1] * percent))

img = img.resize((resized_width, resized_height), Image.LANCZOS)

return img

def crop_image(img, target_size, center):

width, height = img.size

size = target_size

if center == True:

w_start = (width - size) / 2

h_start = (height - size) / 2

else:

w_start = np.random.randint(0, width - size + 1)

h_start = np.random.randint(0, height - size + 1)

w_end = w_start + size

h_end = h_start + size

img = img.crop((w_start, h_start, w_end, h_end))

return img

def random_crop(img, size, scale=[0.08, 1.0], ratio=[3. / 4., 4. / 3.]):

aspect_ratio = math.sqrt(np.random.uniform(*ratio))

w = 1. * aspect_ratio

h = 1. / aspect_ratio

bound = min((float(img.size[0]) / img.size[1]) / (w**2),

(float(img.size[1]) / img.size[0]) / (h**2))

scale_max = min(scale[1], bound)

scale_min = min(scale[0], bound)

target_area = img.size[0] * img.size[1] * np.random.uniform(scale_min,

scale_max)

target_size = math.sqrt(target_area)

w = int(target_size * w)

h = int(target_size * h)

i = np.random.randint(0, img.size[0] - w + 1)

j = np.random.randint(0, img.size[1] - h + 1)

img = img.crop((i, j, i + w, j + h))

img = img.resize((size, size), Image.LANCZOS)

return img

def rotate_image(img):

angle = np.random.randint(-10, 11)

img = img.rotate(angle)

return img

def distort_color(img):

def random_brightness(img, lower=0.5, upper=1.5):

e = np.random.uniform(lower, upper)

return ImageEnhance.Brightness(img).enhance(e)

def random_contrast(img, lower=0.5, upper=1.5):

e = np.random.uniform(lower, upper)

return ImageEnhance.Contrast(img).enhance(e)

def random_color(img, lower=0.5, upper=1.5):

e = np.random.uniform(lower, upper)

return ImageEnhance.Color(img).enhance(e)

ops = [random_brightness, random_contrast, random_color]

np.random.shuffle(ops)

img = ops[0](img)

img = ops[1](img)

img = ops[2](img)

return img

def process_image(sample, mode, color_jitter, rotate):

img_path = sample[0]

img = Image.open(img_path)

if mode == 'train':

if rotate: img = rotate_image(img)

img = random_crop(img, DATA_DIM)

else:

img = resize_short(img, target_size=256)

img = crop_image(img, target_size=DATA_DIM, center=True)

if mode == 'train':

if color_jitter:

img = distort_color(img)

if np.random.randint(0, 2) == 1:

img = img.transpose(Image.FLIP_LEFT_RIGHT)

if img.mode != 'RGB':

img = img.convert('RGB')

img = np.array(img).astype('float32').transpose((2, 0, 1)) / 255

img -= img_mean

img /= img_std

if mode == 'train' or mode == 'val':

return img, sample[1]

elif mode == 'test':

return [img]

def process_batch_data(input_data, mode, color_jitter, rotate):

batch_data = []

for sample in input_data:

batch_data.append(process_image(sample, mode, color_jitter, rotate))

return batch_data

def _reader_creator(file_list,

batch_size,

mode,

shuffle=False,

color_jitter=False,

rotate=False,

data_dir=DATA_DIR,

shuffle_seed=0,

infinite=False):

def reader():

def read_file_list():

with open(file_list) as flist:

full_lines = [line.strip() for line in flist]

if shuffle:

if shuffle_seed is not None:

np.random.seed(shuffle_seed)

np.random.shuffle(full_lines)

batch_data = []

for line in full_lines:

img_path, label = line.split()

img_path = os.path.join(data_dir, img_path)

batch_data.append([img_path, int(label)])

if len(batch_data) == batch_size:

if mode == 'train' or mode == 'val':

yield batch_data

elif mode == 'test':

yield [sample[0] for sample in batch_data]

batch_data = []

return read_file_list

data_reader = reader()

num_trainers = int(os.environ.get('PADDLE_TRAINERS_NUM', 1))

if mode == 'train' and num_trainers > 1:

assert shuffle_seed is not None, \

"If num_trainers > 1, the shuffle_seed must be set, because " \

"the order of batch data generated by reader " \

"must be the same in the respective processes."

data_reader = fluid.contrib.reader.distributed_batch_reader(data_reader)

mapper = functools.partial(

process_batch_data, mode=mode, color_jitter=color_jitter, rotate=rotate)

return paddle.reader.xmap_readers(mapper, data_reader, THREAD, BUF_SIZE)

def train(batch_size, data_dir=DATA_DIR, shuffle_seed=0, infinite=False):

file_list = os.path.join(data_dir, 'train_list.txt')

return _reader_creator(

file_list,

batch_size,

'train',

shuffle=True,

color_jitter=False,

rotate=False,

data_dir=data_dir,

shuffle_seed=shuffle_seed,

infinite=infinite)

def val(batch_size, data_dir=DATA_DIR):

file_list = os.path.join(data_dir, 'val_list.txt')

return _reader_creator(

file_list, batch_size, 'val', shuffle=False, data_dir=data_dir)

def test(batch_size, data_dir=DATA_DIR):

file_list = os.path.join(data_dir, 'val_list.txt')

return _reader_creator(

file_list, batch_size, 'test', shuffle=False, data_dir=data_dir)