134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
import cv2
|
|
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)
|
|
'''
|
|
pred = pred.numpy()[:, 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]
|