operators.py 27.5 KB
Newer Older
F
Felix 已提交
1
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved
F
Felix 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
#
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

G
gaotingquan 已提交
20
from functools import partial
F
flytocc 已提交
21
import io
F
Felix 已提交
22 23 24 25 26
import six
import math
import random
import cv2
import numpy as np
H
HydrogenSulfate 已提交
27
from PIL import Image, ImageOps, __version__ as PILLOW_VERSION
G
gaotingquan 已提交
28
from paddle.vision.transforms import ColorJitter as RawColorJitter
H
HydrogenSulfate 已提交
29
from paddle.vision.transforms import RandomRotation as RawRandomRotation
30 31
from paddle.vision.transforms import ToTensor, Normalize, RandomHorizontalFlip, RandomResizedCrop
from paddle.vision.transforms import functional as F
F
Felix 已提交
32 33
from .autoaugment import ImageNetPolicy
from .functional import augmentations
G
gaotingquan 已提交
34 35 36
from ppcls.utils import logger


G
gaotingquan 已提交
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53
def format_data(func):
    def warpper(self, data):
        if isinstance(data, dict):
            img = data["img"]
            result = func(self, img)
            if not isinstance(result, dict):
                result = {"img": result}
            return { ** data, ** result}
        else:
            result = func(self, data)
            if isinstance(result, dict):
                result = result["img"]
            return result

    return warpper


G
gaotingquan 已提交
54
class UnifiedResize(object):
H
HydrogenSulfate 已提交
55
    def __init__(self, interpolation=None, backend="cv2", return_numpy=True):
G
gaotingquan 已提交
56 57 58 59 60
        _cv2_interp_from_str = {
            'nearest': cv2.INTER_NEAREST,
            'bilinear': cv2.INTER_LINEAR,
            'area': cv2.INTER_AREA,
            'bicubic': cv2.INTER_CUBIC,
61 62
            'lanczos': cv2.INTER_LANCZOS4,
            'random': (cv2.INTER_LINEAR, cv2.INTER_CUBIC)
G
gaotingquan 已提交
63 64 65 66 67 68 69
        }
        _pil_interp_from_str = {
            'nearest': Image.NEAREST,
            'bilinear': Image.BILINEAR,
            'bicubic': Image.BICUBIC,
            'box': Image.BOX,
            'lanczos': Image.LANCZOS,
70 71
            'hamming': Image.HAMMING,
            'random': (Image.BILINEAR, Image.BICUBIC)
G
gaotingquan 已提交
72 73
        }

74 75 76 77 78
        def _cv2_resize(src, size, resample):
            if isinstance(resample, tuple):
                resample = random.choice(resample)
            return cv2.resize(src, size, interpolation=resample)

H
HydrogenSulfate 已提交
79
        def _pil_resize(src, size, resample, return_numpy=True):
80 81
            if isinstance(resample, tuple):
                resample = random.choice(resample)
H
HydrogenSulfate 已提交
82 83
            if isinstance(src, np.ndarray):
                pil_img = Image.fromarray(src)
H
HydrogenSulfate 已提交
84 85
            else:
                pil_img = src
G
gaotingquan 已提交
86
            pil_img = pil_img.resize(size, resample)
H
HydrogenSulfate 已提交
87 88 89
            if return_numpy:
                return np.asarray(pil_img)
            return pil_img
G
gaotingquan 已提交
90 91 92 93

        if backend.lower() == "cv2":
            if isinstance(interpolation, str):
                interpolation = _cv2_interp_from_str[interpolation.lower()]
94
            # compatible with opencv < version 4.4.0
G
gaotingquan 已提交
95
            elif interpolation is None:
96
                interpolation = cv2.INTER_LINEAR
97
            self.resize_func = partial(_cv2_resize, resample=interpolation)
G
gaotingquan 已提交
98 99 100
        elif backend.lower() == "pil":
            if isinstance(interpolation, str):
                interpolation = _pil_interp_from_str[interpolation.lower()]
