title: Copying parts of numpy array to another array with extra dimension
body: <p>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 : <a href="https://www.jeremyjordan.me/semantic-segmentation/" rel="nofollow noreferrer">https://www.jeremyjordan.me/semantic-segmentation/</a>) however the code that I obtain is quite bad and I am certain that numpy has functionalities that could provide a more elegant solution.</p>
  
  <p>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.</p>
  
  <p>I came up with this:</p>
  
  <pre><code>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] &gt; 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] &lt; 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
  </code></pre>
  
  <p><strong>EDIT:</strong> </p>
  
  <p>It actually doesn't run as it should</p>
  
  <p>I will reformulate:</p>
  
  <p>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).</p>
  
  <p>I want to come out of my function with a numpy array with 3 more channels each containing different information:</p>
  
  <ul>
  <li>(88,240,240,0) would have the label_0 pixels with a value of 1 (the rest would be 0)</li>
  <li>(88,240,240,1) would have the label_1 pixels with a value of 1 (the rest would be 0)</li>
  <li>(88,240,240,2) would have the label_2 pixels with a value of 1 (the rest would be 0)</li>
  </ul>
  
  <p>Does anyone has a suggestion ?</p>
  
  <p>Kind regards,</p>
  
  <p>Unic0</p>