File size: 3,013 Bytes
9c93bcf cca303a 690f166 c3ee9a5 b768f7e 690f166 b47dc23 b768f7e c3ee9a5 7777d8a c3ee9a5 b06234a c3ee9a5 b972826 c3ee9a5 b47dc23 c3ee9a5 b39aaad d694199 5b6b34d 8703dc5 c3ee9a5 a393dd3 c3ee9a5 b39aaad 798398e c3ee9a5 8b94057 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import os
os.system('pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html')
os.system("git clone https://github.com/microsoft/unilm.git")
import sys
sys.path.append("unilm")
import cv2
from unilm.dit.object_detection.ditod import add_vit_config
import torch
from detectron2.config import CfgNode as CN
from detectron2.config import get_cfg
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor
import gradio as gr
# Step 1: instantiate config
cfg = get_cfg()
add_vit_config(cfg)
cfg.merge_from_file("cascade_dit_base.yml")
# Step 2: add model weights URL to config
cfg.MODEL.WEIGHTS = "https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_cascade.pth"
# Step 3: set device
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Step 4: define model
predictor = DefaultPredictor(cfg)
def analyze_image(img):
md = MetadataCatalog.get(cfg.DATASETS.TEST[0])
if cfg.DATASETS.TEST[0]=='icdar2019_test':
md.set(thing_classes=["table"])
else:
md.set(thing_classes=["text","title","list","table","figure"])
output = predictor(img)["instances"]
v = Visualizer(img[:, :, ::-1],
md,
scale=1.0,
instance_mode=ColorMode.SEGMENTATION)
result = v.draw_instance_predictions(output.to("cpu"))
result_image = result.get_image()[:, :, ::-1]
return result_image
title = "Interactive demo: Document Layout Analysis with DiT"
description = "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'."
article = "<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>"
examples =[['publaynet_example.jpeg']]
css = ".output-image, .input-image, .image-preview {height: 600px !important}"
iface = gr.Interface(fn=analyze_image,
inputs=gr.inputs.Image(type="numpy", label="document image"),
outputs=gr.outputs.Image(type="numpy", label="annotated document"),
title=title,
description=description,
examples=examples,
article=article,
css=css,
enable_queue=True)
iface.launch(debug=True, cache_examples=True) |