| import os |
| import cv2 |
| import time |
| import random |
| import numpy as np |
|
|
| import torch |
| from torchmetrics import MeanMetric as TorchMeanMetric |
|
|
| from nerv.utils.tensor import to_numpy |
|
|
|
|
| def datetime2str(form='%Y-%m-%d_%H-%M-%S'): |
| datetime = time.strftime(form, time.localtime()) |
| return datetime |
|
|
|
|
| def assert_array_shape(xyz, shapes=()): |
| """Check array shape. |
| |
| Args: |
| xyz (np.ndarray): array |
| shape (tuple of tuple of ints, optional): possible target shapes, |
| -1 means arbitrary. Defaults to ((-1, 3)). |
| """ |
| if not shapes: |
| raise ValueError('"shapes" cannot be empty') |
|
|
| if isinstance(shapes[0], int): |
| shapes = (shapes, ) |
|
|
| flags = {x: True for x in range(len(shapes))} |
| for idx, shape in enumerate(shapes): |
| if len(xyz.shape) != len(shape): |
| flags[idx] = False |
|
|
| for dim, num in enumerate(shape): |
| if num == -1: |
| continue |
| elif xyz.shape[dim] != num: |
| flags[idx] = False |
| if sum(flags.values()) == 0: |
| raise ValueError( |
| f'Input array {xyz.shape} is not in target shapes {shapes}!') |
|
|
|
|
| def array_equal(a, b): |
| """Compare if two arrays are the same. |
| |
| Args: |
| a/b: can be np.ndarray or torch.Tensor. |
| """ |
| if a.shape != b.shape: |
| return False |
| try: |
| assert (a == b).all() |
| return True |
| except: |
| return False |
|
|
|
|
| def array_in_list(array, lst): |
| """Judge whether an array is in a list.""" |
| for v in lst: |
| if array_equal(array, v): |
| return True |
| return False |
|
|
|
|
| def set_seed(seed=1, deterministic=False): |
| """Set the random seed of the environment for reproducibility. |
| |
| Args: |
| seed (int): the random seed to set. |
| deterministic (bool, optional): whether to use deterministic torch |
| backend. Default: False. |
| """ |
| print('Using random seed:', seed) |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| if deterministic: |
| print('Using deterministic pytorch backends') |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
|
|
|
|
| class AverageMeter(object): |
| """Computes and stores the average and current value.""" |
|
|
| def __init__(self, device=torch.device('cpu')): |
| self.reset() |
| self.device = device |
|
|
| def reset(self): |
| self.val = 0. |
| self.avg = 0. |
| self.sum = 0. |
| self.count = 0. |
|
|
| def update(self, val, n=1): |
| self.val = val |
| self.sum += val * n |
| self.count += n |
| self.avg = self.sum / self.count |
|
|
| def compute(self): |
| return torch.tensor(self.avg).to(self.device) |
|
|
| def to(self, device): |
| self.device = device |
| return self |
|
|
|
|
| class MeanMetric: |
| """Compatibility wrapper around `torchmetrics.MeanMetric`. |
| |
| Provides a small API expected by the codebase: `update(value, weight=1.0)`, |
| `compute()` and `to(device)`. Works with torchmetrics 1.x and 2.x. |
| """ |
|
|
| def __init__(self, device, nan_strategy="warn", compute_on_step=True, dist_sync_on_step=False, process_group=None, dist_sync_fn=None): |
| self.device = device |
| |
| self._metric = TorchMeanMetric() |
| try: |
| self._metric.to(device) |
| except Exception: |
| pass |
|
|
| def update(self, value, weight=1.0): |
| |
| |
| try: |
| return self._metric.update(value, weight) |
| except TypeError: |
| return self._metric.update(value) |
|
|
| def compute(self): |
| out = self._metric.compute() |
| if isinstance(out, torch.Tensor): |
| return out.to(self.device) |
| try: |
| return torch.tensor(out).to(self.device) |
| except Exception: |
| return out |
|
|
| def to(self, device): |
| try: |
| self._metric.to(device) |
| except Exception: |
| pass |
| self.device = device |
| return self |
|
|
|
|
| def _convert4save_img(array): |
| """Convert a image array to be saveable.""" |
| |
| if len(array.shape) == 2: |
| array = np.stack([array] * 3, axis=-1) |
| |
| if array.shape[0] == 3: |
| array = array.transpose(1, 2, 0) |
| return np.ascontiguousarray(array) |
|
|
|
|
| def _convert4save_video(array): |
| """Convert a video array to be saveable.""" |
| |
| if len(array.shape) == 3: |
| array = np.stack([array] * 3, axis=-1) |
| |
| if array.shape[1] == 3: |
| array = array.transpose(0, 2, 3, 1) |
| return np.ascontiguousarray(array) |
|
|
|
|
| def convert4save(array, is_video=False): |
| """Check the dtype and value range of input array for save. |
| |
| Need to convert to [(T), H, W, C] with np.uint8 value range [0, 255]. |
| |
| Args: |
| array (np.ndarray or torch.Tensor): array to be converted. |
| is_video (bool, optional): whether the array is a video or image. |
| Default: False (means array is an image). |
| |
| Returns: |
| np.ndarray: the converted array ready for save (image or video). |
| """ |
| array = to_numpy(array) |
| if 'int' in str(array.dtype): |
| assert 0 <= array.min() <= array.max() <= 255 |
| elif 'float' in str(array.dtype): |
| assert 0. <= array.min() <= array.max() <= 1. |
| array = np.round(array * 255.).astype(np.uint8) |
| if is_video: |
| return _convert4save_video(array) |
| else: |
| return _convert4save_img(array) |
|
|
|
|
| def save_video(video, save_path, fps=30, codec='mp4v'): |
| """video: np.ndarray of shape [M, H, W, 3]""" |
| os.makedirs(os.path.dirname(save_path), exist_ok=True) |
| video = convert4save(video, is_video=True) |
| video = video[..., [2, 1, 0]] |
| H, W = video.shape[-3:-1] |
| assert save_path.split('.')[-1] == 'mp4' |
| |
| size = [W, H] |
| out = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*codec), fps, size) |
| for i in range(video.shape[0]): |
| out.write(video[i]) |
| out.release() |
|
|
