81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
import os
|
|
|
|
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
|
|
from gradio.themes.utils import sizes
|
|
|
|
|
|
theme = gr.themes.Default(radius_size=sizes.radius_none).set(
|
|
block_label_text_color = '#4D63FF',
|
|
block_title_text_color = '#4D63FF',
|
|
button_primary_text_color = '#4D63FF',
|
|
button_primary_background_fill='#FFFFFF',
|
|
button_primary_border_color='#4D63FF',
|
|
button_primary_background_fill_hover='#EDEFFF',
|
|
)
|
|
|
|
|
|
# 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
|
|
|
|
|
|
with gr.Blocks(theme=theme, css="footer {visibility: hidden}") as demo:
|
|
gr.Markdown("""
|
|
<div align='center' ><font size='60'>版面分析</font></div>
|
|
""")
|
|
with gr.Row():
|
|
with gr.Column():
|
|
image = gr.Image(label="图片", type="numpy")
|
|
with gr.Row():
|
|
button = gr.Button("提交", variant="primary")
|
|
box2 = gr.Image(label="图片", type="numpy")
|
|
|
|
button.click(fn=analyze_image, inputs=[image], outputs=box2)
|
|
examples = gr.Examples(examples=[['publaynet_example.jpeg']], inputs=[image], label="例子")
|
|
|
|
|
|
demo.launch(server_name = "0.0.0.0")
|