| 21 | from matplotlib import transforms |
| 22 | |
| 23 | def renorm(img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) \ |
| 24 | -> torch.FloatTensor: |
| 25 | # img: tensor(3,H,W) or tensor(B,3,H,W) |
| 26 | # return: same as img |
| 27 | assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim() |
| 28 | if img.dim() == 3: |
| 29 | assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % (img.size(0), str(img.size())) |
| 30 | img_perm = img.permute(1,2,0) |
| 31 | mean = torch.Tensor(mean) |
| 32 | std = torch.Tensor(std) |
| 33 | img_res = img_perm * std + mean |
| 34 | return img_res.permute(2,0,1) |
| 35 | else: # img.dim() == 4 |
| 36 | assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % (img.size(1), str(img.size())) |
| 37 | img_perm = img.permute(0,2,3,1) |
| 38 | mean = torch.Tensor(mean) |
| 39 | std = torch.Tensor(std) |
| 40 | img_res = img_perm * std + mean |
| 41 | return img_res.permute(0,3,1,2) |
| 42 | |
| 43 | class ColorMap(): |
| 44 | def __init__(self, basergb=[255,255,0]): |