H
HydrogenSulfate 已提交
101 102
            self.resize_func = partial(
                _pil_resize, resample=interpolation, return_numpy=return_numpy)
G
gaotingquan 已提交
103 104 105 106 107 108 109
        else:
            logger.warning(
                f"The backend of Resize only support \"cv2\" or \"PIL\". \"f{backend}\" is unavailable. Use \"cv2\" instead."
            )
            self.resize_func = cv2.resize

    def __call__(self, src, size):
H
HydrogenSulfate 已提交
110 111
        if isinstance(size, list):
            size = tuple(size)
G
gaotingquan 已提交
112
        return self.resize_func(src, size)
F
Felix 已提交
113

D
dongshuilong 已提交
114

115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
class RandomInterpolationAugment(object):
    def __init__(self, prob):
        self.prob = prob

    def _aug(self, img):
        img_shape = img.shape
        side_ratio = np.random.uniform(0.2, 1.0)
        small_side = int(side_ratio * img_shape[0])
        interpolation = np.random.choice([
            cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA,
            cv2.INTER_CUBIC, cv2.INTER_LANCZOS4
        ])
        small_img = cv2.resize(
            img, (small_side, small_side), interpolation=interpolation)
        interpolation = np.random.choice([
            cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA,
            cv2.INTER_CUBIC, cv2.INTER_LANCZOS4
        ])
        aug_img = cv2.resize(
            small_img, (img_shape[1], img_shape[0]),
            interpolation=interpolation)
        return aug_img

    def __call__(self, img):
        if np.random.random() < self.prob:
            if isinstance(img, np.ndarray):
                return self._aug(img)
            else:
                pil_img = np.array(img)
                aug_img = self._aug(pil_img)
                img = Image.fromarray(aug_img.astype(np.uint8))
                return img
        else:
            return img


F
Felix 已提交
151 152 153 154 155
class OperatorParamError(ValueError):
    """ OperatorParamError
    """
    pass

D
dongshuilong 已提交
156

F
Felix 已提交
157 158 159
class DecodeImage(object):
    """ decode image """

F
flytocc 已提交
160
    def __init__(self,
Y
Yang Nie 已提交
161
                 to_np=True,
F
flytocc 已提交
162 163
                 to_rgb=True,
                 channel_first=False,
Y
Yang Nie 已提交
164
                 backend="cv2"):
F
Felix 已提交
165
        self.to_np = to_np  # to numpy
Y
Yang Nie 已提交
166
        self.to_rgb = to_rgb  # only enabled when to_np is True
F
Felix 已提交
167 168
        self.channel_first = channel_first  # only enabled when to_np is True

F
flytocc 已提交
169 170
        if backend.lower() not in ["cv2", "pil"]:
            logger.warning(
Y
Yang Nie 已提交
171
                f"The backend of DecodeImage only support \"cv2\" or \"PIL\". \"f{backend}\" is unavailable. Use \"cv2\" instead."
F
flytocc 已提交
172 173 174 175
            )
            backend = "cv2"
        self.backend = backend.lower()

Y
Yang Nie 已提交
176 177 178 179
        if not to_np:
            logger.warning(
                f"\"to_rgb\" and \"channel_first\" are only enabled when to_np is True. \"to_np\" is now {to_np}."
            )
F
flytocc 已提交
180

G
gaotingquan 已提交
181 182
    @format_data
    def __call__(self, img):
Y
Yang Nie 已提交
183
        if isinstance(img, Image.Image):
Y
Yang Nie 已提交
184
            assert self.backend == "pil", "invalid input 'img' in DecodeImage"
Y
Yang Nie 已提交
185
        elif isinstance(img, np.ndarray):
Y
Yang Nie 已提交
186 187
            assert self.backend == "cv2", "invalid input 'img' in DecodeImage"
        elif isinstance(img, bytes):
Y
Yang Nie 已提交
188 189
            if self.backend == "pil":
                data = io.BytesIO(img)
