Python中如何逐行写入EXR图像？低内存EXR网格拼接工具开发

现有代码的问题在于提前将所有EXR合并为完整的numpy数组，当拼接大尺寸网格（如4x4高分辨率EXR）时，会占用大量内存甚至导致溢出。要实现低内存写入，核心思路是避免加载完整的合并图像，逐行构建输出数据并写入EXR文件。

核心思路 #

EXR格式支持增量写入扫描线，我们可以：

• 先确定输出图像的整体尺寸（基于网格行列数和单张EXR的尺寸）
• 初始化EXR文件头，打开输出文件
• 按输出图像的行号，计算对应输入EXR的位置和行号
• 读取对应输入EXR的目标行，拼接成输出行
• 将该行数据写入EXR，写完后释放临时内存

具体实现代码 #

1. 辅助函数：读取单个EXR的指定行 #

这个函数仅加载单张EXR的指定行数据，避免加载整个文件：

import numpy as np
import OpenEXR
import Imath

# 映射numpy类型到EXR类型
_np_to_exr = {
    np.float32: Imath.PixelType(Imath.PixelType.FLOAT),
    np.uint16: Imath.PixelType(Imath.PixelType.HALF),
}

def read_exr_line(exr_path, line_idx, channel_names=['R', 'G', 'B']):
    """读取EXR文件的指定行数据"""
    exr_file = OpenEXR.InputFile(exr_path)
    header = exr_file.header()
    width = header['dataWindow'].max.x + 1
    height = header['dataWindow'].max.y + 1
    
    if line_idx < 0 or line_idx >= height:
        raise ValueError(f"行索引 {line_idx} 超出范围 (0-{height-1})")
    
    # 获取通道数据类型
    channel_type = header['channels'][channel_names[0]].type
    np_dtype = np.float32 if channel_type == Imath.PixelType.FLOAT else np.uint16
    
    # EXR坐标原点在左下角，需转换行号
    exr_y = height - 1 - line_idx
    line_data = {}
    for channel in channel_names:
        channel_bytes = exr_file.readPixels(channel, exr_y, exr_y + 1)
        line_data[channel] = np.frombuffer(channel_bytes, dtype=np_dtype).reshape(1, width)
    
    # 合并为[1, width, channels]格式的数组
    merged_line = np.concatenate([line_data[ch] for ch in channel_names], axis=-1)
    exr_file.close()
    return merged_line

2. 逐行写入EXR的主逻辑 #

修改原有的合并与写入逻辑，改为逐行处理：

def write_exr_row_by_row(output_path, exr_files_grid, channel_names=['R', 'G', 'B']):
    """
    逐行写入拼接后的EXR网格
    Args:
        output_path: 输出EXR路径
        exr_files_grid: 二维列表，代表网格布局，如[[exr1, exr2], [exr3, exr4]]对应2x2网格
        channel_names: 通道名称列表
    """
    if not exr_files_grid or not exr_files_grid[0]:
        raise ValueError("无效的网格布局")
    
    # 获取单张EXR的基准尺寸
    sample_exr = OpenEXR.InputFile(exr_files_grid[0][0])
    sample_header = sample_exr.header()
    single_width = sample_header['dataWindow'].max.x + 1
    single_height = sample_header['dataWindow'].max.y + 1
    sample_exr.close()
    
    # 检查所有输入EXR尺寸一致
    for row in exr_files_grid:
        for exr_path in row:
            exr = OpenEXR.InputFile(exr_path)
            header = exr.header()
            w = header['dataWindow'].max.x + 1
            h = header['dataWindow'].max.y + 1
            if w != single_width or h != single_height:
                raise ValueError(f"EXR {exr_path} 尺寸不匹配: {w}x{h}, 预期 {single_width}x{single_height}")
            exr.close()
    
    # 计算输出图像总尺寸
    grid_rows = len(exr_files_grid)
    grid_cols = len(exr_files_grid[0])
    output_width = grid_cols * single_width
    output_height = grid_rows * single_height
    
    # 初始化EXR文件头
    header = OpenEXR.Header(output_width, output_height)
    sample_exr = OpenEXR.InputFile(exr_files_grid[0][0])
    sample_channel_type = sample_exr.header()['channels'][channel_names[0]].type
    sample_exr.close()
    header['channels'] = {n: Imath.Channel(sample_channel_type) for n in channel_names}
    
    # 逐行写入EXR
    with OpenEXR.OutputFile(output_path, header) as exr_out:
        for output_line_idx in range(output_height):
            # 计算当前行对应的网格位置和单张EXR内的行号
            grid_row_idx = output_line_idx // single_height
            local_line_idx = output_line_idx % single_height
            
            # 拼接当前输出行的所有列数据
            output_line_parts = []
            for grid_col_idx in range(grid_cols):
                exr_path = exr_files_grid[grid_row_idx][grid_col_idx]
                line_data = read_exr_line(exr_path, local_line_idx, channel_names)
                output_line_parts.append(line_data)
            
            # 合并为完整的一行
            full_line = np.concatenate(output_line_parts, axis=1)
            
            # 按通道拆分并写入
            channel_data = {}
            exr_y = output_height - 1 - output_line_idx  # 转换为EXR的左下角坐标
            for idx, ch in enumerate(channel_names):
                channel_data[ch] = full_line[..., idx].tobytes()
            
            exr_out.writePixels(channel_data)
            
            # 显式释放临时内存
            del output_line_parts, full_line, channel_data

# 示例用法：构建2x2网格并写入
exr_grid = [
    ["exr_00.exr", "exr_01.exr"],
    ["exr_10.exr", "exr_11.exr"]
]
write_exr_row_by_row("output_grid.exr", exr_grid, ['R', 'G', 'B'])

关键优化说明 #

• 内存控制：每次仅加载当前需要的一行数据，内存占用仅为单张EXR一行数据的网格列数倍，远低于加载完整合并图像的内存消耗。
• 坐标转换：注意EXR的扫描线坐标原点在左下角，需将从上到下的行号转换为EXR的y坐标。
• 一致性检查：确保所有输入EXR的尺寸、通道数和数据类型一致，避免拼接错误。
• 资源管理：使用with语句管理EXR文件的打开/关闭，避免资源泄漏。

如果逐行写入效率偏低，可改为批量写入多行（如一次写入10行），平衡内存占用和写入速度，只需调整循环逻辑批量处理多行数据即可。

内容的提问来源于stack exchange，提问作者Ahmad Uzzam Masood