高效压缩算法

摘要

医学图像分割掩膜通常是高度稀疏的，这为专门的压缩算法提供了优化空间。本文旨在对比 MedMask 格式与常用的 NIfTI (.nii.gz) 和 NumPy (.npz) 格式在存储二值（Binary）和多标签（Multi-label）掩膜时的压缩效率、编码及解码速度。最后，我们将深入探究 MedMask 的核心技术——PackBits 预处理结合 Zstandard 压缩——为何能针对稀疏掩膜实现卓越的压缩性能。

1. 统一基准测试框架

为了避免代码重复并确保测试的公平性，我们首先定义一个统一的基准测试框架。该框架可以处理不同的数据格式，并返回标准化的性能指标（文件大小、编码时间、解码时间）。

import os
import time
import gzip
import tempfile
import numpy as np
import nibabel as nib
import zstandard as zstd
from pathlib import Path
from medmask import SegmentationMask
from spacetransformer import Space
import pandas as pd

# --- 统一的基准测试函数 ---
def run_benchmark(format_type, data, space=None, label_mapping=None, original_path=None):
    """
    对指定格式运行压缩和解压基准测试。

    Args:
        format_type (str): 'nifti', 'npz', 'medmask'
        data (np.ndarray): 掩膜数据
        space (Space, optional): MedMask/NIfTI所需的空间信息.
        label_mapping (dict, optional): MedMask所需的多标签映射.
        original_path (Path, optional): NIfTI格式的原始路径，用于直接获取文件大小.

    Returns:
        dict: 包含 size, encode_time, decode_time 的字典.
    """
    stats = {'size': 0, 'encode_time': 0, 'decode_time': 0}
    
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_path = Path(tmpdir) / f"tempfile.{format_type.split('_')[0]}"

        # --- 编码过程 ---
        encode_start = time.time()
        if format_type == 'nifti':
            if original_path:
                stats['size'] = original_path.stat().st_size
            else: # 用于创建新的多标签 NIfTI 文件
                affine =  np.eye(4)
                nii_img = nib.Nifti1Image(data.astype(np.uint8), affine)
                tmp_path_nii = tmp_path.with_suffix('.nii.gz')
                nib.save(nii_img, tmp_path_nii)
                stats['size'] = tmp_path_nii.stat().st_size
            # NIfTI 的编码/解码时间通常包含在加载过程中，这里不单独测量以简化对比
            pass 
        elif format_type == 'npz':
            np.savez_compressed(tmp_path, mask=data)
            stats['size'] = tmp_path.stat().st_size
        elif format_type == 'medmask':
            mask_obj = SegmentationMask(data, label_mapping, space=space)
            mask_obj.save(tmp_path)
            stats['size'] = tmp_path.stat().st_size
        stats['encode_time'] = (time.time() - encode_start) * 1000  # ms

        # --- 解码过程 ---
        if stats['size'] > 0: # 仅在成功编码后解码
            decode_start = time.time()
            if format_type == 'nifti':
                 # 为保持一致性，不单独测量 NIfTI 的解码时间
                 pass
            elif format_type == 'npz':
                _ = np.load(tmp_path)['mask']
            elif format_type == 'medmask':
                loaded_mask = SegmentationMask.load(tmp_path)
                # 模拟实际使用，提取所有标签
                if label_mapping:
                    for name in label_mapping.keys():
                        _ = loaded_mask.get_binary_mask_by_names(name)
            stats['decode_time'] = (time.time() - decode_start) * 1000 # ms
            
    return stats

# --- 数据加载 ---
mask_dir = Path('dicube-testdata/mask/s0000')
with open(mask_dir / 'nonzero_masks.txt', 'r') as f:
    valid_files = [line.strip() for line in f.readlines()]
print(f"加载了 {len(valid_files)} 个有效的掩膜文件用于测试。")

2. 二值掩膜（Binary Mask）性能对比

在此部分，我们测试单个器官掩膜的压缩性能。这类掩膜只包含一个标签，是典型的二值稀疏数据。

测试配置与执行

我们选择一系列大小和稀疏度各不相同的掩膜文件进行测试。

# 选择测试文件
binary_test_files = [
    'gluteus_maximus_right.nii.gz', # 大掩膜
    'urinary_bladder.nii.gz',       # 大掩膜
    'colon.nii.gz',                 # 中等掩膜
    'iliopsoas_left.nii.gz',        # 中等掩膜
    'iliac_artery_left.nii.gz',     # 小掩膜
    'small_bowel.nii.gz'            # 极小掩膜
]

