File size: 4,457 Bytes
c147abc |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
"""Simple DocLayout model for inference."""
import json
from pathlib import Path
from typing import Dict, List, Union
import numpy as np
from PIL import Image
from ultralytics import YOLO
class DocLayoutModel:
"""
Document layout detection model.
Examples
--------
>>> model = DocLayoutModel("model.pt")
>>> results = model.predict("document.png")
>>> for det in results:
... print(f"{det['class_name']}: {det['confidence']:.2f}")
"""
# Default class mappings
DOCSTRUCTBENCH_CLASSES = {
0: "title",
1: "plain_text",
2: "abandon",
3: "figure",
4: "figure_caption",
5: "table",
6: "table_caption",
7: "table_footnote",
8: "isolate_formula",
9: "formula_caption",
}
DOCLAYNET_CLASSES = {
0: "Caption",
1: "Footnote",
2: "Formula",
3: "List-item",
4: "Page-footer",
5: "Page-header",
6: "Picture",
7: "Section-header",
8: "Table",
9: "Text",
10: "Title",
}
def __init__(
self,
weights_path: Union[str, Path],
config_path: Union[str, Path, None] = None,
model_type: str = "auto",
):
"""
Initialize model.
Parameters
----------
weights_path : str or Path
Path to model weights (.pt file)
config_path : str or Path, optional
Path to config.json with class names. If None, auto-detects from weights filename.
model_type : str, default="auto"
Model type: "docstructbench", "doclaynet", or "auto" (detect from filename)
"""
self.weights_path = Path(weights_path)
self._model = None
# Load class names from config or auto-detect
if config_path:
with open(config_path) as f:
config = json.load(f)
self.class_names = {i: name for i, name in enumerate(config["class_names"])}
else:
self.class_names = self._get_class_names(model_type)
def _get_class_names(self, model_type: str) -> Dict[int, str]:
"""Get class names based on model type."""
if model_type == "auto":
name = self.weights_path.stem.lower()
if "doclaynet" in name:
return self.DOCLAYNET_CLASSES
return self.DOCSTRUCTBENCH_CLASSES
elif model_type == "doclaynet":
return self.DOCLAYNET_CLASSES
elif model_type == "docstructbench":
return self.DOCSTRUCTBENCH_CLASSES
else:
raise ValueError(f"Unknown model type: {model_type}")
@property
def model(self) -> YOLO:
"""Lazy-load the YOLO model."""
if self._model is None:
self._model = YOLO(str(self.weights_path))
return self._model
def predict(
self,
source: Union[str, Path, Image.Image, np.ndarray],
confidence: float = 0.2,
image_size: int = 1024,
device: str = "cpu",
) -> List[Dict]:
"""
Run inference on an image.
Parameters
----------
source : str, Path, PIL.Image, or np.ndarray
Input image
confidence : float, default=0.2
Confidence threshold
image_size : int, default=1024
Input image size
device : str, default="cpu"
Device to run on ("cpu", "cuda", "mps")
Returns
-------
List[Dict]
List of detections, each with keys:
- class_id: int
- class_name: str
- confidence: float
- bbox: [x1, y1, x2, y2]
"""
results = self.model.predict(
source=str(source) if isinstance(source, Path) else source,
imgsz=image_size,
conf=confidence,
device=device,
save=False,
verbose=False,
)
detections = []
for result in results:
for box in result.boxes:
cls = int(box.cls[0])
detections.append(
{
"class_id": cls,
"class_name": self.class_names.get(cls, f"class_{cls}"),
"confidence": float(box.conf[0]),
"bbox": box.xyxy[0].tolist(),
}
)
return detections
|