如何高效基于标签掩码计算多通道图像的目标像素统计量？

问题背景 #

用tifffile.imread读取得到3D图像矩阵（维度为Y×X×通道，通道数不固定），对应2D标签掩码（0代表背景，1、2……对应不同目标），需要计算每个目标、每个通道的均值、中位数、标准差、最小值和最大值，但原代码运行速度极慢，求优化方案。

原代码：

sample = imread(input_img)
label_mask = np.load(input_mask)

n_stains = sample.shape[2]
n_labels = np.max(label_mask)

# 创建空DataFrame存储强度统计结果
intensity_measurements = pd.DataFrame(columns = ['sample', 'label', 'stain', 'mean', 'median', 'std', 'min', 'max'])

for label in range(1, n_labels+1):
    for stain in range(n_stains):
        # 提取对应通道和标签的像素
        stain_label = sample[:,:,stain][label_mask == label]

        # 计算统计量
        mean = np.mean(stain_label)
        median = np.median(stain_label)
        std = np.std(stain_label)
        min = np.min(stain_label)
        max = np.max(stain_label)

        # 将结果添加到DataFrame
        intensity_measurements = intensity_measurements.append({'sample' : args.input_img, 'label': label, 'stain': stain, 'mean': mean, 'median': median, 'std': std, 'min': min, 'max': max}, ignore_index=True)

原代码性能瓶颈 #

1. 双重循环+重复索引：每次循环都执行布尔索引提取像素，重复计算掩码，冗余操作多。
2. DataFrame.append低效：append每次都会创建新的DataFrame实例，频繁内存拷贝，数据量大时性能骤降。
3. 单统计量多次遍历：每个统计量单独调用numpy函数，多次遍历同一像素数组。

优化方案 #

方案1：预收集数据+一次性生成DataFrame #

先批量收集所有统计结果，最后一次性生成DataFrame，避免频繁的append操作：

import numpy as np
import pandas as pd
from tifffile import imread

sample = imread(input_img)
label_mask = np.load(input_mask)
sample_path = args.input_img

# 获取所有非背景标签
labels = np.unique(label_mask)
labels = labels[labels != 0]
n_stains = sample.shape[2]

# 预存储所有统计结果
stats_list = []

for stain in range(n_stains):
    channel_data = sample[:, :, stain]
    for label in labels:
        pixels = channel_data[label_mask == label]
        stats_list.append({
            'sample': sample_path,
            'label': label,
            'stain': stain,
            'mean': np.mean(pixels),
            'median': np.median(pixels),
            'std': np.std(pixels),
            'min': np.min(pixels),
            'max': np.max(pixels)
        })

# 一次性生成DataFrame
intensity_measurements = pd.DataFrame(stats_list)

方案2：用scipy.stats批量获取统计量 #

scipy.stats.describe可一次性返回均值、方差、最值等统计量，减少函数调用次数：

import numpy as np
import pandas as pd
from tifffile import imread
from scipy import stats

sample = imread(input_img)
label_mask = np.load(input_mask)
sample_path = args.input_img

labels = np.unique(label_mask)
labels = labels[labels != 0]
n_stains = sample.shape[2]

stats_list = []

for stain in range(n_stains):
    channel_data = sample[:, :, stain]
    for label in labels:
        pixels = channel_data[label_mask == label]
        # 一次性获取多统计量（不含中位数，需单独计算）
        desc = stats.describe(pixels, ddof=1)  # ddof=1对应样本标准差
        stats_list.append({
            'sample': sample_path,
            'label': label,
            'stain': stain,
            'mean': desc.mean,
            'median': np.median(pixels),
            'std': np.sqrt(desc.variance),
            'min': desc.minmax[0],
            'max': desc.minmax[1]
        })

intensity_measurements = pd.DataFrame(stats_list)

方案3：用skimage.measure.regionprops批量处理 #

regionprops是专门针对标签掩码的区域统计工具，支持多通道输入，性能最优：

import numpy as np
import pandas as pd
from tifffile import imread
from skimage.measure import regionprops

sample = imread(input_img)
label_mask = np.load(input_mask)
sample_path = args.input_img

# 将图像转置为（通道，Y，X）格式适配regionprops
sample_transposed = np.transpose(sample, (2, 0, 1))
# 批量提取所有区域的统计属性
props = regionprops(label_mask, intensity_image=sample_transposed)

stats_list = []
for prop in props:
    label = prop.label
    for stain in range(sample_transposed.shape[0]):
        # 提取当前通道的像素
        region_pixels = sample_transposed[stain][prop.coords[:, 0], prop.coords[:, 1]]
        stats_list.append({
            'sample': sample_path,
            'label': label,
            'stain': stain,
            'mean': prop.intensity_mean[stain],
            'median': np.median(region_pixels),
            'std': np.std(region_pixels),
            'min': prop.min_intensity[stain],
            'max': prop.max_intensity[stain]
        })

intensity_measurements = pd.DataFrame(stats_list)

方案4：numpy完全向量化（无循环） #

利用bincount实现分组求和，结合pandas分组计算中位数等统计量，适合标签连续的场景：

import numpy as np
import pandas as pd
from tifffile import imread

sample = imread(input_img)
label_mask = np.load(input_mask)
sample_path = args.input_img

n_stains = sample.shape[2]
# 展平图像和掩码
flat_mask = label_mask.flatten()
flat_sample = sample.reshape(-1, n_stains)

# 过滤背景像素
valid_idx = flat_mask != 0
flat_mask_valid = flat_mask[valid_idx]
flat_sample_valid = flat_sample[valid_idx]

# 用pandas分组计算所有统计量
df = pd.DataFrame({
    'label': flat_mask_valid,
    **{f'stain_{s}': flat_sample_valid[:, s] for s in range(n_stains)}
})

grouped = df.groupby('label')
stats_list = []
for label, group in grouped:
    for stain in range(n_stains):
        col = f'stain_{stain}'
        stats_list.append({
            'sample': sample_path,
            'label': label,
            'stain': stain,
            'mean': group[col].mean(),
            'median': group[col].median(),
            'std': group[col].std(),
            'min': group[col].min(),
            'max': group[col].max()
        })

intensity_measurements = pd.DataFrame(stats_list)

性能对比 #

原代码因双重循环+频繁append，在大图像多标签场景下耗时数分钟；优化方案中，regionprops或向量化方法可将速度提升10~100倍，具体取决于图像规模和标签数量。

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