title: How to visualize a matrix of categories as an RGB image?
body: <p>I am using neural network to do semantic segmentation(human parsing), something like taking a photo of people as input and the neural network tells that every pixel is most likely to be head, leg, background or some other parts of human. The algorithm runs smoothly and giving a <code>numpy.ndarray</code> as output . The shape of the array is <code>(1,23,600,400)</code>, where 600*400 is the resolution of the input image and 23 is the number of categories. The 3d matrix looks like a 23-layer stacked 2d matrices, where each layer using a matrix of float to tell the possibility that each pixel is of that category.</p>
  
  <p>To visualize the matrix like the following figure, I used <code>numpy.argmax</code> to squash the 3d matrix into a 2d matrix that holds the index of the most possible category. But I don't have any idea how to proceed to get the visualization I want.</p>
  
  <p><a href="https://i.stack.imgur.com/tgDp2.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/tgDp2.png" alt="The desired visualization effect"></a></p>
  
  <h1>EDIT</h1>
  
  <p>Actually, I can do it in a trivial way. That is, use a for loop to traverse through every pixel and assign a color to it to get a image. However, this is not a vectorized coding, since numpy has built-in way to speed up matrix manipulation. And I need to save CPU cycles for real time segmentation.</p>