基于欧氏距离的掩码生成：PyOpenCL等跨平台实现求助

一、Numba 实现（CPU/AMD GPU 通用） #

Numba通过JIT编译加速，支持CPU和AMD GPU（需安装ROCm），语法接近Numpy，迁移成本低。

1. 二进制掩码生成 #

import numba
import numpy as np

@numba.jit(nopython=True, parallel=True)
def generate_mask_numba(coords, shape, radius):
    mask = np.zeros(shape, dtype=np.bool_)
    ndim = len(shape)
    # 遍历每个坐标点
    for idx in numba.prange(coords.shape[0]):
        coord = coords[idx]
        # 计算当前坐标的有效范围，避免越界
        starts = []
        ends = []
        slices = []
        for d in range(ndim):
            start = max(0, coord[d] - radius)
            end = min(shape[d], coord[d] + radius + 1)
            starts.append(start)
            ends.append(end)
            slices.append(slice(start, end))
        # 生成局部网格并计算距离平方
        grid = np.meshgrid(*[np.arange(s, e) for s, e in zip(starts, ends)], indexing='ij')
        dist_sq = 0.0
        for d in range(ndim):
            dist_sq += (grid[d] - coord[d]) ** 2
        # 标记符合距离条件的区域
        mask[tuple(slices)] |= (dist_sq <= radius ** 2)
    return mask

2. 连通域标记（替代scipy.ndimage.label） #

@numba.jit(nopython=True)
def label_mask_numba(mask):
    ndim = mask.ndim
    labels = np.zeros_like(mask, dtype=np.int32)
    current_label = 1
    # 生成全连通邻域偏移量
    offsets = []
    for d in range(ndim):
        offsets.append(-1)
        offsets.append(0)
        offsets.append(1)
    offsets = np.array(np.meshgrid(*[offsets]*ndim)).reshape(ndim, -1).T
    offsets = offsets[np.any(offsets != 0, axis=1)]

    for idx in np.ndindex(mask.shape):
        if mask[idx] and labels[idx] == 0:
            # BFS遍历连通域
            queue = [idx]
            labels[idx] = current_label
            while queue:
                pos = queue.pop(0)
                for off in offsets:
                    new_pos = tuple(p + o for p, o in zip(pos, off))
                    # 检查边界有效性
                    valid = True
                    for d in range(ndim):
                        if new_pos[d] < 0 or new_pos[d] >= mask.shape[d]:
                            valid = False
                            break
                    if valid and mask[new_pos] and labels[new_pos] == 0:
                        labels[new_pos] = current_label
                        queue.append(new_pos)
            current_label += 1
    return labels, current_label - 1

二、PyOpenCL 实现（跨平台GPU/CPU） #

PyOpenCL支持所有OpenCL兼容设备（AMD、Intel、NVIDIA GPU，甚至CPU），核心是编写OpenCL内核。

1. 二进制掩码生成 #

import pyopencl as cl
import numpy as np

def generate_mask_opencl(coords, shape, radius):
    # 初始化OpenCL环境
    ctx = cl.create_some_context()
    queue = cl.CommandQueue(ctx)
    
    # 准备设备内存
    mask = np.zeros(shape, dtype=np.bool_)
    mf = cl.mem_flags
    mask_buf = cl.Buffer(ctx, mf.WRITE_ONLY, mask.nbytes)
    coords_buf = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=coords.astype(np.int32))
    
    # OpenCL内核代码
    kernel_code = """
    __kernel void generate_mask(__global bool *mask, __global const int *coords, 
                                int num_coords, int ndim, __global const int *shape, int radius) {
        int idx = get_global_id(0);
        int pos[3];
        int temp = idx;
        // 计算全局索引对应的空间坐标
        for (int d = ndim-1; d >=0; d--) {
            pos[d] = temp % shape[d];
            temp /= shape[d];
        }
        
        bool is_in_mask = false;
        for (int c = 0; c < num_coords; c++) {
            int coord[3];
            for (int d=0; d<ndim; d++) {
                coord[d] = coords[c*ndim + d];
            }
            // 计算距离平方并判断是否在半径内
            int dist_sq = 0;
            bool out_of_range = false;
            for (int d=0; d<ndim; d++) {
                int diff = pos[d] - coord[d];
                if (abs(diff) > radius) {
                    out_of_range = true;
                    break;
                }
                dist_sq += diff * diff;
            }
            if (!out_of_range && dist_sq <= radius*radius) {
                is_in_mask = true;
                break;
            }
        }
        mask[idx] = is_in_mask;
    }
    """
    
    # 编译并执行内核
    prg = cl.Program(ctx, kernel_code).build()
    global_size = (np.prod(shape),)
    prg.generate_mask(queue, global_size, None, mask_buf, coords_buf, 
                      np.int32(coords.shape[0]), np.int32(len(shape)), 
                      np.int32(shape), np.int32(radius))
    # 读取结果到主机内存
    cl.enqueue_copy(queue, mask, mask_buf).wait()
    return mask

三、TensorFlow 实现（跨平台，支持CPU/GPU） #

TensorFlow支持所有主流硬件，内置连通域标记函数，适合快速开发。

1. 二进制掩码生成 #

import tensorflow as tf

def generate_mask_tensorflow(coords, shape, radius):
    coords = tf.convert_to_tensor(coords, dtype=tf.int32)
    shape = tf.convert_to_tensor(shape, dtype=tf.int32)
    radius = tf.convert_to_tensor(radius, dtype=tf.int32)
    
    # 生成全局网格坐标
    grids = tf.meshgrid(*[tf.range(s) for s in shape], indexing='ij')
    grid = tf.stack(grids, axis=-1)
    
    # 计算每个网格点到所有坐标点的最小距离平方
    coords_expanded = tf.expand_dims(coords, axis=tuple(range(len(shape))))
    diff = grid - coords_expanded
    dist_sq = tf.reduce_sum(tf.square(diff), axis=-1)
    min_dist_sq = tf.reduce_min(dist_sq, axis=0)
    
    mask = min_dist_sq <= tf.square(radius)
    return mask.numpy()

2. 连通域标记 #

直接使用TensorFlow内置函数：

def label_mask_tensorflow(mask):
    mask = tf.convert_to_tensor(mask, dtype=tf.bool)
    # 返回标签数组（从1开始）和组件数量
    labels, num_components = tf.nn.connected_components(mask)
    return labels.numpy(), num_components.numpy()

方案选择建议 #

• 追求简洁快速开发：优先选TensorFlow，内置函数稳定，跨平台支持完善。
• 追求极致性能且熟悉JIT编译：选Numba，CPU/GPU通用，语法接近Numpy，学习成本低。
• 需要完全控制硬件细节或适配特殊OpenCL设备：选PyOpenCL，但需编写内核代码，调试成本较高。

内容的提问来源于stack exchange，提问作者João Mamede