binary_results = []

for fname in binary_test_files:
    original_path = mask_dir / fname
    nii_img = nib.load(original_path)
    mask_data = nii_img.get_fdata().astype(np.uint8)
    space = Space.from_nifti(nii_img)
    organ_name = fname.replace('.nii.gz', '')
    
    nifti_stats = run_benchmark('nifti', data=mask_data, original_path=original_path)
    npz_stats = run_benchmark('npz', data=mask_data)
    medmask_stats = run_benchmark('medmask', data=mask_data, space=space, label_mapping={organ_name: 1})
    
    binary_results.append({
        '文件名': fname.replace('.nii.gz', ''),
        '非零像素': np.count_nonzero(mask_data),
        'NIfTI (KB)': nifti_stats['size'] / 1024,
        'NPZ (KB)': npz_stats['size'] / 1024,
        'MedMask (KB)': medmask_stats['size'] / 1024,
        'NPZ 编码 (ms)': npz_stats['encode_time'],
        'MedMask 编码 (ms)': medmask_stats['encode_time'],
        'NPZ 解码 (ms)': npz_stats['decode_time'],
        'MedMask 解码 (ms)': medmask_stats['decode_time']
    })

df_binary = pd.DataFrame(binary_results)

# --- 增加平均值行 ---
if not df_binary.empty:
    avg_row = df_binary.select_dtypes(include=np.number).mean()
    avg_row['文件名'] = 'Average'
    avg_row = avg_row.reindex(df_binary.columns, fill_value='-')
    df_binary = pd.concat([df_binary, pd.DataFrame([avg_row])], ignore_index=True)

df_binary['MedMask压缩比 (vs NIfTI)'] = df_binary['NIfTI (KB)'] / df_binary['MedMask (KB)']
df_binary['NPZ压缩比 (vs NIfTI)'] = df_binary['NIfTI (KB)'] / df_binary['NPZ (KB)']

print("--- 二值掩膜压缩性能对比 ---")
display(df_binary.style.format({
    'NIfTI (KB)': '{:.1f}', 'NPZ (KB)': '{:.1f}', 'MedMask (KB)': '{:.1f}',
    'NPZ 编码 (ms)': '{:.1f}', 'MedMask 编码 (ms)': '{:.1f}',
    'NPZ 解码 (ms)': '{:.1f}', 'MedMask 解码 (ms)': '{:.1f}',
    'MedMask压缩比 (vs NIfTI)': '{:.1f}x', 'NPZ压缩比 (vs NIfTI)': '{:.1f}x',
    '非零像素': '{:,.0f}'
}).hide(axis="index"))

二值掩膜结论

从上表可以清晰地看出：

1. 压缩效率: MedMask 的压缩效率显著优于 NPZ 和 NIfTI。如 Average 行所示，MedMask 的平均文件大小远低于其他两者。对于稀疏程度高的小目标（如 iliac_artery_left），MedMask 的压缩比优势尤为突出。数据越稀疏，其优势越明显。
2. 存储大小: 总体而言，MedMask 能有效降低文件大小，通常能达到数倍乃至数十倍的压缩效果，将文件体积从几十KB量级降低到几KB甚至更低。
3. 处理速度: MedMask 的编解码速度与 NPZ 相当，有时甚至更快，完全满足高性能应用的需求。

3. 多标签掩膜（Multi-label Mask）性能对比

在实际应用中，常需要将多个器官掩膜存储在同一个文件中。我们构建一个包含多个标签的数组，然后对比 NIfTI、NPZ 和 MedMask 格式在存储这同一个多标签数组时的性能差异。

测试配置与执行

我们将相关的器官掩膜组合成逻辑分组，并融合成一个多标签数组进行测试。

test_groups = [
    {
        'name': '臀肌群',
        'files': ['gluteus_maximus_left.nii.gz', 'gluteus_maximus_right.nii.gz', 
                  'gluteus_medius_left.nii.gz', 'gluteus_medius_right.nii.gz'],
    },
    {
        'name': '股骨与髋骨',
        'files': ['femur_left.nii.gz', 'femur_right.nii.gz', 
                  'hip_left.nii.gz', 'hip_right.nii.gz'],
    },
    {
        'name': '盆腔器官',
        'files': ['urinary_bladder.nii.gz', 'colon.nii.gz', 'small_bowel.nii.gz'],
    }
]

