如何将包含元组列表的NPZ文件修改后重新保存？

问题背景 #

目标是加载、修改并保存NPZ文件，文件包含5个键：4个为np.array类型，最后一个pair_infos是元组列表，示例元组结构为：

(array([1967, 2929]), 0.23312236286919832, array([888.117, 387.693, 851.922, 124.948]))

尝试将该元组列表转为numpy数组保存时，触发错误：

ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (44601, 3) + inhomogeneous part

用户代码 #

import numpy as np
from numpy import load
import os

directory = '/MY_PATH/scene_info_0.1_0.7'

for filename in os.listdir(directory):
    f_npz = os.path.join(directory,filename)
    data = load(f_npz, allow_pickle = True)
    print(type(data))
    for image_path in data['image_paths']:
        if image_path is not None:
            if 'Undistorted_SfM' in image_path:
                image_path = image_path.replace('Undistorted_SfM', 'phoenix/S6/zl548/MegaDepth_v1')
            else:
                print("Undistorted_SfM NOT IN IMAGE PATH??")
    
    data['pair_infos'] = np.asarray(data['pair_infos'])
    new_file = '/MY_PATH/npz_mess/' + filename
    np.savez(new_file, data['image_paths'], data['depth_paths'], data['intrinsics'],data['poses'],data['pair_infos'])
    print("Saved to ", new_file)

错误原因 #

报错核心是元组内元素形状不一致：每个元组的第一个元素是长度为2的数组，第三个是长度为4的数组，普通numpy数组要求所有元素形状均匀，无法容纳这种非均匀结构，因此触发形状不匹配错误。

解决方案 #

方案1：直接保存元组列表（推荐） #

无需将pair_infos转为numpy数组，直接保存原列表即可，np.savez支持保存Python原生列表，只要加载时开启allow_pickle=True。同时修正代码中image_paths修改无效的问题：

import numpy as np
from numpy import load
import os

directory = '/MY_PATH/scene_info_0.1_0.7'

for filename in os.listdir(directory):
    f_npz = os.path.join(directory,filename)
    data = load(f_npz, allow_pickle = True)
    
    # 修正：通过索引修改数组元素，而非循环变量副本
    for i in range(len(data['image_paths'])):
        image_path = data['image_paths'][i]
        if image_path is not None:
            if 'Undistorted_SfM' in image_path:
                data['image_paths'][i] = image_path.replace('Undistorted_SfM', 'phoenix/S6/zl548/MegaDepth_v1')
            else:
                print("Undistorted_SfM NOT IN IMAGE PATH??")
    
    # 移除转数组的代码，直接保存原列表
    new_file = '/MY_PATH/npz_mess/' + filename
    # 用键值对保存，方便后续加载时通过键访问
    np.savez(new_file, 
             image_paths=data['image_paths'], 
             depth_paths=data['depth_paths'], 
             intrinsics=data['intrinsics'],
             poses=data['poses'],
             pair_infos=data['pair_infos'])
    print("Saved to ", new_file)

方案2：转为object类型numpy数组 #

如果需要将pair_infos转为数组，可以指定dtype=object，让数组容纳任意类型的元素：

import numpy as np
from numpy import load
import os

directory = '/MY_PATH/scene_info_0.1_0.7'

for filename in os.listdir(directory):
    f_npz = os.path.join(directory,filename)
    data = load(f_npz, allow_pickle = True)
    
    # 修正image_paths修改逻辑
    for i in range(len(data['image_paths'])):
        image_path = data['image_paths'][i]
        if image_path is not None:
            if 'Undistorted_SfM' in image_path:
                data['image_paths'][i] = image_path.replace('Undistorted_SfM', 'phoenix/S6/zl548/MegaDepth_v1')
            else:
                print("Undistorted_SfM NOT IN IMAGE PATH??")
    
    # 转为object类型数组
    data['pair_infos'] = np.array(data['pair_infos'], dtype=object)
    new_file = '/MY_PATH/npz_mess/' + filename
    np.savez(new_file, 
             image_paths=data['image_paths'], 
             depth_paths=data['depth_paths'], 
             intrinsics=data['intrinsics'],
             poses=data['poses'],
             pair_infos=data['pair_infos'])
    print("Saved to ", new_file)

加载时仍需使用load(..., allow_pickle=True)才能正确读取object类型数组。

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