393 lines
16 KiB
Python
393 lines
16 KiB
Python
import math
|
|
import cv2
|
|
import numpy as np
|
|
|
|
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 extract_barriers(img_path, regions):
|
|
"""Extracts horizontal and vertical barrier lines using OpenCV."""
|
|
if not img_path:
|
|
return [], []
|
|
|
|
import statistics
|
|
text_widths = [(r['bbox'][2] - r['bbox'][0]) for r in regions if r['type'] == 'text']
|
|
median_col_width = statistics.median(text_widths) if text_widths else 300
|
|
|
|
img = cv2.imread(str(img_path))
|
|
if img is None:
|
|
return [], []
|
|
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
_, thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
|
|
|
|
line_length = int(median_col_width * 0.5)
|
|
h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (line_length, 1))
|
|
h_lines_raw = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, h_kernel, iterations=2)
|
|
|
|
v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, line_length))
|
|
v_lines_raw = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, v_kernel, iterations=2)
|
|
|
|
h_contours_raw, _ = cv2.findContours(h_lines_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
v_contours_raw, _ = cv2.findContours(v_lines_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
h_lines = []
|
|
for cnt in h_contours_raw:
|
|
x, y, w, h = cv2.boundingRect(cnt)
|
|
# Ignore short decorative lines (caption underlines) which are <= 0.6 col
|
|
if h <= 15 and w >= (median_col_width * 0.6):
|
|
h_lines.append((x, y, x + w, y + h))
|
|
|
|
v_lines = []
|
|
for cnt in v_contours_raw:
|
|
x, y, w, h = cv2.boundingRect(cnt)
|
|
if w <= 15: # Thin lines only
|
|
v_lines.append((x, y, x + w, y + h))
|
|
|
|
return h_lines, v_lines
|
|
|
|
def is_blocked_by_horizontal_barrier(top_y, bottom_y, item_left, item_right, h_lines):
|
|
"""
|
|
Checks if a horizontal line sits between top_y and bottom_y and spans across the item.
|
|
"""
|
|
item_width = item_right - item_left
|
|
tolerance = item_width * 0.2
|
|
|
|
for (lx1, ly1, lx2, ly2) in h_lines:
|
|
line_y = ly1
|
|
if top_y < line_y < bottom_y:
|
|
# Does it span across the item?
|
|
if lx1 <= (item_left + tolerance) and lx2 >= (item_right - tolerance):
|
|
return True
|
|
return False
|
|
|
|
def is_blocked_by_vertical_barrier(t_bbox, item_bbox, v_lines):
|
|
"""
|
|
Checks if a vertical line sits between the centers of two regions.
|
|
"""
|
|
t_cx = (t_bbox[0] + t_bbox[2]) / 2
|
|
item_cx = (item_bbox[0] + item_bbox[2]) / 2
|
|
|
|
left_x = min(t_cx, item_cx)
|
|
right_x = max(t_cx, item_cx)
|
|
|
|
# Only check if they are in different columns
|
|
if right_x - left_x < 100:
|
|
return False
|
|
|
|
y_overlap_top = max(t_bbox[1], item_bbox[1])
|
|
y_overlap_bottom = min(t_bbox[3], item_bbox[3])
|
|
|
|
for (vx1, vy1, vx2, vy2) in v_lines:
|
|
line_x = vx1
|
|
if left_x < line_x < right_x:
|
|
# Does the line overlap them vertically?
|
|
# A true vertical column separator will be very tall
|
|
if vy1 < y_overlap_bottom and vy2 > y_overlap_top:
|
|
return True
|
|
return False
|
|
|
|
def group_surya_regions_geometrically(regions, img_path=None):
|
|
if not regions:
|
|
return []
|
|
|
|
# Extract barriers using OpenCV
|
|
h_lines, v_lines = extract_barriers(img_path, regions)
|
|
if img_path:
|
|
print(f" [grouper] Extracted {len(h_lines)} horizontal and {len(v_lines)} vertical OpenCV barriers")
|
|
|
|
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')]
|
|
|
|
region_by_id = {r['id']: r for r in regions}
|
|
|
|
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
|
|
|
|
max_dist_above = 2500 if item['type'] == 'image' else 1500
|
|
candidates_above = []
|
|
for t in titles:
|
|
t_bbox = t['bbox']
|
|
if t_bbox[1] <= item_cy:
|
|
# Barrier Check
|
|
if is_blocked_by_horizontal_barrier(t_bbox[3], item_bbox[1], item_bbox[0], item_bbox[2], h_lines):
|
|
continue
|
|
if is_blocked_by_vertical_barrier(t_bbox, item_bbox, v_lines):
|
|
continue
|
|
|
|
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 < max_dist_above:
|
|
candidates_above.append(t)
|
|
|
|
if candidates_above:
|
|
def _above_score(t):
|
|
vert = item_bbox[1] - t['bbox'][3]
|
|
horiz = abs(item_cx - (t['bbox'][0] + t['bbox'][2]) / 2)
|
|
return vert + horiz * 0.8
|
|
best_title = min(candidates_above, key=_above_score)
|
|
else:
|
|
candidates_below = []
|
|
for t in titles:
|
|
t_bbox = t['bbox']
|
|
if t_bbox[1] > item_cy:
|
|
# Barrier Check
|
|
if is_blocked_by_horizontal_barrier(item_bbox[3], t_bbox[1], item_bbox[0], item_bbox[2], h_lines):
|
|
continue
|
|
if is_blocked_by_vertical_barrier(t_bbox, item_bbox, v_lines):
|
|
continue
|
|
|
|
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:
|
|
candidates_below.append(t)
|
|
|
|
if candidates_below:
|
|
def _below_score(t):
|
|
vert = t['bbox'][1] - item_bbox[3]
|
|
horiz = abs(item_cx - (t['bbox'][0] + t['bbox'][2]) / 2)
|
|
return vert + horiz * 0.8
|
|
best_title = min(candidates_below, key=_below_score)
|
|
else:
|
|
# Absolute fallback
|
|
valid_closest = []
|
|
for t in titles:
|
|
if is_blocked_by_horizontal_barrier(min(item_bbox[3], t['bbox'][3]), max(item_bbox[1], t['bbox'][1]), item_bbox[0], item_bbox[2], h_lines):
|
|
continue
|
|
if is_blocked_by_vertical_barrier(t['bbox'], item_bbox, v_lines):
|
|
continue
|
|
valid_closest.append(t)
|
|
|
|
if valid_closest:
|
|
closest_t = min(valid_closest, key=lambda t: math.hypot(
|
|
item_cx - ((t['bbox'][0] + t['bbox'][2]) / 2),
|
|
item_cy - ((t['bbox'][1] + t['bbox'][3]) / 2)
|
|
))
|
|
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:
|
|
virtual_id = f"orphan_{item['id']}"
|
|
article_groups[virtual_id] = [item]
|
|
|
|
# Format output
|
|
articles = []
|
|
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)
|
|
|
|
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"
|
|
})
|
|
|
|
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']
|
|
|
|
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
|
|
|
|
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:
|
|
# Check barrier before merging orphans
|
|
if not is_blocked_by_horizontal_barrier(current_orphan['bbox'][3], o['bbox'][1], current_orphan['bbox'][0], current_orphan['bbox'][2], h_lines):
|
|
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])
|
|
continue
|
|
|
|
merged_orphans.append(current_orphan)
|
|
current_orphan = o
|
|
|
|
if current_orphan:
|
|
merged_orphans.append(current_orphan)
|
|
|
|
# Orphan Adoption: If an orphan sits directly below a valid article with high overlap and small gap, adopt it (ignoring photo borders)
|
|
adopted_orphans = set()
|
|
for o in merged_orphans:
|
|
best_parent = None
|
|
min_gap = float('inf')
|
|
for a in valid_articles:
|
|
overlap = get_horizontal_overlap(a['bbox'], o['bbox'])
|
|
gap = o['bbox'][1] - a['bbox'][3]
|
|
if overlap > 0 and 0 <= gap < 200:
|
|
width_overlap_ratio = overlap / min((a['bbox'][2] - a['bbox'][0]), (o['bbox'][2] - o['bbox'][0]))
|
|
if width_overlap_ratio > 0.8:
|
|
if gap < min_gap:
|
|
min_gap = gap
|
|
best_parent = a
|
|
if best_parent:
|
|
best_parent['member_regions'].extend(o['member_regions'])
|
|
best_parent['bbox'][0] = min(best_parent['bbox'][0], o['bbox'][0])
|
|
best_parent['bbox'][1] = min(best_parent['bbox'][1], o['bbox'][1])
|
|
best_parent['bbox'][2] = max(best_parent['bbox'][2], o['bbox'][2])
|
|
best_parent['bbox'][3] = max(best_parent['bbox'][3], o['bbox'][3])
|
|
adopted_orphans.add(o['article_id'])
|
|
|
|
final_orphans_pass_1 = [o for o in merged_orphans if o['article_id'] not in adopted_orphans]
|
|
|
|
# Strict Orphan Photo Binding: An image may never stand alone.
|
|
remaining_orphans = []
|
|
for o in final_orphans_pass_1:
|
|
has_image = any(region_by_id[rid]['type'] == 'image' for rid in o['member_regions'] if rid in region_by_id)
|
|
has_title = any(region_by_id[rid]['type'] in ('doc_title', 'paragraph_title') for rid in o['member_regions'] if rid in region_by_id)
|
|
|
|
if has_image and not has_title:
|
|
best_article = None
|
|
min_dist = float('inf')
|
|
|
|
for a in valid_articles:
|
|
dist = min(abs(a['bbox'][3] - o['bbox'][1]), abs(o['bbox'][3] - a['bbox'][1]))
|
|
overlap = get_horizontal_overlap(a['bbox'], o['bbox'])
|
|
if overlap > 0 and dist < min_dist:
|
|
min_dist = dist
|
|
best_article = a
|
|
|
|
if best_article:
|
|
best_article['member_regions'].extend(o['member_regions'])
|
|
best_article['bbox'][0] = min(best_article['bbox'][0], o['bbox'][0])
|
|
best_article['bbox'][1] = min(best_article['bbox'][1], o['bbox'][1])
|
|
best_article['bbox'][2] = max(best_article['bbox'][2], o['bbox'][2])
|
|
best_article['bbox'][3] = max(best_article['bbox'][3], o['bbox'][3])
|
|
continue # successfully bound
|
|
|
|
remaining_orphans.append(o)
|
|
|
|
final_articles = valid_articles + remaining_orphans
|
|
|
|
MAX_SPLIT_HEIGHT = 2500
|
|
MIN_MEMBERS_SPLIT = 5
|
|
MIN_SPLIT_GAP = 30
|
|
MIN_LOWER_HEIGHT = 400
|
|
|
|
split_result = []
|
|
for art in final_articles:
|
|
_, art_t, _, art_b = art['bbox']
|
|
if (art_b - art_t) <= MAX_SPLIT_HEIGHT or len(art['member_regions']) < MIN_MEMBERS_SPLIT:
|
|
split_result.append(art)
|
|
continue
|
|
|
|
members = [region_by_id[rid] for rid in art['member_regions'] if rid in region_by_id]
|
|
members.sort(key=lambda reg: reg['bbox'][1])
|
|
|
|
best_gap = MIN_SPLIT_GAP
|
|
split_idx = -1
|
|
max_bottom_so_far = 0
|
|
for i in range(len(members) - 1):
|
|
cur_bottom = members[i]['bbox'][3]
|
|
max_bottom_so_far = max(max_bottom_so_far, cur_bottom)
|
|
|
|
next_top = members[i + 1]['bbox'][1]
|
|
gap = next_top - max_bottom_so_far
|
|
|
|
# If there is a barrier line crossing this gap, artificially inflate the gap score to force a split here
|
|
is_barrier = is_blocked_by_horizontal_barrier(max_bottom_so_far, next_top, members[i]['bbox'][0], members[i]['bbox'][2], h_lines)
|
|
effective_gap = gap + 10000 if is_barrier else gap
|
|
|
|
if effective_gap > best_gap:
|
|
lower_group = members[i + 1:]
|
|
if not any(reg['type'] == 'image' for reg in lower_group):
|
|
best_gap = effective_gap
|
|
split_idx = i
|
|
|
|
if split_idx < 0:
|
|
split_result.append(art)
|
|
continue
|
|
|
|
upper = members[:split_idx + 1]
|
|
lower = members[split_idx + 1:]
|
|
|
|
lower_bottom = max(reg['bbox'][3] for reg in lower)
|
|
if lower_bottom - lower[0]['bbox'][1] < MIN_LOWER_HEIGHT:
|
|
split_result.append(art)
|
|
continue
|
|
|
|
def _bbox(regs):
|
|
return [min(reg['bbox'][0] for reg in regs), min(reg['bbox'][1] for reg in regs),
|
|
max(reg['bbox'][2] for reg in regs), max(reg['bbox'][3] for reg in regs)]
|
|
|
|
real_gap = members[split_idx+1]['bbox'][1] - max(r['bbox'][3] for r in upper)
|
|
print(f" [split] Article {art['article_id']}: split at gap={real_gap}px after region "
|
|
f"{upper[-1]['id']} (upper {len(upper)} regions, lower {len(lower)} regions)")
|
|
|
|
upper_ids = [reg['id'] for reg in upper]
|
|
lower_ids = [reg['id'] for reg in lower]
|
|
|
|
split_result.append({
|
|
'article_id': art['article_id'],
|
|
'title_region': art['title_region'] if art['title_region'] in upper_ids else None,
|
|
'member_regions': upper_ids,
|
|
'bbox': _bbox(upper),
|
|
'grouped_by': art['grouped_by'],
|
|
})
|
|
split_result.append({
|
|
'article_id': None,
|
|
'title_region': art['title_region'] if art['title_region'] in lower_ids else None,
|
|
'member_regions': lower_ids,
|
|
'bbox': _bbox(lower),
|
|
'grouped_by': 'geometric_split',
|
|
})
|
|
|
|
for idx, a in enumerate(split_result, start=1):
|
|
a['article_id'] = idx
|
|
|
|
return split_result
|