operators.py 16.2 KB
Newer Older
W
WenmuZhou 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
"""
# Copyright (c) 2020 PaddlePaddle Authors. 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.
"""

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

import sys
import six
import cv2
import numpy as np
z37757's avatar
z37757 已提交
26
import math
X
xiaoting 已提交
27
from PIL import Image
W
WenmuZhou 已提交
28 29 30 31 32


class DecodeImage(object):
    """ decode image """

z37757's avatar
z37757 已提交
33 34 35 36 37
    def __init__(self,
                 img_mode='RGB',
                 channel_first=False,
                 ignore_orientation=False,
                 **kwargs):
W
WenmuZhou 已提交
38 39
        self.img_mode = img_mode
        self.channel_first = channel_first
z37757's avatar
z37757 已提交
40
        self.ignore_orientation = ignore_orientation
W
WenmuZhou 已提交
41 42 43 44 45 46 47 48 49 50

    def __call__(self, data):
        img = data['image']
        if six.PY2:
            assert type(img) is str and len(
                img) > 0, "invalid input 'img' in DecodeImage"
        else:
            assert type(img) is bytes and len(
                img) > 0, "invalid input 'img' in DecodeImage"
        img = np.frombuffer(img, dtype='uint8')
z37757's avatar
z37757 已提交
51 52 53 54 55
        if self.ignore_orientation:
            img = cv2.imdecode(img, cv2.IMREAD_IGNORE_ORIENTATION |
                               cv2.IMREAD_COLOR)
        else:
            img = cv2.imdecode(img, 1)
L
LDOUBLEV 已提交
56 57
        if img is None:
            return None
W
WenmuZhou 已提交
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113
        if self.img_mode == 'GRAY':
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        elif self.img_mode == 'RGB':
            assert img.shape[2] == 3, 'invalid shape of image[%s]' % (img.shape)
            img = img[:, :, ::-1]

        if self.channel_first:
            img = img.transpose((2, 0, 1))

        data['image'] = img
        return data


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

    def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
        if isinstance(scale, str):
            scale = eval(scale)
        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
        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]

        shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
        self.mean = np.array(mean).reshape(shape).astype('float32')
        self.std = np.array(std).reshape(shape).astype('float32')

    def __call__(self, data):
        img = data['image']
        from PIL import Image
        if isinstance(img, Image.Image):
            img = np.array(img)
        assert isinstance(img,
                          np.ndarray), "invalid input 'img' in NormalizeImage"
        data['image'] = (
            img.astype('float32') * self.scale - self.mean) / self.std
        return data


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

    def __init__(self, **kwargs):
        pass

    def __call__(self, data):
        img = data['image']
        from PIL import Image
        if isinstance(img, Image.Image):
            img = np.array(img)
        data['image'] = img.transpose((2, 0, 1))
        return data


T
tink2123 已提交
114 115
class Fasttext(object):
    def __init__(self, path="None", **kwargs):
T
tink2123 已提交
116
        import fasttext
T
tink2123 已提交
117 118 119 120 121 122 123 124 125
        self.fast_model = fasttext.load_model(path)

    def __call__(self, data):
        label = data['label']
        fast_label = self.fast_model[label]
        data['fast_label'] = fast_label
        return data


D
dyning 已提交
126
class KeepKeys(object):
W
WenmuZhou 已提交
127 128 129 130 131 132 133 134 135 136
    def __init__(self, keep_keys, **kwargs):
        self.keep_keys = keep_keys

    def __call__(self, data):
        data_list = []
        for key in self.keep_keys:
            data_list.append(data[key])
        return data_list


z37757's avatar
z37757 已提交
137
class Pad(object):
z37757's avatar
z37757 已提交
138 139 140 141 142 143 144
    def __init__(self, size=None, size_div=32, **kwargs):
        if size is not None and not isinstance(size, (int, list, tuple)):
            raise TypeError("Type of target_size is invalid. Now is {}".format(
                type(size)))
        if isinstance(size, int):
            size = [size, size]
        self.size = size
z37757's avatar
z37757 已提交
145 146 147 148 149
        self.size_div = size_div

    def __call__(self, data):

        img = data['image']
z37757's avatar
z37757 已提交
150 151 152 153 154 155 156 157 158 159 160 161 162
        img_h, img_w = img.shape[0], img.shape[1]
        if self.size:
            resize_h2, resize_w2 = self.size
            assert (
                img_h < resize_h2 and img_w < resize_w2
            ), '(h, w) of target size should be greater than (img_h, img_w)'
        else:
            resize_h2 = max(
                int(math.ceil(img.shape[0] / self.size_div) * self.size_div),
                self.size_div)
            resize_w2 = max(
                int(math.ceil(img.shape[1] / self.size_div) * self.size_div),
                self.size_div)
