image.py 11.2 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13
# 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.
D
dangqingqing 已提交
14
"""
D
dangqingqing 已提交
15
This file contains some common interfaces for image preprocess.
D
dangqingqing 已提交
16
Many users are confused about the image layout. We introduce
D
dangqingqing 已提交
17
the image layout as follows.
D
dangqingqing 已提交
18 19

- CHW Layout
L
Luo Tao 已提交
20

D
dangqingqing 已提交
21
  - The abbreviations: C=channel, H=Height, W=Width
D
dangqingqing 已提交
22 23 24
  - The default layout of image opened by cv2 or PIL is HWC.
    PaddlePaddle only supports the CHW layout. And CHW is simply
    a transpose of HWC. It must transpose the input image.
D
dangqingqing 已提交
25 26

- Color format: RGB or BGR
L
Luo Tao 已提交
27

D
dangqingqing 已提交
28
  OpenCV use BGR color format. PIL use RGB color format. Both
D
dangqingqing 已提交
29 30
  formats can be used for training. Noted that, the format should
  be keep consistent between the training and inference peroid.
D
dangqingqing 已提交
31
"""
32 33 34

from __future__ import print_function

L
Luo Tao 已提交
35 36 37 38 39 40 41
import numpy as np
try:
    import cv2
except ImportError:
    cv2 = None
import os
import tarfile
42
import six.moves.cPickle as pickle
L
Luo Tao 已提交
43 44 45 46 47 48

__all__ = [
    "load_image_bytes", "load_image", "resize_short", "to_chw", "center_crop",
    "random_crop", "left_right_flip", "simple_transform", "load_and_transform",
    "batch_images_from_tar"
]
D
dangqingqing 已提交
49 50


51 52 53 54 55 56 57 58 59 60 61 62
def _check_cv2():
    if cv2 is None:
        import sys
        sys.stderr.write(
            '''Warning with paddle image module: opencv-python should be imported,
         or paddle image module could NOT work; please install opencv-python first.'''
        )
        return False
    else:
        return True


63 64 65 66 67 68
def batch_images_from_tar(data_file,
                          dataset_name,
                          img2label,
                          num_per_batch=1024):
    """
    Read images from tar file and batch them into batch file.
L
Luo Tao 已提交
69 70 71 72 73

    :param data_file: path of image tar file
    :type data_file: string
    :param dataset_name: 'train','test' or 'valid'
    :type dataset_name: string
M
minqiyang 已提交
74
    :param img2label: a dic with image file name as key
75
                    and image's label as value
L
Luo Tao 已提交
76 77 78 79 80
    :type img2label: dic
    :param num_per_batch: image number per batch file
    :type num_per_batch: int
    :return: path of list file containing paths of batch file
    :rtype: string
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
    """
    batch_dir = data_file + "_batch"
    out_path = "%s/%s" % (batch_dir, dataset_name)
    meta_file = "%s/%s.txt" % (batch_dir, dataset_name)

    if os.path.exists(out_path):
        return meta_file
    else:
        os.makedirs(out_path)

    tf = tarfile.open(data_file)
    mems = tf.getmembers()
    data = []
    labels = []
    file_id = 0
    for mem in mems:
        if mem.name in img2label:
            data.append(tf.extractfile(mem).read())
            labels.append(img2label[mem.name])
            if len(data) == num_per_batch:
                output = {}
                output['label'] = labels
                output['data'] = data
104
                pickle.dump(
105
                    output,
M
minqiyang 已提交
106
                    open('%s/batch_%d' % (out_path, file_id), 'wb'),
107
                    protocol=pickle.HIGHEST_PROTOCOL)
108 109 110 111 112 113 114
                file_id += 1
                data = []
                labels = []
    if len(data) > 0:
        output = {}
        output['label'] = labels
        output['data'] = data
115
        pickle.dump(
116
            output,
M
minqiyang 已提交
117
            open('%s/batch_%d' % (out_path, file_id), 'wb'),
118
            protocol=pickle.HIGHEST_PROTOCOL)
