MMOCR / tests /test_dataset /test_crop.py
tomofi's picture
Add application file
2366e36
raw
history blame
3.31 kB
# Copyright (c) OpenMMLab. All rights reserved.
import math
from itertools import chain, permutations
import numpy as np
import pytest
from mmocr.datasets.pipelines.box_utils import sort_vertex, sort_vertex8
from mmocr.datasets.pipelines.crop import box_jitter, crop_img, warp_img
def test_order_vertex():
dummy_points_x = [20, 20, 120, 120]
dummy_points_y = [20, 40, 40, 20]
expect_points_x = [20, 120, 120, 20]
expect_points_y = [20, 20, 40, 40]
with pytest.raises(AssertionError):
sort_vertex([], dummy_points_y)
with pytest.raises(AssertionError):
sort_vertex(dummy_points_x, [])
for perm in set(permutations([0, 1, 2, 3])):
points_x = [dummy_points_x[i] for i in perm]
points_y = [dummy_points_y[i] for i in perm]
ordered_points_x, ordered_points_y = sort_vertex(points_x, points_y)
assert np.allclose(ordered_points_x, expect_points_x)
assert np.allclose(ordered_points_y, expect_points_y)
def test_sort_vertex8():
dummy_points_x = [21, 21, 122, 122]
dummy_points_y = [21, 39, 39, 21]
expect_points = [21, 21, 122, 21, 122, 39, 21, 39]
for perm in set(permutations([0, 1, 2, 3])):
points_x = [dummy_points_x[i] for i in perm]
points_y = [dummy_points_y[i] for i in perm]
points = list(chain.from_iterable(zip(points_x, points_y)))
ordered_points = sort_vertex8(points)
assert np.allclose(ordered_points, expect_points)
def test_box_jitter():
dummy_points_x = [20, 120, 120, 20]
dummy_points_y = [20, 20, 40, 40]
kwargs = dict(jitter_ratio_x=0.0, jitter_ratio_y=0.0)
with pytest.raises(AssertionError):
box_jitter([], dummy_points_y)
with pytest.raises(AssertionError):
box_jitter(dummy_points_x, [])
with pytest.raises(AssertionError):
box_jitter(dummy_points_x, dummy_points_y, jitter_ratio_x=1.)
with pytest.raises(AssertionError):
box_jitter(dummy_points_x, dummy_points_y, jitter_ratio_y=1.)
box_jitter(dummy_points_x, dummy_points_y, **kwargs)
assert np.allclose(dummy_points_x, [20, 120, 120, 20])
assert np.allclose(dummy_points_y, [20, 20, 40, 40])
def test_opencv_crop():
dummy_img = np.ones((600, 600, 3), dtype=np.uint8)
dummy_box = [20, 20, 120, 20, 120, 40, 20, 40]
cropped_img = warp_img(dummy_img, dummy_box)
with pytest.raises(AssertionError):
warp_img(dummy_img, [])
with pytest.raises(AssertionError):
warp_img(dummy_img, [20, 40, 40, 20])
assert math.isclose(cropped_img.shape[0], 20)
assert math.isclose(cropped_img.shape[1], 100)
def test_min_rect_crop():
dummy_img = np.ones((600, 600, 3), dtype=np.uint8)
dummy_box = [20, 20, 120, 20, 120, 40, 20, 40]
cropped_img = crop_img(
dummy_img,
dummy_box,
0.,
0.,
)
with pytest.raises(AssertionError):
crop_img(dummy_img, [])
with pytest.raises(AssertionError):
crop_img(dummy_img, [20, 40, 40, 20])
with pytest.raises(AssertionError):
crop_img(dummy_img, dummy_box, 4, 0.2)
with pytest.raises(AssertionError):
crop_img(dummy_img, dummy_box, 0.4, 1.2)
assert math.isclose(cropped_img.shape[0], 20)
assert math.isclose(cropped_img.shape[1], 100)