z37757's avatar
z37757 已提交
163 164 165
        img = cv2.copyMakeBorder(
            img,
            0,
z37757's avatar
z37757 已提交
166
            resize_h2 - img_h,
z37757's avatar
z37757 已提交
167
            0,
z37757's avatar
z37757 已提交
168
            resize_w2 - img_w,
z37757's avatar
z37757 已提交
169 170 171 172 173 174
            cv2.BORDER_CONSTANT,
            value=0)
        data['image'] = img
        return data


L
LDOUBLEV 已提交
175 176 177 178 179 180 181 182 183 184 185 186 187 188
class Resize(object):
    def __init__(self, size=(640, 640), **kwargs):
        self.size = size

    def resize_image(self, img):
        resize_h, resize_w = self.size
        ori_h, ori_w = img.shape[:2]  # (h, w, c)
        ratio_h = float(resize_h) / ori_h
        ratio_w = float(resize_w) / ori_w
        img = cv2.resize(img, (int(resize_w), int(resize_h)))
        return img, [ratio_h, ratio_w]

    def __call__(self, data):
        img = data['image']
189 190
        if 'polys' in data:
            text_polys = data['polys']
L
LDOUBLEV 已提交
191 192

        img_resize, [ratio_h, ratio_w] = self.resize_image(img)
193 194 195 196 197 198 199 200
        if 'polys' in data:
            new_boxes = []
            for box in text_polys:
                new_box = []
                for cord in box:
                    new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
                new_boxes.append(new_box)
            data['polys'] = np.array(new_boxes, dtype=np.float32)
L
LDOUBLEV 已提交
201 202 203 204
        data['image'] = img_resize
        return data


W
WenmuZhou 已提交
205 206 207 208
class DetResizeForTest(object):
    def __init__(self, **kwargs):
        super(DetResizeForTest, self).__init__()
        self.resize_type = 0
W
wangjingyeye 已提交
209
        self.keep_ratio = False
W
WenmuZhou 已提交
210 211 212
        if 'image_shape' in kwargs:
            self.image_shape = kwargs['image_shape']
            self.resize_type = 1
W
wangjingyeye 已提交
213 214
            if 'keep_ratio' in kwargs:
                self.keep_ratio = kwargs['keep_ratio']
文幕地方's avatar
文幕地方 已提交
215
        elif 'limit_side_len' in kwargs:
W
WenmuZhou 已提交
216 217
            self.limit_side_len = kwargs['limit_side_len']
            self.limit_type = kwargs.get('limit_type', 'min')
文幕地方's avatar
文幕地方 已提交
218
        elif 'resize_long' in kwargs:
M
MissPenguin 已提交
219 220
            self.resize_type = 2
            self.resize_long = kwargs.get('resize_long', 960)
W
WenmuZhou 已提交
221 222 223 224 225 226
        else:
            self.limit_side_len = 736
            self.limit_type = 'min'

    def __call__(self, data):
        img = data['image']
M
MissPenguin 已提交
227
        src_h, src_w, _ = img.shape
W
WenmuZhou 已提交
228 229

        if self.resize_type == 0:
M
MissPenguin 已提交
230 231 232 233
            # img, shape = self.resize_image_type0(img)
            img, [ratio_h, ratio_w] = self.resize_image_type0(img)
        elif self.resize_type == 2:
            img, [ratio_h, ratio_w] = self.resize_image_type2(img)
W
WenmuZhou 已提交
234
        else:
M
MissPenguin 已提交
235 236
            # img, shape = self.resize_image_type1(img)
            img, [ratio_h, ratio_w] = self.resize_image_type1(img)
W
WenmuZhou 已提交
237
        data['image'] = img
M
MissPenguin 已提交
238
        data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
W
WenmuZhou 已提交
239 240 241 242 243
        return data

    def resize_image_type1(self, img):
        resize_h, resize_w = self.image_shape
        ori_h, ori_w = img.shape[:2]  # (h, w, c)
W
wangjingyeye 已提交
244
        if self.keep_ratio is True:
W
wangjingyeye 已提交
245 246 247
            resize_w = ori_w * resize_h / ori_h
            N = math.ceil(resize_w / 32)
            resize_w = N * 32
M
MissPenguin 已提交
248 249
        ratio_h = float(resize_h) / ori_h
        ratio_w = float(resize_w) / ori_w
