基于LayoutLM（philschmid/lilt-en-funsd）的收据文本提取问题求助

你的问题核心是模型将同一实体（如商品名、价格）拆分为多个B-ANSWER框，需要合并这些框的文本并按商品项分组提取字段。以下是修改后的代码和关键逻辑说明：

修改后的完整推理函数 #

def run_inference(image_path, model=model, processor=processor, output_image=True):
    # Load image from the path
    image = Image.open(image_path).convert("RGB")

    # get predictions
    encoding = processor(image, return_tensors="pt")
    del encoding["pixel_values"]
    outputs = model(**encoding)
    predictions = outputs.logits.argmax(-1).squeeze().tolist()
    labels = [model.config.id2label[prediction] for prediction in predictions]
    boxes = encoding["bbox"][0].tolist()
    model_name = model.name_or_path.split('/')[-1]

    # 收集所有ANSWER相关的标签和框（包含B-ANSWER和I-ANSWER）
    answer_indices = [i for i, label in enumerate(labels) if label in ("B-ANSWER", "I-ANSWER")]
    answer_boxes = [encoding["bbox"][0][i].detach().numpy() for i in answer_indices]
    answer_labels = [labels[i] for i in answer_indices]

    # 提取每个框对应的文本
    box_texts = extract_text_from_boxes(image, answer_boxes, image_path, model_name)

    # 合并同一实体的连续框：拼接文本并合并坐标
    merged_boxes_and_texts = []
    current_box = None
    current_text = []

    for idx, (box, label, text) in enumerate(zip(answer_boxes, answer_labels, box_texts)):
        if label == "B-ANSWER":
            # 保存上一个已合并的实体
            if current_box is not None:
                merged_text = " ".join(current_text)
                merged_boxes_and_texts.append((current_box, merged_text))
            # 初始化新实体的框和文本
            current_box = box.copy()
            current_text = [text]
        else:  # 处理I-ANSWER，延续当前实体
            if current_box is not None:
                # 合并框坐标：取最小左上、最大右下
                current_box[0] = min(current_box[0], box[0])
                current_box[1] = min(current_box[1], box[1])
                current_box[2] = max(current_box[2], box[2])
                current_box[3] = max(current_box[3], box[3])
                current_text.append(text)
    # 处理最后一个未保存的实体
    if current_box is not None:
        merged_text = " ".join(current_text)
        merged_boxes_and_texts.append((current_box, merged_text))

    # 按行分组，提取商品、数量、价格
    receipt_items = []
    if merged_boxes_and_texts:
        # 先按垂直位置（y坐标）排序，再按水平位置（x坐标）排序
        sorted_items = sorted(merged_boxes_and_texts, key=lambda x: (x[0][1], x[0][0]))
        current_item = []
        for box, text in sorted_items:
            y_center = (box[1] + box[3]) / 2
            if not current_item:
                current_item.append((box, text))
            else:
                last_y_center = (current_item[-1][0][1] + current_item[-1][0][3]) / 2
                # 阈值可根据收据尺寸调整，判断是否属于同一行
                if abs(y_center - last_y_center) < 20:
                    current_item.append((box, text))
                else:
                    # 对当前行的内容按水平位置排序，拆分字段
                    sorted_row = sorted(current_item, key=lambda x: x[0][0])
                    receipt_items.append({
                        "product": sorted_row[0][1].strip() if len(sorted_row)>=1 else "",
                        "quantity": sorted_row[1][1].strip() if len(sorted_row)>=2 else "",
                        "price": sorted_row[2][1].strip() if len(sorted_row)>=3 else ""
                    })
                    current_item = [(box, text)]
        # 处理最后一行的内容
        sorted_row = sorted(current_item, key=lambda x: x[0][0])
        receipt_items.append({
            "product": sorted_row[0][1].strip() if len(sorted_row)>=1 else "",
            "quantity": sorted_row[1][1].strip() if len(sorted_row)>=2 else "",
            "price": sorted_row[2][1].strip() if len(sorted_row)>=3 else ""
        })

    if output_image:
        image_with_boxes = draw_boxes(image, encoding["bbox"][0], labels)
        return image_with_boxes, receipt_items
    else:
        return draw_boxes(image, encoding["bbox"][0], labels), receipt_items

关键逻辑说明 #

1. 合并拆分的实体框：

   • 收集所有B-ANSWER和I-ANSWER标签的框，避免遗漏同一实体的后续部分。
   • 对连续的同实体框，合并坐标范围并拼接文本，解决内容被拆分的问题。如果你的模型输出全是B-ANSWER，可将标签判断条件改为label == "B-ANSWER"，并通过位置距离判断是否属于同一实体。

2. 按行分组提取字段：

   • 根据框的垂直中心坐标判断是否属于同一行，将同一行的文本归为一个商品项。
   • 对每行内容按水平位置排序，默认从左到右依次为商品名、数量、价格，可根据你的收据布局调整字段顺序。

3. 阈值调整：

   • 代码中判断同一行的像素阈值（20）可根据收据的实际尺寸调整，确保同一行的框被正确分组。

内容的提问来源于stack exchange，提问作者Asim