multilabel_results = []

for group in test_groups:
    # 1. 创建多标签数据
    first_img = nib.load(mask_dir / group['files'][0])
    space = Space.from_nifti(first_img)
    multilabel_array = np.zeros(first_img.shape, dtype=np.uint8)
    label_mapping = {}
    
    for i, fname in enumerate(group['files']):
        fpath = mask_dir / fname
        if not fpath.exists(): continue
        
        data = nib.load(fpath).get_fdata().astype(np.uint8)
        
        # 构建多标签数组
        label_value = i + 1
        organ_name = fname.replace('.nii.gz', '')
        multilabel_array[data > 0] = label_value
        label_mapping[organ_name] = label_value

    # 2. 对同一个多标签数组，用不同格式进行基准测试
    nifti_stats = run_benchmark('nifti', data=multilabel_array, space=space)
    npz_stats = run_benchmark('npz', data=multilabel_array)
    medmask_stats = run_benchmark('medmask', data=multilabel_array, space=space, label_mapping=label_mapping)
    
    multilabel_results.append({
        '测试组': group['name'],
        '标签数': len(group['files']),
        'NIfTI (KB)': nifti_stats['size'] / 1024,
        'NPZ (KB)': npz_stats['size'] / 1024,
        'MedMask (KB)': medmask_stats['size'] / 1024,
        'NPZ 编码 (ms)': npz_stats['encode_time'],
        'MedMask 编码 (ms)': medmask_stats['encode_time'],
        'NPZ 解码 (ms)': npz_stats['decode_time'],
        'MedMask 解码 (ms)': medmask_stats['decode_time']
    })

df_multi = pd.DataFrame(multilabel_results)

# --- 增加平均值行 ---
if not df_multi.empty:
    avg_row = df_multi.select_dtypes(include=np.number).mean()
    avg_row['测试组'] = 'Average'
    avg_row = avg_row.reindex(df_multi.columns, fill_value='-')
    df_multi = pd.concat([df_multi, pd.DataFrame([avg_row])], ignore_index=True)

df_multi['MedMask压缩比 (vs NIfTI)'] = df_multi['NIfTI (KB)'] / df_multi['MedMask (KB)']
df_multi['NPZ压缩比 (vs NIfTI)'] = df_multi['NIfTI (KB)'] / df_multi['NPZ (KB)']

print("--- 多标签掩膜压缩性能对比 ---")
display(df_multi.style.format({
    '独立NIfTI (KB)': '{:.1f}', '独立NPZ (KB)': '{:.1f}', 'MedMask合并 (KB)': '{:.1f}',
    'NIfTI (KB)': '{:.1f}', 'NPZ (KB)': '{:.1f}', 'MedMask (KB)': '{:.1f}',
    'NPZ 编码 (ms)': '{:.1f}', 'MedMask 编码 (ms)': '{:.1f}',
    'NPZ 解码 (ms)': '{:.1f}', 'MedMask 解码 (ms)': '{:.1f}',
    'MedMask压缩比 (vs NIfTI)': '{:.1f}x', 'NPZ压缩比 (vs NIfTI)': '{:.1f}x'
}).hide(axis="index"))

多标签掩膜结论

在处理包含多个标签的单一掩膜文件时，MedMask 同样展现出卓越的压缩性能。从上表可以看出，存储相同的多标签数据，MedMask 格式生成的文件体积显著小于 NIfTI 和 NPZ。

4. 核心技术探究：PackBits + Zstandard 的威力

MedMask 的高压缩率源于其针对不同掩膜类型的双层压缩策略。对于二值掩膜（Binary Mask），它采用 PackBits + Zstandard (Zstd) 的组合；而对于多标签掩膜，则直接使用 Zstd。本节将重点探究为何 PackBits 预处理能为二值掩膜带来显著的性能提升。

PackBits 预处理机制

MedMask 中使用的 PackBits 是一种针对二值稀疏数据的位打包（Bit Packing）技术，而非传统的游程编码（Run-Length Encoding）。其核心思想是将多个布尔值（在数组中通常以 uint8 的 0 或 1 存储）压缩到单个字节的位（bit）中。由于一个 uint8 字节包含8个位，该算法可以：

