File size: 5,352 Bytes
1e4dcb8 3f224a9 1e4dcb8 3f224a9 1e4dcb8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | # SPDX-License-Identifier: MIT
# Copyright (c) 2026 Baiyu Chen and the COMODO-AnyMo-Bench contributors
"""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 # A batch of windows with different lengths.
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)})
|