106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
import numpy as np
|
|
from scipy.spatial import distance
|
|
|
|
class ArticleSegmenter:
|
|
def __init__(self, col_gap_ratio=0.02):
|
|
self.col_gap_ratio = col_gap_ratio
|
|
|
|
def _vertical_overlap(self, a, b):
|
|
top = max(a[1], b[1]); bot = min(a[3], b[3])
|
|
return max(0, bot - top)
|
|
|
|
def _horizontal_gap(self, a, b):
|
|
return b[0] - a[2]
|
|
|
|
def segment(self, blocks, page_width, page_height):
|
|
# 1. Remove ads
|
|
content = [b for b in blocks if not b["is_ad"] and b["type"] != "Advertisement"]
|
|
|
|
# 2. Sort reading order: column-aware (top-to-bottom within columns)
|
|
content = self._reading_order(content, page_width)
|
|
|
|
# 3. Seed articles at each Title block
|
|
articles = []
|
|
current = None
|
|
for b in content:
|
|
if b["type"] == "Title":
|
|
if current:
|
|
articles.append(current)
|
|
current = {"blocks": [b], "title_block": b}
|
|
else:
|
|
if current is None:
|
|
current = {"blocks": [b], "title_block": None}
|
|
else:
|
|
if self._belongs(current, b, page_width):
|
|
current["blocks"].append(b)
|
|
else:
|
|
articles.append(current)
|
|
current = {"blocks": [b], "title_block": None}
|
|
if current:
|
|
articles.append(current)
|
|
|
|
# 4. Merge cross-column continuation of same article
|
|
articles = self._merge_columns(articles, page_width)
|
|
|
|
# 5. Compute union bounding box per article
|
|
for a in articles:
|
|
a["bbox"] = self._union_bbox(a["blocks"])
|
|
a["confidence"] = float(np.mean([blk["score"] for blk in a["blocks"]]))
|
|
return articles
|
|
|
|
def _reading_order(self, blocks, page_width, n_cols_guess=None):
|
|
if not blocks:
|
|
return blocks
|
|
xs = np.array([b["bbox"][0] for b in blocks]).reshape(-1, 1)
|
|
# Estimate columns via 1D clustering on x-start
|
|
from sklearn.cluster import KMeans
|
|
k = self._estimate_columns(xs, page_width)
|
|
km = KMeans(n_clusters=k, n_init=5).fit(xs)
|
|
for b, lbl in zip(blocks, km.labels_):
|
|
b["col"] = int(lbl)
|
|
# Order columns left→right, blocks top→bottom
|
|
col_order = np.argsort([xs[km.labels_ == c].mean() for c in range(k)])
|
|
ordered = []
|
|
for c in col_order:
|
|
col_blocks = [b for b in blocks if b["col"] == c]
|
|
col_blocks.sort(key=lambda b: b["bbox"][1])
|
|
ordered.extend(col_blocks)
|
|
return ordered
|
|
|
|
def _estimate_columns(self, xs, page_width):
|
|
from sklearn.metrics import silhouette_score
|
|
best_k, best_score = 1, -1
|
|
for k in range(1, min(7, len(xs))):
|
|
if k == 1:
|
|
continue
|
|
from sklearn.cluster import KMeans
|
|
labels = KMeans(n_clusters=k, n_init=3).fit_predict(xs)
|
|
try:
|
|
s = silhouette_score(xs, labels)
|
|
if s > best_score:
|
|
best_score, best_k = s, k
|
|
except Exception:
|
|
pass
|
|
return max(best_k, 1)
|
|
|
|
def _belongs(self, article, block, page_width):
|
|
"""Block belongs to current article if vertically continuous & aligned."""
|
|
last = article["blocks"][-1]["bbox"]
|
|
cur = block["bbox"]
|
|
same_col = abs(last[0] - cur[0]) < 0.05 * page_width
|
|
gap = cur[1] - last[3]
|
|
return same_col and gap < 0.04 * page_width
|
|
|
|
def _merge_columns(self, articles, page_width):
|
|
# Merge article fragments whose headline spans wide / text flows to next col
|
|
# Simplified: if an article has no title and sits at top of next column
|
|
# right after a titled article, merge. Production: use LayoutLMv3 relations.
|
|
return articles
|
|
|
|
def _union_bbox(self, blocks):
|
|
xs1 = min(b["bbox"][0] for b in blocks)
|
|
ys1 = min(b["bbox"][1] for b in blocks)
|
|
xs2 = max(b["bbox"][2] for b in blocks)
|
|
ys2 = max(b["bbox"][3] for b in blocks)
|
|
return [xs1, ys1, xs2 - xs1, ys2 - ys1]
|