Multiclass Segmentation using u-net, Data format (how to label (mask it)

from future import print_function
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
import os
import glob
import skimage.io as io
import skimage.transform as trans

def adjustData(img,mask,flag_multi_class,num_class):
if(flag_multi_class):
img = img / 255
mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]
new_mask = np.zeros(mask.shape + (num_class,))
for i in range(num_class):
#for one pixel in the image, find the class in mask and convert it into one-hot vector
#index = np.where(mask == i)
#index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)
#new_mask[index_mask] = 1
new_mask[mask == i,i] = 1
new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))
mask = new_mask
elif(np.max(img) > 1):
img = img / 255
mask = mask /255
mask[mask > 0.5] = 1
mask[mask <= 0.5] = 0
return (img,mask)

def trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = “grayscale”,
mask_color_mode = “grayscale”,image_save_prefix = “image”,mask_save_prefix = “mask”,
flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):
‘’’
can generate image and mask at the same time
use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same
if you want to visualize the results of generator, set save_to_dir = “your path”
‘’’
image_datagen = ImageDataGenerator(**aug_dict)
mask_datagen = ImageDataGenerator(**aug_dict)
image_generator = image_datagen.flow_from_directory(
train_path,
classes = [image_folder],
class_mode = None,
color_mode = image_color_mode,
target_size = target_size,
batch_size = batch_size,
save_to_dir = save_to_dir,
save_prefix = image_save_prefix,
seed = seed)
mask_generator = mask_datagen.flow_from_directory(
train_path,
classes = [mask_folder],
class_mode = None,
color_mode = mask_color_mode,
target_size = target_size,
batch_size = batch_size,
save_to_dir = save_to_dir,
save_prefix = mask_save_prefix,
seed = seed)
train_generator = zip(image_generator, mask_generator)
for (img,mask) in train_generator:
img,mask = adjustData(img,mask,flag_multi_class,num_class)
yield (img,mask)

def valGenerator(batch_size,train_path,image_folder,mask_folder,image_color_mode = “grayscale”,
mask_color_mode = “grayscale”,image_save_prefix = “image”,mask_save_prefix = “mask”,
flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):
‘’’
can generate image and mask at the same time
use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same
if you want to visualize the results of generator, set save_to_dir = “your path”
‘’’
image_datagen = ImageDataGenerator()
mask_datagen = ImageDataGenerator()
image_generator = image_datagen.flow_from_directory(
train_path,
classes = [image_folder],
class_mode = None,
color_mode = image_color_mode,
target_size = target_size,
batch_size = batch_size,
save_to_dir = save_to_dir,
save_prefix = image_save_prefix,
seed = seed)
mask_generator = mask_datagen.flow_from_directory(
train_path,
classes = [mask_folder],
class_mode = None,
color_mode = mask_color_mode,
target_size = target_size,
batch_size = batch_size,
save_to_dir = save_to_dir,
save_prefix = mask_save_prefix,
seed = seed)
train_generator = zip(image_generator, mask_generator)
for (img,mask) in train_generator:
img,mask = adjustData(img,mask,flag_multi_class,num_class)
yield (img,mask)

def testGenerator(test_path,num_image = 18,target_size = (256,256),flag_multi_class = False,as_gray = True):
for i in test_path:
img = io.imread(i, as_gray = as_gray)
print(i)

img = img / 255

    img = trans.resize(img,target_size)
    img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img
    img = np.reshape(img,(1,)+img.shape)
    yield img

def testGenerator2(batch_size,test_path,target_size = (256,256), image_color_mode = “grayscale”, num_class = 2,flag_multi_class = False,as_gray = True):
image_datagen = ImageDataGenerator()
image_generator = image_datagen.flow_from_directory(
test_path,
class_mode = None,
color_mode = image_color_mode,
target_size = target_size,
batch_size = batch_size
)
return image_generator

def geneTrainNpy(image_path,mask_path,flag_multi_class = False,num_class = 2,image_prefix = “image”,mask_prefix = “mask”,image_as_gray = True,mask_as_gray = True):
image_name_arr = glob.glob(os.path.join(image_path,"%s*.png"%image_prefix))
image_arr = []
mask_arr = []
for index,item in enumerate(image_name_arr):
img = io.imread(item,as_gray = image_as_gray)
img = np.reshape(img,img.shape + (1,)) if image_as_gray else img
mask = io.imread(item.replace(image_path,mask_path).replace(image_prefix,mask_prefix),as_gray = mask_as_gray)
mask = np.reshape(mask,mask.shape + (1,)) if mask_as_gray else mask
img,mask = adjustData(img,mask,flag_multi_class,num_class)
image_arr.append(img)
mask_arr.append(mask)
image_arr = np.array(image_arr)
mask_arr = np.array(mask_arr)
return image_arr,mask_arr

def labelVisualize(num_class,color_dict,img):
img = img[:,:,0] if len(img.shape) == 3 else img
img_out = np.zeros(img.shape + (3,))
for i in range(num_class):
img_out[img == i,:] = color_dict[i]
return img_out

def labelVisualizeBinary(img):
img[img > 0.5] = 1
img[img<= 0.5] = 0
return img

def saveResult(save_path,results_filename,npyfile,flag_multi_class = False,num_class = 2):
print(results_filename[0])
for i,item in enumerate(npyfile):
img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]
print(i)
io.imsave(os.path.join(save_path,results_filename[i]),img)

my problem

this is my data code , i have problem how to change it to multiclass segmentation,
i divided codes with 3 part , 1 for data(to transform (num classes), 1 for train(like epoch,path_to_train,etc) , and 1 for model (unet)

in my model i was change the loss function to sparse_categorical_crossentropy, and the last layer to softmax,i know my problem in the data codes (how to transform it multiclass (num classes), i has 3 output (black,white, red) that mean my num_classes = 3), but i dont know how to change it? i just change the num_classes but nothing happend , maybe you can solve my problem thank you :slight_smile: (codes up here)

I suggest you to try FastAI 2 unet_leaner, it is mounted above pytorch. Semantic segmentation is very easy with that library

if you have a link or codes can you give it to me? maybe i can try it… (m

thank you i will try it later , but i want to try use my codes use models etc, just has simple problem but dont know how to solve that (num_classes and how to give it to label)

i see the fastai that"s good but that make me confused :), where save_path ,where the classes, how to label (ex: red,black,white) etc , thank you

There is a segmentation data loader that does this for you, the mask is an image with int values each value represents a class. In that reposotory are linked some videos that explain how to use it!!!

You are welcome !