From 2bd4ae8d5a7a7b5c1ac9d26099f2af56edf67919 Mon Sep 17 00:00:00 2001 From: Kai Date: Thu, 6 May 2021 16:23:08 +0200 Subject: [PATCH] add ned-priority option to page2tsv --- tsvtools/cli.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tsvtools/cli.py b/tsvtools/cli.py index d94af03..9fcc215 100644 --- a/tsvtools/cli.py +++ b/tsvtools/cli.py @@ -78,8 +78,9 @@ def annotate_tsv(tsv_file, annotated_tsv_file): @click.option('--ned-threshold', type=float, default=None) @click.option('--min-confidence', type=float, default=None) @click.option('--max-confidence', type=float, default=None) +@click.option('--ned-priority', type=int, default=1) def page2tsv(page_xml_file, tsv_out_file, purpose, image_url, ner_rest_endpoint, ned_rest_endpoint, noproxy, scale_factor, - ned_threshold, min_confidence, max_confidence): + ned_threshold, min_confidence, max_confidence, ned_priority): if purpose == "NERD": out_columns = ['No.', 'TOKEN', 'NE-TAG', 'NE-EMB', 'ID', 'url_id', 'left', 'right', 'top', 'bottom', 'conf'] elif purpose == "OCR": @@ -128,7 +129,8 @@ def page2tsv(page_xml_file, tsv_out_file, purpose, image_url, ner_rest_endpoint, if min_confidence is not None and max_confidence is not None: line_info['ocrconf'] = line_info.conf.map(lambda x: get_conf_color(x, min_confidence, max_confidence)) - tsv = pd.DataFrame(tsv, columns=['rid', 'line', 'hcenter'] + ['TEXT', 'url_id', 'left', 'right', 'top', 'bottom', 'line_id']) + tsv = pd.DataFrame(tsv, columns=['rid', 'line', 'hcenter'] + + ['TEXT', 'url_id', 'left', 'right', 'top', 'bottom', 'line_id']) if len(tsv) == 0: return @@ -173,7 +175,7 @@ def page2tsv(page_xml_file, tsv_out_file, purpose, image_url, ner_rest_endpoint, if ned_rest_endpoint is not None: - tsv, _ = ned(tsv, ner_result, ned_rest_endpoint, threshold=ned_threshold) + tsv, _ = ned(tsv, ner_result, ned_rest_endpoint, threshold=ned_threshold, priority=ned_priority) tsv.to_csv(tsv_out_file, sep="\t", quoting=3, index=False, mode='a', header=False) except requests.HTTPError as e: