vitouphy/document-analysis

1import os2os.system('pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html')3os.system("git clone https://github.com/microsoft/unilm.git")4 5import sys6sys.path.append("unilm")7 8import cv29 10from unilm.dit.object_detection.ditod import add_vit_config11 12import torch13 14from detectron2.config import CfgNode as CN15from detectron2.config import get_cfg16from detectron2.utils.visualizer import ColorMode, Visualizer17from detectron2.data import MetadataCatalog18from detectron2.engine import DefaultPredictor19 20import gradio as gr21 22 23# Step 1: instantiate config24cfg = get_cfg()25add_vit_config(cfg)26cfg.merge_from_file("cascade_dit_base.yml")27 28# Step 2: add model weights URL to config29cfg.MODEL.WEIGHTS = "https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_cascade.pth"30 31# Step 3: set device32cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"33 34# Step 4: define model35predictor = DefaultPredictor(cfg)36 37 38def analyze_image(img):39    md = MetadataCatalog.get(cfg.DATASETS.TEST[0])40    if cfg.DATASETS.TEST[0]=='icdar2019_test':41        md.set(thing_classes=["table"])42    else:43        md.set(thing_classes=["text","title","list","table","figure"])44    45    output = predictor(img)["instances"]46    v = Visualizer(img[:, :, ::-1],47                    md,48                    scale=1.0,49                    instance_mode=ColorMode.SEGMENTATION)50    result = v.draw_instance_predictions(output.to("cpu"))51    result_image = result.get_image()[:, :, ::-1]52    53    return result_image54    55title = "Interactive demo: Document Layout Analysis with DiT"56description = "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'."57article = "<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>"58examples =[['publaynet_example.jpeg']]59css = ".output-image, .input-image, .image-preview {height: 600px !important}"60 61iface = gr.Interface(fn=analyze_image, 62                     inputs=gr.inputs.Image(type="numpy", label="document image"), 63                     outputs=gr.outputs.Image(type="numpy", label="annotated document"),64                     title=title,65                     description=description,66                     examples=examples,67                     article=article,68                     css=css,69                     enable_queue=True)70iface.launch(debug=True, cache_examples=True)