Spaces:
Running on Zero
Running on Zero
| from tactile_prediction.model.common.normalizer import SingleFieldLinearNormalizer | |
| from common.pytorch_util import dict_apply, dict_apply_reduce, dict_apply_split | |
| import numpy as np | |
| import torch | |
| def get_range_normalizer_from_stat(stat, output_max=1, output_min=-1, range_eps=1e-7): | |
| # -1, 1 normalization | |
| input_max = stat['max'] | |
| input_min = stat['min'] | |
| input_range = input_max - input_min | |
| ignore_dim = input_range < range_eps | |
| input_range[ignore_dim] = output_max - output_min | |
| scale = (output_max - output_min) / input_range | |
| offset = output_min - scale * input_min | |
| offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim] | |
| return SingleFieldLinearNormalizer.create_manual( | |
| scale=scale, | |
| offset=offset, | |
| input_stats_dict=stat | |
| ) | |
| def get_image_range_normalizer(): | |
| scale = np.array([2], dtype=np.float32) | |
| offset = np.array([-1], dtype=np.float32) | |
| stat = { | |
| 'min': np.array([0], dtype=np.float32), | |
| 'max': np.array([1], dtype=np.float32), | |
| 'mean': np.array([0.5], dtype=np.float32), | |
| 'std': np.array([np.sqrt(1/12)], dtype=np.float32) | |
| } | |
| return SingleFieldLinearNormalizer.create_manual( | |
| scale=scale, | |
| offset=offset, | |
| input_stats_dict=stat | |
| ) | |
| def get_image_identity_normalizer(): | |
| scale = np.array([1], dtype=np.float32) | |
| offset = np.array([0], dtype=np.float32) | |
| stat = { | |
| 'min': np.array([0], dtype=np.float32), | |
| 'max': np.array([1], dtype=np.float32), | |
| 'mean': np.array([0.5], dtype=np.float32), | |
| 'std': np.array([np.sqrt(1/12)], dtype=np.float32) | |
| } | |
| return SingleFieldLinearNormalizer.create_manual( | |
| scale=scale, | |
| offset=offset, | |
| input_stats_dict=stat | |
| ) | |
| def get_image_imagenet_normalizer(): | |
| scale = np.array([1 / 0.229, 1 / 0.224, 1 / 0.225], dtype=np.float32) | |
| offset = np.array([-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], dtype=np.float32) | |
| stat = { | |
| 'min': np.array([0.0, 0.0, 0.0], dtype=np.float32), | |
| 'max': np.array([1.0, 1.0, 1.0], dtype=np.float32), | |
| 'mean': np.array([0.485, 0.456, 0.406], dtype=np.float32), | |
| 'std': np.array([0.229, 0.224, 0.225], dtype=np.float32), | |
| } | |
| return SingleFieldLinearNormalizer.create_manual( | |
| scale=scale, | |
| offset=offset, | |
| input_stats_dict=stat | |
| ) | |
| def get_identity_normalizer_from_stat(stat): | |
| scale = np.ones_like(stat['min']) | |
| offset = np.zeros_like(stat['min']) | |
| return SingleFieldLinearNormalizer.create_manual( | |
| scale=scale, | |
| offset=offset, | |
| input_stats_dict=stat | |
| ) | |
| def array_to_stats(arr: np.ndarray): | |
| stat = { | |
| 'min': np.min(arr, axis=0), | |
| 'max': np.max(arr, axis=0), | |
| 'mean': np.mean(arr, axis=0), | |
| 'std': np.std(arr, axis=0) | |
| } | |
| return stat | |
| def concatenate_normalizer(normalizers: list): | |
| scale = torch.concatenate([normalizer.params_dict['scale'] for normalizer in normalizers], axis=-1) | |
| offset = torch.concatenate([normalizer.params_dict['offset'] for normalizer in normalizers], axis=-1) | |
| input_stats_dict = dict_apply_reduce( | |
| [normalizer.params_dict['input_stats'] for normalizer in normalizers], | |
| lambda x: torch.concatenate(x,axis=-1)) | |
| return SingleFieldLinearNormalizer.create_manual( | |
| scale=scale, | |
| offset=offset, | |
| input_stats_dict=input_stats_dict | |
| ) |