| # Detectron2 + DocLayNet | |
| Model made for document layout analysis | |
| ## Load the model | |
| First install the required dependencies: | |
| ```bash | |
| pip install -r requirements.txt | |
| ``` | |
| In a `.py` or `.ipynb` file: | |
| ```python | |
| 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() | |
| ``` | |