Spaces:
Runtime error
Runtime error
| import argparse | |
| import cv2 | |
| from ditod import add_vit_config | |
| from detectron2.config import get_cfg | |
| from detectron2.utils.visualizer import ColorMode, Visualizer | |
| from detectron2.data import MetadataCatalog | |
| from detectron2.engine import DefaultPredictor | |
| # Step 1: instantiate config | |
| cfg = get_cfg() | |
| add_vit_config(cfg) | |
| cfg.merge_from_file("cascade_dit_base.yaml") | |
| # Step 2: add model weights URL to config | |
| cfg.MODEL.WEIGHTS = https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_mrcnn.pth | |
| # Step 3: set device | |
| # TODO also support GPU | |
| cfg.MODEL.DEVICE='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 = "This is a demo for Microsoft's Document Image Transformer (DiT)." | |
| examples =[['document.png']] | |
| iface = gr.Interface(fn=analyze_image, | |
| inputs=gr.inputs.Image(type="numpy"), | |
| outputs=gr.outputs.Image(type="numpy", label="analyzed image"), | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| enable_queue=True) | |
| iface.launch(debug=True) |