baseline_generate_dataset.py 文件源码

def process_one(image_filename, output_dir, basename):
    page = PAGE.parse_file(get_page_filename(image_filename))
    text_lines = [tl for tr in page.text_regions for tl in tr.text_lines]
    img = imread(image_filename, mode='RGB')
    gt = np.zeros_like(img)
    cv2.fillPoly(gt, [PAGE.Point.list_to_cv2poly(tl.coords) for tl in text_lines], DRAWING_COLOR)
    save_and_resize(img, os.path.join(output_dir, 'images', '{}.jpg'.format(basename)))
    save_and_resize(gt, os.path.join(output_dir, 'labels', '{}.png'.format(basename)), nearest=True)

    classes = np.stack([(0, 0, 0), DRAWING_COLOR])
    np.savetxt(os.path.join(output_dir, 'classes.txt'), classes, fmt='%d')