W
WenmuZhou 已提交
250
        img = cv2.resize(img, (int(resize_w), int(resize_h)))
M
MissPenguin 已提交
251 252
        # return img, np.array([ori_h, ori_w])
        return img, [ratio_h, ratio_w]
W
WenmuZhou 已提交
253 254 255 256 257 258 259 260 261 262

    def resize_image_type0(self, img):
        """
        resize image to a size multiple of 32 which is required by the network
        args:
            img(array): array with shape [h, w, c]
        return(tuple):
            img, (ratio_h, ratio_w)
        """
        limit_side_len = self.limit_side_len
W
WenmuZhou 已提交
263
        h, w, c = img.shape
W
WenmuZhou 已提交
264 265 266 267 268 269 270 271 272 273

        # limit the max side
        if self.limit_type == 'max':
            if max(h, w) > limit_side_len:
                if h > w:
                    ratio = float(limit_side_len) / h
                else:
                    ratio = float(limit_side_len) / w
            else:
                ratio = 1.
W
WenmuZhou 已提交
274
        elif self.limit_type == 'min':
W
WenmuZhou 已提交
275 276 277 278 279 280 281
            if min(h, w) < limit_side_len:
                if h < w:
                    ratio = float(limit_side_len) / h
                else:
                    ratio = float(limit_side_len) / w
            else:
                ratio = 1.
W
WenmuZhou 已提交
282
        elif self.limit_type == 'resize_long':
L
LDOUBLEV 已提交
283
            ratio = float(limit_side_len) / max(h, w)
W
WenmuZhou 已提交
284 285
        else:
            raise Exception('not support limit type, image ')
W
WenmuZhou 已提交
286 287 288
        resize_h = int(h * ratio)
        resize_w = int(w * ratio)

Z
zhoujun 已提交
289 290
        resize_h = max(int(round(resize_h / 32) * 32), 32)
        resize_w = max(int(round(resize_w / 32) * 32), 32)
W
WenmuZhou 已提交
291 292 293 294 295 296 297 298

        try:
            if int(resize_w) <= 0 or int(resize_h) <= 0:
                return None, (None, None)
            img = cv2.resize(img, (int(resize_w), int(resize_h)))
        except:
            print(img.shape, resize_w, resize_h)
            sys.exit(0)
M
MissPenguin 已提交
299 300 301
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)
        return img, [ratio_h, ratio_w]
L
LDOUBLEV 已提交
302

M
MissPenguin 已提交
303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324
    def resize_image_type2(self, img):
        h, w, _ = img.shape

        resize_w = w
        resize_h = h

        if resize_h > resize_w:
            ratio = float(self.resize_long) / resize_h
        else:
            ratio = float(self.resize_long) / resize_w

        resize_h = int(resize_h * ratio)
        resize_w = int(resize_w * ratio)

        max_stride = 128
        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
        img = cv2.resize(img, (int(resize_w), int(resize_h)))
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)

        return img, [ratio_h, ratio_w]
J
Jethong 已提交
325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347


class E2EResizeForTest(object):
    def __init__(self, **kwargs):
        super(E2EResizeForTest, self).__init__()
        self.max_side_len = kwargs['max_side_len']
        self.valid_set = kwargs['valid_set']

    def __call__(self, data):
        img = data['image']
        src_h, src_w, _ = img.shape
        if self.valid_set == 'totaltext':
            im_resized, [ratio_h, ratio_w] = self.resize_image_for_totaltext(
                img, max_side_len=self.max_side_len)
        else:
            im_resized, (ratio_h, ratio_w) = self.resize_image(
                img, max_side_len=self.max_side_len)
        data['image'] = im_resized
        data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
        return data

    def resize_image_for_totaltext(self, im, max_side_len=512):

348
        h, w, _ = im.shape
J
Jethong 已提交
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
        resize_w = w
        resize_h = h
        ratio = 1.25
        if h * ratio > max_side_len:
            ratio = float(max_side_len) / resize_h
        resize_h = int(resize_h * ratio)
        resize_w = int(resize_w * ratio)

        max_stride = 128
        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
        im = cv2.resize(im, (int(resize_w), int(resize_h)))
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)
        return im, (ratio_h, ratio_w)

    def resize_image(self, im, max_side_len=512):
        """
        resize image to a size multiple of max_stride which is required by the network
        :param im: the resized image
        :param max_side_len: limit of max image size to avoid out of memory in gpu
        :return: the resized image and the resize ratio
        """
        h, w, _ = im.shape

        resize_w = w
        resize_h = h

        # Fix the longer side
        if resize_h > resize_w:
            ratio = float(max_side_len) / resize_h
        else:
            ratio = float(max_side_len) / resize_w

        resize_h = int(resize_h * ratio)
        resize_w = int(resize_w * ratio)

        max_stride = 128
        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
        im = cv2.resize(im, (int(resize_w), int(resize_h)))
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)

        return im, (ratio_h, ratio_w)
