Spaces:
Runtime error
Runtime error
| import os | |
| os.system('!python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.10/index.html') | |
| import cv2 | |
| 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 | |
| def add_vit_config(cfg): | |
| """ | |
| Add config for VIT. | |
| """ | |
| _C = cfg | |
| _C.MODEL.VIT = CN() | |
| # CoaT model name. | |
| _C.MODEL.VIT.NAME = "" | |
| # Output features from CoaT backbone. | |
| _C.MODEL.VIT.OUT_FEATURES = ["layer3", "layer5", "layer7", "layer11"] | |
| _C.MODEL.VIT.IMG_SIZE = [224, 224] | |
| _C.MODEL.VIT.POS_TYPE = "shared_rel" | |
| _C.MODEL.VIT.DROP_PATH = 0. | |
| _C.MODEL.VIT.MODEL_KWARGS = "{}" | |
| _C.SOLVER.OPTIMIZER = "ADAMW" | |
| _C.SOLVER.BACKBONE_MULTIPLIER = 1.0 | |
| _C.AUG = CN() | |
| _C.AUG.DETR = False | |
| # 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 =[['publaynet_example.jpeg']] | |
| 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) |