| 195 | return cmap |
| 196 | |
| 197 | class Colorize(object): |
| 198 | def __init__(self, n=35): |
| 199 | self.cmap = labelcolormap(n) |
| 200 | self.cmap = torch.from_numpy(self.cmap[:n]) |
| 201 | |
| 202 | def __call__(self, gray_image): |
| 203 | size = gray_image.size() |
| 204 | color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0) |
| 205 | |
| 206 | for label in range(0, len(self.cmap)): |
| 207 | mask = (label == gray_image[0]).cpu() |
| 208 | color_image[0][mask] = self.cmap[label][0] |
| 209 | color_image[1][mask] = self.cmap[label][1] |
| 210 | color_image[2][mask] = self.cmap[label][2] |
| 211 | |
| 212 | return color_image |