| |
| |
|
|
| """Preprocessing that matches the original Nymeria COMODO loader and collate.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.nn import functional as F |
| from transformers import BatchFeature, ProcessorMixin |
| from transformers.utils.hub import cached_file |
|
|
|
|
| class ComodoProcessor(ProcessorMixin): |
| attributes = [] |
| model_input_names = ["imu_values"] |
|
|
| def __init__(self, mean, std, sampling_rate=60, sequence_length=512): |
| super().__init__() |
| self.mean = np.asarray(mean, dtype=np.float32) |
| self.std = np.asarray(std, dtype=np.float32) |
| if self.mean.shape != (18,) or self.std.shape != (18,): |
| raise ValueError("COMODO requires normalization statistics for 18 channels") |
| self.std = np.where(self.std < 1e-6, 1.0, self.std).astype(np.float32) |
| self.sampling_rate = sampling_rate |
| self.sequence_length = sequence_length |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
| variant = kwargs.pop("variant", None) |
| hub_kwargs = {key: kwargs[key] for key in |
| ("subfolder", "revision", "token", "cache_dir", "local_files_only", |
| "force_download") if key in kwargs} |
| if variant is not None: |
| if variant not in ("core50", "fine150"): |
| raise ValueError("variant must be 'core50' or 'fine150'") |
| if hub_kwargs.get("subfolder") not in (None, "", variant): |
| raise ValueError("variant conflicts with subfolder") |
| hub_kwargs["subfolder"] = variant |
| elif not hub_kwargs.get("subfolder"): |
| metadata_path = cached_file(str(pretrained_model_name_or_path), "config.json", |
| _raise_exceptions_for_missing_entries=False, **hub_kwargs) |
| metadata = json.loads(Path(metadata_path).read_text()) if metadata_path else {} |
| if metadata.get("is_variant_bundle", False): |
| hub_kwargs["subfolder"] = metadata.get("variant", "core50") |
| config = json.loads(Path(cached_file(str(pretrained_model_name_or_path), |
| "processor_config.json", **hub_kwargs)).read_text()) |
| stats = json.loads(Path(cached_file(str(pretrained_model_name_or_path), |
| "normalization_stats.json", **hub_kwargs)).read_text()) |
| return cls(stats["mean"], stats["std"], config["sampling_rate"], config["sequence_length"]) |
|
|
| def save_pretrained(self, save_directory, **kwargs): |
| target = Path(save_directory) |
| target.mkdir(parents=True, exist_ok=True) |
| config = {"processor_class": "ComodoProcessor", "sampling_rate": self.sampling_rate, |
| "sequence_length": self.sequence_length, |
| "auto_map": {"AutoProcessor": "processing_comodo.ComodoProcessor"}} |
| (target / "processor_config.json").write_text(json.dumps(config, indent=2) + "\n") |
| (target / "normalization_stats.json").write_text(json.dumps({ |
| "mean": self.mean.tolist(), "std": self.std.tolist()}, indent=2) + "\n") |
| source = Path(__file__) |
| if source.resolve() != (target / source.name).resolve(): |
| (target / source.name).write_text(source.read_text()) |
| return [str(target / "processor_config.json")] |
|
|
| def __call__(self, imu, sampling_rate=60, return_tensors="pt", **kwargs): |
| if kwargs: |
| raise TypeError(f"Unknown preprocessing arguments: {sorted(kwargs)}") |
| if sampling_rate != self.sampling_rate: |
| raise ValueError(f"Provide {self.sampling_rate} Hz IMU windows; got {sampling_rate} Hz") |
| if return_tensors != "pt": |
| raise ValueError("Only return_tensors='pt' is supported") |
| if isinstance(imu, torch.Tensor): |
| imu = imu.detach().cpu().numpy() |
| try: |
| values = np.asarray(imu, dtype=np.float32) |
| except ValueError: |
| values = None |
| if values is not None and values.ndim == 2: |
| windows = [values] |
| elif values is not None and values.ndim == 3: |
| windows = list(values) |
| elif values is None and isinstance(imu, (list, tuple)): |
| windows = [np.asarray(row, dtype=np.float32) for row in imu] |
| else: |
| raise ValueError("Expected [T, 18], [B, T, 18], or a list of [T, 18] windows") |
| resized = [] |
| for window in windows: |
| if window.ndim != 2 or window.shape[1] != 18 or window.shape[0] == 0: |
| raise ValueError("Each IMU window must have shape [T, 18] with T > 0") |
| window = np.nan_to_num(window, copy=True) |
| window = (window - self.mean.reshape(1, -1)) / self.std.reshape(1, -1) |
| tensor = torch.from_numpy(np.ascontiguousarray(window.T)).unsqueeze(0) |
| resized.append(F.interpolate(tensor, size=self.sequence_length, mode="linear", |
| align_corners=False).squeeze(0)) |
| if not resized: |
| raise ValueError("The input batch must contain at least one window") |
| return BatchFeature({"imu_values": torch.stack(resized)}) |
|
|
