超大JSON文件加载清洗求助：Snapshot Serengeti数据集处理

问题背景 #

我正在开展一项基于Snapshot Serengeti数据集的图像分类项目，数据集包含一个5GB以上的超大JSON文件，内含多个顶级键，但训练仅需提取其中"images": [{...}, {...}, ...]数组中的数据。该文件体积过大，无法直接打开读取或存入字典。

图像条目格式 #

图像条目的JSON结构如下：

{
"id": "S1/B04/B04_R1/S1_B04_R1_PICT0003",
"file_name": "S1/B04/B04_R1/S1_B04_R1_PICT0003.JPG",
"frame_num": 1,
"seq_id": "SER_S1#B04#1#3",
"width": 2048,
"height": 1536,
"corrupt": false,
"location": "B04",
"seq_num_frames": 1,
"datetime": "2010-07-20 06:14:06"
},

尝试过的方法及问题 #

1. 分块读取失败 #

尝试按100MB分块读取文件，但文件存在格式问题（单引号、NaN值），未预处理直接解析会报错，尝试代码如下：

with open(labels_json) as f:         
    for chunk in iter(lambda: f.read(100*1024*1024), ""):         
        data = json.loads(chunk)

问题：分块内容未必是完整JSON结构，且未处理格式错误直接解析必然失败。

2. 按季节拆分耗尽存储 #

因图像按11个季节分类，我尝试将数据写入11个独立文件以便单独加载，但单个季节未处理完就耗尽云存储，脚本存在写入效率问题，原脚本如下：

import json

labels_json = annotations_directory + "SS_Labels.json"

get_filename = lambda n : f"SS_labels_S{i}.json"

# Define the 11 output files
seasons = {}
started = {}
for i in range(1, 12):
    filename = get_filename(i)
    seasons[i] = open(filename, "w")
    seasons[i].write('[')
    started[i] = False

def seperate_seasons(dir):
    line_num = 0
    decoder = json.JSONDecoder()
    with open(dir, 'r') as labels:
        begin_writing = False
        buffer = []
        id = 1
        for line in labels:
            if not begin_writing: # Begin writing for the line after "images"
                if 'images' in line:
                    begin_writing = True
            else:
                line.replace('NaN', 'null') # clean NaN values
                line.replace("'", '"')      # clean incorrect key values

                buffer.append(line.strip()) # add line to buffer

                getID = lambda l: int(line.split('"')[3].split('/')[0][1])
                if '"id"' in line or "'id'" in line:
                    previous_id = id
                    id = getID(line)        # get id of object

                if line.strip() == '},' or line.strip() == '}': # when the object has finished, write it to the appropriate image folder
                    label = ','.join(buffer)
                    if label[-1] != ',':
                        label += ','

                    if started[id] == False:
                        print(f'Beginning Season {id}')
                        started[id] = True

                        if id != 1:
                            seasons[previous_id].write(']')
                            seasons[previous_id].close()
                            del seasons[previous_id]


                    seasons[id].write(label)                    # add label entry to file

seperate_seasons(labels_json)

# Close all remaining label files
for season in seasons.values():
    season.write(']')
    season.close()

问题：字符串替换未生效（replace返回新字符串但未赋值）、buffer处理冗余导致重复写入、未及时刷新缓冲区、拆分逻辑依赖行格式易出错。

解决方案 #

优化思路 #

1. 使用流式JSON解析库（如ijson）遍历images数组的单个元素，避免加载整个文件到内存。
2. 实时处理格式问题（替换NaN为null、单引号为双引号）。
3. 按季节写入时减少冗余操作，及时刷新缓冲区，避免内存堆积。

优化后代码 #

import ijson
import json

labels_json = annotations_directory + "SS_Labels.json"

# 初始化季节文件句柄和状态
season_files = {}
season_first_entry = {}
for season_num in range(1, 12):
    filename = f"SS_labels_S{season_num}.json"
    f = open(filename, 'w')
    f.write('[')
    season_files[season_num] = f
    season_first_entry[season_num] = True

def clean_json_data(data):
    # 递归处理数据中的NaN值，转换为null
    if isinstance(data, dict):
        return {k: clean_json_data(v) for k, v in data.items()}
    elif isinstance(data, list):
        return [clean_json_data(item) for item in data]
    elif data is float and str(data) == 'nan':
        return None
    else:
        return data

with open(labels_json, 'r', encoding='utf-8') as f:
    # 流式遍历images数组中的每个元素
    for raw_image in ijson.items(f, 'images.item'):
        # 清理数据中的NaN值
        image = clean_json_data(raw_image)
        # 提取季节编号（从id字段的Sx中取x）
        season_num = int(image['id'].split('/')[0][1])
        
        # 转成标准JSON字符串
        image_str = json.dumps(image)
        
        # 写入对应季节文件
        if season_first_entry[season_num]:
            season_files[season_num].write(image_str)
            season_first_entry[season_num] = False
        else:
            season_files[season_num].write(',' + image_str)
        
        # 每写入一次刷新缓冲区，避免缓存占用过多
        season_files[season_num].flush()

# 关闭所有季节文件，补全JSON数组结尾
for f in season_files.values():
    f.write(']')
    f.close()

关键改进点 #

• 用ijson流式解析，仅加载单个图像元素，内存占用极低。
• 直接处理解析后的字典，避免手动拼接JSON字符串的错误。
• 递归清理NaN值，比字符串替换更可靠。
• 控制逗号添加逻辑，避免冗余逗号问题。
• 实时刷新缓冲区，减少磁盘缓存占用，避免云存储突发写入量过大。

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