Copying parts of numpy array to another array with extra dimension

Question

I am trying to use machine learning for semantic segmentation and I managed to find a way to get the proper one hot encoding (using this : https://www.jeremyjordan.me/semantic-segmentation/) however the code that I obtain is quite bad and I am certain that numpy has functionalities that could provide a more elegant solution.

The idea is the following: from a label array (88,240,240) creating a new array (88,240,240,3) with the proper values in each channel.

I came up with this:

def data_reshape(train_image_list, train_label_list, img_size):
    temp = np.empty(shape=[train_label_list.shape[0], img_size[1], img_size[0], 3])
    temp[:,:,:,0] = train_label_list
    temp[temp[:,:,:,0] > 0] = 2
    temp[temp[:,:,:,0] == 0] = 1
    temp[temp[:,:,:,0] == 2] = 0

    temp[:,:,:,1] = train_label_list
    temp[temp[:,:,:,1] == 2] = 0

    temp[:,:,:,2] = train_label_list
    temp[temp[:,:,:,2] < 2] = 0
    temp[temp[:,:,:,2] == 2] = 1
    train_image_list = np.reshape(train_image_list, newshape=[-1, img_size[1], img_size[0], 1])
    train_label_list = np.reshape(temp, newshape=[-1, img_size[1], img_size[0], 3])

    return train_image_list, train_label_list

EDIT:

It actually doesn't run as it should

I will reformulate:

I have a numpy array : (88,240,240) which contains information for 3 different labels on each of the 88 images (0 for pixel of label_0, 1 for pixel of label_1,2 for pixels of the label_2).

I want to come out of my function with a numpy array with 3 more channels each containing different information:

• (88,240,240,0) would have the label_0 pixels with a value of 1 (the rest would be 0)
• (88,240,240,1) would have the label_1 pixels with a value of 1 (the rest would be 0)
• (88,240,240,2) would have the label_2 pixels with a value of 1 (the rest would be 0)

Does anyone has a suggestion ?

Kind regards,

Unic0

Answer 1

train_label_list has values 0,1,2 and you want to expand it to 3 channels. Is that right?

temp = np.zeros(shape=[train_label_list.shape[0], img_size[1], img_size[0], 3])
temp[:, :, :, 0] = train_label_list == 0
temp[:, :, :, 1] = train_label_list == 1
temp[:, :, :, 2] = train_label_list == 2

This should do the trick.