diff --git a/ppocr/data/imaug/fce_aug.py b/ppocr/data/imaug/fce_aug.py index e1668d77bb793c47132289455108eef2ccec0bd5..7ef00de33bad3076d5b6d585cb9aa14df17a015d 100644 --- a/ppocr/data/imaug/fce_aug.py +++ b/ppocr/data/imaug/fce_aug.py @@ -18,7 +18,7 @@ https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/transform import numpy as np from PIL import Image, ImageDraw import cv2 -import Polygon as plg +from shapely.geometry import Polygon import math from ppocr.utils.poly_nms import poly_intersection @@ -129,16 +129,16 @@ class RandomCropFlip: pts = np.stack([[xmin, xmax, xmax, xmin], [ymin, ymin, ymax, ymax]]).T.astype(np.int32) - pp = plg.Polygon(pts) + pp = Polygon(pts) fail_flag = False for polygon, ignore_tag in zip(polygons, ignore_tags): - ppi = plg.Polygon(polygon.reshape(-1, 2)) + ppi = Polygon(polygon.reshape(-1, 2)) ppiou, _ = poly_intersection(ppi, pp) - if np.abs(ppiou - float(ppi.area())) > self.epsilon and \ + if np.abs(ppiou - float(ppi.area)) > self.epsilon and \ np.abs(ppiou) > self.epsilon: fail_flag = True break - elif np.abs(ppiou - float(ppi.area())) < self.epsilon: + elif np.abs(ppiou - float(ppi.area)) < self.epsilon: polys_new.append(polygon) ignore_tags_new.append(ignore_tag) else: diff --git a/ppocr/utils/poly_nms.py b/ppocr/utils/poly_nms.py index 2ee4ac0e0b0082203173a8a3e174680912a63c08..31eb6642aca1c29ab937ada7c044d51f922c7730 100644 --- a/ppocr/utils/poly_nms.py +++ b/ppocr/utils/poly_nms.py @@ -13,7 +13,7 @@ # limitations under the License. import numpy as np -import Polygon as plg +from shapely.geometry import Polygon def points2polygon(points): @@ -33,7 +33,7 @@ def points2polygon(points): assert (points.size % 2 == 0) and (points.size >= 8) point_mat = points.reshape([-1, 2]) - return plg.Polygon(point_mat) + return Polygon(point_mat) def poly_intersection(poly_det, poly_gt): @@ -46,13 +46,11 @@ def poly_intersection(poly_det, poly_gt): Returns: intersection_area (float): The intersection area between two polygons. """ - assert isinstance(poly_det, plg.Polygon) - assert isinstance(poly_gt, plg.Polygon) + assert isinstance(poly_det, Polygon) + assert isinstance(poly_gt, Polygon) - poly_inter = poly_det & poly_gt - if len(poly_inter) == 0: - return 0, poly_inter - return poly_inter.area(), poly_inter + poly_inter = poly_det.buffer(0.001) & poly_gt.buffer(0.001) + return poly_inter.area, poly_inter def poly_union(poly_det, poly_gt): @@ -65,11 +63,11 @@ def poly_union(poly_det, poly_gt): Returns: union_area (float): The union area between two polygons. """ - assert isinstance(poly_det, plg.Polygon) - assert isinstance(poly_gt, plg.Polygon) + assert isinstance(poly_det, Polygon) + assert isinstance(poly_gt, Polygon) - area_det = poly_det.area() - area_gt = poly_gt.area() + area_det = poly_det.area + area_gt = poly_gt.area area_inters, _ = poly_intersection(poly_det, poly_gt) return area_det + area_gt - area_inters @@ -114,8 +112,8 @@ def poly_iou(poly_det, poly_gt): Returns: iou (float): The IOU between two polygons. """ - assert isinstance(poly_det, plg.Polygon) - assert isinstance(poly_gt, plg.Polygon) + assert isinstance(poly_det, Polygon) + assert isinstance(poly_gt, Polygon) area_inters, _ = poly_intersection(poly_det, poly_gt) area_union = poly_union(poly_det, poly_gt) if area_union == 0: @@ -142,4 +140,4 @@ def poly_nms(polygons, threshold): remove_index = np.where(iou_list > threshold) index = np.delete(index, remove_index) - return keep_poly \ No newline at end of file + return keep_poly