1. 将8个连续的 uint8 类型的掩膜像素值（每个占用1字节）打包成一个 uint8（占用1字节）。
2. 每个原始像素（0 或 1）映射到新字节中的一个位（0 或 1）。

工作示例： 假设有8个连续的像素值 [0, 1, 0, 0, 0, 0, 0, 1]。在内存中，它们占用8个字节。经过 PackBits 处理后，它们会被编码成单个字节。该字节的二进制表示为 10000010（注：位的顺序取决于具体实现），在十进制中为 130。这样，仅此一步就实现了理论上接近8:1的数据压缩。

这种机制决定了 PackBits 预处理仅适用于二值数据。对于多标签掩膜，像素值可以大于1（例如 2, 3, 4, ...），无法用单个位来表示，因此 MedMask 会跳过此步骤，直接对原始多标签数组应用 Zstd 压缩。这也解释了为何 MedMask 在处理二值掩膜时的压缩比通常优于处理多标签掩膜。

为了量化 PackBits 预处理的有效性，我们对比纯 Zstd 压缩和 MedMask (PackBits + Zstd) 在处理同一个二值掩膜时的效果。

# 加载中等大小的掩膜数据
mask_path = mask_dir / 'urinary_bladder.nii.gz'
img = nib.load(mask_path)
data = img.get_fdata()>0
data_bytes = data.tobytes()
data_packbit_bytes = np.packbits(data).tobytes()

# 1. 原始大小
raw_size = len(data_bytes)
packbit_size = len(data_packbit_bytes)

# 2. 纯 Gzip 压缩 (NIfTI/NPZ 使用)
gzip_size = len(gzip.compress(data_bytes))

# 3. 纯 Zstandard 压缩
zstd_size = len(zstd.ZstdCompressor().compress(data_bytes))

# 4. MedMask (PackBits + Zstd) 压缩
packbit_zstd_size = len(zstd.ZstdCompressor().compress(data_packbit_bytes))


# 结果汇总
packbits_analysis = {
    "方法": ["原始数据 (Bytes)","Packbits数据 (Bytes)", "纯 Gzip", "纯 Zstandard", "MedMask (PackBits + Zstd)"],
    "大小 (KB)": [raw_size / 1024, packbit_size/1024, gzip_size / 1024, zstd_size / 1024, packbit_zstd_size / 1024],
    "压缩比 (vs 原始)": [1.0, raw_size/packbit_size, raw_size / gzip_size, raw_size / zstd_size, raw_size / packbit_zstd_size]
}
df_packbits = pd.DataFrame(packbits_analysis)

print("--- 压缩策略对稀疏数据的效果对比 ---")
print(f"测试对象: {mask_path.name}, 稀疏度: {1 - np.count_nonzero(data) / data.size:.2%}")
display(df_packbits.style.format({
    "大小 (KB)": "{:.2f}",
    "压缩比 (vs 原始)": "{:.1f}x"
}).hide(axis="index"))

结论

综合以上分析，MedMask 之所以能成为一种高效的医学分割掩膜存储格式，其核心在于它根据数据特性采用的自适应双层压缩策略。

1. 对于二值掩膜，MedMask 采用 PackBits 预处理 + Zstandard 压缩 的模式。
   • PackBits 作为一种位打包预处理，首先消除了用整个字节（uint8）存储单个布尔值所带来的结构性冗余，实现了第一重压缩。
   • 经过预处理的紧凑数据流，再由高效的 Zstandard 算法进行第二重压缩。这种针对性的组合策略，使其压缩比远超仅依赖 Gzip (DEFLATE) 的 NIfTI 和 NPZ 格式，且性能与数据稀疏度正相关。
2. 对于多标签掩膜，由于无法进行位打包，MedMask 直接应用 Zstandard 进行压缩。虽然缺少了 PackBits 带来的巨额增益，但 Zstandard 本身的性能依然优于传统的 DEFLATE 算法，因此 MedMask 在存储多标签数据时仍比 NIfTI 和 NPZ 更具空间效率。

总而言之，MedMask 的设计精髓在于它并非简单应用通用压缩，而是通过领域特定的预处理（位打包）来最大化后续通用压缩算法（Zstd）的效率。这种设计使其在显著降低存储空间的同时，保持了高性能的读写速度，为医学图像分析工作流提供了切实的优化。