|
import cv2
|
|
import numpy as np
|
|
import pyclipper
|
|
from shapely.geometry import Polygon
|
|
from collections import namedtuple
|
|
import warnings
|
|
import torch
|
|
warnings.filterwarnings('ignore')
|
|
|
|
|
|
def iou_rotate(box_a, box_b, method='union'):
|
|
rect_a = cv2.minAreaRect(box_a)
|
|
rect_b = cv2.minAreaRect(box_b)
|
|
r1 = cv2.rotatedRectangleIntersection(rect_a, rect_b)
|
|
if r1[0] == 0:
|
|
return 0
|
|
else:
|
|
inter_area = cv2.contourArea(r1[1])
|
|
area_a = cv2.contourArea(box_a)
|
|
area_b = cv2.contourArea(box_b)
|
|
union_area = area_a + area_b - inter_area
|
|
if union_area == 0 or inter_area == 0:
|
|
return 0
|
|
if method == 'union':
|
|
iou = inter_area / union_area
|
|
elif method == 'intersection':
|
|
iou = inter_area / min(area_a, area_b)
|
|
else:
|
|
raise NotImplementedError
|
|
return iou
|
|
|
|
class SegDetectorRepresenter():
|
|
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, batch, pred, is_output_polygon=False, height=None, width=None):
|
|
'''
|
|
batch: (image, polygons, ignore_tags
|
|
batch: a dict produced by dataloaders.
|
|
image: tensor of shape (N, C, H, W).
|
|
polygons: tensor of shape (N, K, 4, 2), the polygons of objective regions.
|
|
ignore_tags: tensor of shape (N, K), indicates whether a region is ignorable or not.
|
|
shape: the original shape of images.
|
|
filename: the original filenames of images.
|
|
pred:
|
|
binary: text region segmentation map, with shape (N, H, W)
|
|
thresh: [if exists] thresh hold prediction with shape (N, H, W)
|
|
thresh_binary: [if exists] binarized with threshold, (N, H, W)
|
|
'''
|
|
pred = pred[:, 0, :, :]
|
|
segmentation = self.binarize(pred)
|
|
boxes_batch = []
|
|
scores_batch = []
|
|
|
|
batch_size = pred.size(0) if isinstance(pred, torch.Tensor) else pred.shape[0]
|
|
|
|
if height is None:
|
|
height = pred.shape[1]
|
|
if width is None:
|
|
width = pred.shape[2]
|
|
|
|
for batch_index in range(batch_size):
|
|
if is_output_polygon:
|
|
boxes, scores = self.polygons_from_bitmap(pred[batch_index], segmentation[batch_index], width, height)
|
|
else:
|
|
boxes, scores = self.boxes_from_bitmap(pred[batch_index], segmentation[batch_index], width, height)
|
|
boxes_batch.append(boxes)
|
|
scores_batch.append(scores)
|
|
return boxes_batch, scores_batch
|
|
|
|
def binarize(self, pred) -> np.ndarray:
|
|
return pred > self.thresh
|
|
|
|
def polygons_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
|
|
'''
|
|
_bitmap: single map with shape (H, W),
|
|
whose values are binarized as {0, 1}
|
|
'''
|
|
|
|
assert len(_bitmap.shape) == 2
|
|
bitmap = _bitmap.cpu().numpy()
|
|
pred = pred.cpu().detach().numpy()
|
|
height, width = bitmap.shape
|
|
boxes = []
|
|
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, contour.squeeze(1))
|
|
if self.box_thresh > score:
|
|
continue
|
|
|
|
if points.shape[0] > 2:
|
|
box = self.unclip(points, unclip_ratio=self.unclip_ratio)
|
|
if len(box) > 1:
|
|
continue
|
|
else:
|
|
continue
|
|
box = box.reshape(-1, 2)
|
|
_, sside = self.get_mini_boxes(box.reshape((-1, 1, 2)))
|
|
if sside < self.min_size + 2:
|
|
continue
|
|
|
|
if not isinstance(dest_width, int):
|
|
dest_width = dest_width.item()
|
|
dest_height = dest_height.item()
|
|
|
|
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)
|
|
scores.append(score)
|
|
return boxes, scores
|
|
|
|
def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
|
|
'''
|
|
_bitmap: single map with shape (H, W),
|
|
whose values are binarized as {0, 1}
|
|
'''
|
|
|
|
assert len(_bitmap.shape) == 2
|
|
if isinstance(pred, torch.Tensor):
|
|
bitmap = _bitmap.cpu().numpy()
|
|
pred = pred.cpu().detach().numpy()
|
|
else:
|
|
bitmap = _bitmap
|
|
|
|
height, width = bitmap.shape
|
|
contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
|
|
num_contours = min(len(contours), self.max_candidates)
|
|
boxes = np.zeros((num_contours, 4, 2), dtype=np.int16)
|
|
scores = np.zeros((num_contours,), dtype=np.float32)
|
|
|
|
for index in range(num_contours):
|
|
contour = contours[index].squeeze(1)
|
|
points, sside = self.get_mini_boxes(contour)
|
|
|
|
|
|
if sside < 2:
|
|
continue
|
|
points = np.array(points)
|
|
score = self.box_score_fast(pred, contour)
|
|
|
|
|
|
|
|
box = self.unclip(points, unclip_ratio=self.unclip_ratio).reshape(-1, 1, 2)
|
|
box, sside = self.get_mini_boxes(box)
|
|
|
|
|
|
box = np.array(box)
|
|
if not isinstance(dest_width, int):
|
|
dest_width = dest_width.item()
|
|
dest_height = dest_height.item()
|
|
|
|
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[index, :, :] = box.astype(np.int16)
|
|
scores[index] = score
|
|
return boxes, 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.int32), 0, w - 1)
|
|
xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int32), 0, w - 1)
|
|
ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int32), 0, h - 1)
|
|
ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int32), 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)
|
|
if bitmap.dtype == np.float16:
|
|
bitmap = bitmap.astype(np.float32)
|
|
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
|
|
|
|
class AverageMeter(object):
|
|
"""Computes and stores the average and current value"""
|
|
|
|
def __init__(self):
|
|
self.reset()
|
|
|
|
def reset(self):
|
|
self.val = 0
|
|
self.avg = 0
|
|
self.sum = 0
|
|
self.count = 0
|
|
|
|
def update(self, val, n=1):
|
|
self.val = val
|
|
self.sum += val * n
|
|
self.count += n
|
|
self.avg = self.sum / self.count
|
|
return self
|
|
|
|
|
|
class DetectionIoUEvaluator(object):
|
|
def __init__(self, is_output_polygon=False, iou_constraint=0.5, area_precision_constraint=0.5):
|
|
self.is_output_polygon = is_output_polygon
|
|
self.iou_constraint = iou_constraint
|
|
self.area_precision_constraint = area_precision_constraint
|
|
|
|
def evaluate_image(self, gt, pred):
|
|
|
|
def get_union(pD, pG):
|
|
return Polygon(pD).union(Polygon(pG)).area
|
|
|
|
def get_intersection_over_union(pD, pG):
|
|
return get_intersection(pD, pG) / get_union(pD, pG)
|
|
|
|
def get_intersection(pD, pG):
|
|
return Polygon(pD).intersection(Polygon(pG)).area
|
|
|
|
def compute_ap(confList, matchList, numGtCare):
|
|
correct = 0
|
|
AP = 0
|
|
if len(confList) > 0:
|
|
confList = np.array(confList)
|
|
matchList = np.array(matchList)
|
|
sorted_ind = np.argsort(-confList)
|
|
confList = confList[sorted_ind]
|
|
matchList = matchList[sorted_ind]
|
|
for n in range(len(confList)):
|
|
match = matchList[n]
|
|
if match:
|
|
correct += 1
|
|
AP += float(correct) / (n + 1)
|
|
|
|
if numGtCare > 0:
|
|
AP /= numGtCare
|
|
|
|
return AP
|
|
|
|
perSampleMetrics = {}
|
|
|
|
matchedSum = 0
|
|
|
|
Rectangle = namedtuple('Rectangle', 'xmin ymin xmax ymax')
|
|
|
|
numGlobalCareGt = 0
|
|
numGlobalCareDet = 0
|
|
|
|
arrGlobalConfidences = []
|
|
arrGlobalMatches = []
|
|
|
|
recall = 0
|
|
precision = 0
|
|
hmean = 0
|
|
|
|
detMatched = 0
|
|
|
|
iouMat = np.empty([1, 1])
|
|
|
|
gtPols = []
|
|
detPols = []
|
|
|
|
gtPolPoints = []
|
|
detPolPoints = []
|
|
|
|
|
|
gtDontCarePolsNum = []
|
|
|
|
detDontCarePolsNum = []
|
|
|
|
pairs = []
|
|
detMatchedNums = []
|
|
|
|
arrSampleConfidences = []
|
|
arrSampleMatch = []
|
|
|
|
evaluationLog = ""
|
|
|
|
for n in range(len(gt)):
|
|
points = gt[n]['points']
|
|
|
|
dontCare = gt[n]['ignore']
|
|
|
|
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
|
continue
|
|
|
|
gtPol = points
|
|
gtPols.append(gtPol)
|
|
gtPolPoints.append(points)
|
|
if dontCare:
|
|
gtDontCarePolsNum.append(len(gtPols) - 1)
|
|
|
|
evaluationLog += "GT polygons: " + str(len(gtPols)) + (" (" + str(len(
|
|
gtDontCarePolsNum)) + " don't care)\n" if len(gtDontCarePolsNum) > 0 else "\n")
|
|
|
|
for n in range(len(pred)):
|
|
points = pred[n]['points']
|
|
if not Polygon(points).is_valid or not Polygon(points).is_simple:
|
|
continue
|
|
|
|
detPol = points
|
|
detPols.append(detPol)
|
|
detPolPoints.append(points)
|
|
if len(gtDontCarePolsNum) > 0:
|
|
for dontCarePol in gtDontCarePolsNum:
|
|
dontCarePol = gtPols[dontCarePol]
|
|
intersected_area = get_intersection(dontCarePol, detPol)
|
|
pdDimensions = Polygon(detPol).area
|
|
precision = 0 if pdDimensions == 0 else intersected_area / pdDimensions
|
|
if (precision > self.area_precision_constraint):
|
|
detDontCarePolsNum.append(len(detPols) - 1)
|
|
break
|
|
|
|
evaluationLog += "DET polygons: " + str(len(detPols)) + (" (" + str(len(
|
|
detDontCarePolsNum)) + " don't care)\n" if len(detDontCarePolsNum) > 0 else "\n")
|
|
|
|
if len(gtPols) > 0 and len(detPols) > 0:
|
|
|
|
outputShape = [len(gtPols), len(detPols)]
|
|
iouMat = np.empty(outputShape)
|
|
gtRectMat = np.zeros(len(gtPols), np.int8)
|
|
detRectMat = np.zeros(len(detPols), np.int8)
|
|
if self.is_output_polygon:
|
|
for gtNum in range(len(gtPols)):
|
|
for detNum in range(len(detPols)):
|
|
pG = gtPols[gtNum]
|
|
pD = detPols[detNum]
|
|
iouMat[gtNum, detNum] = get_intersection_over_union(pD, pG)
|
|
else:
|
|
|
|
|
|
for gtNum in range(len(gtPols)):
|
|
for detNum in range(len(detPols)):
|
|
pG = np.float32(gtPols[gtNum])
|
|
pD = np.float32(detPols[detNum])
|
|
iouMat[gtNum, detNum] = iou_rotate(pD, pG)
|
|
for gtNum in range(len(gtPols)):
|
|
for detNum in range(len(detPols)):
|
|
if gtRectMat[gtNum] == 0 and detRectMat[
|
|
detNum] == 0 and gtNum not in gtDontCarePolsNum and detNum not in detDontCarePolsNum:
|
|
if iouMat[gtNum, detNum] > self.iou_constraint:
|
|
gtRectMat[gtNum] = 1
|
|
detRectMat[detNum] = 1
|
|
detMatched += 1
|
|
pairs.append({'gt': gtNum, 'det': detNum})
|
|
detMatchedNums.append(detNum)
|
|
evaluationLog += "Match GT #" + \
|
|
str(gtNum) + " with Det #" + str(detNum) + "\n"
|
|
|
|
numGtCare = (len(gtPols) - len(gtDontCarePolsNum))
|
|
numDetCare = (len(detPols) - len(detDontCarePolsNum))
|
|
if numGtCare == 0:
|
|
recall = float(1)
|
|
precision = float(0) if numDetCare > 0 else float(1)
|
|
else:
|
|
recall = float(detMatched) / numGtCare
|
|
precision = 0 if numDetCare == 0 else float(
|
|
detMatched) / numDetCare
|
|
|
|
hmean = 0 if (precision + recall) == 0 else 2.0 * \
|
|
precision * recall / (precision + recall)
|
|
|
|
matchedSum += detMatched
|
|
numGlobalCareGt += numGtCare
|
|
numGlobalCareDet += numDetCare
|
|
|
|
perSampleMetrics = {
|
|
'precision': precision,
|
|
'recall': recall,
|
|
'hmean': hmean,
|
|
'pairs': pairs,
|
|
'iouMat': [] if len(detPols) > 100 else iouMat.tolist(),
|
|
'gtPolPoints': gtPolPoints,
|
|
'detPolPoints': detPolPoints,
|
|
'gtCare': numGtCare,
|
|
'detCare': numDetCare,
|
|
'gtDontCare': gtDontCarePolsNum,
|
|
'detDontCare': detDontCarePolsNum,
|
|
'detMatched': detMatched,
|
|
'evaluationLog': evaluationLog
|
|
}
|
|
|
|
return perSampleMetrics
|
|
|
|
def combine_results(self, results):
|
|
numGlobalCareGt = 0
|
|
numGlobalCareDet = 0
|
|
matchedSum = 0
|
|
for result in results:
|
|
numGlobalCareGt += result['gtCare']
|
|
numGlobalCareDet += result['detCare']
|
|
matchedSum += result['detMatched']
|
|
|
|
methodRecall = 0 if numGlobalCareGt == 0 else float(
|
|
matchedSum) / numGlobalCareGt
|
|
methodPrecision = 0 if numGlobalCareDet == 0 else float(
|
|
matchedSum) / numGlobalCareDet
|
|
methodHmean = 0 if methodRecall + methodPrecision == 0 else 2 * \
|
|
methodRecall * methodPrecision / (
|
|
methodRecall + methodPrecision)
|
|
|
|
methodMetrics = {'precision': methodPrecision,
|
|
'recall': methodRecall, 'hmean': methodHmean}
|
|
|
|
return methodMetrics
|
|
|
|
class QuadMetric():
|
|
def __init__(self, is_output_polygon=False):
|
|
self.is_output_polygon = is_output_polygon
|
|
self.evaluator = DetectionIoUEvaluator(is_output_polygon=is_output_polygon)
|
|
|
|
def measure(self, batch, output, box_thresh=0.6):
|
|
'''
|
|
batch: (image, polygons, ignore_tags
|
|
batch: a dict produced by dataloaders.
|
|
image: tensor of shape (N, C, H, W).
|
|
polygons: tensor of shape (N, K, 4, 2), the polygons of objective regions.
|
|
ignore_tags: tensor of shape (N, K), indicates whether a region is ignorable or not.
|
|
shape: the original shape of images.
|
|
filename: the original filenames of images.
|
|
output: (polygons, ...)
|
|
'''
|
|
results = []
|
|
gt_polyons_batch = batch['text_polys']
|
|
ignore_tags_batch = batch['ignore_tags']
|
|
pred_polygons_batch = np.array(output[0])
|
|
pred_scores_batch = np.array(output[1])
|
|
for polygons, pred_polygons, pred_scores, ignore_tags in zip(gt_polyons_batch, pred_polygons_batch, pred_scores_batch, ignore_tags_batch):
|
|
gt = [dict(points=np.int64(polygons[i]), ignore=ignore_tags[i]) for i in range(len(polygons))]
|
|
if self.is_output_polygon:
|
|
pred = [dict(points=pred_polygons[i]) for i in range(len(pred_polygons))]
|
|
else:
|
|
pred = []
|
|
|
|
for i in range(pred_polygons.shape[0]):
|
|
if pred_scores[i] >= box_thresh:
|
|
|
|
pred.append(dict(points=pred_polygons[i, :, :].astype(np.int32)))
|
|
|
|
results.append(self.evaluator.evaluate_image(gt, pred))
|
|
return results
|
|
|
|
def validate_measure(self, batch, output, box_thresh=0.6):
|
|
return self.measure(batch, output, box_thresh)
|
|
|
|
def evaluate_measure(self, batch, output):
|
|
return self.measure(batch, output), np.linspace(0, batch['image'].shape[0]).tolist()
|
|
|
|
def gather_measure(self, raw_metrics):
|
|
raw_metrics = [image_metrics
|
|
for batch_metrics in raw_metrics
|
|
for image_metrics in batch_metrics]
|
|
|
|
result = self.evaluator.combine_results(raw_metrics)
|
|
|
|
precision = AverageMeter()
|
|
recall = AverageMeter()
|
|
fmeasure = AverageMeter()
|
|
|
|
precision.update(result['precision'], n=len(raw_metrics))
|
|
recall.update(result['recall'], n=len(raw_metrics))
|
|
fmeasure_score = 2 * precision.val * recall.val / (precision.val + recall.val + 1e-8)
|
|
fmeasure.update(fmeasure_score)
|
|
|
|
return {
|
|
'precision': precision,
|
|
'recall': recall,
|
|
'fmeasure': fmeasure
|
|
}
|
|
|
|
def shrink_polygon_py(polygon, shrink_ratio):
|
|
"""
|
|
对框进行缩放，返回去的比例为1/shrink_ratio 即可
|
|
"""
|
|
cx = polygon[:, 0].mean()
|
|
cy = polygon[:, 1].mean()
|
|
polygon[:, 0] = cx + (polygon[:, 0] - cx) * shrink_ratio
|
|
polygon[:, 1] = cy + (polygon[:, 1] - cy) * shrink_ratio
|
|
return polygon
|
|
|
|
|
|
def shrink_polygon_pyclipper(polygon, shrink_ratio):
|
|
from shapely.geometry import Polygon
|
|
import pyclipper
|
|
polygon_shape = Polygon(polygon)
|
|
distance = polygon_shape.area * (1 - np.power(shrink_ratio, 2)) / polygon_shape.length
|
|
subject = [tuple(l) for l in polygon]
|
|
padding = pyclipper.PyclipperOffset()
|
|
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
|
shrunk = padding.Execute(-distance)
|
|
if shrunk == []:
|
|
shrunk = np.array(shrunk)
|
|
else:
|
|
shrunk = np.array(shrunk[0]).reshape(-1, 2)
|
|
return shrunk
|
|
|
|
class MakeShrinkMap():
|
|
r'''
|
|
Making binary mask from detection data with ICDAR format.
|
|
Typically following the process of class `MakeICDARData`.
|
|
'''
|
|
|
|
def __init__(self, min_text_size=4, shrink_ratio=0.4, shrink_type='pyclipper'):
|
|
shrink_func_dict = {'py': shrink_polygon_py, 'pyclipper': shrink_polygon_pyclipper}
|
|
self.shrink_func = shrink_func_dict[shrink_type]
|
|
self.min_text_size = min_text_size
|
|
self.shrink_ratio = shrink_ratio
|
|
|
|
def __call__(self, data: dict) -> dict:
|
|
"""
|
|
从scales中随机选择一个尺度，对图片和文本框进行缩放
|
|
:param data: {'imgs':,'text_polys':,'texts':,'ignore_tags':}
|
|
:return:
|
|
"""
|
|
image = data['imgs']
|
|
text_polys = data['text_polys']
|
|
ignore_tags = data['ignore_tags']
|
|
|
|
h, w = image.shape[:2]
|
|
text_polys, ignore_tags = self.validate_polygons(text_polys, ignore_tags, h, w)
|
|
gt = np.zeros((h, w), dtype=np.float32)
|
|
mask = np.ones((h, w), dtype=np.float32)
|
|
for i in range(len(text_polys)):
|
|
polygon = text_polys[i]
|
|
height = max(polygon[:, 1]) - min(polygon[:, 1])
|
|
width = max(polygon[:, 0]) - min(polygon[:, 0])
|
|
if ignore_tags[i] or min(height, width) < self.min_text_size:
|
|
cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
|
|
ignore_tags[i] = True
|
|
else:
|
|
shrunk = self.shrink_func(polygon, self.shrink_ratio)
|
|
if shrunk.size == 0:
|
|
cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
|
|
ignore_tags[i] = True
|
|
continue
|
|
cv2.fillPoly(gt, [shrunk.astype(np.int32)], 1)
|
|
|
|
data['shrink_map'] = gt
|
|
data['shrink_mask'] = mask
|
|
return data
|
|
|
|
def validate_polygons(self, polygons, ignore_tags, h, w):
|
|
'''
|
|
polygons (numpy.array, required): of shape (num_instances, num_points, 2)
|
|
'''
|
|
if len(polygons) == 0:
|
|
return polygons, ignore_tags
|
|
assert len(polygons) == len(ignore_tags)
|
|
for polygon in polygons:
|
|
polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1)
|
|
polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1)
|
|
|
|
for i in range(len(polygons)):
|
|
area = self.polygon_area(polygons[i])
|
|
if abs(area) < 1:
|
|
ignore_tags[i] = True
|
|
if area > 0:
|
|
polygons[i] = polygons[i][::-1, :]
|
|
return polygons, ignore_tags
|
|
|
|
def polygon_area(self, polygon):
|
|
return cv2.contourArea(polygon)
|
|
|
|
|
|
class MakeBorderMap():
|
|
def __init__(self, shrink_ratio=0.4, thresh_min=0.3, thresh_max=0.7):
|
|
self.shrink_ratio = shrink_ratio
|
|
self.thresh_min = thresh_min
|
|
self.thresh_max = thresh_max
|
|
|
|
def __call__(self, data: dict) -> dict:
|
|
"""
|
|
从scales中随机选择一个尺度，对图片和文本框进行缩放
|
|
:param data: {'imgs':,'text_polys':,'texts':,'ignore_tags':}
|
|
:return:
|
|
"""
|
|
im = data['imgs']
|
|
text_polys = data['text_polys']
|
|
ignore_tags = data['ignore_tags']
|
|
|
|
canvas = np.zeros(im.shape[:2], dtype=np.float32)
|
|
mask = np.zeros(im.shape[:2], dtype=np.float32)
|
|
|
|
for i in range(len(text_polys)):
|
|
if ignore_tags[i]:
|
|
continue
|
|
self.draw_border_map(text_polys[i], canvas, mask=mask)
|
|
canvas = canvas * (self.thresh_max - self.thresh_min) + self.thresh_min
|
|
|
|
data['threshold_map'] = canvas
|
|
data['threshold_mask'] = mask
|
|
return data
|
|
|
|
def draw_border_map(self, polygon, canvas, mask):
|
|
polygon = np.array(polygon)
|
|
assert polygon.ndim == 2
|
|
assert polygon.shape[1] == 2
|
|
|
|
polygon_shape = Polygon(polygon)
|
|
if polygon_shape.area <= 0:
|
|
return
|
|
distance = polygon_shape.area * (1 - np.power(self.shrink_ratio, 2)) / polygon_shape.length
|
|
subject = [tuple(l) for l in polygon]
|
|
padding = pyclipper.PyclipperOffset()
|
|
padding.AddPath(subject, pyclipper.JT_ROUND,
|
|
pyclipper.ET_CLOSEDPOLYGON)
|
|
|
|
padded_polygon = np.array(padding.Execute(distance)[0])
|
|
cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)
|
|
|
|
xmin = padded_polygon[:, 0].min()
|
|
xmax = padded_polygon[:, 0].max()
|
|
ymin = padded_polygon[:, 1].min()
|
|
ymax = padded_polygon[:, 1].max()
|
|
width = xmax - xmin + 1
|
|
height = ymax - ymin + 1
|
|
|
|
polygon[:, 0] = polygon[:, 0] - xmin
|
|
polygon[:, 1] = polygon[:, 1] - ymin
|
|
|
|
xs = np.broadcast_to(
|
|
np.linspace(0, width - 1, num=width).reshape(1, width), (height, width))
|
|
ys = np.broadcast_to(
|
|
np.linspace(0, height - 1, num=height).reshape(height, 1), (height, width))
|
|
|
|
distance_map = np.zeros(
|
|
(polygon.shape[0], height, width), dtype=np.float32)
|
|
for i in range(polygon.shape[0]):
|
|
j = (i + 1) % polygon.shape[0]
|
|
absolute_distance = self.distance(xs, ys, polygon[i], polygon[j])
|
|
distance_map[i] = np.clip(absolute_distance / distance, 0, 1)
|
|
distance_map = distance_map.min(axis=0)
|
|
|
|
xmin_valid = min(max(0, xmin), canvas.shape[1] - 1)
|
|
xmax_valid = min(max(0, xmax), canvas.shape[1] - 1)
|
|
ymin_valid = min(max(0, ymin), canvas.shape[0] - 1)
|
|
ymax_valid = min(max(0, ymax), canvas.shape[0] - 1)
|
|
canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1] = np.fmax(
|
|
1 - distance_map[
|
|
ymin_valid - ymin:ymax_valid - ymax + height,
|
|
xmin_valid - xmin:xmax_valid - xmax + width],
|
|
canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1])
|
|
|
|
def distance(self, xs, ys, point_1, point_2):
|
|
'''
|
|
compute the distance from point to a line
|
|
ys: coordinates in the first axis
|
|
xs: coordinates in the second axis
|
|
point_1, point_2: (x, y), the end of the line
|
|
'''
|
|
height, width = xs.shape[:2]
|
|
square_distance_1 = np.square(xs - point_1[0]) + np.square(ys - point_1[1])
|
|
square_distance_2 = np.square(xs - point_2[0]) + np.square(ys - point_2[1])
|
|
square_distance = np.square(point_1[0] - point_2[0]) + np.square(point_1[1] - point_2[1])
|
|
|
|
cosin = (square_distance - square_distance_1 - square_distance_2) / (2 * np.sqrt(square_distance_1 * square_distance_2))
|
|
square_sin = 1 - np.square(cosin)
|
|
square_sin = np.nan_to_num(square_sin)
|
|
|
|
result = np.sqrt(square_distance_1 * square_distance_2 * square_sin / square_distance)
|
|
result[cosin < 0] = np.sqrt(np.fmin(square_distance_1, square_distance_2))[cosin < 0]
|
|
return result
|
|
|
|
def extend_line(self, point_1, point_2, result):
|
|
ex_point_1 = (int(round(point_1[0] + (point_1[0] - point_2[0]) * (1 + self.shrink_ratio))),
|
|
int(round(point_1[1] + (point_1[1] - point_2[1]) * (1 + self.shrink_ratio))))
|
|
cv2.line(result, tuple(ex_point_1), tuple(point_1), 4096.0, 1, lineType=cv2.LINE_AA, shift=0)
|
|
ex_point_2 = (int(round(point_2[0] + (point_2[0] - point_1[0]) * (1 + self.shrink_ratio))),
|
|
int(round(point_2[1] + (point_2[1] - point_1[1]) * (1 + self.shrink_ratio))))
|
|
cv2.line(result, tuple(ex_point_2), tuple(point_2), 4096.0, 1, lineType=cv2.LINE_AA, shift=0)
|
|
return ex_point_1, ex_point_2 |