119 120 121 122 123 124 125

    with open(meta_file, 'a') as meta:
        for file in os.listdir(out_path):
            meta.write(os.path.abspath("%s/%s" % (out_path, file)) + "\n")
    return meta_file


126 127 128 129 130
def load_image_bytes(bytes, is_color=True):
    """
    Load an color or gray image from bytes array.

    Example usage:
M
minqiyang 已提交
131

132
    .. code-block:: python
L
Luo Tao 已提交
133

134
        with open('cat.jpg') as f:
135
            im = load_image_bytes(f.read())
136 137

    :param bytes: the input image bytes array.
L
Luo Tao 已提交
138
    :type bytes: str
139 140 141
    :param is_color: If set is_color True, it will load and
                     return a color image. Otherwise, it will
                     load and return a gray image.
L
Luo Tao 已提交
142
    :type is_color: bool
143
    """
144
    assert _check_cv2() is True
M
minqiyang 已提交
145

146 147 148 149 150 151
    flag = 1 if is_color else 0
    file_bytes = np.asarray(bytearray(bytes), dtype=np.uint8)
    img = cv2.imdecode(file_bytes, flag)
    return img


D
dangqingqing 已提交
152 153 154 155 156
def load_image(file, is_color=True):
    """
    Load an color or gray image from the file path.

    Example usage:
M
minqiyang 已提交
157

D
dangqingqing 已提交
158
    .. code-block:: python
L
Luo Tao 已提交
159

D
dangqingqing 已提交
160 161 162 163 164 165 166
        im = load_image('cat.jpg')

    :param file: the input image path.
    :type file: string
    :param is_color: If set is_color True, it will load and
                     return a color image. Otherwise, it will
                     load and return a gray image.
L
Luo Tao 已提交
167
    :type is_color: bool
D
dangqingqing 已提交
168
    """
169
    assert _check_cv2() is True
M
minqiyang 已提交
170

D
dangqingqing 已提交
171 172 173 174 175 176 177
    # cv2.IMAGE_COLOR for OpenCV3
    # cv2.CV_LOAD_IMAGE_COLOR for older OpenCV Version
    # cv2.IMAGE_GRAYSCALE for OpenCV3
    # cv2.CV_LOAD_IMAGE_GRAYSCALE for older OpenCV Version
    # Here, use constant 1 and 0
    # 1: COLOR, 0: GRAYSCALE
    flag = 1 if is_color else 0
D
dangqingqing 已提交
178 179 180 181 182
    im = cv2.imread(file, flag)
    return im


def resize_short(im, size):
M
minqiyang 已提交
183
    """
D
dangqingqing 已提交
184 185 186
    Resize an image so that the length of shorter edge is size.

    Example usage:
M
minqiyang 已提交
187

D
dangqingqing 已提交
188
    .. code-block:: python
L
Luo Tao 已提交
189

D
dangqingqing 已提交
190 191
        im = load_image('cat.jpg')
        im = resize_short(im, 256)
M
minqiyang 已提交
192

D
dangqingqing 已提交
193 194 195 196 197
    :param im: the input image with HWC layout.
    :type im: ndarray
    :param size: the shorter edge size of image after resizing.
    :type size: int
    """
198
    assert _check_cv2() is True
M
minqiyang 已提交
199

D
dangqingqing 已提交
200 201 202
    h, w = im.shape[:2]
    h_new, w_new = size, size
    if h > w:
M
minqiyang 已提交
203
        h_new = size * h // w
D
dangqingqing 已提交
204
    else:
M
minqiyang 已提交
205
        w_new = size * w // h
206
    im = cv2.resize(im, (w_new, h_new), interpolation=cv2.INTER_CUBIC)
D
dangqingqing 已提交
207 208 209 210 211 212
    return im


