Args: tensor: Tensor of shape (..., 3, h, w) in range [-1, 1] or [0, 1] denormalize_zero_one: If True, denormalizes image from range [0, 1] otherwise from [-1, 1] to [0, 255] Returns: Numpy array of shape (h, w, 3) in range [0, 255] if tensor shape is (3,
(tensor, denormalize_zero_one=False)
| 26 | |
| 27 | |
| 28 | def tensor2im(tensor, denormalize_zero_one=False): |
| 29 | """ |
| 30 | Args: |
| 31 | tensor: Tensor of shape (..., 3, h, w) in range [-1, 1] or [0, 1] |
| 32 | denormalize_zero_one: If True, denormalizes image from range [0, 1] otherwise |
| 33 | from [-1, 1] to [0, 255] |
| 34 | Returns: |
| 35 | Numpy array of shape (h, w, 3) in range [0, 255] if tensor shape is (3, h, w) |
| 36 | or (1, 3, h, w). Otherwise, returns array of shape (..., h, w, 3). |
| 37 | """ |
| 38 | if len(tensor.shape) == 4 and tensor.shape[0] == 1: |
| 39 | tensor = tensor[0] |
| 40 | if isinstance(tensor, torch.Tensor): |
| 41 | tensor = tensor.detach().cpu().numpy() |
| 42 | im = einops.rearrange(tensor, '... c h w -> ... h w c') |
| 43 | if denormalize_zero_one: |
| 44 | im = im * 255. |
| 45 | else: |
| 46 | im = (im + 1.) * 127.5 |
| 47 | im = np.clip(im, 0, 255).astype(np.uint8) |
| 48 | return im |
| 49 | |
| 50 | |
| 51 | class ImageVisualizer: |