57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import layoutparser as lp
|
|
import numpy as np
|
|
|
|
CLASS_MAP = {
|
|
0: "Text", 1: "Title", 2: "List",
|
|
3: "Table", 4: "Figure", 5: "Advertisement",
|
|
}
|
|
|
|
class LayoutDetector:
|
|
def __init__(self):
|
|
try:
|
|
# Detectron2 PubLayNet — swap with custom DocLayNet/YOLOv11 weights
|
|
self.model = lp.Detectron2LayoutModel(
|
|
config_path="lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config",
|
|
extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.55],
|
|
label_map={0:"Text",1:"Title",2:"List",3:"Table",4:"Figure"},
|
|
)
|
|
except Exception:
|
|
self.model = None
|
|
# Optional secondary ad classifier (custom-trained)
|
|
self.ad_classifier = AdClassifier()
|
|
|
|
def detect(self, image: np.ndarray):
|
|
if self.model is None:
|
|
# Fallback to single block if detectron2 isn't installed
|
|
h, w = image.shape[:2]
|
|
return [{
|
|
"type": "Text",
|
|
"bbox": [0, 0, w, h],
|
|
"score": 1.0,
|
|
"is_ad": False
|
|
}]
|
|
|
|
layout = self.model.detect(image)
|
|
blocks = []
|
|
for b in layout:
|
|
box = b.coordinates # (x1,y1,x2,y2)
|
|
block = {
|
|
"type": b.type,
|
|
"bbox": [int(c) for c in box],
|
|
"score": float(b.score),
|
|
}
|
|
block["is_ad"] = self.ad_classifier.predict(image, box)
|
|
blocks.append(block)
|
|
return blocks
|
|
|
|
|
|
class AdClassifier:
|
|
"""
|
|
Ads rarely match column grids and contain logos/price patterns.
|
|
Train a lightweight CNN (MobileNet) on labeled ad/non-ad crops,
|
|
OR use heuristics + LLM verification. Stub below.
|
|
"""
|
|
def predict(self, image, box) -> bool:
|
|
# Heuristic placeholder: replace with trained model
|
|
return False
|