(
tensor: th.Tensor,
x_max: Optional[float] = 1.0,
x_min: Optional[float] = 0.0,
mode: str = "rgb",
mask: Optional[th.Tensor] = None,
label: Optional[str] = None,
)
| 68 | |
| 69 | |
| 70 | def tensor2image( |
| 71 | tensor: th.Tensor, |
| 72 | x_max: Optional[float] = 1.0, |
| 73 | x_min: Optional[float] = 0.0, |
| 74 | mode: str = "rgb", |
| 75 | mask: Optional[th.Tensor] = None, |
| 76 | label: Optional[str] = None, |
| 77 | ) -> np.ndarray: |
| 78 | |
| 79 | tensor = tensor.detach() |
| 80 | |
| 81 | # Apply mask |
| 82 | if mask is not None: |
| 83 | tensor = tensor * mask |
| 84 | |
| 85 | if len(tensor.size()) == 2: |
| 86 | tensor = tensor[None] |
| 87 | |
| 88 | # Make three channel image |
| 89 | assert len(tensor.size()) == 3, tensor.size() |
| 90 | n_channels = tensor.shape[0] |
| 91 | if n_channels == 1: |
| 92 | tensor = tensor.repeat(3, 1, 1) |
| 93 | elif n_channels != 3: |
| 94 | raise ValueError(f"Unsupported number of channels {n_channels}.") |
| 95 | |
| 96 | # Convert to display format |
| 97 | img = tensor.permute(1, 2, 0) |
| 98 | |
| 99 | if mode == "rgb": |
| 100 | img = tensor2rgb(img, x_max=x_max, x_min=x_min) |
| 101 | elif mode == "jet": |
| 102 | # `cv2.applyColorMap` assumes input format in BGR |
| 103 | img[:, :, :3] = img[:, :, [2, 1, 0]] |
| 104 | img = tensor2rgbjet(img, x_max=x_max, x_min=x_min) |
| 105 | # convert back to rgb |
| 106 | img[:, :, :3] = img[:, :, [2, 1, 0]] |
| 107 | else: |
| 108 | raise ValueError(f"Unsupported mode {mode}.") |
| 109 | |
| 110 | if label is not None: |
| 111 | img = add_label_centered(img, label) |
| 112 | |
| 113 | return img |
| 114 | |
| 115 | # d: b x 1 x H x W |
| 116 | # screenCoords: b x 2 x H X W |
no test coverage detected