def test_db_encoder():
db_encoder = DBEncoder(cfg)
contours = [
np.array([[0, 0], [40, 0], [80, 40], [80, 80], [0, 80]]),
np.array([[100, 100], [140, 140], [140, 100]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, True])
num_polygons = len(polygons)
db_encoder.init_encoder()
for idx in range(num_polygons):
db_encoder.calculate_encoder(polygons.get(idx))
target = db_encoder.create_target()
plt.figure("img")
plt.figure(figsize=(12, 12))
plt.imshow(target['db_probability_map'][0])
plt.show()
plt.figure("img")
plt.figure(figsize=(12, 12))
plt.imshow(target['db_threshold_map'][0])
plt.show()
plt.figure("img")
plt.figure(figsize=(12, 12))
plt.imshow(target['db_threshold_mask'][0])
plt.show()
def test_get(self):
contours = [
np.array([[2., 1.], [1., 3.], [3., 5.], [4., 2.]]),
np.array([[0.5, 0.5], [10.5, 0.5], [10.5, 10.5]]),
np.array([[10, 0], [10, 20], [0, 10], [20, 10], [20, 20]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, False, True])
polygon = polygons.get(1)
np.testing.assert_equal(
polygon['contour'],
np.array([[0.5, 0.5], [10.5, 0.5], [10.5, 10.5]]))
self.assertEqual(polygon['fields']['training_tag'], False)
self.assertEqual(polygon['fields']['num_points'], 3)
polygon['fields']['training_tag'] = True
polygon['contour'] = np.array([[10, 0], [10, 20], [0, 10], [20, 10],
[20, 20]])
np.testing.assert_equal(polygons.get_contours(), [
np.array([[2., 1.], [1., 3.], [3., 5.], [4., 2.]]),
np.array([[0.5, 0.5], [10.5, 0.5], [10.5, 10.5]]),
np.array([[10, 0], [10, 20], [0, 10], [20, 10], [20, 20]])
])
self.assertEqual(polygons.get_field('training_tag'),
[True, False, True])
def test_get_field_numpy(self):
contours = [np.zeros([4, 2])]
polygons = Polygons(contours)
np_field = [np.zeros([4])]
polygons.add_field('np_field', np_field)
np_field_modified = polygons.get_field('np_field')
np_field_modified[0][0] = 1
np.testing.assert_equal(np_field_modified, np.array([[1, 0, 0, 0]]))
np_field_original = polygons.get_field('np_field')
np.testing.assert_equal(np_field_original, np.array([[0, 0, 0, 0]]))
def test_delete(self):
contours = [
np.array([[2., 1.], [1., 3.], [3., 5.], [4., 2.]]),
np.array([[0.5, 0.5], [10.5, 0.5], [10.5, 10.5]]),
np.array([[10, 0], [10, 20], [0, 10], [20, 10], [20, 20]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, False, True])
polygons.delete([0, 2])
np.testing.assert_equal(polygons.get_contours(), [contours[1]])
def test_fliter_small_with_exist(self):
contours = [
np.array([[2., 1.], [1., 3.], [3., 5.], [4., 2.]]),
np.array([[0.5, 0.5], [10.5, 0.5], [10.5, 10.5]]),
np.array([[10, 0], [10, 20], [0, 10], [20, 10], [20, 20]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, False, True])
polygons.fliter_small(10)
np.testing.assert_equal(polygons.get_field('training_tag'),
np.array([False, False, True]))
np.testing.assert_equal(polygons.get_contours(), contours)
def test_fliter_out_with_exist(self):
contours = [
np.array([[20, 10], [10, 30], [30, 50], [40, 20]]),
np.array([[0, 0], [10, 0], [10, 10]]),
np.array([[10, 0], [10, 19], [0, 10], [19, 10], [19, 19]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, False, True])
window = np.array([[0, 0], [20, 20]])
polygons.fliter_out(window)
np.testing.assert_equal(polygons.get_field('training_tag'),
np.array([False, False, True]))
np.testing.assert_equal(polygons.get_contours(), contours)
def test_copy(self):
contours = [np.zeros([4, 2]), np.zeros([4, 2]), np.zeros([4, 2])]
polygons = Polygons(contours)
field = ['s', 's', 's']
polygons.add_field('field', field)
p = polygons.copy()
np.testing.assert_equal(p.get_contours(), contours)
self.assertEqual(p.get_field('num_points'), [4, 4, 4])
self.assertEqual(p.get_field('field'), ['s', 's', 's'])
p.set_contours([np.ones([4, 2]), np.ones([4, 2]), np.ones([4, 2])])
p.set_field('field', ['w', 'w', 'w'])
np.testing.assert_equal(polygons.get_contours(), contours)
self.assertEqual(polygons.get_field('field'), ['s', 's', 's'])
def test_pse_encoder():
pse_encoder = PSEEncoder(cfg)
contours = [
np.array([[0, 0], [40, 0], [80, 40], [80, 80], [0, 80]]),
np.array([[100, 100], [140, 140], [140, 100]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, True])
num_polygons = len(polygons)
pse_encoder.init_encoder()
for idx in range(num_polygons):
pse_encoder.calculate_encoder(polygons.get(idx))
target = pse_encoder.create_target()
plt.figure("img")
plt.figure(figsize=(12, 12))
plt.imshow(np.sum(target['pse_map'], 0) / 6)
plt.show()
def test_east_encoder():
east_encoder = EASTEncoder(cfg)
contours = [
np.array([[200, 200], [240, 200], [280, 240], [280, 280], [200, 280]]),
np.array([[100, 100], [140, 140], [140, 100]])
]
polygons = Polygons(contours)
polygons.add_field('training_tag', [True, True])
num_polygons = len(polygons)
east_encoder.init_encoder()
for idx in range(num_polygons):
east_encoder.calculate_encoder(polygons.get(idx))
target = east_encoder.create_target()
plt.figure("img")
plt.figure(figsize=(12, 12))
plt.imshow(target['east_score_map'][0])
plt.show()
canvas = np.zeros([512, 512])
scores_map = target['east_score_map'].transpose([1, 2, 0])
distance_map = target['east_distance_map'].transpose([1, 2, 0])
rotation_map = target['east_rotation_map'].transpose([1, 2, 0])
box_map_index = np.where(scores_map[:, :, 0] > 0)
box_map_index = np.stack(box_map_index, -1)
for index in range(box_map_index.shape[0]):
(row, col) = box_map_index[index]
rotation_angle = rotation_map[row][col]
d_top, d_right, d_bottom, d_left = distance_map[row][col]
print(d_top, d_right, d_bottom, d_left)
x0, y0, x1, y1, x2, y2, x3, y3 = rbox2poly(
[d_top, d_right, d_bottom, d_left], (col * 4, row * 4),
rotation_angle)
print(x0, y0, x1, y1, x2, y2, x3, y3)
cv2.line(canvas, (x0, y0), (x1, y1), 1, 1)
cv2.line(canvas, (x1, y1), (x2, y2), 1, 1)
cv2.line(canvas, (x2, y2), (x3, y3), 1, 1)
cv2.line(canvas, (x3, y3), (x0, y0), 1, 1)
plt.figure("img")
plt.figure(figsize=(12, 12))
plt.imshow(canvas)
plt.show()
def test_add_field_alread_in(self):
contours = [np.zeros([4, 2]), np.zeros([4, 2]), np.zeros([4, 2])]
polygons = Polygons(contours)
field = ['s', 's', 's']
with self.assertRaises(ValueError):
polygons.add_field('num_points', field)
def test_add_field_fifferent_type(self):
contours = [np.zeros([4, 2]), np.zeros([4, 2]), np.zeros([4, 2])]
polygons = Polygons(contours)
field = [0, 's', 0.001]
with self.assertRaises(TypeError):
polygons.add_field('field', field)
def test_add_field_fifferent_len(self):
contours = [np.zeros([4, 2]), np.zeros([4, 2]), np.zeros([4, 2])]
polygons = Polygons(contours)
field = ['s', 's']
with self.assertRaises(ValueError):
polygons.add_field('field', field)