# Copyright 2026 Toyota Research Institute. All rights reserved. import torch import random from anydata.utils.types import is_list, is_double_list, is_tuple, is_tensor, is_dict, is_seq, is_str from anydata.utils.decorators import iterate12 import torch.nn.functional as tfn def flatten(lst): """Flatten a list of lists into a list""" return [l for ls in lst for l in ls] if is_double_list(lst) else lst def keys_with(dic, string, without=()): """Return keys from a dictionary that contain a certain string""" return [key for key in dic if string in key and not any(w in key for w in make_list(without))] def keys_in(dic, keys): """Return only keys in a dictionary""" return [key for key in keys if key in dic] def first_key(data): """Returns the first key of a dictionary""" if data is None: return None return list(data.keys())[0] def first_value(data): """Returns the first value of a dictionary""" if data is None: return None return data[first_key(data)] def make_list(var, n=None): """Wraps the input into a list, and optionally repeats it to be size n""" var = var if is_list(var) or is_tuple(var) else [var] if n is None: return var else: assert len(var) == 1 or len(var) == n, 'Wrong list length for make_list' return var * n if len(var) == 1 else var def to_namespace(data): """Convert to namespace""" for key in data.keys(): if isinstance(data[key], dict): data[key] = to_namespace(data[key]) return data def modrem(v, n): """Return round division and remainder""" return v // n, v % n def same_shape(shape1, shape2): """Checks if two shapes are the same""" if len(shape1) != len(shape2): return False for i in range(len(shape1)): if shape1[i] != shape2[i]: return False return True @iterate12 def interleave_dict(val1, val2): """Interleave two dictionaries (A B -> A0 B0 A1 B1 A2 B2 ...)""" if val1.dim() == 4 and val2.dim() == 4: b, c, h, w = val1.shape val1 = val1.permute(0, 2, 3, 1).reshape(b, -1, 3) val2 = val2.permute(0, 2, 3, 1).reshape(b, -1, 3) else: b, c, n = val1.shape val1 = val1.permute(0, 2, 1).reshape(b, -1, 3) val2 = val2.permute(0, 2, 1).reshape(b, -1, 3) out = torch.zeros((b, val1.shape[1] + val2.shape[1], 3), device=val1.device, dtype=val1.dtype) out[:, 0::2], out[:, 1::2] = val1, val2 return out def cat_channel_ones(tensor, n=1): """Concatenate tensor with an extra channel of ones""" # Get tensor shape with 1 channel shape = list(tensor.shape) shape[n] = 1 # Return concatenation of tensor with ones return torch.cat([tensor, torch.ones(shape, device=tensor.device, dtype=tensor.dtype)], n) def norm_pixel_grid(grid, hw=None, in_place=False): """Normalize a pixel grid from [W,H] to [-1,+1].""" if hw is None: hw = grid.shape[-2:] if not in_place: grid = grid.clone() # if align_corners(): grid[:, 0] = 2.0 * grid[:, 0] / (hw[1] - 1) - 1.0 grid[:, 1] = 2.0 * grid[:, 1] / (hw[0] - 1) - 1.0 # else: # grid[:, 0] = 2.0 * grid[:, 0] / hw[1] - 1.0 # grid[:, 1] = 2.0 * grid[:, 1] / hw[0] - 1.0 return grid def unnorm_pixel_grid(grid: torch.Tensor, hw=None, in_place=False): """Unnormalize a pixel grid from [-1,+1] to [W,H].""" if hw is None: hw = grid.shape[-2:] if not in_place: grid = grid.clone() # if align_corners(): grid[:, 0] = 0.5 * (hw[1] - 1) * (grid[:, 0] + 1) grid[:, 1] = 0.5 * (hw[0] - 1) * (grid[:, 1] + 1) # else: # grid[:, 0] = 0.5 * hw[1] * (grid[:, 0] + 1) # grid[:, 1] = 0.5 * hw[0] * (grid[:, 1] + 1) return grid def pixel_grid(hw, b=None, with_ones=False, device=None, normalize=False): """Returns a pixel grid for a certain hw resolution""" if is_tensor(hw): b, hw, device = hw.shape[0], hw.shape[-2:], hw.device if is_tensor(device): device = device.device # if align_corners(): hi, hf = 0, hw[0] - 1 wi, wf = 0, hw[1] - 1 # else: # hi, hf = 0.5, hw[0] - 0.5 # wi, wf = 0.5, hw[1] - 0.5 yy, xx = torch.meshgrid([torch.linspace(hi, hf, hw[0], device=device), torch.linspace(wi, wf, hw[1], device=device)], indexing='ij') if with_ones: # Add ones in third column grid = torch.stack([xx, yy, torch.ones(hw, device=device)], 0) else: grid = torch.stack([xx, yy], 0) if b is not None: # Repeat for batch size grid = grid.unsqueeze(0).repeat(b, 1, 1, 1) if normalize: # Normalize pixel grid grid = norm_pixel_grid(grid) return grid def shuffle_dict(data, n=None, interval=None): """Shuffle dictionary for random sampling""" keys = list(data.keys()) # Decide which keys to keep in the shuffling if interval is not None: if not is_list(interval): # Absolute distance keys = [key for key in keys if abs(key) <= interval] else: # Get parameters mode, interval = interval[0], interval[1:] # Only near keys if mode == 'near': if len(interval) == 1: keys = [key for key in keys if abs(key) <= interval[0]] elif len(interval) == 2: keys = [key for key in keys if interval[0] <= abs(key) <= interval[1]] # Only far keys elif mode == 'far': if len(interval) == 1: keys = [key for key in keys if abs(key) > interval[0]] elif len(interval) == 2: keys = [key for key in keys if key < interval[0] or key > interval[1]] # Shuffle keys random.shuffle(keys) # Sampling strategy if n is not None: if is_str(n): if '-' in n: n = [int(v) for v in n.split('-')] n = random.choice(list(range(n[0], n[1] + 1))) keys = keys[:n] elif '/' in n: n = [int(v) for v in n.split('/')] previous = [key for key in keys if key < 0] forwards = [key for key in keys if key > 0] previous, forwards = previous[:n[0]], forwards[:n[1]] keys = previous + forwards else: raise ValueError('Invalid context sampling strategy') else: keys = keys[:n] # Return shuffled dict return {key: data[key] for key in keys} def merge_dicts(metadata: dict, base: dict): """Merge base (metadata_shared) into metadata (per-episode). Per-episode values take priority. For dict values, sub-keys are merged. 'cameras' and 'labels' are skipped: in metadata_shared they represent the union over all episodes, not valid defaults for an individual episode. """ if base is None: return metadata if metadata is None: return base _episode_only = {'cameras', 'labels', 'cameras_available', 'labels_available'} for key1, val1 in base.items(): if key1 in _episode_only: continue if isinstance(val1, dict): if key1 not in metadata: metadata[key1] = dict() for key2, val2 in val1.items(): if key2 not in metadata[key1]: metadata[key1][key2] = val2 else: if key1 not in metadata: metadata[key1] = val1 return metadata def interpolate(tensor, size, scale_factor, mode): """Generic interpolation function""" if size is None and scale_factor is None: return tensor if is_tensor(size): size = size.shape[-2:] return tfn.interpolate( tensor, size=size, scale_factor=scale_factor, recompute_scale_factor=False, mode=mode, align_corners=None, ) def interpolate_image(image, shape=None, scale_factor=None, mode='bilinear', recompute_scale_factor=False): """Interpolate an image to a different resolution""" assert shape is not None or scale_factor is not None, 'Invalid option for interpolate_image' # Take last two dimensions as shape if shape is not None: if is_tensor(shape): shape = shape.shape if len(shape) > 2: shape = shape[-2:] # If the shapes are the same, do nothing if same_shape(image.shape[-2:], shape): return image # Interpolate image to match the shape return interpolate(image, size=shape, scale_factor=scale_factor, mode=mode)