File size: 1,058 Bytes
42cca4c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
# 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()
```
|