Y
Yang Nie 已提交
190
                img = Image.open(data)
Y
Yang Nie 已提交
191
            else:
Y
Yang Nie 已提交
192
                data = np.frombuffer(img, dtype="uint8")
Y
Yang Nie 已提交
193
                img = cv2.imdecode(data, 1)
Y
Yang Nie 已提交
194 195 196 197 198 199 200
        else:
            raise ValueError("invalid input 'img' in DecodeImage")

        if self.to_np:
            if self.backend == "pil":
                assert img.mode == "RGB", f"invalid shape of image[{img.shape}]"
                img = np.asarray(img)[:, :, ::-1]  # BRG
Y
Yang Nie 已提交
201 202

            if self.to_rgb:
H
HydrogenSulfate 已提交
203 204
                assert img.shape[
                    2] == 3, f"invalid shape of image[{img.shape}]"
Y
Yang Nie 已提交
205 206 207 208
                img = img[:, :, ::-1]

            if self.channel_first:
                img = img.transpose((2, 0, 1))
G
gaotingquan 已提交
209
        return img
F
Felix 已提交
210 211 212 213 214


class ResizeImage(object):
    """ resize image """

G
gaotingquan 已提交
215 216 217 218
    def __init__(self,
                 size=None,
                 resize_short=None,
                 interpolation=None,
H
HydrogenSulfate 已提交
219 220
                 backend="cv2",
                 return_numpy=True):
F
Felix 已提交
221 222 223 224 225 226 227 228 229 230 231 232
        if resize_short is not None and resize_short > 0:
            self.resize_short = resize_short
            self.w = None
            self.h = None
        elif size is not None:
            self.resize_short = None
            self.w = size if type(size) is int else size[0]
            self.h = size if type(size) is int else size[1]
        else:
            raise OperatorParamError("invalid params for ReisizeImage for '\
                'both 'size' and 'resize_short' are None")

G
gaotingquan 已提交
233
        self._resize_func = UnifiedResize(
H
HydrogenSulfate 已提交
234 235 236
            interpolation=interpolation,
            backend=backend,
            return_numpy=return_numpy)
G
gaotingquan 已提交
237

F
Felix 已提交
238
    def __call__(self, img):
H
HydrogenSulfate 已提交
239 240 241 242 243
        if isinstance(img, np.ndarray):
            img_h, img_w = img.shape[:2]
        else:
            img_w, img_h = img.size

F
Felix 已提交
244 245 246 247 248 249 250
        if self.resize_short is not None:
            percent = float(self.resize_short) / min(img_w, img_h)
            w = int(round(img_w * percent))
            h = int(round(img_h * percent))
        else:
            w = self.w
            h = self.h
G
gaotingquan 已提交
251
        return self._resize_func(img, (w, h))
F
Felix 已提交
252 253


254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299
class CropWithPadding(RandomResizedCrop):
    """
    crop image and padding to original size
    """

    def __init__(self,
                 prob=1,
                 padding_num=0,
                 size=224,
                 scale=(0.08, 1.0),
                 ratio=(3. / 4, 4. / 3),
                 interpolation='bilinear',
                 key=None):
        super().__init__(size, scale, ratio, interpolation, key)
        self.prob = prob
        self.padding_num = padding_num

    def __call__(self, img):
        is_cv2_img = False
        if isinstance(img, np.ndarray):
            flag = True
        if np.random.random() < self.prob:
            # RandomResizedCrop augmentation
            new = np.zeros_like(np.array(img)) + self.padding_num
            #  orig_W, orig_H = F._get_image_size(sample)
            orig_W, orig_H = self._get_image_size(img)
            i, j, h, w = self._get_param(img)
            cropped = F.crop(img, i, j, h, w)
            new[i:i + h, j:j + w, :] = np.array(cropped)
            if not isinstance:
                new = Image.fromarray(new.astype(np.uint8))
            return new
        else:
            return img

    def _get_image_size(self, img):
        if F._is_pil_image(img):
            return img.size
        elif F._is_numpy_image(img):
            return img.shape[:2][::-1]
        elif F._is_tensor_image(img):
            return img.shape[1:][::-1]  # chw
        else:
            raise TypeError("Unexpected type {}".format(type(img)))


