fendiprime/dit-document-layout-analysis

1import os2os.system('git clone https://github.com/facebookresearch/detectron2.git')3os.system('pip install -e detectron2')4os.system("git clone https://github.com/microsoft/unilm.git")5os.system("sed -i 's/from collections import Iterable/from collections.abc import Iterable/' unilm/dit/object_detection/ditod/table_evaluation/data_structure.py")6os.system("curl -LJ -o publaynet_dit-b_cascade.pth 'https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_cascade.pth?sv=2022-11-02&ss=b&srt=o&sp=r&se=2033-06-08T16:48:15Z&st=2023-06-08T08:48:15Z&spr=https&sig=a9VXrihTzbWyVfaIDlIT1Z0FoR1073VB0RLQUMuudD4%3D'")7 8import sys9sys.path.append("unilm")10sys.path.append("detectron2")11 12import cv213 14import torch15 16from collections.abc import Iterable as Iterable17from detectron2.config import CfgNode as CN18from detectron2.config import get_cfg19from detectron2.utils.visualizer import ColorMode, Visualizer20from detectron2.data import MetadataCatalog21from detectron2.engine import DefaultPredictor22 23from unilm.dit.object_detection.ditod import add_vit_config24 25import gradio as gr26 27 28# Step 1: instantiate config29cfg = get_cfg()30add_vit_config(cfg)31cfg.merge_from_file("cascade_dit_base.yml")32 33# Step 2: add model weights URL to config34cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth"35 36# Step 3: set device37cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"38 39# Step 4: define model40predictor = DefaultPredictor(cfg)41 42 43def analyze_image(img):44    md = MetadataCatalog.get(cfg.DATASETS.TEST[0])45    if cfg.DATASETS.TEST[0]=='icdar2019_test':46        md.set(thing_classes=["table"])47    else:48        md.set(thing_classes=["text","title","list","table","figure"])49    50    output = predictor(img)["instances"]51    v = Visualizer(img[:, :, ::-1],52                    md,53                    scale=1.0,54                    instance_mode=ColorMode.SEGMENTATION)55    result = v.draw_instance_predictions(output.to("cpu"))56    result_image = result.get_image()[:, :, ::-1]57    58    return result_image59    60title = "Interactive demo: Document Layout Analysis with DiT"61description = "Demo for Microsoft's DiT, the Document Image Transformer for state-of-the-art document understanding tasks. This particular model is fine-tuned on PubLayNet, a large dataset for document layout analysis (read more at the links below). To use it, simply upload an image or use the example image below and click 'Submit'. Results will show up in a few seconds. If you want to make the output bigger, right-click on it and select 'Open image in new tab'."62article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2203.02378' target='_blank'>Paper</a> | <a href='https://github.com/microsoft/unilm/tree/master/dit' target='_blank'>Github Repo</a></p> | <a href='https://huggingface.co/docs/transformers/master/en/model_doc/dit' target='_blank'>HuggingFace doc</a></p>"63examples =[['publaynet_example.jpeg']]64css = ".output-image, .input-image, .image-preview {height: 600px !important}"65 66iface = gr.Interface(fn=analyze_image, 67                     inputs=gr.inputs.Image(type="numpy", label="document image"), 68                     outputs=gr.outputs.Image(type="numpy", label="annotated document"),69                     title=title,70                     description=description,71                     examples=examples,72                     article=article,73                     css=css,74                     enable_queue=True)75iface.launch(debug=True, cache_examples=True)