fastai分类模型批量推理遇Pillow像素超限问题求解决方案

核心思路 #

Pillow的这个错误是因为图像解压后的总像素数超过默认限制（178956970，约1.7亿像素）。我们可以在实际加载图像像素数据前先读取图像的尺寸元数据，判断是否超限，若超限则按比例缩小到限制范围内再加载，这样既不用修改Pillow的全局限制，也能处理大图像。

具体实现代码 #

1. 单图像处理函数 #

先写一个通用的图像加载函数，处理超限情况：

from PIL import Image
import math

MAX_PIXELS = 178956970  # Pillow默认限制值

def load_safe_image(image_path):
    # 先打开图像但不加载像素数据
    with Image.open(image_path) as img:
        # 获取图像原始尺寸
        width, height = img.size
        total_pixels = width * height
        
        if total_pixels <= MAX_PIXELS:
            # 像素数合规，直接返回加载后的图像
            return Image.open(image_path)
        else:
            # 计算缩放比例，确保缩放后总像素不超过限制
            scale_factor = math.sqrt(MAX_PIXELS / total_pixels)
            new_width = int(width * scale_factor)
            new_height = int(height * scale_factor)
            
            # 重新打开并缩放图像（用LANCZOS插值保证画质）
            return Image.open(image_path).resize((new_width, new_height), Image.Resampling.LANCZOS)

2. 集成到fastai推理流程 #

如果是用fastai的批量推理管道，可以把这个函数替换默认的图像加载逻辑：

from fastai.vision.all import *

# 自定义图像加载器
def safe_image_loader(path):
    img = load_safe_image(path)
    # 转换为fastai需要的RGB格式
    return img.convert('RGB')

# 构建推理用的DataLoader
dls = ImageDataLoaders.from_folder(
    path='your_inference_images_folder',
    bs=32,  # 你的批次大小
    item_tfms=Resize(224),  # 保持训练时的预处理逻辑
    batch_tfms=Normalize.from_stats(*imagenet_stats),
    loader=safe_image_loader  # 使用自定义安全加载器
)

# 加载训练好的模型并执行批量推理
learn = load_learner('your_model.pkl')
preds, _, probs = learn.get_preds(dl=dls.test_dl(dls.items))

额外说明 #

• 缩放时用Image.Resampling.LANCZOS（Pillow 9.1.0+）或Image.LANCZOS（旧版本），这是画质损失最小的插值算法，适配图像分类场景需求。
• 单张图像推理时，直接调用load_safe_image处理后传入模型即可：

  img = load_safe_image('problem_image.jpg')
  pred, idx, prob = learn.predict(img)
  

• 该方法仅在内存中处理图像数据，不会修改原始文件。

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