Spaces:
Runtime error
Runtime error
| import numpy as np | |
| from shapely.geometry import Polygon | |
| def points2polygon(points): | |
| """Convert k points to 1 polygon. | |
| Args: | |
| points (ndarray or list): A ndarray or a list of shape (2k) | |
| that indicates k points. | |
| Returns: | |
| polygon (Polygon): A polygon object. | |
| """ | |
| if isinstance(points, list): | |
| points = np.array(points) | |
| assert isinstance(points, np.ndarray) | |
| assert (points.size % 2 == 0) and (points.size >= 8) | |
| point_mat = points.reshape([-1, 2]) | |
| return Polygon(point_mat) | |
| def poly_intersection(poly_det, poly_gt, buffer=0.0001): | |
| """Calculate the intersection area between two polygon. | |
| Args: | |
| poly_det (Polygon): A polygon predicted by detector. | |
| poly_gt (Polygon): A gt polygon. | |
| Returns: | |
| intersection_area (float): The intersection area between two polygons. | |
| """ | |
| assert isinstance(poly_det, Polygon) | |
| assert isinstance(poly_gt, Polygon) | |
| if buffer == 0: | |
| poly_inter = poly_det & poly_gt | |
| else: | |
| poly_inter = poly_det.buffer(buffer) & poly_gt.buffer(buffer) | |
| return poly_inter.area, poly_inter | |
| def poly_union(poly_det, poly_gt): | |
| """Calculate the union area between two polygon. | |
| Args: | |
| poly_det (Polygon): A polygon predicted by detector. | |
| poly_gt (Polygon): A gt polygon. | |
| Returns: | |
| union_area (float): The union area between two polygons. | |
| """ | |
| assert isinstance(poly_det, Polygon) | |
| assert isinstance(poly_gt, Polygon) | |
| 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 | |
| def valid_boundary(x, with_score=True): | |
| num = len(x) | |
| if num < 8: | |
| return False | |
| if num % 2 == 0 and (not with_score): | |
| return True | |
| if num % 2 == 1 and with_score: | |
| return True | |
| return False | |
| def boundary_iou(src, target): | |
| """Calculate the IOU between two boundaries. | |
| Args: | |
| src (list): Source boundary. | |
| target (list): Target boundary. | |
| Returns: | |
| iou (float): The iou between two boundaries. | |
| """ | |
| assert valid_boundary(src, False) | |
| assert valid_boundary(target, False) | |
| src_poly = points2polygon(src) | |
| target_poly = points2polygon(target) | |
| return poly_iou(src_poly, target_poly) | |
| def poly_iou(poly_det, poly_gt): | |
| """Calculate the IOU between two polygons. | |
| Args: | |
| poly_det (Polygon): A polygon predicted by detector. | |
| poly_gt (Polygon): A gt polygon. | |
| Returns: | |
| iou (float): The IOU between two polygons. | |
| """ | |
| 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: | |
| return 0.0 | |
| return area_inters / area_union | |
| def poly_nms(polygons, threshold): | |
| assert isinstance(polygons, list) | |
| polygons = np.array(sorted(polygons, key=lambda x: x[-1])) | |
| keep_poly = [] | |
| index = [i for i in range(polygons.shape[0])] | |
| while len(index) > 0: | |
| keep_poly.append(polygons[index[-1]].tolist()) | |
| A = polygons[index[-1]][:-1] | |
| index = np.delete(index, -1) | |
| iou_list = np.zeros((len(index),)) | |
| for i in range(len(index)): | |
| B = polygons[index[i]][:-1] | |
| iou_list[i] = boundary_iou(A, B) | |
| remove_index = np.where(iou_list > threshold) | |
| index = np.delete(index, remove_index) | |
| return keep_poly | |