F
Felix 已提交
300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319
class CropImage(object):
    """ crop image """

    def __init__(self, size):
        if type(size) is int:
            self.size = (size, size)
        else:
            self.size = size  # (h, w)

    def __call__(self, img):
        w, h = self.size
        img_h, img_w = img.shape[:2]
        w_start = (img_w - w) // 2
        h_start = (img_h - h) // 2

        w_end = w_start + w
        h_end = h_start + h
        return img[h_start:h_end, w_start:w_end, :]


Z
zhiboniu 已提交
320
class Padv2(object):
Z
zhiboniu 已提交
321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415
    def __init__(self,
                 size=None,
                 size_divisor=32,
                 pad_mode=0,
                 offsets=None,
                 fill_value=(127.5, 127.5, 127.5)):
        """
        Pad image to a specified size or multiple of size_divisor.
        Args:
            size (int, list): image target size, if None, pad to multiple of size_divisor, default None
            size_divisor (int): size divisor, default 32
            pad_mode (int): pad mode, currently only supports four modes [-1, 0, 1, 2]. if -1, use specified offsets
                if 0, only pad to right and bottom. if 1, pad according to center. if 2, only pad left and top
            offsets (list): [offset_x, offset_y], specify offset while padding, only supported pad_mode=-1
            fill_value (bool): rgb value of pad area, default (127.5, 127.5, 127.5)
        """

        if not isinstance(size, (int, list)):
            raise TypeError(
                "Type of target_size is invalid when random_size is True. \
                            Must be List, now is {}".format(type(size)))

        if isinstance(size, int):
            size = [size, size]

        assert pad_mode in [
            -1, 0, 1, 2
        ], 'currently only supports four modes [-1, 0, 1, 2]'
        if pad_mode == -1:
            assert offsets, 'if pad_mode is -1, offsets should not be None'

        self.size = size
        self.size_divisor = size_divisor
        self.pad_mode = pad_mode
        self.fill_value = fill_value
        self.offsets = offsets

    def apply_image(self, image, offsets, im_size, size):
        x, y = offsets
        im_h, im_w = im_size
        h, w = size
        canvas = np.ones((h, w, 3), dtype=np.float32)
        canvas *= np.array(self.fill_value, dtype=np.float32)
        canvas[y:y + im_h, x:x + im_w, :] = image.astype(np.float32)
        return canvas

    def __call__(self, img):
        im_h, im_w = img.shape[:2]
        if self.size:
            w, h = self.size
            assert (
                im_h <= h and im_w <= w
            ), '(h, w) of target size should be greater than (im_h, im_w)'
        else:
            h = int(np.ceil(im_h / self.size_divisor) * self.size_divisor)
            w = int(np.ceil(im_w / self.size_divisor) * self.size_divisor)

        if h == im_h and w == im_w:
            return img.astype(np.float32)

        if self.pad_mode == -1:
            offset_x, offset_y = self.offsets
        elif self.pad_mode == 0:
            offset_y, offset_x = 0, 0
        elif self.pad_mode == 1:
            offset_y, offset_x = (h - im_h) // 2, (w - im_w) // 2
        else:
            offset_y, offset_x = h - im_h, w - im_w

        offsets, im_size, size = [offset_x, offset_y], [im_h, im_w], [h, w]

        return self.apply_image(img, offsets, im_size, size)


class RandomCropImage(object):
    """Random crop image only
    """

    def __init__(self, size):
        super(RandomCropImage, self).__init__()
        if isinstance(size, int):
            size = [size, size]
        self.size = size

    def __call__(self, img):

        h, w = img.shape[:2]
        tw, th = self.size
        i = random.randint(0, h - th)
        j = random.randint(0, w - tw)

        img = img[i:i + th, j:j + tw, :]
        return img


