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()
```