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 | |