F
Felix 已提交
416 417 418
class RandCropImage(object):
    """ random crop image """

G
gaotingquan 已提交
419 420 421 422 423 424
    def __init__(self,
                 size,
                 scale=None,
                 ratio=None,
                 interpolation=None,
                 backend="cv2"):
F
Felix 已提交
425 426 427 428 429 430 431 432
        if type(size) is int:
            self.size = (size, size)  # (h, w)
        else:
            self.size = size

        self.scale = [0.08, 1.0] if scale is None else scale
        self.ratio = [3. / 4., 4. / 3.] if ratio is None else ratio

G
gaotingquan 已提交
433 434 435
        self._resize_func = UnifiedResize(
            interpolation=interpolation, backend=backend)

G
gaotingquan 已提交
436 437
    @format_data
    def __call__(self, img):
F
Felix 已提交
438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460
        size = self.size
        scale = self.scale
        ratio = self.ratio

        aspect_ratio = math.sqrt(random.uniform(*ratio))
        w = 1. * aspect_ratio
        h = 1. / aspect_ratio

        img_h, img_w = img.shape[:2]

        bound = min((float(img_w) / img_h) / (w**2),
                    (float(img_h) / img_w) / (h**2))
        scale_max = min(scale[1], bound)
        scale_min = min(scale[0], bound)

        target_area = img_w * img_h * random.uniform(scale_min, scale_max)
        target_size = math.sqrt(target_area)
        w = int(target_size * w)
        h = int(target_size * h)

        i = random.randint(0, img_w - w)
        j = random.randint(0, img_h - h)

461
        img = self._resize_func(img[j:j + h, i:i + w, :], size)
G
gaotingquan 已提交
462
        return img
F
Felix 已提交
463 464


H
HydrogenSulfate 已提交
465 466 467 468 469 470 471 472 473 474 475 476 477
class RandCropImageV2(object):
    """ RandCropImageV2 is different from RandCropImage,
    it will Select a cutting position randomly in a uniform distribution way,
    and cut according to the given size without resize at last."""

    def __init__(self, size):
        if type(size) is int:
            self.size = (size, size)  # (h, w)
        else:
            self.size = size

    def __call__(self, img):
        if isinstance(img, np.ndarray):
H
HydrogenSulfate 已提交
478
            img_h, img_w = img.shape[0], img.shape[1]
H
HydrogenSulfate 已提交
479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498
        else:
            img_w, img_h = img.size
        tw, th = self.size

        if img_h + 1 < th or img_w + 1 < tw:
            raise ValueError(
                "Required crop size {} is larger then input image size {}".
                format((th, tw), (img_h, img_w)))

        if img_w == tw and img_h == th:
            return img

        top = random.randint(0, img_h - th + 1)
        left = random.randint(0, img_w - tw + 1)
        if isinstance(img, np.ndarray):
            return img[top:top + th, left:left + tw, :]
        else:
            return img.crop((left, top, left + tw, top + th))


F
Felix 已提交
499 500 501 502 503 504 505 506 507 508 509 510 511 512 513
class RandFlipImage(object):
    """ random flip image
        flip_code:
            1: Flipped Horizontally
            0: Flipped Vertically
            -1: Flipped Horizontally & Vertically
    """

    def __init__(self, flip_code=1):
        assert flip_code in [-1, 0, 1
                             ], "flip_code should be a value in [-1, 0, 1]"
        self.flip_code = flip_code

    def __call__(self, img):
        if random.randint(0, 1) == 1:
H
HydrogenSulfate 已提交
514 515 516
            if isinstance(img, np.ndarray):
                return cv2.flip(img, self.flip_code)
            else:
H
HydrogenSulfate 已提交
517 518 519 520 521 522 523
                if self.flip_code == 1:
                    return img.transpose(Image.FLIP_LEFT_RIGHT)
                elif self.flip_code == 0:
                    return img.transpose(Image.FLIP_TOP_BOTTOM)
                else:
                    return img.transpose(Image.FLIP_LEFT_RIGHT).transpose(
                        Image.FLIP_LEFT_RIGHT)
