Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| """ | |
| PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation | |
| Official implementation of the paper: | |
| "PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation" | |
| by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis | |
| Licensed under a modified MIT license | |
| """ | |
| """``chumpy.ch`` namespace expected by legacy SMAL pickles.""" | |
| import numpy as np | |
| class Ch: | |
| """Minimal stand-in for ``chumpy.ch.Ch`` (unpickling only).""" | |
| def __init__(self, *args, **kwargs): | |
| self._data = None | |
| if args: | |
| self._data = np.asarray(args[0]) | |
| def _resolve(self) -> np.ndarray: | |
| # Real chumpy Ch instances store the underlying ndarray on attribute ``x``; | |
| # legacy pickles unpickle by restoring ``__dict__`` without calling ``__init__``, | |
| # so try common attribute names before falling back to ``_data``. | |
| r = self.__dict__.get("r") | |
| if isinstance(r, np.ndarray): | |
| return np.asarray(r) | |
| for attr in ("x", "_x", "_data"): | |
| val = self.__dict__.get(attr) | |
| if val is not None: | |
| return np.asarray(val) | |
| data = self.__dict__.get("_data") | |
| if data is not None: | |
| return np.asarray(data) | |
| return np.zeros((), dtype=np.float32) | |
| def r(self) -> np.ndarray: | |
| """Match real chumpy API (``ch.r()`` returns the underlying array).""" | |
| return self._resolve() | |
| def __array__(self, dtype=None): | |
| arr = self._resolve() | |
| if dtype is not None: | |
| arr = arr.astype(dtype, copy=False) | |
| return arr | |
| class ChArray(np.ndarray): | |
| """Minimal stand-in for ``chumpy.ch.ChArray``.""" | |
| def _unwrap_ch(value, dtype=np.float32) -> np.ndarray: | |
| """Resolve a chumpy ``Ch`` (method, property, or unpickled ``r``/``x`` attrs) to ndarray.""" | |
| if isinstance(value, Ch): | |
| return np.asarray(value._resolve(), dtype=dtype) | |
| r = getattr(value, "r", None) | |
| if isinstance(r, np.ndarray): | |
| return np.asarray(r, dtype=dtype) | |
| if callable(r): | |
| return np.asarray(r(), dtype=dtype) | |
| for attr in ("x", "_x", "_data"): | |
| val = getattr(value, attr, None) | |
| if val is not None: | |
| return np.asarray(val, dtype=dtype) | |
| raise TypeError(f"Cannot materialize chumpy-like object: {type(value)!r}") | |
| def materialize(value, dtype=np.float32) -> np.ndarray: | |
| """Recursively unwrap ``Ch`` / object arrays from legacy SMAL pickles.""" | |
| if isinstance(value, Ch) or ( | |
| type(value).__name__ == "Ch" | |
| and hasattr(value, "r") | |
| and not isinstance(value, (np.ndarray, list, tuple, dict, str, bytes)) | |
| ): | |
| return _unwrap_ch(value, dtype=dtype) | |
| if isinstance(value, np.ndarray): | |
| if value.dtype == object: | |
| flat = [materialize(x, dtype=dtype) for x in value.ravel()] | |
| return np.stack(flat).reshape(value.shape) | |
| return np.asarray(value, dtype=dtype) | |
| if isinstance(value, (list, tuple)): | |
| return np.asarray([materialize(x, dtype=dtype) for x in value], dtype=dtype) | |
| return np.asarray(value, dtype=dtype) | |
| __all__ = ["Ch", "ChArray", "materialize"] | |