sundeep-news-scan/xy_cut.py

import numpy as np

class Box:
    def __init__(self, region):
        self.id = region['id']
        self.type = region['type']
        self.bbox = region['bbox']
        self.x1, self.y1, self.x2, self.y2 = self.bbox
        self.region = region

def recursive_xy_cut(boxes, min_x_gap=20, min_y_gap=20):
    if not boxes:
        return []
        
    if len(boxes) == 1:
        return [boxes]
        
    # Find bounding box of all boxes
    min_x = min(b.x1 for b in boxes)
    max_x = max(b.x2 for b in boxes)
    min_y = min(b.y1 for b in boxes)
    max_y = max(b.y2 for b in boxes)
    
    # Projection profiles
    x_profile = np.zeros(max_x - min_x + 1, dtype=int)
    y_profile = np.zeros(max_y - min_y + 1, dtype=int)
    
    for b in boxes:
        x_profile[max(0, b.x1 - min_x):b.x2 - min_x + 1] = 1
        y_profile[max(0, b.y1 - min_y):b.y2 - min_y + 1] = 1
        
    # Find gaps
    def find_gaps(profile):
        gaps = []
        in_gap = profile[0] == 0
        start = 0 if in_gap else -1
        for i, val in enumerate(profile):
            if val == 0 and not in_gap:
                in_gap = True
                start = i
            elif val == 1 and in_gap:
                in_gap = False
                gaps.append((start, i))
        if in_gap:
            gaps.append((start, len(profile)))
        return gaps
        
    x_gaps = find_gaps(x_profile)
    y_gaps = find_gaps(y_profile)
    
    # Filter gaps by threshold
    valid_x_gaps = [g for g in x_gaps if g[1] - g[0] >= min_x_gap]
    valid_y_gaps = [g for g in y_gaps if g[1] - g[0] >= min_y_gap]
    
    # Choose best cut (prefer Y cuts to separate vertical sections, then X cuts for columns)
    # Actually, in newspapers, Y cut (horizontal line) separates main stories or headlines from columns
    # X cut (vertical line) separates columns.
    # Standard XY cut alternates or picks the widest gap.
    best_cut = None
    axis = None
    
    if valid_y_gaps:
        # Pick the widest Y gap
        widest = max(valid_y_gaps, key=lambda g: g[1] - g[0])
        best_cut = widest[0] + (widest[1] - widest[0]) // 2 + min_y
        axis = 'y'
    elif valid_x_gaps:
        # Pick the widest X gap
        widest = max(valid_x_gaps, key=lambda g: g[1] - g[0])
        best_cut = widest[0] + (widest[1] - widest[0]) // 2 + min_x
        axis = 'x'
        
    if best_cut is None:
        # No cuts can be made, this is a leaf node
        return [boxes]
        
    # Split boxes
    left_top = []
    right_bottom = []
    
    if axis == 'y':
        for b in boxes:
            if b.y1 + (b.y2 - b.y1)//2 < best_cut:
                left_top.append(b)
            else:
                right_bottom.append(b)
    else:
        for b in boxes:
            if b.x1 + (b.x2 - b.x1)//2 < best_cut:
                left_top.append(b)
            else:
                right_bottom.append(b)
                
    # Recurse
    # If a cut didn't actually split anything (due to bounding box overlap centers), stop
    if not left_top or not right_bottom:
        return [boxes]
        
    return recursive_xy_cut(left_top, min_x_gap, min_y_gap) + recursive_xy_cut(right_bottom, min_x_gap, min_y_gap)

def group_surya_regions_xycut(regions):
    """
    Groups regions by using XY-Cut to establish reading order, 
    then grouping by paragraph_title / doc_title.
    """
    if not regions:
        return []
        
    boxes = [Box(r) for r in regions]
    
    # Step 1: XY-Cut
    leaves = recursive_xy_cut(boxes, min_x_gap=30, min_y_gap=25)
    
    # Step 2: Establish Reading order
    # Leaves are already in a roughly top-to-bottom, left-to-right order because 
    # recursive_xy_cut processes Y cuts then X cuts and appends left_top before right_bottom.
    
    # Flatten boxes in reading order
    ordered_boxes = []
    for leaf in leaves:
        # Sort within leaf just in case
        leaf.sort(key=lambda b: (b.y1, b.x1))
        ordered_boxes.extend(leaf)
        
    # Step 3: Group into articles
    articles = []
    current_article = []
    current_title = None
    
    for box in ordered_boxes:
        # Start a new article if we hit a title
        if box.type in ('paragraph_title', 'doc_title'):
            # If the current article ONLY contains images, these images likely belong to this new title!
            # Or if it contains images at the very end, wait, let's just check if it's ONLY images.
            if current_article and all(b.type == 'image' for b in current_article):
                current_article.append(box)
            else:
                if current_article:
                    articles.append(current_article)
                current_article = [box]
        else:
            current_article.append(box)
            
    if current_article:
        articles.append(current_article)
        
    # Format output
    output = []
    for i, art_boxes in enumerate(articles, start=1):
        member_ids = [b.id for b in art_boxes]
        
        # Calculate union bbox
        min_x = min(b.x1 for b in art_boxes)
        min_y = min(b.y1 for b in art_boxes)
        max_x = max(b.x2 for b in art_boxes)
        max_y = max(b.y2 for b in art_boxes)
        
        output.append({
            "article_id": i,
            "title_region": member_ids[0] if art_boxes[0].type in ('paragraph_title', 'doc_title') else None,
            "member_regions": member_ids,
            "bbox": [min_x, min_y, max_x, max_y],
            "dateline": None,
            "grouped_by": "xy_cut"
        })
        
    return output