| 7 | |
| 8 | |
| 9 | def tensor2im(img, imtype=np.uint8, unnormalize=True, idx=0, nrows=None): |
| 10 | # select a sample or create grid if img is a batch |
| 11 | if len(img.shape) == 4: |
| 12 | nrows = nrows if nrows is not None else int(math.sqrt(img.size(0))) |
| 13 | img = img[idx] if idx >= 0 else torchvision.utils.make_grid(img, nrows) |
| 14 | |
| 15 | img = img.cpu().float() |
| 16 | if unnormalize: |
| 17 | mean = [0.5, 0.5, 0.5] |
| 18 | std = [0.5, 0.5, 0.5] |
| 19 | |
| 20 | for i, m, s in zip(img, mean, std): |
| 21 | i.mul_(s).add_(m) |
| 22 | |
| 23 | image_numpy = img.numpy() |
| 24 | image_numpy_t = np.transpose(image_numpy, (1, 2, 0)) |
| 25 | image_numpy_t = image_numpy_t*254.0 |
| 26 | |
| 27 | return image_numpy_t.astype(imtype) |
| 28 | |
| 29 | def tensor2maskim(mask, imtype=np.uint8, idx=0, nrows=1): |
| 30 | im = tensor2im(mask, imtype=imtype, idx=idx, unnormalize=False, nrows=nrows) |