def ReadPdfContent():
return "ReadPdfContent" >> TransformAndCount(
beam.Map(
lambda pdf_url: {
DataProps.SOURCE_FILENAME: pdf_url,
DataProps.PDF_CONTENT: read_all_from_path(pdf_url)
}), MetricCounters.FILES)
def ReadFileContent():
return "ReadFileContent" >> TransformAndCount(
beam.Map(
lambda file_url: {
DataProps.SOURCE_FILENAME: file_url,
DataProps.FILENAME: file_url,
DataProps.CONTENT: read_all_from_path(file_url)
}), MetricCounters.FILES)
def process_file(file_url: str,
simple_runner: SimplePipelineRunner,
get_output_file_for_source_url: Callable[[str], str],
session: requests.Session,
request_args: MultiDict = None):
output_file_url = get_output_file_for_source_url(file_url)
file_content = read_all_from_path(file_url)
LOGGER.info('read source content: %s (%s)', file_url,
format_size(len(file_content)))
data_type = guess_type(file_url)[0]
LOGGER.debug('data_type: %s', data_type)
LOGGER.debug('session: %s', session)
context = {RequestsPipelineStep.REQUESTS_SESSION_KEY: session}
if request_args:
context['request_args'] = request_args
result = simple_runner.convert(file_content,
file_url,
data_type,
context=context)
LOGGER.debug('result.keys: %s', result.keys())
output_content = encode_if_text_type(result[DataProps.CONTENT])
save_file_content(output_file_url, output_content)
LOGGER.info('saved output to: %s (%s)', output_file_url,
format_size(len(output_content)))
def fix_source_file_to(source_url: str,
target_url: str,
encoding: str = 'utf-8'):
xml_bytes = read_all_from_path(source_url)
fixed_xml_bytes = get_fixed_xml_bytes(xml_bytes, encoding=encoding)
save_file_content(target_url, fixed_xml_bytes)
def load_annotation_image(path, color_map):
get_logger().debug('loading annotation image: %s', path)
return AnnotatedImage(
np.asarray(Image.open(BytesIO(read_all_from_path(
path, 'rb'))).convert('RGB'),
dtype=np.uint8), color_map)
def fix_source_file_to(source_url: str, target_url: str):
source_data = read_all_from_path(source_url)
data = source_data
if b'</content>' in data and b'<content>' not in data:
data = data.replace(b'</content>', b'')
save_file_content(target_url, data)
def configure_pipeline(p, opt):
image_size = ((opt.image_width, opt.image_height)
if opt.image_width and opt.image_height else None)
page_range = opt.pages
first_page = page_range[0] if page_range else 1
xml_mapping = parse_xml_mapping(opt.xml_mapping_path)
if opt.lxml_path:
lxml_xml_file_pairs = (
p | beam.Create(
[[
join_if_relative_path(opt.base_data_path, s)
for s in [opt.lxml_path, opt.xml_path]
]]) | "FindFilePairs" >> TransformAndLog(
beam.FlatMap(lambda patterns: islice(
find_file_pairs_grouped_by_parent_directory_or_name(
patterns), opt.limit)),
log_prefix='file pairs: ',
log_level='debug') | PreventFusion()
| "ReadFileContent" >> beam.Map(
lambda filenames: {
'source_filename': filenames[0],
'xml_filename': filenames[1],
'lxml_content': read_all_from_path(filenames[0]),
'xml_content': read_all_from_path(filenames[1])
}))
elif opt.pdf_path or opt.pdf_xml_file_list:
if opt.pdf_xml_file_list:
pdf_xml_url_pairs = (
p | "ReadFilePairUrls" >> ReadDictCsv(opt.pdf_xml_file_list,
limit=opt.limit)
| "TranslateFilePairUrls" >>
beam.Map(lambda row: (row['source_url'], row['xml_url'])))
else:
pdf_xml_url_pairs = (p | beam.Create([[
join_if_relative_path(opt.base_data_path, s)
for s in [opt.pdf_path, opt.xml_path]
]]) | "FindFilePairs" >> TransformAndLog(
beam.FlatMap(lambda patterns: islice(
find_file_pairs_grouped_by_parent_directory_or_name(
patterns), opt.limit)),
log_prefix='file pairs: ',
log_level='debug'))
pdf_xml_file_pairs = (
pdf_xml_url_pairs | PreventFusion()
| "ReadFileContent" >> TransformAndCount(
beam.Map(
lambda filenames: {
'source_filename': filenames[0],
'xml_filename': filenames[1],
'pdf_content': read_all_from_path(filenames[0]),
'xml_content': read_all_from_path(filenames[1])
}), MetricCounters.FILE_PAIR))
lxml_xml_file_pairs = (
pdf_xml_file_pairs | "ConvertPdfToLxml" >> MapOrLog(
lambda v: remove_keys_from_dict(
extend_dict(
v, {
'lxml_content':
convert_pdf_bytes_to_lxml(v['pdf_content'],
path=v['source_filename'
],
page_range=page_range)
}),
# we don't need the pdf_content unless we are writing tf_records
None if opt.save_tfrecords else {'pdf_content'}),
log_fn=lambda e, v: (get_logger().warning(
'caught exception (ignoring item): %s, pdf: %s, xml: %s',
e,
v['source_filename'],
v['xml_filename'],
exc_info=e)),
error_count=MetricCounters.CONVERT_PDF_TO_LXML_ERROR))
else:
raise RuntimeError('either lxml-path or pdf-path required')
if opt.save_png or opt.save_tfrecords:
with_pdf_png_pages = (
(lxml_xml_file_pairs if opt.save_tfrecords else pdf_xml_file_pairs)
| "ConvertPdfToPng" >> MapOrLog(
lambda v: remove_keys_from_dict(
extend_dict(
v, {
'pdf_png_pages':
list(
pdf_bytes_to_png_pages(v['pdf_content'],
dpi=opt.png_dpi,
image_size=image_size,
page_range=page_range))
}),
{'pdf_content'} # we no longer need the pdf_content
),
error_count=MetricCounters.CONVERT_PDF_TO_PNG_ERROR))
if opt.save_png:
_ = (with_pdf_png_pages | "SavePdfToPng" >> TransformAndLog(
beam.Map(lambda v: save_pages(
FileSystems.join(
opt.output_path,
change_ext(
relative_path(opt.base_data_path, v[
'source_filename']), None, '.png.zip')),
'.png', v['pdf_png_pages'])),
log_fn=lambda x: get_logger().info('saved result: %s', x)))
if opt.save_lxml:
_ = (lxml_xml_file_pairs
| "SaveLxml" >> TransformAndLog(
beam.Map(lambda v: save_file_content(
FileSystems.join(
opt.output_path,
change_ext(
relative_path(opt.base_data_path, v[
'source_filename']), None, '.lxml.gz')), v[
'lxml_content'])),
log_fn=lambda x: get_logger().info('saved lxml: %s', x)))
annotation_results = ((
with_pdf_png_pages if opt.save_tfrecords else lxml_xml_file_pairs
) | "ConvertLxmlToSvgAndAnnotate" >> TransformAndCount(
MapOrLog(
lambda v: remove_keys_from_dict(
extend_dict(
v, {
'svg_pages':
list(
convert_and_annotate_lxml_content(
v['lxml_content'],
v['xml_content'],
xml_mapping,
name=v['source_filename']))
}),
# Won't need the XML anymore
{'lxml_content', 'xml_content'}),
log_fn=lambda e, v: (get_logger().warning(
'caught exception (ignoring item): %s, source: %s, xml: %s',
e,
v['source_filename'],
v['xml_filename'],
exc_info=e)),
error_count=MetricCounters.CONVERT_LXML_TO_SVG_ANNOT_ERROR),
MetricCounters.PAGE,
lambda v: len(v['svg_pages'])))
if opt.save_svg:
_ = (annotation_results | "SaveSvgPages" >> TransformAndLog(
beam.Map(lambda v: save_svg_roots(
FileSystems.join(
opt.output_path,
change_ext(
relative_path(opt.base_data_path, v['source_filename']
), None, '.svg.zip')), v['svg_pages'])),
log_fn=lambda x: get_logger().info('saved result: %s', x)))
if opt.annotation_evaluation_csv or opt.min_annotation_percentage:
annotation_evaluation_results = (
annotation_results | "EvaluateAnnotations" >> TransformAndLog(
beam.Map(lambda v: remove_keys_from_dict(
extend_dict(
v, {
'annotation_evaluation':
evaluate_document_by_page(
SvgStructuredDocument(v['svg_pages']))
}), None
if opt.min_annotation_percentage else {'svg_pages'})),
log_fn=lambda x: get_logger().info(
'annotation evaluation result: %s: %s', x[
'source_filename'], x['annotation_evaluation'])))
if opt.save_block_png or opt.save_tfrecords:
color_map = parse_color_map_from_file(opt.color_map)
with_block_png_pages = (
(annotation_evaluation_results
if opt.min_annotation_percentage else annotation_results)
| "GenerateBlockPng" >> beam.Map(lambda v: remove_keys_from_dict(
extend_dict(
v, {
'block_png_pages': [
svg_page_to_blockified_png_bytes(
svg_page, color_map, image_size=image_size)
for svg_page in v['svg_pages']
]
}), {'svg_pages'})))
if opt.save_block_png:
_ = (with_block_png_pages | "SaveBlockPng" >> TransformAndLog(
beam.Map(lambda v: save_pages(
FileSystems.join(
opt.output_path,
change_ext(
relative_path(opt.base_data_path, v[
'source_filename']), None, '.block-png.zip')),
'.png', v['block_png_pages'])),
log_fn=lambda x: get_logger().info('saved result: %s', x)))
if opt.save_tfrecords:
if opt.min_annotation_percentage:
filtered_pages = (
with_block_png_pages | "FilterPages" >> TransformAndCount(
beam.Map(lambda v: filter_list_props_by_indices(
v,
get_page_indices_with_min_annotation_percentage(
v['annotation_evaluation'], opt.
min_annotation_percentage),
{'pdf_png_pages', 'block_png_pages'})),
MetricCounters.FILTERED_PAGE,
lambda v: len(v['block_png_pages'])))
else:
filtered_pages = with_block_png_pages
_ = (filtered_pages | "WriteTFRecords" >> WritePropsToTFRecord(
FileSystems.join(opt.output_path, 'data'), lambda v: ({
'input_uri':
v['source_filename'] + '#page%d' % (first_page + i),
'input_image':
pdf_png_page,
'annotation_uri':
(v['source_filename'] + '.annot' + '#page%d' %
(first_page + i)),
'annotation_image':
block_png_page,
'page_no':
first_page + i
} for i, pdf_png_page, block_png_page in zip(
range(len(v['pdf_png_pages'])), v['pdf_png_pages'], v[
'block_png_pages']))))
if opt.annotation_evaluation_csv:
annotation_evaluation_csv_name, annotation_evaluation_ext = (
os.path.splitext(opt.annotation_evaluation_csv))
_ = ( # flake8: noqa
annotation_evaluation_results | "FlattenAnotationEvaluationResults"
>> beam.FlatMap(lambda v: to_annotation_evaluation_csv_dict_rows(
v['annotation_evaluation'],
document=basename(v['source_filename'])))
| "WriteAnnotationEvaluationToCsv" >> WriteDictCsv(
join_if_relative_path(opt.output_path,
annotation_evaluation_csv_name),
file_name_suffix=annotation_evaluation_ext,
columns=DEFAULT_EVALUATION_COLUMNS))