求助：为显微镜图像栈语义分割代码实现多核心并行处理

针对单CPU处理图像分割耗时过长的问题，下面用Ray多进程框架改造代码，核心是并行处理每个图像补丁的预测任务，同时通过Ray Actor避免重复加载模型的开销：

步骤1：安装Ray #

先确保安装Ray库：

pip install ray

改造后的完整代码 #

import numpy as np
from patchify import patchify, unpatchify
import os
import cv2
from tqdm import tqdm
from tensorflow import keras
from tensorflow.keras.utils import normalize
import natsort
import ray

# 初始化Ray，自动使用全部CPU核心，也可手动指定num_cpus参数
ray.init(num_cpus=os.cpu_count())

# 创建Ray Actor，每个Worker进程仅加载一次模型，避免重复加载开销
@ray.remote(num_cpus=1)
class ModelPredictor:
    def __init__(self, model_path):
        self.model = keras.models.load_model(model_path, compile=False)
    
    def predict_patch(self, patch):
        # 补丁预处理逻辑和原代码一致
        patch_norm = normalize(np.array(patch), axis=1)
        patch_input = np.stack((patch_norm,)*3, axis=-1)
        patch_input = np.expand_dims(patch_input, 0)
        # 预测并二值化结果
        prediction = (self.model.predict(patch_input)[0,:,:,0] > 0.5).astype(np.uint8)
        return prediction

# 初始化模型预测Actor
model_path = "C:/mymodel.h5"
predictor = ModelPredictor.remote(model_path)

# 创建重建图像存储目录
recon_image_directory = "C:/Users/recon"
if not os.path.exists(recon_image_directory):
    os.makedirs(recon_image_directory)

large_image_path = "C:/original_images/" 
check_images = natsort.natsorted(os.listdir(large_image_path)) 

for num, large_image_name in tqdm(enumerate(check_images), total=len(check_images)):
    if large_image_name.split('.')[1] == "tif":
        img = cv2.imread(os.path.join(large_image_path, large_image_name), 0)
        patches = patchify(img, (256, 256), step=256)
        
        # 批量提交补丁预测的异步任务
        prediction_tasks = []
        for i in range(patches.shape[0]):
            for j in range(patches.shape[1]):
                single_patch = patches[i,j,:,:]
                task = predictor.predict_patch.remote(single_patch)
                prediction_tasks.append(task)
        
        # 等待所有任务完成，按原顺序获取预测结果
        predicted_patches = ray.get(prediction_tasks)
        predicted_patches = np.array(predicted_patches)
        
        # 重建完整分割图像
        predicted_patches_reshaped = np.reshape(predicted_patches, (patches.shape[0], patches.shape[1], 256,256))
        reconstructed_image = unpatchify(predicted_patches_reshaped, img.shape)
        
        # 保存结果
        save_path = os.path.join(recon_image_directory, f"recon_{num}.tif")
        cv2.imwrite(save_path, reconstructed_image)

# 关闭Ray资源
ray.shutdown()

核心改动说明 #

• Ray Actor管理模型：通过Actor类让每个Worker进程只加载一次模型，既节省重复加载的时间，也解决TensorFlow模型在多进程中的兼容性问题。
• 并行处理补丁：将每个补丁的预测任务异步提交，批量获取结果时保持原补丁顺序，充分利用多CPU核心的计算能力。
• 路径优化：用os.path.join替代字符串拼接，避免不同系统下的路径分隔符问题。

注意事项 #

• 若模型体积较大，可适当降低ray.init的num_cpus参数（比如设为CPU核心数的一半），避免内存占用过高。
• 若不需要GPU加速，可在Actor的__init__方法中添加tf.config.set_visible_devices([], 'GPU')，减少多进程下的GPU资源冲突。

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