pipeline.py

"""
Table-layout cleaning utilities (in-memory version of process_table_pipeline.py).

Exposes:
  - run_detection(model, image_np, conf, iou, imgsz) -> detections dict
  - create_labelme_json(...)  : build labelme JSON + apply expand/dedupe logic
  - clean_labelme(...)        : fix row widths & column heights
  - process_image(model, pil_image, conf, iou, imgsz) -> (labelme_json, shapes)

These functions work on numpy/PIL data so they can be called from a Gradio app
without touching the filesystem.
"""

import logging

import numpy as np

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Geometry helpers
# ---------------------------------------------------------------------------

def calculate_iou(box1, box2):
    """IoU between two 4-point polygons."""
    def coords(points):
        a = np.array(points)
        return np.min(a[:, 0]), np.min(a[:, 1]), np.max(a[:, 0]), np.max(a[:, 1])

    x1_min, y1_min, x1_max, y1_max = coords(box1)
    x2_min, y2_min, x2_max, y2_max = coords(box2)

    ixmin = max(x1_min, x2_min)
    iymin = max(y1_min, y2_min)
    ixmax = min(x1_max, x2_max)
    iymax = min(y1_max, y2_max)

    if ixmax < ixmin or iymax < iymin:
        return 0.0

    inter = (ixmax - ixmin) * (iymax - iymin)
    a1 = (x1_max - x1_min) * (y1_max - y1_min)
    a2 = (x2_max - x2_min) * (y2_max - y2_min)
    union = a1 + a2 - inter
    return inter / union if union else 0.0


def expand_rows_to_column_width(boxes):
    """Expand row/header width to match leftmost/rightmost column x-coordinates."""
    columns = [b for b in boxes if 'column' in b['class_name'].lower()]
    if not columns:
        return boxes

    all_x = []
    for col in columns:
        all_x.extend(np.array(col['points'])[:, 0])
    if not all_x:
        return boxes

    col_min_x, col_max_x = min(all_x), max(all_x)

    out = []
    for box in boxes:
        if any(k in box['class_name'].lower() for k in ['row', 'header']):
            a = np.array(box['points'])
            min_y, max_y = np.min(a[:, 1]), np.max(a[:, 1])
            box_copy = box.copy()
            box_copy['points'] = [
                [float(col_min_x), float(min_y)],
                [float(col_max_x), float(min_y)],
                [float(col_max_x), float(max_y)],
                [float(col_min_x), float(max_y)],
            ]
            out.append(box_copy)
        else:
            out.append(box)
    return out


def remove_duplicate_boxes(boxes, iou_threshold=0.3):
    """Keep only the highest-confidence box among overlapping ones."""
    if not boxes:
        return boxes
    sorted_boxes = sorted(boxes, key=lambda b: b['confidence'], reverse=True)
    unique = []
    for cur in sorted_boxes:
        if all(calculate_iou(cur['points'], k['points']) < iou_threshold for k in unique):
            unique.append(cur)
    return unique


def expand_line_width(points, image_width):
    if len(points) != 4:
        return points
    a = np.array(points)
    min_y, max_y = np.min(a[:, 1]), np.max(a[:, 1])
    return [
        [0.0, min_y], [float(image_width), min_y],
        [float(image_width), max_y], [0.0, max_y],
    ]


def expand_column_height(points, image_height):
    if len(points) != 4:
        return points
    a = np.array(points)
    min_x, max_x = np.min(a[:, 0]), np.max(a[:, 0])
    return [
        [min_x, 0.0], [max_x, 0.0],
        [max_x, float(image_height)], [min_x, float(image_height)],
    ]


# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------

def run_detection(model, image_np, conf=0.1, iou=0.1, imgsz=640):
    """Run YOLO on a numpy image (BGR or RGB) and return a detections dict."""
    results = model.predict(image_np, save=False, verbose=False,
                            conf=conf, iou=iou, imgsz=imgsz)

    detections = {'boxes': []}
    if not results:
        return detections

    result = results[0]
    if result.boxes is None or len(result.boxes) == 0:
        return detections

    bboxes = result.boxes.xyxy
    confs = result.boxes.conf
    classes = result.boxes.cls

    for i, bbox in enumerate(bboxes):
        try:
            if hasattr(bbox, 'cpu'):
                bbox_np = bbox.cpu().numpy()
            elif hasattr(bbox, 'numpy'):
                bbox_np = bbox.numpy()
            else:
                bbox_np = np.asarray(bbox)

            x1, y1, x2, y2 = bbox_np
            points = [
                [float(x1), float(y1)], [float(x2), float(y1)],
                [float(x2), float(y2)], [float(x1), float(y2)],
            ]
            conf_v = float(confs[i].item() if hasattr(confs[i], 'item') else confs[i])
            cls = int(classes[i].item() if hasattr(classes[i], 'item') else classes[i])
            class_name = model.names.get(cls, f"unknown_class_{cls}")

            detections['boxes'].append({
                'points': points,
                'confidence': conf_v,
                'class_name': class_name,
                'class_id': cls,
            })
        except Exception as e:
            logger.error(f"Error processing bbox {i}: {e}")

    return detections


