1"""
2# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15"""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20from __future__ import unicode_literals
21
22import sys
23import cv2
24import numpy as np
25import math
26from PIL import Image
27from paddle import get_device
28
29
30class DecodeImage(object):
31 """decode image"""
32
33 def __init__(
34 self, img_mode="RGB", channel_first=False, ignore_orientation=False, **kwargs
35 ):
36 self.img_mode = img_mode
37 self.channel_first = channel_first
38 self.ignore_orientation = ignore_orientation
39
40 def __call__(self, data):
41 img = data["image"]
42 assert type(img) is bytes and len(img) > 0, "invalid input 'img' in DecodeImage"
43 img = np.frombuffer(img, dtype="uint8")
44 if self.img_mode == "GRAY":
45 # For GRAY mode, decode directly to a single-channel grayscale image.
46 decode_flag = cv2.IMREAD_GRAYSCALE
47 else:
48 # For RGB mode, decode to a 3-channel color image.
49 decode_flag = cv2.IMREAD_COLOR
50
51 if self.ignore_orientation:
52 decode_flag |= cv2.IMREAD_IGNORE_ORIENTATION
53
54 img = cv2.imdecode(img, decode_flag)
55
56 if img is None:
57 return None
58 if self.img_mode == "GRAY":
59 img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
60 elif self.img_mode == "RGB":
61 assert img.shape[2] == 3, "invalid shape of image[%s]" % (img.shape)
62 img = img[:, :, ::-1]
63
64 if self.channel_first:
65 img = img.transpose((2, 0, 1))
66
67 data["image"] = img
68 return data
69
70
71class NormalizeImage(object):
72 """normalize image such as subtract mean, divide std"""
73
74 def __init__(self, scale=None, mean=None, std=None, order="chw", **kwargs):
75 if isinstance(scale, str):
76 scale = eval(scale)
77 self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
78 mean = mean if mean is not None else [0.485, 0.456, 0.406]
79 std = std if std is not None else [0.229, 0.224, 0.225]
80
81 shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
82 self.mean = np.array(mean).reshape(shape).astype("float32")
83 self.std = np.array(std).reshape(shape).astype("float32")
84
85 def __call__(self, data):
86 img = data["image"]
87 from PIL import Image
88
89 if isinstance(img, Image.Image):
90 img = np.array(img)
91 assert isinstance(img, np.ndarray), "invalid input 'img' in NormalizeImage"
92 data["image"] = (img.astype("float32") * self.scale - self.mean) / self.std
93 return data
94
95
96class ToCHWImage(object):
97 """convert hwc image to chw image"""
98
99 def __init__(self, **kwargs):
100 pass
101
102 def __call__(self, data):
103 img = data["image"]
104 from PIL import Image
105
106 if isinstance(img, Image.Image):
107 img = np.array(img)
108 data["image"] = img.transpose((2, 0, 1))
109 return data
110
111
112class Fasttext(object):
113 def __init__(self, path="None", **kwargs):
114 import fasttext
115
116 self.fast_model = fasttext.load_model(path)
117
118 def __call__(self, data):
119 label = data["label"]
120 fast_label = self.fast_model[label]
121 data["fast_label"] = fast_label
122 return data
123
124
125class KeepKeys(object):
126 def __init__(self, keep_keys, **kwargs):
127 self.keep_keys = keep_keys
128
129 def __call__(self, data):
130 data_list = []
131 for key in self.keep_keys:
132 data_list.append(data[key])
133 return data_list
134
135
136class Pad(object):
137 def __init__(self, size=None, size_div=32, **kwargs):
138 if size is not None and not isinstance(size, (int, list, tuple)):
139 raise TypeError(
140 "Type of target_size is invalid. Now is {}".format(type(size))
141 )
142 if isinstance(size, int):
143 size = [size, size]
144 self.size = size
145 self.size_div = size_div
146
147 def __call__(self, data):
148 img = data["image"]
149 img_h, img_w = img.shape[0], img.shape[1]
150 if self.size:
151 resize_h2, resize_w2 = self.size
152 assert (
153 img_h < resize_h2 and img_w < resize_w2
154 ), "(h, w) of target size should be greater than (img_h, img_w)"
155 else:
156 resize_h2 = max(
157 int(math.ceil(img.shape[0] / self.size_div) * self.size_div),
158 self.size_div,
159 )
160 resize_w2 = max(
161 int(math.ceil(img.shape[1] / self.size_div) * self.size_div),
162 self.size_div,
163 )
164 img = cv2.copyMakeBorder(
165 img,
166 0,
167 resize_h2 - img_h,
168 0,
169 resize_w2 - img_w,
170 cv2.BORDER_CONSTANT,
171 value=0,
172 )
173 data["image"] = img
174 return data
175
176
177class Resize(object):
178 def __init__(self, size=(640, 640), **kwargs):
179 self.size = size
180
181 def resize_image(self, img):
182 resize_h, resize_w = self.size
183 ori_h, ori_w = img.shape[:2] # (h, w, c)
184 ratio_h = float(resize_h) / ori_h
185 ratio_w = float(resize_w) / ori_w
186 img = cv2.resize(img, (int(resize_w), int(resize_h)))
187 return img, [ratio_h, ratio_w]
188
189 def __call__(self, data):
190 img = data["image"]
191 if "polys" in data:
192 text_polys = data["polys"]
193
194 img_resize, [ratio_h, ratio_w] = self.resize_image(img)
195 if "polys" in data:
196 new_boxes = []
197 for box in text_polys:
198 new_box = []
199 for cord in box:
200 new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
201 new_boxes.append(new_box)
202 data["polys"] = np.array(new_boxes, dtype=np.float32)
203 data["image"] = img_resize
204 return data
205
206
207class DetResizeForTest(object):
208 def __init__(self, **kwargs):
209 super(DetResizeForTest, self).__init__()
210 self.resize_type = 0
211 self.keep_ratio = False
212 if "image_shape" in kwargs:
213 self.image_shape = kwargs["image_shape"]
214 self.resize_type = 1
215 if "keep_ratio" in kwargs:
216 self.keep_ratio = kwargs["keep_ratio"]
217 elif "limit_side_len" in kwargs:
218 self.limit_side_len = kwargs["limit_side_len"]
219 self.limit_type = kwargs.get("limit_type", "min")
220 elif "resize_long" in kwargs:
221 self.resize_type = 2
222 self.resize_long = kwargs.get("resize_long", 960)
223 else:
224 self.limit_side_len = 736
225 self.limit_type = "min"
226
227 def __call__(self, data):
228 img = data["image"]
229 src_h, src_w, _ = img.shape
230 if sum([src_h, src_w]) < 64:
231 img = self.image_padding(img)
232
233 if self.resize_type == 0:
234 # img, shape = self.resize_image_type0(img)
235 img, [ratio_h, ratio_w] = self.resize_image_type0(img)
236 elif self.resize_type == 2:
237 img, [ratio_h, ratio_w] = self.resize_image_type2(img)
238 else:
239 # img, shape = self.resize_image_type1(img)
240 img, [ratio_h, ratio_w] = self.resize_image_type1(img)
241 data["image"] = img
242 data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
243 if "iluvatar_gpu" in get_device():
244 data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w]).astype(np.float32)
245 return data
246
247 def image_padding(self, im, value=0):
248 h, w, c = im.shape
249 im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
250 im_pad[:h, :w, :] = im
251 return im_pad
252
253 def resize_image_type1(self, img):
254 resize_h, resize_w = self.image_shape
255 ori_h, ori_w = img.shape[:2] # (h, w, c)
256 if self.keep_ratio is True:
257 resize_w = ori_w * resize_h / ori_h
258 N = math.ceil(resize_w / 32)
259 resize_w = N * 32
260 ratio_h = float(resize_h) / ori_h
261 ratio_w = float(resize_w) / ori_w
262 img = cv2.resize(img, (int(resize_w), int(resize_h)))
263 # return img, np.array([ori_h, ori_w])
264 return img, [ratio_h, ratio_w]
265
266 def resize_image_type0(self, img):
267 """
268 resize image to a size multiple of 32 which is required by the network
269 args:
270 img(array): array with shape [h, w, c]
271 return(tuple):
272 img, (ratio_h, ratio_w)
273 """
274 limit_side_len = self.limit_side_len
275 h, w, c = img.shape
276
277 # limit the max side
278 if self.limit_type == "max":
279 if max(h, w) > limit_side_len:
280 if h > w:
281 ratio = float(limit_side_len) / h
282 else:
283 ratio = float(limit_side_len) / w
284 else:
285 ratio = 1.0
286 elif self.limit_type == "min":
287 if min(h, w) < limit_side_len:
288 if h < w:
289 ratio = float(limit_side_len) / h
290 else:
291 ratio = float(limit_side_len) / w
292 else:
293 ratio = 1.0
294 elif self.limit_type == "resize_long":
295 ratio = float(limit_side_len) / max(h, w)
296 else:
297 raise Exception("not support limit type, image ")
298 resize_h = int(h * ratio)
299 resize_w = int(w * ratio)
300
301 resize_h = max(int(round(resize_h / 32) * 32), 32)
302 resize_w = max(int(round(resize_w / 32) * 32), 32)
303
304 try:
305 if int(resize_w) <= 0 or int(resize_h) <= 0:
306 return None, (None, None)
307 img = cv2.resize(img, (int(resize_w), int(resize_h)))
308 except:
309 print(img.shape, resize_w, resize_h)
310 sys.exit(0)
311 ratio_h = resize_h / float(h)
312 ratio_w = resize_w / float(w)
313 return img, [ratio_h, ratio_w]
314
315 def resize_image_type2(self, img):
316 h, w, _ = img.shape
317
318 resize_w = w
319 resize_h = h
320
321 if resize_h > resize_w:
322 ratio = float(self.resize_long) / resize_h
323 else:
324 ratio = float(self.resize_long) / resize_w
325
326 resize_h = int(resize_h * ratio)
327 resize_w = int(resize_w * ratio)
328
329 max_stride = 128
330 resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
331 resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
332 img = cv2.resize(img, (int(resize_w), int(resize_h)))
333 ratio_h = resize_h / float(h)
334 ratio_w = resize_w / float(w)
335
336 return img, [ratio_h, ratio_w]
337
338
339class E2EResizeForTest(object):
340 def __init__(self, **kwargs):
341 super(E2EResizeForTest, self).__init__()
342 self.max_side_len = kwargs["max_side_len"]
343 self.valid_set = kwargs["valid_set"]
344
345 def __call__(self, data):
346 img = data["image"]
347 src_h, src_w, _ = img.shape
348 if self.valid_set == "totaltext":
349 im_resized, [ratio_h, ratio_w] = self.resize_image_for_totaltext(
350 img, max_side_len=self.max_side_len
351 )
352 else:
353 im_resized, (ratio_h, ratio_w) = self.resize_image(
354 img, max_side_len=self.max_side_len
355 )
356 data["image"] = im_resized
357 data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
358 return data
359
360 def resize_image_for_totaltext(self, im, max_side_len=512):
361 h, w, _ = im.shape
362 resize_w = w
363 resize_h = h
364 ratio = 1.25
365 if h * ratio > max_side_len:
366 ratio = float(max_side_len) / resize_h
367 resize_h = int(resize_h * ratio)
368 resize_w = int(resize_w * ratio)
369
370 max_stride = 128
371 resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
372 resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
373 im = cv2.resize(im, (int(resize_w), int(resize_h)))
374 ratio_h = resize_h / float(h)
375 ratio_w = resize_w / float(w)
376 return im, (ratio_h, ratio_w)
377
378 def resize_image(self, im, max_side_len=512):
379 """
380 resize image to a size multiple of max_stride which is required by the network
381 :param im: the resized image
382 :param max_side_len: limit of max image size to avoid out of memory in gpu
383 :return: the resized image and the resize ratio
384 """
385 h, w, _ = im.shape
386
387 resize_w = w
388 resize_h = h
389
390 # Fix the longer side
391 if resize_h > resize_w:
392 ratio = float(max_side_len) / resize_h
393 else:
394 ratio = float(max_side_len) / resize_w
395
396 resize_h = int(resize_h * ratio)
397 resize_w = int(resize_w * ratio)
398
399 max_stride = 128
400 resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
401 resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
402 im = cv2.resize(im, (int(resize_w), int(resize_h)))
403 ratio_h = resize_h / float(h)
404 ratio_w = resize_w / float(w)
405
406 return im, (ratio_h, ratio_w)
407
408
409class KieResize(object):
410 def __init__(self, **kwargs):
411 super(KieResize, self).__init__()
412 self.max_side, self.min_side = kwargs["img_scale"][0], kwargs["img_scale"][1]
413
414 def __call__(self, data):
415 img = data["image"]
416 points = data["points"]
417 src_h, src_w, _ = img.shape
418 (
419 im_resized,
420 scale_factor,
421 [ratio_h, ratio_w],
422 [new_h, new_w],
423 ) = self.resize_image(img)
424 resize_points = self.resize_boxes(img, points, scale_factor)
425 data["ori_image"] = img
426 data["ori_boxes"] = points
427 data["points"] = resize_points
428 data["image"] = im_resized
429 data["shape"] = np.array([new_h, new_w])
430 return data
431
432 def resize_image(self, img):
433 norm_img = np.zeros([1024, 1024, 3], dtype="float32")
434 scale = [512, 1024]
435 h, w = img.shape[:2]
436 max_long_edge = max(scale)
437 max_short_edge = min(scale)
438 scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
439 resize_w, resize_h = int(w * float(scale_factor) + 0.5), int(
440 h * float(scale_factor) + 0.5
441 )
442 max_stride = 32
443 resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
444 resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
445 im = cv2.resize(img, (resize_w, resize_h))
446 new_h, new_w = im.shape[:2]
447 w_scale = new_w / w
448 h_scale = new_h / h
449 scale_factor = np.array([w_scale, h_scale, w_scale, h_scale], dtype=np.float32)
450 norm_img[:new_h, :new_w, :] = im
451 return norm_img, scale_factor, [h_scale, w_scale], [new_h, new_w]
452
453 def resize_boxes(self, im, points, scale_factor):
454 points = points * scale_factor
455 img_shape = im.shape[:2]
456 points[:, 0::2] = np.clip(points[:, 0::2], 0, img_shape[1])
457 points[:, 1::2] = np.clip(points[:, 1::2], 0, img_shape[0])
458 return points
459
460
461class SRResize(object):
462 def __init__(
463 self,
464 imgH=32,
465 imgW=128,
466 down_sample_scale=4,
467 keep_ratio=False,
468 min_ratio=1,
469 mask=False,
470 infer_mode=False,
471 **kwargs,
472 ):
473 self.imgH = imgH
474 self.imgW = imgW
475 self.keep_ratio = keep_ratio
476 self.min_ratio = min_ratio
477 self.down_sample_scale = down_sample_scale
478 self.mask = mask
479 self.infer_mode = infer_mode
480
481 def __call__(self, data):
482 imgH = self.imgH
483 imgW = self.imgW
484 images_lr = data["image_lr"]
485 transform2 = ResizeNormalize(
486 (imgW // self.down_sample_scale, imgH // self.down_sample_scale)
487 )
488 images_lr = transform2(images_lr)
489 data["img_lr"] = images_lr
490 if self.infer_mode:
491 return data
492
493 images_HR = data["image_hr"]
494 label_strs = data["label"]
495 transform = ResizeNormalize((imgW, imgH))
496 images_HR = transform(images_HR)
497 data["img_hr"] = images_HR
498 return data
499
500
501class ResizeNormalize(object):
502 def __init__(self, size, interpolation=Image.BICUBIC):
503 self.size = size
504 self.interpolation = interpolation
505
506 def __call__(self, img):
507 img = img.resize(self.size, self.interpolation)
508 img_numpy = np.array(img).astype("float32")
509 img_numpy = img_numpy.transpose((2, 0, 1)) / 255
510 return img_numpy
511
512
513class GrayImageChannelFormat(object):
514 """
515 format gray scale image's channel: (3,h,w) -> (1,h,w)
516 Args:
517 inverse: inverse gray image
518 """
519
520 def __init__(self, inverse=False, **kwargs):
521 self.inverse = inverse
522
523 def __call__(self, data):
524 img = data["image"]
525 img_single_channel = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
526 img_expanded = np.expand_dims(img_single_channel, 0)
527
528 if self.inverse:
529 data["image"] = np.abs(img_expanded - 1)
530 else:
531 data["image"] = img_expanded
532
533 data["src_image"] = img
534 return data
535