如何在同脚本中将函数输出作为另一函数输入？以预处理图像为例

现有代码中，get_preprocessed_image可以生成预处理后的PIL图像，要让infer_text直接使用这个图像作为输入，无需先保存到文件再读取，可通过以下两种方式实现：

方案一：修改infer_text支持直接接收PIL图像（推荐） #

直接调整infer_text的参数和内部逻辑，让它既能接收文件路径，也能接收预处理后的PIL图像对象，避免多余的磁盘IO操作。

修改后的infer_text函数 #

from io import BytesIO
import pathlib
import PIL.Image
import pyabbyy
import graphanno
import geometric
import numpy as np
from tqdm import tqdm

def infer_text(input_source: pathlib.Path | PIL.Image.Image) -> graphanno.GraphAnnotation:
    # 处理输入源：路径则读取文件字节，PIL图像则转为字节流
    if isinstance(input_source, pathlib.Path):
        with input_source.open("rb") as img:
            img_bytes = img.read()
    else:
        # 将PIL图像转为字节流，选用PNG格式保证兼容性
        buffer = BytesIO()
        input_source.save(buffer, format="PNG")
        img_bytes = buffer.getvalue()
    
    # 已提前完成预处理，此处关闭自动预处理
    words = pyabbyy.read_text(img_bytes, preprocess=False)
    
    nodes = []
    for word in words:
        box = geometric.Box(
            origin_x=word["origin_x"],
            origin_y=word["origin_y"],
            width=word["max_x"] - word["origin_x"],
            height=word["max_y"] - word["origin_y"],
        )
        nodes.append(graphanno.Node(text=word["text"], box=box))
    
    num_nodes = len(nodes)
    return graphanno.GraphAnnotation(
        tuple(nodes),
        graphanno.Adjacency(np.zeros((num_nodes, num_nodes))),
        graphanno.Adjacency(np.zeros((num_nodes, num_nodes))),
        graphanno.Adjacency(np.zeros((num_nodes, num_nodes))),
    )

调用方式 #

• 单文件处理：

image_folder = pathlib.Path("/home/Tasks/NM_spanish/Invoices_pdf")
sample_path = next(image_folder.rglob("*.pdf"))  # 获取一个示例文件
preprocessed_img = get_preprocessed_image(sample_path)
annotation = infer_text(preprocessed_img)

• 批量处理（替代原有的先保存再读取逻辑）：

def batch_process(image_folder: pathlib.Path) -> list[graphanno.GraphAnnotation]:
    annotations = []
    for image_path in tqdm.tqdm(list(image_folder.rglob("*.pdf"))):
        try:
            preprocessed_img = get_preprocessed_image(image_path)
            annotations.append(infer_text(preprocessed_img))
        except (RuntimeError, AttributeError):
            print(f"处理失败：{image_path}")
    return annotations

if __name__ == '__main__':
    all_annotations = batch_process(image_folder)

方案二：临时保存预处理图像再传路径（不推荐） #

如果不想修改infer_text的原有逻辑，可以把预处理后的图像临时保存到文件，再将临时文件路径传入函数：

from tempfile import NamedTemporaryFile

image_path = pathlib.Path("/home/Tasks/NM_spanish/Invoices_pdf/example.pdf")
preprocessed_img = get_preprocessed_image(image_path)

# 创建临时文件保存预处理图像
with NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
    preprocessed_img.save(tmp_file, format="PNG")
    tmp_path = pathlib.Path(tmp_file.name)

# 调用原有infer_text函数
annotation = infer_text(tmp_path)

# 清理临时文件
tmp_path.unlink()

传入参数说明 #

• 方案一中，修改后的infer_text支持两种参数：
  1. pathlib.Path对象：兼容原有逻辑，传入文件路径
  2. PIL.Image.Image对象：直接传入get_preprocessed_image生成的预处理图像

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