| from unilm.dit.object_detection.ditod import add_vit_config | |
| import torch | |
| 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 | |
| cfg = get_cfg() | |
| add_vit_config(cfg) | |
| cfg.merge_from_file("cascade_dit_base.yml") | |
| cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth" | |
| cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| 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 = " Table Detection with DiT" | |
| css = ".output-image, .input-image, .image-preview {height: 600px !important}" | |
| iface = gr.Interface( | |
| fn=analyze_image, | |
| inputs=[gr.Image(type="numpy", label="document image")], | |
| outputs=[gr.Image(type="numpy", label="detected tables")], | |
| title=title, | |
| css=css, | |
| ) | |
| iface.launch(debug=True, share=True) | |