如何高效为numpy二维灰度图像的同值非零区域分配连续整数标记？

Great question! When dealing with hundreds of thousands of images, avoiding Python-level loops is critical for performance. Let's break down the best approaches to map each unique non-zero pixel value to a sequential integer label (starting at 1) efficiently.

Core Idea #

Each non-zero unique pixel value represents a distinct object. We need to:

1. Identify all unique non-zero values in an image
2. Map each value to a unique integer label (1, 2, ..., n where n is the number of unique non-zero values)
3. Apply this mapping to the entire image (and scale this to thousands of images without slow Python loops)

1. Single Image Processing (Fully Vectorized) #

For a single 2D NumPy array (img), we can use NumPy's built-in np.unique with the return_inverse parameter—this is a C-optimized operation that avoids all Python loops:

import numpy as np

def label_single_image(img):
    # Create a mask for non-zero pixels
    non_zero_mask = img != 0
    
    # Handle edge case: all-zero image
    if not np.any(non_zero_mask):
        return np.zeros_like(img, dtype=np.int32)
    
    # Get unique non-zero values and their inverse indices (for mapping)
    unique_vals, inverse_indices = np.unique(img[non_zero_mask], return_inverse=True)
    
    # Initialize labeled image with zeros (background)
    labeled_img = np.zeros_like(img, dtype=np.int32)
    
    # Assign labels: inverse indices start at 0, so add 1 to shift to 1-based
    labeled_img[non_zero_mask] = inverse_indices + 1
    
    return labeled_img

This method is extremely fast because all operations are executed in optimized C code, no Python-level iteration involved.

2. Batch Processing for Hundreds of Thousands of Images #

For large datasets, we need to avoid looping through each image in Python. The best options are:

Option A: Numba JIT Compilation (Parallelized)

Numba compiles Python code to machine code, and with parallelization, it can handle tens of thousands of images efficiently. This is ideal if you can load all images into a 3D NumPy array (shape=(num_images, height, width)):

import numba

@numba.jit(nopython=True, parallel=True)
def batch_label_images(imgs):
    num_imgs, height, width = imgs.shape
    labeled_imgs = np.zeros((num_imgs, height, width), dtype=np.int32)
    
    # Parallelize over images (uses all CPU cores)
    for i in numba.prange(num_imgs):
        img = imgs[i]
        non_zero = img[img != 0]
        
        if non_zero.size == 0:
            continue  # Skip all-zero images
        
        # Get unique values and create a mapping dictionary
        unique_vals = np.unique(non_zero)
        val_to_label = {v: idx + 1 for idx, v in enumerate(unique_vals)}
        
        # Assign labels (compiled to machine code, no Python overhead)
        for y in range(height):
            for x in range(width):
                val = img[y, x]
                if val != 0:
                    labeled_imgs[i, y, x] = val_to_label[val]
    
    return labeled_imgs

• parallel=True uses all available CPU cores to process images simultaneously.
• The nested loops over pixels are compiled to machine code, so they're just as fast as native C loops.

Option B: Dask for Out-of-Core Processing (If Memory is Limited)

If you can't fit all images into RAM, use Dask to process images in chunks. Dask mimics NumPy's API but works with datasets larger than memory:

import dask.array as da

# Assume `dask_imgs` is a Dask array of shape (num_images, height, width)
def dask_label_image(img):
    non_zero_mask = img != 0
    unique_vals, inverse = da.unique(img[non_zero_mask], return_inverse=True)
    labeled = da.zeros_like(img, dtype=np.int32)
    labeled = da.where(non_zero_mask, inverse + 1, labeled)
    return labeled

# Apply the function to each image in the Dask array
labeled_dask_imgs = dask_imgs.map_blocks(dask_label_image, dtype=np.int32)

# Compute results (or save directly to disk)
labeled_imgs = labeled_dask_imgs.compute()

Key Performance Notes #

• Avoid Python loops at all costs: Python-level iteration over images or pixels is orders of magnitude slower than optimized C/LLVM code.
• Use appropriate dtypes: For labels, int32 is sufficient (supports up to 2 billion unique labels, way more than any single image will have).
• Handle edge cases: Always check for all-zero images to avoid errors in np.unique.

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