table_detector.py 2.49 KB
import cv2
import numpy as np
import os

def detect_tables(image_path):
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"Không đọc được ảnh: {image_path}")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (3, 3), 0)

    # Edge detection
    edges = cv2.Canny(blur, 50, 150, apertureSize=3)

    # --- Horizontal lines ---
    lines_h = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=120,
                              minLineLength=int(img.shape[1] * 0.6), maxLineGap=20)
    ys_candidates, line_segments = [], []
    if lines_h is not None:
        for l in lines_h:
            x1, y1, x2, y2 = l[0]
            if abs(y1 - y2) <= 3:  # ngang
                y_mid = int(round((y1 + y2) / 2))
                ys_candidates.append(y_mid)
                line_segments.append((x1, x2, y_mid))

    # gom nhóm các y
    ys, tol_y = [], 10
    for y in sorted(ys_candidates):
        if not ys or abs(y - ys[-1]) > tol_y:
            ys.append(y)

    total_rows = max(0, len(ys) - 1)

    # --- Vertical lines ---
    lines_v = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100,
                              minLineLength=int(img.shape[0] * 0.5), maxLineGap=20)
    xs = []
    if lines_v is not None:
        for l in lines_v:
            x1, y1, x2, y2 = l[0]
            if abs(x1 - x2) <= 3:
                xs.append(int(round((x1 + x2) / 2)))

    # gom nhóm cột
    x_pos, tol_v = [], 10
    for v in sorted(xs):
        if not x_pos or v - x_pos[-1] > tol_v:
            x_pos.append(v)

    total_cols = max(0, len(x_pos) - 1)

    tables = []
    if len(ys) >= 3 and line_segments:
        y_min, y_max = ys[0], ys[-1]
        min_x = min(seg[0] for seg in line_segments)
        max_x = max(seg[1] for seg in line_segments)
        table_box = (min_x, y_min, max_x, y_max)

        rows = []
        for i in range(len(ys) - 1):
            row_box = (min_x, ys[i], max_x, ys[i+1])
            rows.append({"row": tuple(int(v) for v in row_box)})
            cv2.rectangle(img, (row_box[0], row_box[1]), (row_box[2], row_box[3]), (0, 255, 255), 2)

        tables.append({
            "total_rows": int(total_rows),
            "total_cols": int(total_cols),
            "table_box": tuple(int(v) for v in table_box),
            "rows_box": rows
        })

        cv2.rectangle(img, (min_x, y_min), (max_x, y_max), (255, 0, 0), 3)

    debug_path = os.path.splitext(image_path)[0] + "_debug.jpg"
    cv2.imwrite(debug_path, img)

    return tables