Can someone help me using 'ray' or other multiprocessing libraries for my semantic segmentation?

Question

I'm doing semantic segmentation for microscope image stacks. My code works fine but the thing is it only uses one core of my CPU which makes me wait a long time to get the segmented images.

I recently knew that there are ways to use multicore processing with other python libraries, but I don't know how to implement it.

So someone can help me edit my code with one of the multiprocessing libraries? My code is in the below.

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

model = keras.models.load_model("C:/mymodel.h5", compile=False)

  

#creating recon image directory
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(large_image_path + large_image_name, 0)
        patches = patchify(img, (256, 256), step=256)
        
        predicted_patches = []
        for i in range(patches.shape[0]):
            for j in range(patches.shape[1]):
                
                single_patch = patches[i,j,:,:] #(256, 256)
                single_patch_norm = normalize(np.array(single_patch), axis=1) 
                single_patch_input = np.stack((single_patch_norm,)*3, axis=-1) # (256, 256, 3) 
                single_patch_input = np.expand_dims(single_patch_input, 0) #(1,256,256,3)
                
                single_patch_prediction = (model.predict(single_patch_input)[0,:,:,0]>0.5).astype(np.uint8) 
               
                predicted_patches.append(single_patch_prediction)
                
        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)
        
        cv2.imwrite(recon_image_directory + "/recon"+'_' + str(num) +  ".tif", reconstructed_image)

Answer 1

Does this snippet work? It should run each prediction in a separate process.

@ray.remote
def predict(large_image_name: str) -> None:
    img = cv2.imread(large_image_path + large_image_name, 0)
    patches = patchify(img, (256, 256), step=256)
    
    predicted_patches = []
    for i in range(patches.shape[0]):
        for j in range(patches.shape[1]):
            
            single_patch = patches[i,j,:,:] #(256, 256)
            single_patch_norm = normalize(np.array(single_patch), axis=1) 
            single_patch_input = np.stack((single_patch_norm,)*3, axis=-1) # (256, 256, 3) 
            single_patch_input = np.expand_dims(single_patch_input, 0) #(1,256,256,3)
            
            single_patch_prediction = (model.predict(single_patch_input)[0,:,:,0]>0.5).astype(np.uint8) 
            
            predicted_patches.append(single_patch_prediction)
            
    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)
    
    cv2.imwrite(recon_image_directory + "/recon"+'_' + str(num) +  ".tif", reconstructed_image)
    
futures = []
for num, large_image_name in tqdm(enumerate(check_images), total=len(check_images)):
    if (large_image_name.split('.')[1] == "tif"):
        futures.append(predict.remote(large_image_name))

ray.get(futures)

We also have a high-level abstraction for doing this sort of thing. If you're interested, you should check out the Ray AI Runtime (AIR).