# ---------------------------------------------------------------------------
# Labelme JSON build + clean
# ---------------------------------------------------------------------------

def create_labelme_json(image_name, detections, image_height, image_width,
                        iou_threshold=0.3):
    expanded = expand_rows_to_column_width(detections['boxes'])
    filtered = remove_duplicate_boxes(expanded, iou_threshold)

    shapes = []
    for box_data in filtered:
        points = box_data['points']
        label = box_data['class_name']
        confidence = box_data['confidence']

        if 'line' in label.lower():
            points = expand_line_width(points, image_width)
        elif 'column' in label.lower():
            points = expand_column_height(points, image_height)

        a = np.array(points)
        rect_points = [
            [float(np.min(a[:, 0])), float(np.min(a[:, 1]))],
            [float(np.max(a[:, 0])), float(np.max(a[:, 1]))],
        ]
        shapes.append({
            "label": label,
            "points": rect_points,
            "group_id": None,
            "description": f"confidence: {confidence:.2f}",
            "shape_type": "rectangle",
            "flags": {},
        })

    return {
        "version": "5.0.1",
        "flags": {},
        "shapes": shapes,
        "imagePath": image_name,
        "imageData": None,
        "imageHeight": image_height,
        "imageWidth": image_width,
    }


def _x_bounds(shapes):
    columns = [s for s in shapes if 'column' in s['label'].lower()]
    if not columns:
        return None, None
    xs = []
    for c in columns:
        xs.extend(np.array(c['points'])[:, 0])
    return (min(xs), max(xs)) if xs else (None, None)


def _y_bounds(shapes):
    headers = [s for s in shapes if 'header' in s['label'].lower()]
    rows = [s for s in shapes if 'row' in s['label'].lower()]
    if not headers and not rows:
        return None, None

    top_y = None
    src = headers if headers else rows
    if src:
        ys = []
        for s in src:
            ys.extend(np.array(s['points'])[:, 1])
        top_y = min(ys) if ys else None

    bottom_y = None
    if rows:
        ys = []
        for r in rows:
            ys.extend(np.array(r['points'])[:, 1])
        bottom_y = max(ys) if ys else None

    return top_y, bottom_y


def fix_row_widths(labelme_data):
    shapes = labelme_data.get('shapes', [])
    rows = [s for s in shapes if 'row' in s['label'].lower()]
    if not rows:
        return labelme_data
    left_x, right_x = _x_bounds(shapes)
    if left_x is None:
        return labelme_data
    for r in rows:
        a = np.array(r['points'])
        min_y, max_y = np.min(a[:, 1]), np.max(a[:, 1])
        r['points'] = [[float(left_x), float(min_y)], [float(right_x), float(max_y)]]
    return labelme_data


def fix_column_heights(labelme_data):
    shapes = labelme_data.get('shapes', [])
    columns = [s for s in shapes if 'column' in s['label'].lower()]
    if not columns:
        return labelme_data
    top_y, bottom_y = _y_bounds(shapes)
    if top_y is None or bottom_y is None:
        return labelme_data
    for c in columns:
        a = np.array(c['points'])
        min_x, max_x = np.min(a[:, 0]), np.max(a[:, 0])
        c['points'] = [[float(min_x), float(top_y)], [float(max_x), float(bottom_y)]]
    return labelme_data


def clean_labelme(labelme_data):
    """Fix row widths then column heights (in place)."""
    labelme_data = fix_row_widths(labelme_data)
    labelme_data = fix_column_heights(labelme_data)
    return labelme_data


# ---------------------------------------------------------------------------
# Top-level convenience
# ---------------------------------------------------------------------------

def process_image(model, image_np, image_name="image.png",
                  conf=0.1, iou=0.1, imgsz=640, dedupe_iou=0.3):
    """
    Full pipeline on a single numpy image (RGB or BGR).
    Returns (cleaned_labelme_json, shapes_list).
    """
    height, width = image_np.shape[:2]
    detections = run_detection(model, image_np, conf=conf, iou=iou, imgsz=imgsz)
    labelme_json = create_labelme_json(image_name, detections, height, width, dedupe_iou)
    labelme_json = clean_labelme(labelme_json)
    return labelme_json, labelme_json["shapes"]