ocr_verify/table_line_generator/editor/drawing.py

"""
表格线绘制功能
"""

import streamlit as st
from PIL import Image, ImageDraw, ImageFont
import json
import hashlib


def draw_table_lines_with_numbers(image, structure, line_width=2, show_numbers=True):
    """
    绘制带编号的表格线（使用线坐标列表）
    
    Args:
        image: PIL Image 对象
        structure: 表格结构字典（包含 horizontal_lines 和 vertical_lines）
        line_width: 线条宽度
        show_numbers: 是否显示编号
    
    Returns:
        绘制了表格线和编号的图片
    """
    img_with_lines = image.copy()
    draw = ImageDraw.Draw(img_with_lines)
    
    # 尝试加载字体
    try:
        font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 20)
    except:
        font = ImageFont.load_default()
    
    # 使用线坐标列表
    horizontal_lines = structure.get('horizontal_lines', [])
    vertical_lines = structure.get('vertical_lines', [])
    modified_h_lines = structure.get('modified_h_lines', set())
    modified_v_lines = structure.get('modified_v_lines', set())
    
    # 计算绘制范围
    x_start = vertical_lines[0] if vertical_lines else 0
    x_end = vertical_lines[-1] if vertical_lines else img_with_lines.width
    y_start = horizontal_lines[0] if horizontal_lines else 0
    y_end = horizontal_lines[-1] if horizontal_lines else img_with_lines.height
    
    # 绘制横线
    for idx, y in enumerate(horizontal_lines):
        color = (255, 0, 0) if idx in modified_h_lines else (0, 0, 255)
        draw.line([(x_start, y), (x_end, y)], fill=color, width=line_width)
        
        # 绘制行编号
        if show_numbers:
            text = f"R{idx+1}"
            bbox = draw.textbbox((x_start - 35, y - 10), text, font=font)
            draw.rectangle(bbox, fill='white', outline='black')
            draw.text((x_start - 35, y - 10), text, fill=color, font=font)
    
    # 绘制竖线
    for idx, x in enumerate(vertical_lines):
        color = (255, 0, 0) if idx in modified_v_lines else (0, 0, 255)
        draw.line([(x, y_start), (x, y_end)], fill=color, width=line_width)
        
        # 绘制列编号
        if show_numbers:
            text = f"C{idx+1}"
            bbox = draw.textbbox((x - 10, y_start - 25), text, font=font)
            draw.rectangle(bbox, fill='white', outline='black')
            draw.text((x - 10, y_start - 25), text, fill=color, font=font)
            bbox = draw.textbbox((x - 10, y_end + 25), text, font=font)
            draw.rectangle(bbox, fill='white', outline='black')
            draw.text((x - 10, y_end + 25), text, fill=color, font=font)
    
    return img_with_lines


def draw_clean_table_lines(image, structure, line_width=2, line_color=(0, 0, 0)):
    """
    绘制纯净的表格线（用于保存）
    - 所有线用统一颜色
    - 不显示编号
    
    Args:
        image: PIL Image 对象
        structure: 表格结构字典
        line_width: 线条宽度
        line_color: 线条颜色，默认黑色 (0, 0, 0)
    
    Returns:
        绘制了纯净表格线的图片
    """
    img_with_lines = image.copy()
    draw = ImageDraw.Draw(img_with_lines)
    
    horizontal_lines = structure.get('horizontal_lines', [])
    vertical_lines = structure.get('vertical_lines', [])
    
    if not horizontal_lines or not vertical_lines:
        return img_with_lines
    
    # 计算绘制范围
    x_start = vertical_lines[0]
    x_end = vertical_lines[-1]
    y_start = horizontal_lines[0]
    y_end = horizontal_lines[-1]
    
    # 绘制横线
    for y in horizontal_lines:
        draw.line([(x_start, y), (x_end, y)], fill=line_color, width=line_width)
    
    # 绘制竖线
    for x in vertical_lines:
        draw.line([(x, y_start), (x, y_end)], fill=line_color, width=line_width)
    
    return img_with_lines


def get_structure_hash(structure, line_width, show_numbers):
    """生成结构的哈希值，用于判断是否需要重新绘制"""
    key_data = {
        'horizontal_lines': structure.get('horizontal_lines', []),
        'vertical_lines': structure.get('vertical_lines', []),
        'modified_h_lines': sorted(list(structure.get('modified_h_lines', set()))),
        'modified_v_lines': sorted(list(structure.get('modified_v_lines', set()))),
        'line_width': line_width,
        'show_numbers': show_numbers
    }
    
    key_str = json.dumps(key_data, sort_keys=True)
    return hashlib.md5(key_str.encode()).hexdigest()


def get_cached_table_lines_image(image, structure, line_width, show_numbers):
    """
    获取缓存的表格线图片，如果缓存不存在或失效则重新绘制
    
    Args:
        image: PIL Image 对象
        structure: 表格结构字典
        line_width: 线条宽度
        show_numbers: 是否显示编号
    
    Returns:
        绘制了表格线和编号的图片
    """
    # 初始化缓存
    if 'cached_table_image' not in st.session_state:
        st.session_state.cached_table_image = None
    if 'cached_table_hash' not in st.session_state:
        st.session_state.cached_table_hash = None
    
    # 计算当前结构的哈希
    current_hash = get_structure_hash(structure, line_width, show_numbers)
    
    # 检查缓存是否有效
    if (st.session_state.cached_table_hash == current_hash and 
        st.session_state.cached_table_image is not None):
        return st.session_state.cached_table_image
    
    # 缓存失效，重新绘制
    img_with_lines = draw_table_lines_with_numbers(
        image, 
        structure, 
        line_width=line_width,
        show_numbers=show_numbers
    )
    
    # 更新缓存
    st.session_state.cached_table_image = img_with_lines
    st.session_state.cached_table_hash = current_hash
    
    return img_with_lines


def clear_table_image_cache():
    """清除表格图片缓存"""
    if 'cached_table_image' in st.session_state:
        st.session_state.cached_table_image = None
    if 'cached_table_hash' in st.session_state:
        st.session_state.cached_table_hash = None