171 lines
7.4 KiB
Python
171 lines
7.4 KiB
Python
import math
|
|
|
|
def get_horizontal_overlap(bbox1, bbox2):
|
|
x_left = max(bbox1[0], bbox2[0])
|
|
x_right = min(bbox1[2], bbox2[2])
|
|
return max(0, x_right - x_left)
|
|
|
|
def group_surya_regions_geometrically(regions):
|
|
"""
|
|
Groups regions spatially.
|
|
Instead of a linear reading order, this associates text and images
|
|
with the headline that is directly above them in the same column.
|
|
"""
|
|
if not regions:
|
|
return []
|
|
|
|
titles = [r for r in regions if r['type'] in ('doc_title', 'paragraph_title')]
|
|
non_titles = [r for r in regions if r['type'] not in ('doc_title', 'paragraph_title')]
|
|
|
|
# If no titles exist, return everything as one article
|
|
if not titles:
|
|
member_ids = [r['id'] for r in regions]
|
|
min_x = min(r['bbox'][0] for r in regions)
|
|
min_y = min(r['bbox'][1] for r in regions)
|
|
max_x = max(r['bbox'][2] for r in regions)
|
|
max_y = max(r['bbox'][3] for r in regions)
|
|
return [{
|
|
"article_id": 1,
|
|
"title_region": None,
|
|
"member_regions": member_ids,
|
|
"bbox": [min_x, min_y, max_x, max_y],
|
|
"grouped_by": "geometric_fallback"
|
|
}]
|
|
|
|
article_groups = {t['id']: [t] for t in titles}
|
|
|
|
for item in non_titles:
|
|
item_bbox = item['bbox']
|
|
item_cy = (item_bbox[1] + item_bbox[3]) / 2
|
|
item_cx = (item_bbox[0] + item_bbox[2]) / 2
|
|
item_width = item_bbox[2] - item_bbox[0]
|
|
|
|
best_title = None
|
|
|
|
# -------------------------------------------------------------
|
|
# Rule 1: Find titles ABOVE the item that share the same column
|
|
# -------------------------------------------------------------
|
|
candidates_above = []
|
|
for t in titles:
|
|
t_bbox = t['bbox']
|
|
if t_bbox[1] <= item_cy:
|
|
overlap = get_horizontal_overlap(item_bbox, t_bbox)
|
|
t_width = t_bbox[2] - t_bbox[0]
|
|
t_cx = (t_bbox[0] + t_bbox[2]) / 2
|
|
|
|
if overlap > (item_width * 0.1) or overlap > (t_width * 0.1) or abs(item_cx - t_cx) < 150:
|
|
distance = item_bbox[1] - t_bbox[3]
|
|
if distance < 1500: # MAX DISTANCE ABOVE
|
|
candidates_above.append(t)
|
|
|
|
if candidates_above:
|
|
best_title = min(candidates_above, key=lambda t: abs(item_bbox[1] - t['bbox'][3]))
|
|
else:
|
|
# -------------------------------------------------------------
|
|
# Rule 2: If no title is above, find titles BELOW the item
|
|
# -------------------------------------------------------------
|
|
candidates_below = []
|
|
for t in titles:
|
|
t_bbox = t['bbox']
|
|
if t_bbox[1] > item_cy:
|
|
overlap = get_horizontal_overlap(item_bbox, t_bbox)
|
|
t_width = t_bbox[2] - t_bbox[0]
|
|
t_cx = (t_bbox[0] + t_bbox[2]) / 2
|
|
|
|
if overlap > (item_width * 0.1) or overlap > (t_width * 0.1) or abs(item_cx - t_cx) < 150:
|
|
distance = t_bbox[1] - item_bbox[3]
|
|
if distance < 800: # MAX DISTANCE BELOW (stricter because text usually flows down, not up)
|
|
candidates_below.append(t)
|
|
|
|
if candidates_below:
|
|
best_title = min(candidates_below, key=lambda t: abs(t['bbox'][1] - item_bbox[3]))
|
|
else:
|
|
# -------------------------------------------------------------
|
|
# Rule 3: Absolute fallback. No column overlap at all.
|
|
# Find the absolute closest title via Euclidean distance.
|
|
# -------------------------------------------------------------
|
|
closest_t = min(titles, key=lambda t: math.hypot(
|
|
item_cx - ((t['bbox'][0] + t['bbox'][2]) / 2),
|
|
item_cy - ((t['bbox'][1] + t['bbox'][3]) / 2)
|
|
))
|
|
# Only attach if it's reasonably close, else leave as orphan
|
|
dist = math.hypot(item_cx - ((closest_t['bbox'][0] + closest_t['bbox'][2]) / 2),
|
|
item_cy - ((closest_t['bbox'][1] + closest_t['bbox'][3]) / 2))
|
|
if dist < 1500:
|
|
best_title = closest_t
|
|
|
|
if best_title:
|
|
article_groups[best_title['id']].append(item)
|
|
else:
|
|
# Treat as an orphan (create a new article group for it, or group orphans together)
|
|
# For simplicity, we'll assign it a virtual title ID based on its own ID so it becomes a standalone article
|
|
virtual_id = f"orphan_{item['id']}"
|
|
article_groups[virtual_id] = [item]
|
|
|
|
# Format output
|
|
articles = []
|
|
# Sort groups top-to-bottom, left-to-right based on title position
|
|
sorted_titles = sorted(titles, key=lambda t: (t['bbox'][1], t['bbox'][0]))
|
|
|
|
for idx, (title_id, group_regions) in enumerate(article_groups.items(), start=1):
|
|
if not group_regions:
|
|
continue
|
|
|
|
min_x = min(r['bbox'][0] for r in group_regions)
|
|
min_y = min(r['bbox'][1] for r in group_regions)
|
|
max_x = max(r['bbox'][2] for r in group_regions)
|
|
max_y = max(r['bbox'][3] for r in group_regions)
|
|
|
|
# Determine if it's an orphan group
|
|
is_orphan = str(title_id).startswith("orphan_")
|
|
|
|
articles.append({
|
|
"article_id": idx,
|
|
"title_region": None if is_orphan else title_id,
|
|
"member_regions": [r['id'] for r in group_regions],
|
|
"bbox": [min_x, min_y, max_x, max_y],
|
|
"grouped_by": "geometric_orphan" if is_orphan else "geometric"
|
|
})
|
|
|
|
# We should merge orphans that are vertically adjacent into the same orphan group to prevent shattering
|
|
# Simple vertical merge for orphans
|
|
orphan_articles = [a for a in articles if a['grouped_by'] == 'geometric_orphan']
|
|
valid_articles = [a for a in articles if a['grouped_by'] == 'geometric']
|
|
|
|
# Sort orphans top-to-bottom
|
|
orphan_articles.sort(key=lambda a: a['bbox'][1])
|
|
|
|
merged_orphans = []
|
|
current_orphan = None
|
|
|
|
for o in orphan_articles:
|
|
if not current_orphan:
|
|
current_orphan = o
|
|
continue
|
|
|
|
# Check if they overlap horizontally and are close vertically
|
|
overlap = get_horizontal_overlap(current_orphan['bbox'], o['bbox'])
|
|
vertical_dist = o['bbox'][1] - current_orphan['bbox'][3]
|
|
|
|
if overlap > 0 and vertical_dist < 400:
|
|
# Merge
|
|
current_orphan['member_regions'].extend(o['member_regions'])
|
|
current_orphan['bbox'][0] = min(current_orphan['bbox'][0], o['bbox'][0])
|
|
current_orphan['bbox'][1] = min(current_orphan['bbox'][1], o['bbox'][1])
|
|
current_orphan['bbox'][2] = max(current_orphan['bbox'][2], o['bbox'][2])
|
|
current_orphan['bbox'][3] = max(current_orphan['bbox'][3], o['bbox'][3])
|
|
else:
|
|
merged_orphans.append(current_orphan)
|
|
current_orphan = o
|
|
|
|
if current_orphan:
|
|
merged_orphans.append(current_orphan)
|
|
|
|
final_articles = valid_articles + merged_orphans
|
|
|
|
# Reassign IDs
|
|
for idx, a in enumerate(final_articles, start=1):
|
|
a['article_id'] = idx
|
|
|
|
return final_articles
|