Spaces:
Running on Zero
Running on Zero
| import glob | |
| import os | |
| import re | |
| from pathlib import Path | |
| import torch | |
| from .scaler import MinMaxScaler | |
| def increment_path(path, exist_ok=False, sep="", mkdir=False): | |
| # Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc. | |
| path = Path(path) # os-agnostic | |
| if path.exists() and not exist_ok: | |
| suffix = path.suffix | |
| path = path.with_suffix("") | |
| dirs = glob.glob(f"{path}{sep}*") # similar paths | |
| matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs] | |
| i = [int(m.groups()[0]) for m in matches if m] # indices | |
| n = max(i) + 1 if i else 2 # increment number | |
| path = Path(f"{path}{sep}{n}{suffix}") # update path | |
| dir = path if path.suffix == "" else path.parent # directory | |
| if not dir.exists() and mkdir: | |
| dir.mkdir(parents=True, exist_ok=True) # make directory | |
| return path | |
| class Normalizer: | |
| def __init__(self, data, symmetric_indices=()): | |
| flat = data.reshape(-1, data.shape[-1]) | |
| self.scaler = MinMaxScaler((-1, 1), clip=True) | |
| self.scaler.fit(flat) | |
| self.symmetric_indices = tuple(int(index) for index in symmetric_indices) | |
| if self.symmetric_indices: | |
| self.make_symmetric(self.symmetric_indices) | |
| def make_symmetric(self, indices): | |
| """Use zero-centred symmetric MinMax ranges for selected channels.""" | |
| scaler = self.scaler | |
| indices = torch.as_tensor(indices, dtype=torch.long) | |
| max_abs = torch.maximum( | |
| scaler.data_min_[indices].abs(), scaler.data_max_[indices].abs() | |
| ) | |
| max_abs = max_abs.clamp_min(torch.finfo(max_abs.dtype).eps) | |
| scaler.data_min_[indices] = -max_abs | |
| scaler.data_max_[indices] = max_abs | |
| scaler.data_range_[indices] = 2.0 * max_abs | |
| feature_min, feature_max = scaler.feature_range | |
| scaler.scale_[indices] = (feature_max - feature_min) / (2.0 * max_abs) | |
| scaler.min_[indices] = 0.0 | |
| self.symmetric_indices = tuple(int(index) for index in indices.tolist()) | |
| def normalize(self, x): | |
| batch, seq, ch = x.shape | |
| x = x.reshape(-1, ch) | |
| return self.scaler.transform(x).reshape((batch, seq, ch)) | |
| def unnormalize(self, x): | |
| batch, seq, ch = x.shape | |
| x = x.reshape(-1, ch) | |
| x = torch.clip(x, -1, 1) # clip to force compatibility | |
| return self.scaler.inverse_transform(x).reshape((batch, seq, ch)) | |
| def vectorize_many(data): | |
| # given a list of batch x seqlen x joints? x channels, flatten all to batch x seqlen x -1, concatenate | |
| batch_size = data[0].shape[0] | |
| seq_len = data[0].shape[1] | |
| out = [x.reshape(batch_size, seq_len, -1).contiguous() for x in data] | |
| global_pose_vec_gt = torch.cat(out, dim=2) | |
| return global_pose_vec_gt | |