L
add kie  
LDOUBLEV 已提交
394 395 396 397 398 399 400 401 402 403 404 405


class KieResize(object):
    def __init__(self, **kwargs):
        super(KieResize, self).__init__()
        self.max_side, self.min_side = kwargs['img_scale'][0], kwargs[
            'img_scale'][1]

    def __call__(self, data):
        img = data['image']
        points = data['points']
        src_h, src_w, _ = img.shape
L
debug  
LDOUBLEV 已提交
406 407
        im_resized, scale_factor, [ratio_h, ratio_w
                                   ], [new_h, new_w] = self.resize_image(img)
L
add kie  
LDOUBLEV 已提交
408 409 410 411 412
        resize_points = self.resize_boxes(img, points, scale_factor)
        data['ori_image'] = img
        data['ori_boxes'] = points
        data['points'] = resize_points
        data['image'] = im_resized
L
debug  
LDOUBLEV 已提交
413
        data['shape'] = np.array([new_h, new_w])
L
add kie  
LDOUBLEV 已提交
414 415 416
        return data

    def resize_image(self, img):
L
debug  
LDOUBLEV 已提交
417
        norm_img = np.zeros([1024, 1024, 3], dtype='float32')
L
add kie  
LDOUBLEV 已提交
418 419 420 421 422 423
        scale = [512, 1024]
        h, w = img.shape[:2]
        max_long_edge = max(scale)
        max_short_edge = min(scale)
        scale_factor = min(max_long_edge / max(h, w),
                           max_short_edge / min(h, w))
L
debug  
LDOUBLEV 已提交
424 425 426 427 428 429
        resize_w, resize_h = int(w * float(scale_factor) + 0.5), int(h * float(
            scale_factor) + 0.5)
        max_stride = 32
        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
        im = cv2.resize(img, (resize_w, resize_h))
L
add kie  
LDOUBLEV 已提交
430 431 432 433 434 435
        new_h, new_w = im.shape[:2]
        w_scale = new_w / w
        h_scale = new_h / h
        scale_factor = np.array(
            [w_scale, h_scale, w_scale, h_scale], dtype=np.float32)
        norm_img[:new_h, :new_w, :] = im
L
debug  
LDOUBLEV 已提交
436
        return norm_img, scale_factor, [h_scale, w_scale], [new_h, new_w]
L
add kie  
LDOUBLEV 已提交
437 438 439 440 441 442 443

    def resize_boxes(self, im, points, scale_factor):
        points = points * scale_factor
        img_shape = im.shape[:2]
        points[:, 0::2] = np.clip(points[:, 0::2], 0, img_shape[1])
        points[:, 1::2] = np.clip(points[:, 1::2], 0, img_shape[0])
        return points
X
xiaoting 已提交
444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492


class SRResize(object):
    def __init__(self,
                 imgH=32,
                 imgW=128,
                 down_sample_scale=4,
                 keep_ratio=False,
                 min_ratio=1,
                 mask=False,
                 infer_mode=False,
                 **kwargs):
        self.imgH = imgH
        self.imgW = imgW
        self.keep_ratio = keep_ratio
        self.min_ratio = min_ratio
        self.down_sample_scale = down_sample_scale
        self.mask = mask
        self.infer_mode = infer_mode

    def __call__(self, data):
        imgH = self.imgH
        imgW = self.imgW
        images_lr = data["image_lr"]
        transform2 = ResizeNormalize(
            (imgW // self.down_sample_scale, imgH // self.down_sample_scale))
        images_lr = transform2(images_lr)
        data["img_lr"] = images_lr
        if self.infer_mode:
            return data

        images_HR = data["image_hr"]
        label_strs = data["label"]
        transform = ResizeNormalize((imgW, imgH))
        images_HR = transform(images_HR)
        data["img_hr"] = images_HR
        return data


class ResizeNormalize(object):
    def __init__(self, size, interpolation=Image.BICUBIC):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img):
        img = img.resize(self.size, self.interpolation)
        img_numpy = np.array(img).astype("float32")
        img_numpy = img_numpy.transpose((2, 0, 1)) / 255
        return img_numpy