mirror of https://github.com/open-mmlab/mmocr.git
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from mmocr.utils import point_distance, points_center
|
|
|
|
|
|
class TestPointDistance(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.p1_list = [1, 2]
|
|
self.p2_list = [2, 2]
|
|
self.p1_array = np.array([1, 2])
|
|
self.p2_array = np.array([2, 2])
|
|
self.p1_tensor = torch.Tensor([1, 2])
|
|
self.p2_tensor = torch.Tensor([2, 2])
|
|
self.invalid_p = [1, 2, 3]
|
|
|
|
def test_point_distance(self):
|
|
# list
|
|
self.assertEqual(point_distance(self.p1_list, self.p2_list), 1)
|
|
self.assertEqual(point_distance(self.p1_list, self.p1_list), 0)
|
|
# array
|
|
self.assertEqual(point_distance(self.p1_array, self.p2_array), 1)
|
|
self.assertEqual(point_distance(self.p1_array, self.p1_array), 0)
|
|
# tensor
|
|
self.assertEqual(point_distance(self.p1_tensor, self.p2_tensor), 1)
|
|
self.assertEqual(point_distance(self.p1_tensor, self.p1_tensor), 0)
|
|
with self.assertRaises(AssertionError):
|
|
point_distance(self.invalid_p, self.invalid_p)
|
|
|
|
|
|
class TestPointCenter(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.point_list = [1, 2, 3, 4]
|
|
self.point_nparray = np.array([1, 2, 3, 4])
|
|
self.point_tensor = torch.Tensor([1, 2, 3, 4])
|
|
self.incorrect_input = [1, 3, 4]
|
|
self.gt = np.array([2, 3])
|
|
|
|
def test_point_center(self):
|
|
# list
|
|
self.assertTrue(
|
|
np.array_equal(points_center(self.point_list), self.gt))
|
|
# array
|
|
self.assertTrue(
|
|
np.array_equal(points_center(self.point_nparray), self.gt))
|
|
# tensor
|
|
self.assertTrue(
|
|
np.array_equal(points_center(self.point_tensor), self.gt))
|
|
with self.assertRaises(AssertionError):
|
|
points_center(self.incorrect_input)
|