| 260 | |
| 261 | |
| 262 | class Colorize(object): |
| 263 | def __init__(self, n=35): |
| 264 | self.cmap = labelcolormap(n) |
| 265 | self.cmap = torch.from_numpy(self.cmap[:n]) |
| 266 | |
| 267 | def __call__(self, gray_image): |
| 268 | size = gray_image.size() |
| 269 | color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0) |
| 270 | |
| 271 | for label in range(0, len(self.cmap)): |
| 272 | mask = (label == gray_image[0]).cpu() |
| 273 | color_image[0][mask] = self.cmap[label][0] |
| 274 | color_image[1][mask] = self.cmap[label][1] |
| 275 | color_image[2][mask] = self.cmap[label][2] |
| 276 | |
| 277 | return color_image |