diff --git a/ppocr/postprocess/db_postprocess_torch.py b/ppocr/postprocess/db_postprocess_torch.py deleted file mode 100644 index d1466327f13aeca74057e7593c890e7d3fff8c6b..0000000000000000000000000000000000000000 --- a/ppocr/postprocess/db_postprocess_torch.py +++ /dev/null @@ -1,136 +0,0 @@ -import cv2 -import paddle -import numpy as np -import pyclipper -from shapely.geometry import Polygon - - -class DBPostProcess(): - def __init__(self, - thresh=0.3, - box_thresh=0.7, - max_candidates=1000, - unclip_ratio=1.5): - self.min_size = 3 - self.thresh = thresh - self.box_thresh = box_thresh - self.max_candidates = max_candidates - self.unclip_ratio = unclip_ratio - - def __call__(self, pred, shape_list, is_output_polygon=False): - ''' - batch: (image, polygons, ignore_tags - h_w_list: 包含[h,w]的数组 - pred: - binary: text region segmentation map, with shape (N, 1,H, W) - ''' - if isinstance(pred, paddle.Tensor): - pred = pred.numpy() - pred = pred[:, 0, :, :] - segmentation = self.binarize(pred) - batch_out = [] - for batch_index in range(pred.shape[0]): - height, width = shape_list[batch_index] - boxes, scores = self.post_p( - pred[batch_index], - segmentation[batch_index], - width, - height, - is_output_polygon=is_output_polygon) - batch_out.append({"points": boxes}) - return batch_out - - def binarize(self, pred): - return pred > self.thresh - - def post_p(self, - pred, - bitmap, - dest_width, - dest_height, - is_output_polygon=True): - ''' - _bitmap: single map with shape (H, W), - whose values are binarized as {0, 1} - ''' - height, width = pred.shape - boxes = [] - new_scores = [] - contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), - cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) - for contour in contours[:self.max_candidates]: - epsilon = 0.005 * cv2.arcLength(contour, True) - approx = cv2.approxPolyDP(contour, epsilon, True) - points = approx.reshape((-1, 2)) - if points.shape[0] < 4: - continue - score = self.box_score_fast(pred, points.reshape(-1, 2)) - if self.box_thresh > score: - continue - - if points.shape[0] > 2: - box = self.unclip(points, unclip_ratio=self.unclip_ratio) - if len(box) > 1 or len(box) == 0: - continue - else: - continue - four_point_box, sside = self.get_mini_boxes(box.reshape((-1, 1, 2))) - if sside < self.min_size + 2: - continue - - if not is_output_polygon: - box = np.array(four_point_box) - else: - box = box.reshape(-1, 2) - 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) - new_scores.append(score) - return boxes, new_scores - - def unclip(self, box, unclip_ratio=1.5): - 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): - 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] \ No newline at end of file