db_utils.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701
  1. import cv2
  2. import numpy as np
  3. import pyclipper
  4. from shapely.geometry import Polygon
  5. from collections import namedtuple
  6. import torch
  7. import warnings
  8. warnings.filterwarnings('ignore')
  9. def iou_rotate(box_a, box_b, method='union'):
  10. rect_a = cv2.minAreaRect(box_a)
  11. rect_b = cv2.minAreaRect(box_b)
  12. r1 = cv2.rotatedRectangleIntersection(rect_a, rect_b)
  13. if r1[0] == 0:
  14. return 0
  15. else:
  16. inter_area = cv2.contourArea(r1[1])
  17. area_a = cv2.contourArea(box_a)
  18. area_b = cv2.contourArea(box_b)
  19. union_area = area_a + area_b - inter_area
  20. if union_area == 0 or inter_area == 0:
  21. return 0
  22. if method == 'union':
  23. iou = inter_area / union_area
  24. elif method == 'intersection':
  25. iou = inter_area / min(area_a, area_b)
  26. else:
  27. raise NotImplementedError
  28. return iou
  29. class SegDetectorRepresenter():
  30. def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000, unclip_ratio=1.5):
  31. self.min_size = 3
  32. self.thresh = thresh
  33. self.box_thresh = box_thresh
  34. self.max_candidates = max_candidates
  35. self.unclip_ratio = unclip_ratio
  36. def __call__(self, batch, pred, is_output_polygon=False):
  37. '''
  38. batch: (image, polygons, ignore_tags
  39. batch: a dict produced by dataloaders.
  40. image: tensor of shape (N, C, H, W).
  41. polygons: tensor of shape (N, K, 4, 2), the polygons of objective regions.
  42. ignore_tags: tensor of shape (N, K), indicates whether a region is ignorable or not.
  43. shape: the original shape of images.
  44. filename: the original filenames of images.
  45. pred:
  46. binary: text region segmentation map, with shape (N, H, W)
  47. thresh: [if exists] thresh hold prediction with shape (N, H, W)
  48. thresh_binary: [if exists] binarized with threshold, (N, H, W)
  49. '''
  50. pred = pred[:, 0, :, :]
  51. segmentation = self.binarize(pred)
  52. boxes_batch = []
  53. scores_batch = []
  54. # print(pred.size())
  55. batch_size = pred.size(0) if isinstance(pred, torch.Tensor) else pred.shape[0]
  56. for batch_index in range(batch_size):
  57. # height, width = batch['shape'][batch_index]
  58. height, width = pred.shape[1], pred.shape[2]
  59. if is_output_polygon:
  60. boxes, scores = self.polygons_from_bitmap(pred[batch_index], segmentation[batch_index], width, height)
  61. else:
  62. boxes, scores = self.boxes_from_bitmap(pred[batch_index], segmentation[batch_index], width, height)
  63. boxes_batch.append(boxes)
  64. scores_batch.append(scores)
  65. return boxes_batch, scores_batch
  66. def binarize(self, pred):
  67. return pred > self.thresh
  68. def polygons_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
  69. '''
  70. _bitmap: single map with shape (H, W),
  71. whose values are binarized as {0, 1}
  72. '''
  73. assert len(_bitmap.shape) == 2
  74. bitmap = _bitmap.cpu().numpy() # The first channel
  75. pred = pred.cpu().detach().numpy()
  76. height, width = bitmap.shape
  77. boxes = []
  78. scores = []
  79. contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
  80. for contour in contours[:self.max_candidates]:
  81. epsilon = 0.005 * cv2.arcLength(contour, True)
  82. approx = cv2.approxPolyDP(contour, epsilon, True)
  83. points = approx.reshape((-1, 2))
  84. if points.shape[0] < 4:
  85. continue
  86. # _, sside = self.get_mini_boxes(contour)
  87. # if sside < self.min_size:
  88. # continue
  89. score = self.box_score_fast(pred, contour.squeeze(1))
  90. if self.box_thresh > score:
  91. continue
  92. if points.shape[0] > 2:
  93. box = self.unclip(points, unclip_ratio=self.unclip_ratio)
  94. if len(box) > 1:
  95. continue
  96. else:
  97. continue
  98. box = box.reshape(-1, 2)
  99. _, sside = self.get_mini_boxes(box.reshape((-1, 1, 2)))
  100. if sside < self.min_size + 2:
  101. continue
  102. if not isinstance(dest_width, int):
  103. dest_width = dest_width.item()
  104. dest_height = dest_height.item()
  105. box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
  106. box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
  107. boxes.append(box)
  108. scores.append(score)
  109. return boxes, scores
  110. def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
  111. '''
  112. _bitmap: single map with shape (H, W),
  113. whose values are binarized as {0, 1}
  114. '''
  115. assert len(_bitmap.shape) == 2
  116. if isinstance(pred, torch.Tensor):
  117. bitmap = _bitmap.cpu().numpy() # The first channel
  118. pred = pred.cpu().detach().numpy()
  119. else:
  120. bitmap = _bitmap
  121. height, width = bitmap.shape
  122. contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
  123. num_contours = min(len(contours), self.max_candidates)
  124. boxes = np.zeros((num_contours, 4, 2), dtype=np.int16)
  125. scores = np.zeros((num_contours,), dtype=np.float32)
  126. for index in range(num_contours):
  127. contour = contours[index].squeeze(1)
  128. points, sside = self.get_mini_boxes(contour)
  129. # if sside < self.min_size:
  130. # continue
  131. if sside < 2:
  132. continue
  133. points = np.array(points)
  134. score = self.box_score_fast(pred, contour)
  135. # if self.box_thresh > score:
  136. # continue
  137. box = self.unclip(points, unclip_ratio=self.unclip_ratio).reshape(-1, 1, 2)
  138. box, sside = self.get_mini_boxes(box)
  139. # if sside < 5:
  140. # continue
  141. box = np.array(box)
  142. if not isinstance(dest_width, int):
  143. dest_width = dest_width.item()
  144. dest_height = dest_height.item()
  145. box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
  146. box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
  147. boxes[index, :, :] = box.astype(np.int16)
  148. scores[index] = score
  149. return boxes, scores
  150. def unclip(self, box, unclip_ratio=1.5):
  151. poly = Polygon(box)
  152. distance = poly.area * unclip_ratio / poly.length
  153. offset = pyclipper.PyclipperOffset()
  154. offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
  155. expanded = np.array(offset.Execute(distance))
  156. return expanded
  157. def get_mini_boxes(self, contour):
  158. bounding_box = cv2.minAreaRect(contour)
  159. points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
  160. index_1, index_2, index_3, index_4 = 0, 1, 2, 3
  161. if points[1][1] > points[0][1]:
  162. index_1 = 0
  163. index_4 = 1
  164. else:
  165. index_1 = 1
  166. index_4 = 0
  167. if points[3][1] > points[2][1]:
  168. index_2 = 2
  169. index_3 = 3
  170. else:
  171. index_2 = 3
  172. index_3 = 2
  173. box = [points[index_1], points[index_2], points[index_3], points[index_4]]
  174. return box, min(bounding_box[1])
  175. def box_score_fast(self, bitmap, _box):
  176. h, w = bitmap.shape[:2]
  177. box = _box.copy()
  178. xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int64), 0, w - 1)
  179. xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int64), 0, w - 1)
  180. ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int64), 0, h - 1)
  181. ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int64), 0, h - 1)
  182. mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
  183. box[:, 0] = box[:, 0] - xmin
  184. box[:, 1] = box[:, 1] - ymin
  185. cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
  186. if bitmap.dtype == np.float16:
  187. bitmap = bitmap.astype(np.float32)
  188. return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
  189. class AverageMeter(object):
  190. """Computes and stores the average and current value"""
  191. def __init__(self):
  192. self.reset()
  193. def reset(self):
  194. self.val = 0
  195. self.avg = 0
  196. self.sum = 0
  197. self.count = 0
  198. def update(self, val, n=1):
  199. self.val = val
  200. self.sum += val * n
  201. self.count += n
  202. self.avg = self.sum / self.count
  203. return self
  204. class DetectionIoUEvaluator(object):
  205. def __init__(self, is_output_polygon=False, iou_constraint=0.5, area_precision_constraint=0.5):
  206. self.is_output_polygon = is_output_polygon
  207. self.iou_constraint = iou_constraint
  208. self.area_precision_constraint = area_precision_constraint
  209. def evaluate_image(self, gt, pred):
  210. def get_union(pD, pG):
  211. return Polygon(pD).union(Polygon(pG)).area
  212. def get_intersection_over_union(pD, pG):
  213. return get_intersection(pD, pG) / get_union(pD, pG)
  214. def get_intersection(pD, pG):
  215. return Polygon(pD).intersection(Polygon(pG)).area
  216. def compute_ap(confList, matchList, numGtCare):
  217. correct = 0
  218. AP = 0
  219. if len(confList) > 0:
  220. confList = np.array(confList)
  221. matchList = np.array(matchList)
  222. sorted_ind = np.argsort(-confList)
  223. confList = confList[sorted_ind]
  224. matchList = matchList[sorted_ind]
  225. for n in range(len(confList)):
  226. match = matchList[n]
  227. if match:
  228. correct += 1
  229. AP += float(correct) / (n + 1)
  230. if numGtCare > 0:
  231. AP /= numGtCare
  232. return AP
  233. perSampleMetrics = {}
  234. matchedSum = 0
  235. Rectangle = namedtuple('Rectangle', 'xmin ymin xmax ymax')
  236. numGlobalCareGt = 0
  237. numGlobalCareDet = 0
  238. arrGlobalConfidences = []
  239. arrGlobalMatches = []
  240. recall = 0
  241. precision = 0
  242. hmean = 0
  243. detMatched = 0
  244. iouMat = np.empty([1, 1])
  245. gtPols = []
  246. detPols = []
  247. gtPolPoints = []
  248. detPolPoints = []
  249. # Array of Ground Truth Polygons' keys marked as don't Care
  250. gtDontCarePolsNum = []
  251. # Array of Detected Polygons' matched with a don't Care GT
  252. detDontCarePolsNum = []
  253. pairs = []
  254. detMatchedNums = []
  255. arrSampleConfidences = []
  256. arrSampleMatch = []
  257. evaluationLog = ""
  258. for n in range(len(gt)):
  259. points = gt[n]['points']
  260. # transcription = gt[n]['text']
  261. dontCare = gt[n]['ignore']
  262. if not Polygon(points).is_valid or not Polygon(points).is_simple:
  263. continue
  264. gtPol = points
  265. gtPols.append(gtPol)
  266. gtPolPoints.append(points)
  267. if dontCare:
  268. gtDontCarePolsNum.append(len(gtPols) - 1)
  269. evaluationLog += "GT polygons: " + str(len(gtPols)) + (" (" + str(len(
  270. gtDontCarePolsNum)) + " don't care)\n" if len(gtDontCarePolsNum) > 0 else "\n")
  271. for n in range(len(pred)):
  272. points = pred[n]['points']
  273. if not Polygon(points).is_valid or not Polygon(points).is_simple:
  274. continue
  275. detPol = points
  276. detPols.append(detPol)
  277. detPolPoints.append(points)
  278. if len(gtDontCarePolsNum) > 0:
  279. for dontCarePol in gtDontCarePolsNum:
  280. dontCarePol = gtPols[dontCarePol]
  281. intersected_area = get_intersection(dontCarePol, detPol)
  282. pdDimensions = Polygon(detPol).area
  283. precision = 0 if pdDimensions == 0 else intersected_area / pdDimensions
  284. if (precision > self.area_precision_constraint):
  285. detDontCarePolsNum.append(len(detPols) - 1)
  286. break
  287. evaluationLog += "DET polygons: " + str(len(detPols)) + (" (" + str(len(
  288. detDontCarePolsNum)) + " don't care)\n" if len(detDontCarePolsNum) > 0 else "\n")
  289. if len(gtPols) > 0 and len(detPols) > 0:
  290. # Calculate IoU and precision matrixs
  291. outputShape = [len(gtPols), len(detPols)]
  292. iouMat = np.empty(outputShape)
  293. gtRectMat = np.zeros(len(gtPols), np.int8)
  294. detRectMat = np.zeros(len(detPols), np.int8)
  295. if self.is_output_polygon:
  296. for gtNum in range(len(gtPols)):
  297. for detNum in range(len(detPols)):
  298. pG = gtPols[gtNum]
  299. pD = detPols[detNum]
  300. iouMat[gtNum, detNum] = get_intersection_over_union(pD, pG)
  301. else:
  302. # gtPols = np.float32(gtPols)
  303. # detPols = np.float32(detPols)
  304. for gtNum in range(len(gtPols)):
  305. for detNum in range(len(detPols)):
  306. pG = np.float32(gtPols[gtNum])
  307. pD = np.float32(detPols[detNum])
  308. iouMat[gtNum, detNum] = iou_rotate(pD, pG)
  309. for gtNum in range(len(gtPols)):
  310. for detNum in range(len(detPols)):
  311. if gtRectMat[gtNum] == 0 and detRectMat[
  312. detNum] == 0 and gtNum not in gtDontCarePolsNum and detNum not in detDontCarePolsNum:
  313. if iouMat[gtNum, detNum] > self.iou_constraint:
  314. gtRectMat[gtNum] = 1
  315. detRectMat[detNum] = 1
  316. detMatched += 1
  317. pairs.append({'gt': gtNum, 'det': detNum})
  318. detMatchedNums.append(detNum)
  319. evaluationLog += "Match GT #" + \
  320. str(gtNum) + " with Det #" + str(detNum) + "\n"
  321. numGtCare = (len(gtPols) - len(gtDontCarePolsNum))
  322. numDetCare = (len(detPols) - len(detDontCarePolsNum))
  323. if numGtCare == 0:
  324. recall = float(1)
  325. precision = float(0) if numDetCare > 0 else float(1)
  326. else:
  327. recall = float(detMatched) / numGtCare
  328. precision = 0 if numDetCare == 0 else float(
  329. detMatched) / numDetCare
  330. hmean = 0 if (precision + recall) == 0 else 2.0 * \
  331. precision * recall / (precision + recall)
  332. matchedSum += detMatched
  333. numGlobalCareGt += numGtCare
  334. numGlobalCareDet += numDetCare
  335. perSampleMetrics = {
  336. 'precision': precision,
  337. 'recall': recall,
  338. 'hmean': hmean,
  339. 'pairs': pairs,
  340. 'iouMat': [] if len(detPols) > 100 else iouMat.tolist(),
  341. 'gtPolPoints': gtPolPoints,
  342. 'detPolPoints': detPolPoints,
  343. 'gtCare': numGtCare,
  344. 'detCare': numDetCare,
  345. 'gtDontCare': gtDontCarePolsNum,
  346. 'detDontCare': detDontCarePolsNum,
  347. 'detMatched': detMatched,
  348. 'evaluationLog': evaluationLog
  349. }
  350. return perSampleMetrics
  351. def combine_results(self, results):
  352. numGlobalCareGt = 0
  353. numGlobalCareDet = 0
  354. matchedSum = 0
  355. for result in results:
  356. numGlobalCareGt += result['gtCare']
  357. numGlobalCareDet += result['detCare']
  358. matchedSum += result['detMatched']
  359. methodRecall = 0 if numGlobalCareGt == 0 else float(
  360. matchedSum) / numGlobalCareGt
  361. methodPrecision = 0 if numGlobalCareDet == 0 else float(
  362. matchedSum) / numGlobalCareDet
  363. methodHmean = 0 if methodRecall + methodPrecision == 0 else 2 * \
  364. methodRecall * methodPrecision / (
  365. methodRecall + methodPrecision)
  366. methodMetrics = {'precision': methodPrecision,
  367. 'recall': methodRecall, 'hmean': methodHmean}
  368. return methodMetrics
  369. class QuadMetric():
  370. def __init__(self, is_output_polygon=False):
  371. self.is_output_polygon = is_output_polygon
  372. self.evaluator = DetectionIoUEvaluator(is_output_polygon=is_output_polygon)
  373. def measure(self, batch, output, box_thresh=0.6):
  374. '''
  375. batch: (image, polygons, ignore_tags
  376. batch: a dict produced by dataloaders.
  377. image: tensor of shape (N, C, H, W).
  378. polygons: tensor of shape (N, K, 4, 2), the polygons of objective regions.
  379. ignore_tags: tensor of shape (N, K), indicates whether a region is ignorable or not.
  380. shape: the original shape of images.
  381. filename: the original filenames of images.
  382. output: (polygons, ...)
  383. '''
  384. results = []
  385. gt_polyons_batch = batch['text_polys']
  386. ignore_tags_batch = batch['ignore_tags']
  387. pred_polygons_batch = np.array(output[0])
  388. pred_scores_batch = np.array(output[1])
  389. for polygons, pred_polygons, pred_scores, ignore_tags in zip(gt_polyons_batch, pred_polygons_batch, pred_scores_batch, ignore_tags_batch):
  390. gt = [dict(points=np.int64(polygons[i]), ignore=ignore_tags[i]) for i in range(len(polygons))]
  391. if self.is_output_polygon:
  392. pred = [dict(points=pred_polygons[i]) for i in range(len(pred_polygons))]
  393. else:
  394. pred = []
  395. # print(pred_polygons.shape)
  396. for i in range(pred_polygons.shape[0]):
  397. if pred_scores[i] >= box_thresh:
  398. # print(pred_polygons[i,:,:].tolist())
  399. pred.append(dict(points=pred_polygons[i, :, :].astype(np.int64)))
  400. # pred = [dict(points=pred_polygons[i,:,:].tolist()) if pred_scores[i] >= box_thresh for i in range(pred_polygons.shape[0])]
  401. results.append(self.evaluator.evaluate_image(gt, pred))
  402. return results
  403. def validate_measure(self, batch, output, box_thresh=0.6):
  404. return self.measure(batch, output, box_thresh)
  405. def evaluate_measure(self, batch, output):
  406. return self.measure(batch, output), np.linspace(0, batch['image'].shape[0]).tolist()
  407. def gather_measure(self, raw_metrics):
  408. raw_metrics = [image_metrics
  409. for batch_metrics in raw_metrics
  410. for image_metrics in batch_metrics]
  411. result = self.evaluator.combine_results(raw_metrics)
  412. precision = AverageMeter()
  413. recall = AverageMeter()
  414. fmeasure = AverageMeter()
  415. precision.update(result['precision'], n=len(raw_metrics))
  416. recall.update(result['recall'], n=len(raw_metrics))
  417. fmeasure_score = 2 * precision.val * recall.val / (precision.val + recall.val + 1e-8)
  418. fmeasure.update(fmeasure_score)
  419. return {
  420. 'precision': precision,
  421. 'recall': recall,
  422. 'fmeasure': fmeasure
  423. }
  424. def shrink_polygon_py(polygon, shrink_ratio):
  425. """
  426. 对框进行缩放,返回去的比例为1/shrink_ratio 即可
  427. """
  428. cx = polygon[:, 0].mean()
  429. cy = polygon[:, 1].mean()
  430. polygon[:, 0] = cx + (polygon[:, 0] - cx) * shrink_ratio
  431. polygon[:, 1] = cy + (polygon[:, 1] - cy) * shrink_ratio
  432. return polygon
  433. def shrink_polygon_pyclipper(polygon, shrink_ratio):
  434. from shapely.geometry import Polygon
  435. import pyclipper
  436. polygon_shape = Polygon(polygon)
  437. distance = polygon_shape.area * (1 - np.power(shrink_ratio, 2)) / polygon_shape.length
  438. subject = [tuple(l) for l in polygon]
  439. padding = pyclipper.PyclipperOffset()
  440. padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
  441. shrunk = padding.Execute(-distance)
  442. if shrunk == []:
  443. shrunk = np.array(shrunk)
  444. else:
  445. shrunk = np.array(shrunk[0]).reshape(-1, 2)
  446. return shrunk
  447. class MakeShrinkMap():
  448. r'''
  449. Making binary mask from detection data with ICDAR format.
  450. Typically following the process of class `MakeICDARData`.
  451. '''
  452. def __init__(self, min_text_size=4, shrink_ratio=0.4, shrink_type='pyclipper'):
  453. shrink_func_dict = {'py': shrink_polygon_py, 'pyclipper': shrink_polygon_pyclipper}
  454. self.shrink_func = shrink_func_dict[shrink_type]
  455. self.min_text_size = min_text_size
  456. self.shrink_ratio = shrink_ratio
  457. def __call__(self, data: dict) -> dict:
  458. """
  459. 从scales中随机选择一个尺度,对图片和文本框进行缩放
  460. :param data: {'imgs':,'text_polys':,'texts':,'ignore_tags':}
  461. :return:
  462. """
  463. image = data['imgs']
  464. text_polys = data['text_polys']
  465. ignore_tags = data['ignore_tags']
  466. h, w = image.shape[:2]
  467. text_polys, ignore_tags = self.validate_polygons(text_polys, ignore_tags, h, w)
  468. gt = np.zeros((h, w), dtype=np.float32)
  469. mask = np.ones((h, w), dtype=np.float32)
  470. for i in range(len(text_polys)):
  471. polygon = text_polys[i]
  472. height = max(polygon[:, 1]) - min(polygon[:, 1])
  473. width = max(polygon[:, 0]) - min(polygon[:, 0])
  474. if ignore_tags[i] or min(height, width) < self.min_text_size:
  475. cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
  476. ignore_tags[i] = True
  477. else:
  478. shrunk = self.shrink_func(polygon, self.shrink_ratio)
  479. if shrunk.size == 0:
  480. cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
  481. ignore_tags[i] = True
  482. continue
  483. cv2.fillPoly(gt, [shrunk.astype(np.int32)], 1)
  484. data['shrink_map'] = gt
  485. data['shrink_mask'] = mask
  486. return data
  487. def validate_polygons(self, polygons, ignore_tags, h, w):
  488. '''
  489. polygons (numpy.array, required): of shape (num_instances, num_points, 2)
  490. '''
  491. if len(polygons) == 0:
  492. return polygons, ignore_tags
  493. assert len(polygons) == len(ignore_tags)
  494. for polygon in polygons:
  495. polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1)
  496. polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1)
  497. for i in range(len(polygons)):
  498. area = self.polygon_area(polygons[i])
  499. if abs(area) < 1:
  500. ignore_tags[i] = True
  501. if area > 0:
  502. polygons[i] = polygons[i][::-1, :]
  503. return polygons, ignore_tags
  504. def polygon_area(self, polygon):
  505. return cv2.contourArea(polygon)
  506. class MakeBorderMap():
  507. def __init__(self, shrink_ratio=0.4, thresh_min=0.3, thresh_max=0.7):
  508. self.shrink_ratio = shrink_ratio
  509. self.thresh_min = thresh_min
  510. self.thresh_max = thresh_max
  511. def __call__(self, data: dict) -> dict:
  512. """
  513. 从scales中随机选择一个尺度,对图片和文本框进行缩放
  514. :param data: {'imgs':,'text_polys':,'texts':,'ignore_tags':}
  515. :return:
  516. """
  517. im = data['imgs']
  518. text_polys = data['text_polys']
  519. ignore_tags = data['ignore_tags']
  520. canvas = np.zeros(im.shape[:2], dtype=np.float32)
  521. mask = np.zeros(im.shape[:2], dtype=np.float32)
  522. for i in range(len(text_polys)):
  523. if ignore_tags[i]:
  524. continue
  525. self.draw_border_map(text_polys[i], canvas, mask=mask)
  526. canvas = canvas * (self.thresh_max - self.thresh_min) + self.thresh_min
  527. data['threshold_map'] = canvas
  528. data['threshold_mask'] = mask
  529. return data
  530. def draw_border_map(self, polygon, canvas, mask):
  531. polygon = np.array(polygon)
  532. assert polygon.ndim == 2
  533. assert polygon.shape[1] == 2
  534. polygon_shape = Polygon(polygon)
  535. if polygon_shape.area <= 0:
  536. return
  537. distance = polygon_shape.area * (1 - np.power(self.shrink_ratio, 2)) / polygon_shape.length
  538. subject = [tuple(l) for l in polygon]
  539. padding = pyclipper.PyclipperOffset()
  540. padding.AddPath(subject, pyclipper.JT_ROUND,
  541. pyclipper.ET_CLOSEDPOLYGON)
  542. padded_polygon = np.array(padding.Execute(distance)[0])
  543. cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)
  544. xmin = padded_polygon[:, 0].min()
  545. xmax = padded_polygon[:, 0].max()
  546. ymin = padded_polygon[:, 1].min()
  547. ymax = padded_polygon[:, 1].max()
  548. width = xmax - xmin + 1
  549. height = ymax - ymin + 1
  550. polygon[:, 0] = polygon[:, 0] - xmin
  551. polygon[:, 1] = polygon[:, 1] - ymin
  552. xs = np.broadcast_to(
  553. np.linspace(0, width - 1, num=width).reshape(1, width), (height, width))
  554. ys = np.broadcast_to(
  555. np.linspace(0, height - 1, num=height).reshape(height, 1), (height, width))
  556. distance_map = np.zeros(
  557. (polygon.shape[0], height, width), dtype=np.float32)
  558. for i in range(polygon.shape[0]):
  559. j = (i + 1) % polygon.shape[0]
  560. absolute_distance = self.distance(xs, ys, polygon[i], polygon[j])
  561. distance_map[i] = np.clip(absolute_distance / distance, 0, 1)
  562. distance_map = distance_map.min(axis=0)
  563. xmin_valid = min(max(0, xmin), canvas.shape[1] - 1)
  564. xmax_valid = min(max(0, xmax), canvas.shape[1] - 1)
  565. ymin_valid = min(max(0, ymin), canvas.shape[0] - 1)
  566. ymax_valid = min(max(0, ymax), canvas.shape[0] - 1)
  567. canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1] = np.fmax(
  568. 1 - distance_map[
  569. ymin_valid - ymin:ymax_valid - ymax + height,
  570. xmin_valid - xmin:xmax_valid - xmax + width],
  571. canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1])
  572. def distance(self, xs, ys, point_1, point_2):
  573. '''
  574. compute the distance from point to a line
  575. ys: coordinates in the first axis
  576. xs: coordinates in the second axis
  577. point_1, point_2: (x, y), the end of the line
  578. '''
  579. height, width = xs.shape[:2]
  580. square_distance_1 = np.square(xs - point_1[0]) + np.square(ys - point_1[1])
  581. square_distance_2 = np.square(xs - point_2[0]) + np.square(ys - point_2[1])
  582. square_distance = np.square(point_1[0] - point_2[0]) + np.square(point_1[1] - point_2[1])
  583. cosin = (square_distance - square_distance_1 - square_distance_2) / (2 * np.sqrt(square_distance_1 * square_distance_2))
  584. square_sin = 1 - np.square(cosin)
  585. square_sin = np.nan_to_num(square_sin)
  586. result = np.sqrt(square_distance_1 * square_distance_2 * square_sin / square_distance)
  587. result[cosin < 0] = np.sqrt(np.fmin(square_distance_1, square_distance_2))[cosin < 0]
  588. return result
  589. def extend_line(self, point_1, point_2, result):
  590. ex_point_1 = (int(round(point_1[0] + (point_1[0] - point_2[0]) * (1 + self.shrink_ratio))),
  591. int(round(point_1[1] + (point_1[1] - point_2[1]) * (1 + self.shrink_ratio))))
  592. cv2.line(result, tuple(ex_point_1), tuple(point_1), 4096.0, 1, lineType=cv2.LINE_AA, shift=0)
  593. ex_point_2 = (int(round(point_2[0] + (point_2[0] - point_1[0]) * (1 + self.shrink_ratio))),
  594. int(round(point_2[1] + (point_2[1] - point_1[1]) * (1 + self.shrink_ratio))))
  595. cv2.line(result, tuple(ex_point_2), tuple(point_2), 4096.0, 1, lineType=cv2.LINE_AA, shift=0)
  596. return ex_point_1, ex_point_2