F
Felix 已提交
524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543
        else:
            return img


class AutoAugment(object):
    def __init__(self):
        self.policy = ImageNetPolicy()

    def __call__(self, img):
        from PIL import Image
        img = np.ascontiguousarray(img)
        img = Image.fromarray(img)
        img = self.policy(img)
        img = np.asarray(img)


class NormalizeImage(object):
    """ normalize image such as substract mean, divide std
    """

littletomatodonkey's avatar
littletomatodonkey 已提交
544 545 546 547 548 549 550
    def __init__(self,
                 scale=None,
                 mean=None,
                 std=None,
                 order='chw',
                 output_fp16=False,
                 channel_num=3):
F
Felix 已提交
551 552
        if isinstance(scale, str):
            scale = eval(scale)
littletomatodonkey's avatar
littletomatodonkey 已提交
553 554 555 556 557
        assert channel_num in [
            3, 4
        ], "channel number of input image should be set to 3 or 4."
        self.channel_num = channel_num
        self.output_dtype = 'float16' if output_fp16 else 'float32'
F
Felix 已提交
558
        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
littletomatodonkey's avatar
littletomatodonkey 已提交
559
        self.order = order
F
Felix 已提交
560 561 562
        mean = mean if mean is not None else [0.485, 0.456, 0.406]
        std = std if std is not None else [0.229, 0.224, 0.225]

littletomatodonkey's avatar
littletomatodonkey 已提交
563
        shape = (3, 1, 1) if self.order == 'chw' else (1, 1, 3)
F
Felix 已提交
564 565 566
        self.mean = np.array(mean).reshape(shape).astype('float32')
        self.std = np.array(std).reshape(shape).astype('float32')

G
gaotingquan 已提交
567 568
    @format_data
    def __call__(self, img):
F
Felix 已提交
569 570 571 572 573 574
        from PIL import Image
        if isinstance(img, Image.Image):
            img = np.array(img)

        assert isinstance(img,
                          np.ndarray), "invalid input 'img' in NormalizeImage"
littletomatodonkey's avatar
littletomatodonkey 已提交
575 576 577 578 579 580 581 582 583 584 585 586 587

        img = (img.astype('float32') * self.scale - self.mean) / self.std

        if self.channel_num == 4:
            img_h = img.shape[1] if self.order == 'chw' else img.shape[0]
            img_w = img.shape[2] if self.order == 'chw' else img.shape[1]
            pad_zeros = np.zeros(
                (1, img_h, img_w)) if self.order == 'chw' else np.zeros(
                    (img_h, img_w, 1))
            img = (np.concatenate(
                (img, pad_zeros), axis=0)
                   if self.order == 'chw' else np.concatenate(
                       (img, pad_zeros), axis=2))
588 589

        img = img.astype(self.output_dtype)
G
gaotingquan 已提交
590
        return img
F
Felix 已提交
591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611


class ToCHWImage(object):
    """ convert hwc image to chw image
    """

    def __init__(self):
        pass

    def __call__(self, img):
        from PIL import Image
        if isinstance(img, Image.Image):
            img = np.array(img)

        return img.transpose((2, 0, 1))


class AugMix(object):
    """ Perform AugMix augmentation and compute mixture.
    """

D
dongshuilong 已提交
612 613 614 615 616 617
    def __init__(self,
                 prob=0.5,
                 aug_prob_coeff=0.1,
                 mixture_width=3,
                 mixture_depth=1,
                 aug_severity=1):
