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