|
import os |
|
import importlib.metadata |
|
from os import getcwd, path, environ |
|
from dotenv import load_dotenv |
|
import json |
|
|
|
|
|
def check_additional_requirements(): |
|
if importlib.util.find_spec("detectron2") is None: |
|
os.system('pip install detectron2@git+https://github.com/facebookresearch/detectron2.git') |
|
if importlib.util.find_spec("gradio") is not None: |
|
if importlib.metadata.version("gradio")!="3.44.3": |
|
os.system("pip uninstall -y gradio") |
|
os.system("pip install gradio==3.44.3") |
|
else: |
|
os.system("pip install gradio==3.44.3") |
|
return |
|
|
|
|
|
load_dotenv() |
|
check_additional_requirements() |
|
|
|
import deepdoctection as dd |
|
from deepdoctection.dataflow.serialize import DataFromList |
|
import time |
|
import gradio as gr |
|
from botocore.config import Config |
|
|
|
|
|
|
|
_DD_ONE = "conf_dd_one.yaml" |
|
|
|
dd.ModelCatalog.register("xrf_layout/model_final_inf_only.pt",dd.ModelProfile( |
|
name="xrf_layout/model_final_inf_only.pt", |
|
description="layout_detection/morning-dragon-114", |
|
config="xrf_dd/layout/CASCADE_RCNN_R_50_FPN_GN.yaml", |
|
size=[274632215], |
|
tp_model=False, |
|
hf_repo_id=environ.get("HF_REPO_LAYOUT"), |
|
hf_model_name="model_final_inf_only.pt", |
|
hf_config_file=["Base-RCNN-FPN.yaml", "CASCADE_RCNN_R_50_FPN_GN.yaml"], |
|
categories={"1": dd.LayoutType.text, |
|
"2": dd.LayoutType.title, |
|
"3": dd.LayoutType.list, |
|
"4": dd.LayoutType.table, |
|
"5": dd.LayoutType.figure}, |
|
model_wrapper="D2FrcnnDetector", |
|
)) |
|
|
|
dd.ModelCatalog.register("xrf_cell/model_final_inf_only.pt", dd.ModelProfile( |
|
name="xrf_cell/model_final_inf_only.pt", |
|
description="cell_detection/restful-eon-6", |
|
config="xrf_dd/cell/CASCADE_RCNN_R_50_FPN_GN.yaml", |
|
size=[274583063], |
|
tp_model=False, |
|
hf_repo_id=environ.get("HF_REPO_CELL"), |
|
hf_model_name="model_final_inf_only.pt", |
|
hf_config_file=["Base-RCNN-FPN.yaml", "CASCADE_RCNN_R_50_FPN_GN.yaml"], |
|
categories={"1": dd.LayoutType.cell}, |
|
model_wrapper="D2FrcnnDetector", |
|
)) |
|
|
|
dd.ModelCatalog.register("xrf_item/model_final_inf_only.pt", dd.ModelProfile( |
|
name="xrf_item/model_final_inf_only.pt", |
|
description="item_detection/firm_plasma_14", |
|
config="xrf_dd/item/CASCADE_RCNN_R_50_FPN_GN.yaml", |
|
size=[274595351], |
|
tp_model=False, |
|
hf_repo_id=environ.get("HF_REPO_ITEM"), |
|
hf_model_name="model_final_inf_only.pt", |
|
hf_config_file=["Base-RCNN-FPN.yaml", "CASCADE_RCNN_R_50_FPN_GN.yaml"], |
|
categories={"1": dd.LayoutType.row, "2": dd.LayoutType.column}, |
|
model_wrapper="D2FrcnnDetector", |
|
)) |
|
|
|
|
|
|
|
cfg = dd.set_config_by_yaml(path.join(getcwd(),_DD_ONE)) |
|
cfg.freeze(freezed=False) |
|
cfg.DEVICE = "cpu" |
|
cfg.freeze() |
|
|
|
|
|
layout_config_path = dd.ModelCatalog.get_full_path_configs(cfg.CONFIG.D2LAYOUT) |
|
layout_weights_path = dd.ModelDownloadManager.maybe_download_weights_and_configs(cfg.WEIGHTS.D2LAYOUT) |
|
categories_layout = dd.ModelCatalog.get_profile(cfg.WEIGHTS.D2LAYOUT).categories |
|
assert categories_layout is not None |
|
assert layout_weights_path is not None |
|
d_layout = dd.D2FrcnnDetector(layout_config_path, layout_weights_path, categories_layout, device=cfg.DEVICE) |
|
|
|
|
|
cell_config_path = dd.ModelCatalog.get_full_path_configs(cfg.CONFIG.D2CELL) |
|
cell_weights_path = dd.ModelDownloadManager.maybe_download_weights_and_configs(cfg.WEIGHTS.D2CELL) |
|
categories_cell = dd.ModelCatalog.get_profile(cfg.WEIGHTS.D2CELL).categories |
|
assert categories_cell is not None |
|
d_cell = dd.D2FrcnnDetector(cell_config_path, cell_weights_path, categories_cell, device=cfg.DEVICE) |
|
|
|
|
|
item_config_path = dd.ModelCatalog.get_full_path_configs(cfg.CONFIG.D2ITEM) |
|
item_weights_path = dd.ModelDownloadManager.maybe_download_weights_and_configs(cfg.WEIGHTS.D2ITEM) |
|
categories_item = dd.ModelCatalog.get_profile(cfg.WEIGHTS.D2ITEM).categories |
|
assert categories_item is not None |
|
d_item = dd.D2FrcnnDetector(item_config_path, item_weights_path, categories_item, device=cfg.DEVICE) |
|
|
|
|
|
|
|
credentials_kwargs={"aws_access_key_id": os.environ["ACCESS_KEY"], |
|
"aws_secret_access_key": os.environ["SECRET_KEY"], |
|
"config": Config(region_name=os.environ["REGION"])} |
|
tex_text = dd.TextractOcrDetector(**credentials_kwargs) |
|
|
|
|
|
def build_gradio_analyzer(): |
|
"""Building the Detectron2/DocTr analyzer based on the given config""" |
|
|
|
cfg.freeze(freezed=False) |
|
cfg.TAB = True |
|
cfg.TAB_REF = True |
|
cfg.OCR = True |
|
cfg.freeze() |
|
|
|
pipe_component_list = [] |
|
layout = dd.ImageLayoutService(d_layout, to_image=True, crop_image=True) |
|
pipe_component_list.append(layout) |
|
|
|
nms_service = dd.AnnotationNmsService(nms_pairs=cfg.LAYOUT_NMS_PAIRS.COMBINATIONS, |
|
thresholds=cfg.LAYOUT_NMS_PAIRS.THRESHOLDS) |
|
pipe_component_list.append(nms_service) |
|
|
|
if cfg.TAB: |
|
|
|
detect_result_generator = dd.DetectResultGenerator(categories_cell) |
|
cell = dd.SubImageLayoutService(d_cell, dd.LayoutType.table, {1: 6}, detect_result_generator) |
|
pipe_component_list.append(cell) |
|
|
|
detect_result_generator = dd.DetectResultGenerator(categories_item) |
|
item = dd.SubImageLayoutService(d_item, dd.LayoutType.table, {1: 7, 2: 8}, detect_result_generator) |
|
pipe_component_list.append(item) |
|
|
|
table_segmentation = dd.TableSegmentationService( |
|
cfg.SEGMENTATION.ASSIGNMENT_RULE, |
|
cfg.SEGMENTATION.THRESHOLD_ROWS, |
|
cfg.SEGMENTATION.THRESHOLD_COLS, |
|
cfg.SEGMENTATION.FULL_TABLE_TILING, |
|
cfg.SEGMENTATION.REMOVE_IOU_THRESHOLD_ROWS, |
|
cfg.SEGMENTATION.REMOVE_IOU_THRESHOLD_COLS, |
|
dd.LayoutType.table, |
|
[dd.CellType.header, dd.CellType.body, dd.LayoutType.cell], |
|
[dd.LayoutType.row, dd.LayoutType.column], |
|
[dd.CellType.row_number, dd.CellType.column_number], |
|
cfg.SEGMENTATION.STRETCH_RULE |
|
) |
|
pipe_component_list.append(table_segmentation) |
|
|
|
if cfg.TAB_REF: |
|
table_segmentation_refinement = dd.TableSegmentationRefinementService() |
|
pipe_component_list.append(table_segmentation_refinement) |
|
|
|
if cfg.OCR: |
|
|
|
t_text = dd.TextExtractionService(tex_text) |
|
pipe_component_list.append(t_text) |
|
|
|
match_words = dd.MatchingService( |
|
parent_categories=cfg.WORD_MATCHING.PARENTAL_CATEGORIES, |
|
child_categories=cfg.WORD_MATCHING.CHILD_CATEGORIES, |
|
matching_rule=cfg.WORD_MATCHING.RULE, |
|
threshold=cfg.WORD_MATCHING.THRESHOLD, |
|
max_parent_only=cfg.WORD_MATCHING.MAX_PARENT_ONLY |
|
) |
|
pipe_component_list.append(match_words) |
|
|
|
order = dd.TextOrderService( |
|
text_container=cfg.TEXT_ORDERING.TEXT_CONTAINER, |
|
floating_text_block_categories=cfg.TEXT_ORDERING.FLOATING_TEXT_BLOCK, |
|
text_block_categories=cfg.TEXT_ORDERING.TEXT_BLOCK, |
|
include_residual_text_container=cfg.TEXT_ORDERING.TEXT_CONTAINER_TO_TEXT_BLOCK) |
|
pipe_component_list.append(order) |
|
|
|
pipe = dd.DoctectionPipe(pipeline_component_list=pipe_component_list) |
|
|
|
return pipe |
|
|
|
|
|
def analyze_image(img, pdf, max_datapoints): |
|
|
|
|
|
analyzer = build_gradio_analyzer() |
|
|
|
if img is not None: |
|
image = dd.Image(file_name=str(time.time()).replace(".","") + ".png", location="") |
|
image.image = img[:, :, ::-1] |
|
|
|
df = DataFromList(lst=[image]) |
|
df = analyzer.analyze(dataset_dataflow=df) |
|
elif pdf: |
|
df = analyzer.analyze(path=pdf.name, max_datapoints=max_datapoints) |
|
else: |
|
raise ValueError |
|
|
|
df.reset_state() |
|
|
|
layout_items_str = "" |
|
jsonl_out = [] |
|
dpts = [] |
|
html_list = [] |
|
|
|
for dp in df: |
|
dpts.append(dp) |
|
out = dp.as_dict() |
|
jsonl_out.append(out) |
|
out.pop("_image") |
|
layout_items = [layout for layout in dp.layouts if layout.reading_order is not None] |
|
layout_items.sort(key=lambda x: x.reading_order) |
|
layout_items_str += f"\n\n -------- PAGE NUMBER: {dp.page_number+1} ------------- \n" |
|
for item in layout_items: |
|
layout_items_str += f"\n {item.category_name}: {item.text}" |
|
html_list.extend([table.html for table in dp.tables]) |
|
if html_list: |
|
html = ("<br /><br /><br />").join(html_list) |
|
else: |
|
html = None |
|
json_object = json.dumps(jsonl_out, indent = 4) |
|
return [dp.viz(show_cells=False) for dp in dpts], layout_items_str, html, json_object |
|
|
|
|
|
demo = gr.Blocks(css="scrollbar.css") |
|
|
|
|
|
with demo: |
|
with gr.Box(): |
|
gr.Markdown("<h1><center>deepdoctection - A Document AI Package</center></h1>") |
|
gr.Markdown("<strong>deep</strong>doctection is a Python library that orchestrates document extraction" |
|
" and document layout analysis tasks using deep learning models. It does not implement models" |
|
" but enables you to build pipelines using highly acknowledged libraries for object detection," |
|
" OCR and selected NLP tasks and provides an integrated frameworks for fine-tuning, evaluating" |
|
" and running models.<br />" |
|
"This pipeline consists of a stack of models powered by <strong>Detectron2" |
|
"</strong> for layout analysis and table recognition. OCR will be provided as well. You can process" |
|
"an image or even a PDF-document. Up to nine pages can be processed. <br />") |
|
gr.Markdown("<center><strong>Please note:</strong> The models for layout detection and table recognition are not open sourced. " |
|
"When you start using deepdoctection you will get models that have been trained on less diversified data and that will perform worse. " |
|
"OCR isn't open sourced either: It uses AWS Textract, which is a commercial service. Keep this in mind, before you get started with " |
|
"your installation and observe dissapointing results. Thanks. </center>") |
|
gr.Markdown("[https://github.com/deepdoctection/deepdoctection](https://github.com/deepdoctection/deepdoctection)") |
|
with gr.Box(): |
|
gr.Markdown("<h2><center>Upload a document and choose setting</center></h2>") |
|
with gr.Row(): |
|
with gr.Column(): |
|
with gr.Tab("Image upload"): |
|
with gr.Column(): |
|
inputs = gr.Image(type='numpy', label="Original Image") |
|
with gr.Tab("PDF upload (only first image will be processed) *"): |
|
with gr.Column(): |
|
inputs_pdf = gr.File(label="PDF") |
|
gr.Markdown("<sup>* If an image is cached in tab, remove it first</sup>") |
|
with gr.Column(): |
|
gr.Examples( |
|
examples=[path.join(getcwd(), "sample_1.jpg"), path.join(getcwd(), "sample_2.png")], |
|
inputs = inputs) |
|
gr.Examples(examples=[path.join(getcwd(), "sample_3.pdf")], inputs = inputs_pdf) |
|
|
|
with gr.Row(): |
|
max_imgs = gr.Slider(1, 8, value=2, step=1, label="Number of pages in multi page PDF", |
|
info="Will stop after 9 pages") |
|
|
|
with gr.Row(): |
|
btn = gr.Button("Run model", variant="primary") |
|
|
|
with gr.Box(): |
|
gr.Markdown("<h2><center>Outputs</center></h2>") |
|
with gr.Row(): |
|
with gr.Column(): |
|
with gr.Box(): |
|
gr.Markdown("<center><strong>Contiguous text</strong></center>") |
|
image_text = gr.Textbox() |
|
with gr.Column(): |
|
with gr.Box(): |
|
gr.Markdown("<center><strong>Layout detection</strong></center>") |
|
gallery = gr.Gallery( |
|
label="Output images", show_label=False, elem_id="gallery" |
|
).style(grid=2) |
|
with gr.Row(): |
|
with gr.Box(): |
|
gr.Markdown("<center><strong>Table</strong></center>") |
|
html = gr.HTML() |
|
|
|
with gr.Row(): |
|
with gr.Box(): |
|
gr.Markdown("<center><strong>JSON</strong></center>") |
|
json_output = gr.JSON() |
|
|
|
btn.click(fn=analyze_image, inputs=[inputs, inputs_pdf, max_imgs], |
|
outputs=[gallery, image_text, html, json_output]) |
|
|
|
demo.launch() |
|
|