F
Felix 已提交
618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645
        """
        Args:
            prob: Probability of taking augmix
            aug_prob_coeff: Probability distribution coefficients.
            mixture_width: Number of augmentation chains to mix per augmented example.
            mixture_depth: Depth of augmentation chains. -1 denotes stochastic depth in [1, 3]'
            aug_severity: Severity of underlying augmentation operators (between 1 to 10).
        """
        # fmt: off
        self.prob = prob
        self.aug_prob_coeff = aug_prob_coeff
        self.mixture_width = mixture_width
        self.mixture_depth = mixture_depth
        self.aug_severity = aug_severity
        self.augmentations = augmentations
        # fmt: on

    def __call__(self, image):
        """Perform AugMix augmentations and compute mixture.
        Returns:
          mixed: Augmented and mixed image.
        """
        if random.random() > self.prob:
            # Avoid the warning: the given NumPy array is not writeable
            return np.asarray(image).copy()

        ws = np.float32(
            np.random.dirichlet([self.aug_prob_coeff] * self.mixture_width))
D
dongshuilong 已提交
646 647
        m = np.float32(
            np.random.beta(self.aug_prob_coeff, self.aug_prob_coeff))
F
Felix 已提交
648 649

        # image = Image.fromarray(image)
D
dongshuilong 已提交
650
        mix = np.zeros(image.shape)
F
Felix 已提交
651 652 653
        for i in range(self.mixture_width):
            image_aug = image.copy()
            image_aug = Image.fromarray(image_aug)
D
dongshuilong 已提交
654 655
            depth = self.mixture_depth if self.mixture_depth > 0 else np.random.randint(
                1, 4)
F
Felix 已提交
656 657 658 659 660 661 662
            for _ in range(depth):
                op = np.random.choice(self.augmentations)
                image_aug = op(image_aug, self.aug_severity)
            mix += ws[i] * np.asarray(image_aug)

        mixed = (1 - m) * image + m * mix
        return mixed.astype(np.uint8)
G
gaotingquan 已提交
663 664 665 666 667 668


class ColorJitter(RawColorJitter):
    """ColorJitter.
    """

669
    def __init__(self, prob=2, *args, **kwargs):
G
gaotingquan 已提交
670
        super().__init__(*args, **kwargs)
671
        self.prob = prob
G
gaotingquan 已提交
672 673

    def __call__(self, img):
674 675 676 677 678 679 680
        if np.random.random() < self.prob:
            if not isinstance(img, Image.Image):
                img = np.ascontiguousarray(img)
                img = Image.fromarray(img)
            img = super()._apply_image(img)
            if isinstance(img, Image.Image):
                img = np.asarray(img)
G
gaotingquan 已提交
681
        return img
H
HydrogenSulfate 已提交
682 683


