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