feat: 添加表格检测与OCR功能

实现表格检测、结构识别和OCR文本提取功能
- 新增表格检测模型加载和预处理工具
- 添加表格结构识别和单元格坐标提取功能
- 集成EasyOCR进行文本识别并输出CSV格式
- 重构前端使用Streamlit展示检测结果
- 更新依赖项以支持新功能

app.py

@@ -17,6 +17,19 @@ import logging
 from PIL.ExifTags import TAGS
 from gradio_client import Client, handle_file
 import requests
+import streamlit as st
+from PIL import Image
+from utils.model import load_detection_model, load_structure_model
+from utils.preprocessing import prepare_image, prepare_cropped_image
+from utils.detection import (
+    detect_tables,
+    detect_cells,
+    outputs_to_objects,
+    objects_to_crops,
+    get_cell_coordinates_by_row,
+)
+from utils.visualization import visualize_detected_tables, plot_results
+from utils.ocr import apply_ocr, save_csv
 
 # 设置日志
 logging.basicConfig(level=logging.INFO)
@@ -319,52 +332,6 @@ async def upload_text(file: UploadFile = File(...)):
 
     - **file**: 需要识别的图片文件 (jpg, png, etc.)
     """
-    # 检查文件类型
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="上传的文件必须是图片格式。")
-
-    try:
-        # 保存上传的文件到临时目录
-        temp_file_path = f"/tmp/{file.filename}"
-        with open(temp_file_path, "wb") as f:
-            f.write(await file.read())
-
-        # 调用千帆模型进行OCR识别
-        client = Client("baidu/Qianfan-OCR-Demo", hf_token=hf_token)
-        result = client.predict(
-            file_paths=[handle_file(temp_file_path)],
-            prompt="Please extract the text from the image.",
-            layout_as_thought=False,
-            max_new_tokens=2048,
-            api_name="/run_inference",
-        )
-
-        # 清理临时文件
-        if os.path.exists(temp_file_path):
-            os.remove(temp_file_path)
-
-        # 检查结果类型并进行适当处理
-        logger.info(f"OCR识别结果类型: {type(result)}")
-        logger.info(f"OCR识别结果: {result}")
-
-        # 确保返回的是有效的JSON
-        if isinstance(result, str):
-            return {"result": result}
-        elif isinstance(result, dict):
-            return {"result": result}
-        else:
-            return {"result": str(result)}
-    except Exception as e:
-        # 清理临时文件
-        if os.path.exists(temp_file_path):
-            os.remove(temp_file_path)
-        logger.error(f"OCR识别错误: {e}")
-        import traceback
-
-        logger.error(f"错误堆栈: {traceback.format_exc()}")
-        raise HTTPException(
-            status_code=500, detail=f"服务器内部错误: {str(e)} {traceback.format_exc()}"
-        )
 
 
 from fastapi.staticfiles import StaticFiles

requirements.txt

@@ -7,4 +7,11 @@ opencv-python
 numpy
 python-multipart==0.0.9
 gradio_client
-requests
+requests
+streamlit
+transformers
+huggingface_hub
+matplotlib
+easyocr
+tqdm
+pandas

utils/detection.py

@@ -0,0 +1,98 @@
+import torch
+
+def detect_tables(model, pixel_values):
+    with torch.no_grad():
+        outputs = model(pixel_values)
+    return outputs
+
+def detect_cells(model, pixel_values):
+    with torch.no_grad():
+        outputs = model(pixel_values)
+    return outputs
+
+def outputs_to_objects(outputs, img_size, id2label):
+    def box_cxcywh_to_xyxy(x):
+        x_c, y_c, w, h = x.unbind(-1)
+        b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
+        return torch.stack(b, dim=1)
+
+    def rescale_bboxes(out_bbox, size):
+        img_w, img_h = size
+        b = box_cxcywh_to_xyxy(out_bbox)
+        b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
+        return b
+
+    # Add "no object" to id2label if not present
+    if len(id2label) not in id2label:
+        id2label[len(id2label)] = "no object"
+
+    m = outputs.logits.softmax(-1).max(-1)
+    pred_labels = list(m.indices.detach().cpu().numpy())[0]
+    pred_scores = list(m.values.detach().cpu().numpy())[0]
+    pred_bboxes = outputs['pred_boxes'].detach().cpu()[0]
+    pred_bboxes = [elem.tolist() for elem in rescale_bboxes(pred_bboxes, img_size)]
+
+    print(f"Predicted labels: {pred_labels}")
+    print(f"id2label: {id2label}")
+
+    objects = []
+    for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):
+        try:
+            class_label = id2label[int(label)]
+        except KeyError:
+            print(f"Label {label} not found in id2label. Skipping.")
+            continue
+        if not class_label == 'no object':
+            objects.append({'label': class_label, 'score': float(score),
+                            'bbox': [float(elem) for elem in bbox]})
+
+    return objects
+
+def objects_to_crops(img, tokens, objects, class_thresholds, padding=10):
+    table_crops = []
+    for obj in objects:
+        if obj['score'] < class_thresholds[obj['label']]:
+            continue
+
+        cropped_table = {}
+        bbox = obj['bbox']
+        bbox = [bbox[0] - padding, bbox[1] - padding, bbox[2] + padding, bbox[3] + padding]
+        cropped_img = img.crop(bbox)
+
+        table_tokens = [token for token in tokens if iob(token['bbox'], bbox) >= 0.5]
+        for token in table_tokens:
+            token['bbox'] = [token['bbox'][0] - bbox[0], token['bbox'][1] - bbox[1], token['bbox'][2] - bbox[0], token['bbox'][3] - bbox[1]]
+
+        if obj['label'] == 'table rotated':
+            cropped_img = cropped_img.rotate(270, expand=True)
+            for token in table_tokens:
+                bbox = token['bbox']
+                bbox = [cropped_img.size[0] - bbox[3] - 1, bbox[0], cropped_img.size[0] - bbox[1] - 1, bbox[2]]
+                token['bbox'] = bbox
+
+        cropped_table['image'] = cropped_img
+        cropped_table['tokens'] = table_tokens
+        table_crops.append(cropped_table)
+
+    return table_crops
+
+def get_cell_coordinates_by_row(table_data):
+    rows = [entry for entry in table_data if entry['label'] == 'table row']
+    columns = [entry for entry in table_data if entry['label'] == 'table column']
+    rows.sort(key=lambda x: x['bbox'][1])
+    columns.sort(key=lambda x: x['bbox'][0])
+
+    def find_cell_coordinates(row, column):
+        cell_bbox = [column['bbox'][0], row['bbox'][1], column['bbox'][2], row['bbox'][3]]
+        return cell_bbox
+
+    cell_coordinates = []
+    for row in rows:
+        row_cells = []
+        for column in columns:
+            cell_bbox = find_cell_coordinates(row, column)
+            row_cells.append({'column': column['bbox'], 'cell': cell_bbox})
+        row_cells.sort(key=lambda x: x['column'][0])
+        cell_coordinates.append({'row': row['bbox'], 'cells': row_cells, 'cell_count': len(row_cells)})
+    cell_coordinates.sort(key=lambda x: x['row'][1])
+    return cell_coordinates

utils/model.py

@@ -0,0 +1,13 @@
+import torch
+from transformers import AutoModelForObjectDetection, TableTransformerForObjectDetection
+
+def load_detection_model():
+    model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+    return model, device
+
+def load_structure_model(device):
+    model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-structure-recognition-v1.1-all")
+    model.to(device)
+    return model

utils/ocr.py

@@ -0,0 +1,34 @@
+import numpy as np
+import easyocr
+from tqdm.auto import tqdm
+import csv
+
+reader = easyocr.Reader(['en'])  # this needs to run only once to load the model into memory
+
+def apply_ocr(cell_coordinates, cropped_table):
+    data = dict()
+    max_num_columns = 0
+    for idx, row in enumerate(tqdm(cell_coordinates)):
+        row_text = []
+        for cell in row["cells"]:
+            cell_image = np.array(cropped_table.crop(cell["cell"]))
+            result = reader.readtext(np.array(cell_image))
+            if len(result) > 0:
+                text = " ".join([x[1] for x in result])
+                row_text.append(text)
+        if len(row_text) > max_num_columns:
+            max_num_columns = len(row_text)
+        data[idx] = row_text
+
+    print("Max number of columns:", max_num_columns)
+    for row, row_data in data.copy().items():
+        if len(row_data) != max_num_columns:
+            row_data = row_data + ["" for _ in range(max_num_columns - len(row_data))]
+        data[row] = row_data
+    return data
+
+def save_csv(data):
+    with open('output.csv', 'w') as result_file:
+        wr = csv.writer(result_file, dialect='excel')
+        for row, row_text in data.items():
+            wr.writerow(row_text)

utils/preprocessing.py

@@ -0,0 +1,35 @@
+from torchvision import transforms
+from PIL import Image
+
+class MaxResize(object):
+    def __init__(self, max_size=800):
+        self.max_size = max_size
+
+    def __call__(self, image):
+        width, height = image.size
+        current_max_size = max(width, height)
+        scale = self.max_size / current_max_size
+        resized_image = image.resize((int(round(scale*width)), int(round(scale*height))))
+        return resized_image
+
+detection_transform = transforms.Compose([
+    MaxResize(800),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+structure_transform = transforms.Compose([
+    MaxResize(1000),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def prepare_image(image, device):
+    pixel_values = detection_transform(image).unsqueeze(0)
+    pixel_values = pixel_values.to(device)
+    return pixel_values
+
+def prepare_cropped_image(cropped_image, device):
+    pixel_values = structure_transform(cropped_image).unsqueeze(0)
+    pixel_values = pixel_values.to(device)
+    return pixel_values

utils/visualization.py

@@ -0,0 +1,63 @@
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.patches import Patch
+
+def visualize_detected_tables(img, det_tables):
+    plt.imshow(img, interpolation="lanczos")
+    fig = plt.gcf()
+    fig.set_size_inches(20, 20)
+    ax = plt.gca()
+
+    for det_table in det_tables:
+        bbox = det_table['bbox']
+        if det_table['label'] == 'table':
+            facecolor = (1, 0, 0.45)
+            edgecolor = (1, 0, 0.45)
+            alpha = 0.3
+            linewidth = 2
+            hatch = '//////'
+        elif det_table['label'] == 'table rotated':
+            facecolor = (0.95, 0.6, 0.1)
+            edgecolor = (0.95, 0.6, 0.1)
+            alpha = 0.3
+            linewidth = 2
+            hatch = '//////'
+        else:
+            continue
+
+        rect = patches.Rectangle(bbox[:2], bbox[2] - bbox[0], bbox[3] - bbox[1], linewidth=linewidth,
+                                 edgecolor='none', facecolor=facecolor, alpha=0.1)
+        ax.add_patch(rect)
+        rect = patches.Rectangle(bbox[:2], bbox[2] - bbox[0], bbox[3] - bbox[1], linewidth=linewidth,
+                                 edgecolor=edgecolor, facecolor='none', linestyle='-', alpha=alpha)
+        ax.add_patch(rect)
+        rect = patches.Rectangle(bbox[:2], bbox[2] - bbox[0], bbox[3] - bbox[1], linewidth=0,
+                                 edgecolor=edgecolor, facecolor='none', linestyle='-', hatch=hatch, alpha=0.2)
+        ax.add_patch(rect)
+
+    plt.xticks([], [])
+    plt.yticks([], [])
+    legend_elements = [Patch(facecolor=(1, 0, 0.45), edgecolor=(1, 0, 0.45), label='Table', hatch='//////', alpha=0.3),
+                       Patch(facecolor=(0.95, 0.6, 0.1), edgecolor=(0.95, 0.6, 0.1), label='Table (rotated)', hatch='//////', alpha=0.3)]
+    plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.02), loc='upper center', borderaxespad=0,
+               fontsize=10, ncol=2)
+    plt.gcf().set_size_inches(10, 10)
+    plt.axis('off')
+    return fig
+
+def plot_results(cropped_table, cells, class_to_visualize):
+    plt.figure(figsize=(16, 10))
+    plt.imshow(cropped_table)
+    ax = plt.gca()
+
+    for cell in cells:
+        score = cell["score"]
+        bbox = cell["bbox"]
+        label = cell["label"]
+        if label == class_to_visualize:
+            xmin, ymin, xmax, ymax = tuple(bbox)
+            ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color="red", linewidth=3))
+            text = f'{cell["label"]}: {score:0.2f}'
+            ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
+    plt.axis('off')
+    return plt.gcf()