NumPy + PyTorch Tensor assignment

lets assume we have a tensor representing an image of the shape (910, 270, 1) which assigned a number (some index) to each pixel with width=910 and height=270.

We also have a numpy array of size (N, 3) which maps a 3-tuple to an index.

I now want to create a new numpy array of shape (920, 270, 3) which has a 3-tuple based on the original tensor index and the mapping-3-tuple-numpy array. How do I do this assignment without for loops and other consuming iterations?

This would look simething like:

color_image = np.zeros((self._w, self._h, 3), dtype=np.int32)
self._colors = np.array(N,3) # this is already present
indexed_image = torch.tensor(920,270,1) # this is already present

#how do I assign it to this numpy array?
color_image[indexed_image.w, indexed_image.h] = self._colors[indexed_image.flatten()]

Answer

Assuming you have _colors, and indexed_image. Something that ressembles to:

>>> indexed_image = torch.randint(0, 10, (920, 270, 1))
>>> _colors  = np.random.randint(0, 255, (N, 3))

A common way of converting a dense map to a RGB map is to loop over the label set:

>>> _colors = torch.FloatTensor(_colors)

>>> rgb = torch.zeros(indexed_image.shape[:-1] + (3,))
>>> for lbl in range(N):
...     rgb[lbl == indexed_image[...,0]] = _colors[lbl]