mirrors/mmocr
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmocr.git synced 2025-06-03 21:54:47 +08:00
Code Issues Projects Releases Wiki Activity

Fix crop_polygon, Resize and poly2shapely

Browse Source
This commit is contained in:
gaotongxiao 2022-05-27 14:32:44 +00:00
parent 00f821315e
commit 0f0f68baf1
3 changed files with 21 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mmocr/datasets/pipelines/processing.py
View File

@ -170,7 +170,8 @@ class Resize(MMCV_Resize):
if polygon is not None: if polygon is not None:
polygons_resize.append(polygon.astype(np.float32)) polygons_resize.append(polygon.astype(np.float32))
else: else:
polygons_resize.append(np.zeros_like(polygons[idx])) polygons_resize.append(
np.zeros_like(polygons[idx], dtype=np.float32))
results['gt_polygons'] = polygons_resize results['gt_polygons'] = polygons_resize
def transform(self, results: dict) -> dict: def transform(self, results: dict) -> dict:

18
mmocr/utils/polygon_utils.py
View File

@ -100,8 +100,8 @@ def poly2shapely(polygon: ArrayLike) -> Polygon:
Returns: Returns:
polygon (Polygon): A polygon object. polygon (Polygon): A polygon object.
""" """
assert len(polygon) % 2 == 0 and len(polygon) >= 8
polygon = np.array(polygon, dtype=np.float32) polygon = np.array(polygon, dtype=np.float32)
assert polygon.size % 2 == 0 and polygon.size >= 6
polygon = polygon.reshape([-1, 2]) polygon = polygon.reshape([-1, 2])
return Polygon(polygon) return Polygon(polygon)
@ -119,7 +119,8 @@ def polys2shapely(polygons: Sequence[ArrayLike]) -> Sequence[Polygon]:
return [poly2shapely(polygon) for polygon in polygons] return [poly2shapely(polygon) for polygon in polygons]
def crop_polygon(polygon: ArrayLike, crop_box: np.ndarray) -> np.ndarray: def crop_polygon(polygon: ArrayLike,
crop_box: np.ndarray) -> Union[np.ndarray, None]:
"""Crop polygon to be within a box region. """Crop polygon to be within a box region.
Args: Args:
@ -127,19 +128,18 @@ def crop_polygon(polygon: ArrayLike, crop_box: np.ndarray) -> np.ndarray:
crop_box (ndarray): target box region in shape (4, ). crop_box (ndarray): target box region in shape (4, ).
Returns: Returns:
np.array or None: Cropped polygon. np.array or None: Cropped polygon. If the polygon is not within the
crop box, return None.
""" """
polygon = np.asarray(polygon, dtype=np.float32) poly = poly2shapely(polygon)
crop_box = np.asarray(crop_box, dtype=np.float32) crop_poly = poly2shapely(bbox2poly(crop_box))
poly = Polygon(polygon.reshape(-1, 2))
crop_poly = Polygon(bbox2poly(crop_box).reshape(-1, 2))
poly_cropped = poly.intersection(crop_poly) poly_cropped = poly.intersection(crop_poly)
if poly_cropped.area == 0.: if poly_cropped.area == 0.:
# If polygon is outside crop_box region, return None. # If polygon is outside crop_box region, return None.
return None return None
else: else:
poly_cropped = np.array(poly_cropped.boundary.xy)[:, :-1] poly_cropped = np.array(poly_cropped.boundary.xy, dtype=np.float32)
return poly_cropped.reshape(-1) return poly_cropped[:, :-1].T.reshape(-1)
def poly_make_valid(poly: Polygon) -> Polygon: def poly_make_valid(poly: Polygon) -> Polygon:

17
tests/test_utils/test_polygon_utils.py
View File

@ -17,22 +17,25 @@ class TestCropPolygon(unittest.TestCase):
# polygon cross box # polygon cross box
polygon = np.array([20., -10., 40., 10., 10., 40., -10., 20.]) polygon = np.array([20., -10., 40., 10., 10., 40., -10., 20.])
crop_box = np.array([0., 0., 60., 60.]) crop_box = np.array([0., 0., 60., 60.])
target_poly_cropped = np.array([[10., 40., 30., 10., 0., 0., 10.], target_poly_cropped = np.array(
[40., 10., 0., 0., 10., 30., 40.]]) [10, 40, 0, 30, 0, 10, 10, 0, 30, 0, 40, 10])
poly_cropped = crop_polygon(polygon, crop_box) poly_cropped = crop_polygon(polygon, crop_box)
self.assertTrue(target_poly_cropped.all() == poly_cropped.all()) self.assertTrue(
poly2shapely(poly_cropped).equals(
poly2shapely(target_poly_cropped)))
# polygon inside box # polygon inside box
polygon = np.array([0., 0., 30., 0., 30., 30., 0., 30.]).reshape(-1, 2) polygon = np.array([0., 0., 30., 0., 30., 30., 0., 30.]).reshape(-1, 2)
crop_box = np.array([0., 0., 60., 60.]) crop_box = np.array([0., 0., 60., 60.])
target_poly_cropped = polygon target_poly_cropped = polygon
poly_cropped = crop_polygon(polygon, crop_box) poly_cropped = crop_polygon(polygon, crop_box)
self.assertTrue(target_poly_cropped.all() == poly_cropped.all()) self.assertTrue(
poly2shapely(poly_cropped).equals(
poly2shapely(target_poly_cropped)))
# polygon outside box # polygon outside box
polygon = np.array([0., 0., 30., 0., 30., 30., 0., 30.]).reshape(-1, 2) polygon = np.array([0., 0., 30., 0., 30., 30., 0., 30.]).reshape(-1, 2)
crop_box = np.array([80., 80., 90., 90.]) crop_box = np.array([80., 80., 90., 90.])
target_poly_cropped = polygon
poly_cropped = crop_polygon(polygon, crop_box) poly_cropped = crop_polygon(polygon, crop_box)
self.assertEqual(poly_cropped, None) self.assertEqual(poly_cropped, None)
@ -133,10 +136,10 @@ class TestPolygonUtils(unittest.TestCase):
self.assertEqual(polys2shapely(polys), polygons) self.assertEqual(polys2shapely(polys), polygons)
# test invalid # test invalid
polys = [0, 0, 1] polys = [0, 0, 1]
with self.assertRaises(TypeError): with self.assertRaises(AssertionError):
polys2shapely(polys) polys2shapely(polys)
polys = [0, 0, 1, 0, 1, 1, 0, 1, 1] polys = [0, 0, 1, 0, 1, 1, 0, 1, 1]
with self.assertRaises(TypeError): with self.assertRaises(AssertionError):
polys2shapely(polys) polys2shapely(polys)
def test_poly_make_valid(self): def test_poly_make_valid(self):