def to_chw(im, order=(2, 0, 1)):
    """
    Transpose the input image order. The image layout is HWC format
D
dangqingqing 已提交
213 214
    opened by cv2 or PIL. Transpose the input image to CHW layout
    according the order (2,0,1).
D
dangqingqing 已提交
215 216

    Example usage:
M
minqiyang 已提交
217

D
dangqingqing 已提交
218
    .. code-block:: python
L
Luo Tao 已提交
219

D
dangqingqing 已提交
220 221 222
        im = load_image('cat.jpg')
        im = resize_short(im, 256)
        im = to_chw(im)
M
minqiyang 已提交
223

D
dangqingqing 已提交
224 225 226
    :param im: the input image with HWC layout.
    :type im: ndarray
    :param order: the transposed order.
M
minqiyang 已提交
227
    :type order: tuple|list
D
dangqingqing 已提交
228 229 230 231 232 233 234 235 236 237 238
    """
    assert len(im.shape) == len(order)
    im = im.transpose(order)
    return im


def center_crop(im, size, is_color=True):
    """
    Crop the center of image with size.

    Example usage:
M
minqiyang 已提交
239

D
dangqingqing 已提交
240
    .. code-block:: python
L
Luo Tao 已提交
241

D
dangqingqing 已提交
242
        im = center_crop(im, 224)
M
minqiyang 已提交
243

D
dangqingqing 已提交
244 245
    :param im: the input image with HWC layout.
    :type im: ndarray
D
dangqingqing 已提交
246
    :param size: the cropping size.
D
dangqingqing 已提交
247 248 249 250 251
    :type size: int
    :param is_color: whether the image is color or not.
    :type is_color: bool
    """
    h, w = im.shape[:2]
M
minqiyang 已提交
252 253
    h_start = (h - size) // 2
    w_start = (w - size) // 2
D
dangqingqing 已提交
254 255 256 257 258 259 260 261 262 263 264 265 266
    h_end, w_end = h_start + size, w_start + size
    if is_color:
        im = im[h_start:h_end, w_start:w_end, :]
    else:
        im = im[h_start:h_end, w_start:w_end]
    return im


def random_crop(im, size, is_color=True):
    """
    Randomly crop input image with size.

    Example usage:
M
minqiyang 已提交
267

D
dangqingqing 已提交
268
    .. code-block:: python
L
Luo Tao 已提交
269

D
dangqingqing 已提交
270
        im = random_crop(im, 224)
M
minqiyang 已提交
271

D
dangqingqing 已提交
272 273
    :param im: the input image with HWC layout.
    :type im: ndarray
D
dangqingqing 已提交
274
    :param size: the cropping size.
D
dangqingqing 已提交
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
    :type size: int
    :param is_color: whether the image is color or not.
    :type is_color: bool
    """
    h, w = im.shape[:2]
    h_start = np.random.randint(0, h - size + 1)
    w_start = np.random.randint(0, w - size + 1)
    h_end, w_end = h_start + size, w_start + size
    if is_color:
        im = im[h_start:h_end, w_start:w_end, :]
    else:
        im = im[h_start:h_end, w_start:w_end]
    return im


290
def left_right_flip(im, is_color=True):
D
dangqingqing 已提交
291 292 293 294 295
    """
    Flip an image along the horizontal direction.
    Return the flipped image.

    Example usage:
M
minqiyang 已提交
296

D
dangqingqing 已提交
297
    .. code-block:: python
L
Luo Tao 已提交
298

D
dangqingqing 已提交
299
        im = left_right_flip(im)
M
minqiyang 已提交
300

Q
qingqing01 已提交
301
    :param im: input image with HWC layout or HW layout for gray image
D
dangqingqing 已提交
302
    :type im: ndarray
Q
qingqing01 已提交
303
    :param is_color: whether input image is color or not
304
    :type is_color: bool
D
dangqingqing 已提交
305
    """
306
    if len(im.shape) == 3 and is_color:
D
dangqingqing 已提交
307 308
        return im[:, ::-1, :]
    else:
309
        return im[:, ::-1]
D
dangqingqing 已提交
310 311


