(label_tensor, n_label, imtype=np.uint8, tile=False)
| 97 | |
| 98 | # Converts a one-hot tensor into a colorful label map |
| 99 | def tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False): |
| 100 | if label_tensor.dim() == 4: |
| 101 | # transform each image in the batch |
| 102 | images_np = [] |
| 103 | for b in range(label_tensor.size(0)): |
| 104 | one_image = label_tensor[b] |
| 105 | one_image_np = tensor2label(one_image, n_label, imtype) |
| 106 | images_np.append(one_image_np.reshape(1, *one_image_np.shape)) |
| 107 | images_np = np.concatenate(images_np, axis=0) |
| 108 | if tile: |
| 109 | images_tiled = tile_images(images_np) |
| 110 | return images_tiled |
| 111 | else: |
| 112 | images_np = images_np[0] |
| 113 | return images_np |
| 114 | |
| 115 | if label_tensor.dim() == 1: |
| 116 | return np.zeros((64, 64, 3), dtype=np.uint8) |
| 117 | if n_label == 0: |
| 118 | return tensor2im(label_tensor, imtype) |
| 119 | label_tensor = label_tensor.cpu().float() |
| 120 | if label_tensor.size()[0] > 1: |
| 121 | label_tensor = label_tensor.max(0, keepdim=True)[1] |
| 122 | label_tensor = Colorize(n_label)(label_tensor) |
| 123 | label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0)) |
| 124 | result = label_numpy.astype(imtype) |
| 125 | return result |
| 126 | |
| 127 | |
| 128 | def save_image(image_numpy, image_path, create_dir=False): |
nothing calls this directly
no test coverage detected