Detectron2 + DocLayNet

Model made for document layout analysis

Load the model

First install the required dependencies:

pip install -r requirements.txt

In a .py or .ipynb file:

import cv2
import json
import matplotlib.pyplot as plt
from detectron2.utils.visualizer import Visualizer
from detectron2.data import Metadata
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor

cfg = get_cfg()
cfg.merge_from_file("config.yml")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set the testing threshold for this model

with open("metadata.json", "r") as f:
    metadata_dict = json.load(f)

predictor = DefaultPredictor(cfg)
metadata = Metadata()
metadata.set(thing_classes=metadata_dict["thing_classes"])
im = cv2.imread("image.jpg")
output = predictor(im)
v = Visualizer(im[:, :, ::-1], metadata=metadata, scale=0.8)
v = v.draw_instance_predictions(output["instances"].to("cpu"))
plt.figure(figsize=(14,10))
plt.imshow(cv2.cvtColor(v.get_image()[:, :, ::-1], cv2.COLOR_BGR2RGB))
plt.show()