312 313 314 315 316 317
def simple_transform(im,
                     resize_size,
                     crop_size,
                     is_train,
                     is_color=True,
                     mean=None):
D
dangqingqing 已提交
318
    """
D
dangqingqing 已提交
319
    Simply data argumentation for training. These operations include
D
dangqingqing 已提交
320 321
    resizing, croping and flipping.

D
dangqingqing 已提交
322
    Example usage:
M
minqiyang 已提交
323

D
dangqingqing 已提交
324
    .. code-block:: python
L
Luo Tao 已提交
325

D
dangqingqing 已提交
326 327
        im = simple_transform(im, 256, 224, True)

D
dangqingqing 已提交
328 329 330 331 332 333 334 335
    :param im: The input image with HWC layout.
    :type im: ndarray
    :param resize_size: The shorter edge length of the resized image.
    :type resize_size: int
    :param crop_size: The cropping size.
    :type crop_size: int
    :param is_train: Whether it is training or not.
    :type is_train: bool
L
Luo Tao 已提交
336 337
    :param is_color: whether the image is color or not.
    :type is_color: bool
M
minqiyang 已提交
338
    :param mean: the mean values, which can be element-wise mean values or
L
Luo Tao 已提交
339 340
                 mean values per channel.
    :type mean: numpy array | list
D
dangqingqing 已提交
341 342 343
    """
    im = resize_short(im, resize_size)
    if is_train:
Y
yeyupiaoling 已提交
344
        im = random_crop(im, crop_size, is_color=is_color)
D
dangqingqing 已提交
345
        if np.random.randint(2) == 0:
346
            im = left_right_flip(im, is_color)
D
dangqingqing 已提交
347
    else:
Y
yeyupiaoling 已提交
348
        im = center_crop(im, crop_size, is_color=is_color)
349 350 351 352 353 354
    if len(im.shape) == 3:
        im = to_chw(im)

    im = im.astype('float32')
    if mean is not None:
        mean = np.array(mean, dtype=np.float32)
M
minqiyang 已提交
355
        # mean value, may be one value per channel
356
        if mean.ndim == 1 and is_color:
357
            mean = mean[:, np.newaxis, np.newaxis]
358 359
        elif mean.ndim == 1:
            mean = mean
360 361 362 363
        else:
            # elementwise mean
            assert len(mean.shape) == len(im)
        im -= mean
D
dangqingqing 已提交
364 365 366 367 368 369 370 371

    return im


def load_and_transform(filename,
                       resize_size,
                       crop_size,
                       is_train,
372 373
                       is_color=True,
                       mean=None):
D
dangqingqing 已提交
374 375
    """
    Load image from the input file `filename` and transform image for
D
dangqingqing 已提交
376 377
    data argumentation. Please refer to the `simple_transform` interface
    for the transform operations.
D
dangqingqing 已提交
378

D
dangqingqing 已提交
379
    Example usage:
M
minqiyang 已提交
380

D
dangqingqing 已提交
381
    .. code-block:: python
L
Luo Tao 已提交
382

D
dangqingqing 已提交
383 384
        im = load_and_transform('cat.jpg', 256, 224, True)

D
dangqingqing 已提交
385 386 387 388 389 390 391 392
    :param filename: The file name of input image.
    :type filename: string
    :param resize_size: The shorter edge length of the resized image.
    :type resize_size: int
    :param crop_size: The cropping size.
    :type crop_size: int
    :param is_train: Whether it is training or not.
    :type is_train: bool
L
Luo Tao 已提交
393 394
    :param is_color: whether the image is color or not.
    :type is_color: bool
M
minqiyang 已提交
395
    :param mean: the mean values, which can be element-wise mean values or
L
Luo Tao 已提交
396 397
                 mean values per channel.
    :type mean: numpy array | list
D
dangqingqing 已提交
398
    """
399
    im = load_image(filename, is_color)
400
    im = simple_transform(im, resize_size, crop_size, is_train, is_color, mean)
D
dangqingqing 已提交
401
    return im
反馈
建议
客服 返回
顶部