# -*- coding:utf-8 -*- from __future__ import absolute_import from __future__ import division from __future__ import print_function import sys import cv2 import numpy as np import paddle import pyclipper from PIL import Image from PIL import ImageDraw from PIL import ImageFont from shapely.geometry import Polygon class DBProcessTest(object): """ DB pre-process for Test mode """ def __init__(self, params): super(DBProcessTest, self).__init__() self.resize_type = 0 if 'test_image_shape' in params: self.image_shape = params['test_image_shape'] self.resize_type = 1 if 'max_side_len' in params: self.max_side_len = params['max_side_len'] else: self.max_side_len = 2400 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.max_side_len h, w, _ = img.shape # limit the max side 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. resize_h = int(h * ratio) resize_w = int(w * ratio) resize_h = max(int(round(resize_h / 32) * 32), 32) resize_w = max(int(round(resize_w / 32) * 32), 32) 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: sys.exit(0) ratio_h = resize_h / float(h) ratio_w = resize_w / float(w) # return img, np.array([h, w]) return img, [ratio_h, ratio_w] def resize_image_type1(self, im): resize_h, resize_w = self.image_shape ori_h, ori_w = im.shape[:2] # (h, w, c) im = cv2.resize(im, (int(resize_w), int(resize_h))) ratio_h = float(resize_h) / ori_h ratio_w = float(resize_w) / ori_w return im, (ratio_h, ratio_w) def normalize(self, im): img_mean = [0.485, 0.456, 0.406] img_std = [0.229, 0.224, 0.225] im = im.astype(np.float32, copy=False) im = im / 255 im[:, :, 0] -= img_mean[0] im[:, :, 1] -= img_mean[1] im[:, :, 2] -= img_mean[2] im[:, :, 0] /= img_std[0] im[:, :, 1] /= img_std[1] im[:, :, 2] /= img_std[2] channel_swap = (2, 0, 1) im = im.transpose(channel_swap) return im def __call__(self, im): src_h, src_w, _ = im.shape if self.resize_type == 0: im, (ratio_h, ratio_w) = self.resize_image_type0(im) else: im, (ratio_h, ratio_w) = self.resize_image_type1(im) im = self.normalize(im) im = im[np.newaxis, :] return [im, (src_h, src_w, ratio_h, ratio_w)] class DBPostProcess(object): """ The post process for Differentiable Binarization (DB). """ def __init__(self, params): self.thresh = params['thresh'] self.box_thresh = params['box_thresh'] self.max_candidates = params['max_candidates'] self.unclip_ratio = params['unclip_ratio'] self.min_size = 3 self.dilation_kernel = None self.score_mode = 'fast' def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height): ''' _bitmap: single map with shape (1, H, W), whose values are binarized as {0, 1} ''' bitmap = _bitmap height, width = bitmap.shape outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) if len(outs) == 3: img, contours, _ = outs[0], outs[1], outs[2] elif len(outs) == 2: contours, _ = outs[0], outs[1] num_contours = min(len(contours), self.max_candidates) boxes = [] scores = [] for index in range(num_contours): contour = contours[index] points, sside = self.get_mini_boxes(contour) if sside < self.min_size: continue points = np.array(points) if self.score_mode == "fast": score = self.box_score_fast(pred, points.reshape(-1, 2)) else: score = self.box_score_slow(pred, contour) if self.box_thresh > score: continue box = self.unclip(points).reshape(-1, 1, 2) box, sside = self.get_mini_boxes(box) if sside < self.min_size + 2: continue box = np.array(box) box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width) box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height) boxes.append(box.astype(np.int16)) scores.append(score) return np.array(boxes, dtype=np.int16), scores def unclip(self, box): unclip_ratio = self.unclip_ratio poly = Polygon(box) distance = poly.area * unclip_ratio / poly.length offset = pyclipper.PyclipperOffset() offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON) expanded = np.array(offset.Execute(distance)) return expanded def get_mini_boxes(self, contour): bounding_box = cv2.minAreaRect(contour) points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0]) index_1, index_2, index_3, index_4 = 0, 1, 2, 3 if points[1][1] > points[0][1]: index_1 = 0 index_4 = 1 else: index_1 = 1 index_4 = 0 if points[3][1] > points[2][1]: index_2 = 2 index_3 = 3 else: index_2 = 3 index_3 = 2 box = [points[index_1], points[index_2], points[index_3], points[index_4]] return box, min(bounding_box[1]) def box_score_fast(self, bitmap, _box): ''' box_score_fast: use bbox mean score as the mean score ''' h, w = bitmap.shape[:2] box = _box.copy() xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1) xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int), 0, w - 1) ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int), 0, h - 1) ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int), 0, h - 1) mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8) box[:, 0] = box[:, 0] - xmin box[:, 1] = box[:, 1] - ymin cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1) return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] def box_score_slow(self, bitmap, contour): ''' box_score_slow: use polyon mean score as the mean score ''' h, w = bitmap.shape[:2] contour = contour.copy() contour = np.reshape(contour, (-1, 2)) xmin = np.clip(np.min(contour[:, 0]), 0, w - 1) xmax = np.clip(np.max(contour[:, 0]), 0, w - 1) ymin = np.clip(np.min(contour[:, 1]), 0, h - 1) ymax = np.clip(np.max(contour[:, 1]), 0, h - 1) mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8) contour[:, 0] = contour[:, 0] - xmin contour[:, 1] = contour[:, 1] - ymin cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype(np.int32), 1) return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] def __call__(self, outs_dict, shape_list): pred = outs_dict['maps'] pred = pred[:, 0, :, :] segmentation = pred > self.thresh boxes_batch = [] for batch_index in range(pred.shape[0]): src_h, src_w, ratio_h, ratio_w = shape_list[batch_index] mask = segmentation[batch_index] tmp_boxes, tmp_scores = self.boxes_from_bitmap(pred[batch_index], mask, src_w, src_h) boxes_batch.append(tmp_boxes) return boxes_batch def draw_boxes(image, boxes, scores=None, drop_score=0.5): img = image.copy() draw = ImageDraw.Draw(img) if scores is None: scores = [1] * len(boxes) for (box, score) in zip(boxes, scores): if score < drop_score: continue draw.line([(box[0][0], box[0][1]), (box[1][0], box[1][1])], fill='red') draw.line([(box[1][0], box[1][1]), (box[2][0], box[2][1])], fill='red') draw.line([(box[2][0], box[2][1]), (box[3][0], box[3][1])], fill='red') draw.line([(box[3][0], box[3][1]), (box[0][0], box[0][1])], fill='red') draw.line([(box[0][0] - 1, box[0][1] + 1), (box[1][0] - 1, box[1][1] + 1)], fill='red') draw.line([(box[1][0] - 1, box[1][1] + 1), (box[2][0] - 1, box[2][1] + 1)], fill='red') draw.line([(box[2][0] - 1, box[2][1] + 1), (box[3][0] - 1, box[3][1] + 1)], fill='red') draw.line([(box[3][0] - 1, box[3][1] + 1), (box[0][0] - 1, box[0][1] + 1)], fill='red') return img def get_image_ext(image): if image.shape[2] == 4: return ".png" return ".jpg"