H
HydrogenSulfate 已提交
684 685 686 687 688 689 690 691 692 693 694 695 696 697
class RandomRotation(RawRandomRotation):
    """RandomRotation.
    """

    def __init__(self, prob=0.5, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.prob = prob

    def __call__(self, img):
        if np.random.random() < self.prob:
            img = super()._apply_image(img)
        return img


H
HydrogenSulfate 已提交
698 699 700 701 702 703
class Pad(object):
    """
    Pads the given PIL.Image on all sides with specified padding mode and fill value.
    adapted from: https://pytorch.org/vision/stable/_modules/torchvision/transforms/transforms.html#Pad
    """

H
HydrogenSulfate 已提交
704 705 706 707 708
    def __init__(self,
                 padding: int,
                 fill: int=0,
                 padding_mode: str="constant",
                 backend: str="pil"):
H
HydrogenSulfate 已提交
709 710 711
        self.padding = padding
        self.fill = fill
        self.padding_mode = padding_mode
H
HydrogenSulfate 已提交
712 713 714 715
        self.backend = backend
        assert backend in [
            "pil", "cv2"
        ], f"backend must in ['pil', 'cv2'], but got {backend}"
H
HydrogenSulfate 已提交
716 717 718 719 720

    def _parse_fill(self, fill, img, min_pil_version, name="fillcolor"):
        # Process fill color for affine transforms
        major_found, minor_found = (int(v)
                                    for v in PILLOW_VERSION.split('.')[:2])
C
cuicheng01 已提交
721 722
        major_required, minor_required = (
            int(v) for v in min_pil_version.split('.')[:2])
H
HydrogenSulfate 已提交
723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749
        if major_found < major_required or (major_found == major_required and
                                            minor_found < minor_required):
            if fill is None:
                return {}
            else:
                msg = (
                    "The option to fill background area of the transformed image, "
                    "requires pillow>={}")
                raise RuntimeError(msg.format(min_pil_version))

        num_bands = len(img.getbands())
        if fill is None:
            fill = 0
        if isinstance(fill, (int, float)) and num_bands > 1:
            fill = tuple([fill] * num_bands)
        if isinstance(fill, (list, tuple)):
            if len(fill) != num_bands:
                msg = (
                    "The number of elements in 'fill' does not match the number of "
                    "bands of the image ({} != {})")
                raise ValueError(msg.format(len(fill), num_bands))

            fill = tuple(fill)

        return {name: fill}

    def __call__(self, img):
H
HydrogenSulfate 已提交
750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766
        if self.backend == "pil":
            opts = self._parse_fill(self.fill, img, "2.3.0", name="fill")
            if img.mode == "P":
                palette = img.getpalette()
                img = ImageOps.expand(img, border=self.padding, **opts)
                img.putpalette(palette)
                return img
            return ImageOps.expand(img, border=self.padding, **opts)
        else:
            img = cv2.copyMakeBorder(
                img,
                self.padding,
                self.padding,
                self.padding,
                self.padding,
                cv2.BORDER_CONSTANT,
                value=(self.fill, self.fill, self.fill))
H
HydrogenSulfate 已提交
767
            return img
768 769


770
# TODO(gaotingquan): integrate into RandomRotation
771 772 773 774 775 776 777
class RandomRot90(object):
    """RandomRot90
    """

    def __init__(self):
        pass

G
gaotingquan 已提交
778 779
    @format_data
    def __call__(self, img):
780 781 782
        orientation = random.choice([0, 1, 2, 3])
        if orientation:
            img = np.rot90(img, orientation)
G
gaotingquan 已提交
783
        return {"img": img, "random_rot90_orientation": orientation}
C
cuicheng01 已提交
784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834


class BlurImage(object):
    """BlurImage
    """

    def __init__(self,
                 ratio=0.5,
                 motion_max_ksize=12,
                 motion_max_angle=45,
                 gaussian_max_ksize=12):
        self.ratio = ratio
        self.motion_max_ksize = motion_max_ksize
        self.motion_max_angle = motion_max_angle
        self.gaussian_max_ksize = gaussian_max_ksize

    def _gaussian_blur(self, img, max_ksize=12):
        ksize = (np.random.choice(np.arange(5, max_ksize, 2)),
                 np.random.choice(np.arange(5, max_ksize, 2)))
        img = cv2.GaussianBlur(img, ksize, 0)
        return img

    def _motion_blur(self, img, max_ksize=12, max_angle=45):
        degree = np.random.choice(np.arange(5, max_ksize, 2))
        angle = np.random.choice(np.arange(-1 * max_angle, max_angle))

        M = cv2.getRotationMatrix2D((degree / 2, degree / 2), angle, 1)
        motion_blur_kernel = np.diag(np.ones(degree))
        motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M,
                                            (degree, degree))

        motion_blur_kernel = motion_blur_kernel / degree
        blurred = cv2.filter2D(img, -1, motion_blur_kernel)

        cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
        img = np.array(blurred, dtype=np.uint8)
        return img

    @format_data
    def __call__(self, img):
        if random.random() > self.ratio:
            label = 0
        else:
            method = random.choice(["gaussian", "motion"])
            if method == "gaussian":
                img = self._gaussian_blur(img, self.gaussian_max_ksize)
            else:
                img = self._motion_blur(img, self.motion_max_ksize,
                                        self.motion_max_angle)
            label = 1
        return {"img": img, "blur_image": label}