删除 db torch后处理
This commit is contained in:
WenmuZhou
parent
33d9688014
commit
367c49dffd
|
|
@ -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]
|
||||
Loading…
Reference in New Issue