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_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_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 draw_polygons(image: np.ndarray,
polygons: Polygons,
thickness=2):
canvas = image.copy().astype(np.uint8)
num = len(polygons)
if num == 0:
return canvas
for idx in range(num):
item_dict = polygons.get(idx)
points = item_dict['contour'].astype(np.int32)
if polygons.has_field('training_tag'):
if not item_dict['fields']['training_tag']:
lines_color = COLORS[-1]
if not polygons.has_field('training_tag') or item_dict['fields']['training_tag']:
lines_color = COLORS[0]
for i in range(len(points)):
cv2.line(canvas, (points[i][0], points[i][1]), (points[i-1][0], points[i-1][1]), lines_color[0], thickness)
return canvas
