layout_generate_dataset.py 文件源码

def process_one(image_dir, page_dir, output_dir, basename, colormap, color_labels):
    image_filename = os.path.join(image_dir, "{}.jpg".format(basename))
    page_filename = os.path.join(page_dir, "{}.xml".format(basename))

    page = PAGE.parse_file(page_filename)
    text_lines = [tl for tr in page.text_regions for tl in tr.text_lines]
    graphic_regions = page.graphic_regions
    img = imread(image_filename, mode='RGB')

    gt = np.zeros_like(img[:, :, 0])
    mask1 = cv2.fillPoly(gt.copy(), [PAGE.Point.list_to_cv2poly(tl.coords)
                                     for tl in text_lines if 'comment' in tl.id], 1)
    mask2 = cv2.fillPoly(gt.copy(), [PAGE.Point.list_to_cv2poly(tl.coords)
                                     for tl in text_lines if not 'comment' in tl.id], 1)
    mask3 = cv2.fillPoly(gt.copy(), [PAGE.Point.list_to_cv2poly(tl.coords)
                                     for tl in graphic_regions], 1)
    arr = np.dstack([mask1, mask2, mask3])

    gt_img = convert_array_masks(arr, colormap, color_labels)
    save_and_resize(img, os.path.join(output_dir, 'images', '{}.jpg'.format(basename)))
    save_and_resize(gt_img, os.path.join(output_dir, 'labels', '{}.png'.format(basename)), nearest=True)