| |
| from typing import Union |
|
|
| import mmcv |
| import numpy as np |
| import torch |
| from torch import Tensor |
|
|
|
|
| def imrenormalize(img: Union[Tensor, np.ndarray], img_norm_cfg: dict, |
| new_img_norm_cfg: dict) -> Union[Tensor, np.ndarray]: |
| """Re-normalize the image. |
| |
| Args: |
| img (Tensor | ndarray): Input image. If the input is a Tensor, the |
| shape is (1, C, H, W). If the input is a ndarray, the shape |
| is (H, W, C). |
| img_norm_cfg (dict): Original configuration for the normalization. |
| new_img_norm_cfg (dict): New configuration for the normalization. |
| |
| Returns: |
| Tensor | ndarray: Output image with the same type and shape of |
| the input. |
| """ |
| if isinstance(img, torch.Tensor): |
| assert img.ndim == 4 and img.shape[0] == 1 |
| new_img = img.squeeze(0).cpu().numpy().transpose(1, 2, 0) |
| new_img = _imrenormalize(new_img, img_norm_cfg, new_img_norm_cfg) |
| new_img = new_img.transpose(2, 0, 1)[None] |
| return torch.from_numpy(new_img).to(img) |
| else: |
| return _imrenormalize(img, img_norm_cfg, new_img_norm_cfg) |
|
|
|
|
| def _imrenormalize(img: Union[Tensor, np.ndarray], img_norm_cfg: dict, |
| new_img_norm_cfg: dict) -> Union[Tensor, np.ndarray]: |
| """Re-normalize the image.""" |
| img_norm_cfg = img_norm_cfg.copy() |
| new_img_norm_cfg = new_img_norm_cfg.copy() |
| for k, v in img_norm_cfg.items(): |
| if (k == 'mean' or k == 'std') and not isinstance(v, np.ndarray): |
| img_norm_cfg[k] = np.array(v, dtype=img.dtype) |
| |
| if 'bgr_to_rgb' in img_norm_cfg: |
| img_norm_cfg['rgb_to_bgr'] = img_norm_cfg['bgr_to_rgb'] |
| img_norm_cfg.pop('bgr_to_rgb') |
| for k, v in new_img_norm_cfg.items(): |
| if (k == 'mean' or k == 'std') and not isinstance(v, np.ndarray): |
| new_img_norm_cfg[k] = np.array(v, dtype=img.dtype) |
| img = mmcv.imdenormalize(img, **img_norm_cfg) |
| img = mmcv.imnormalize(img, **new_img_norm_cfg) |
| return img |
|
|
