Class NumpyArrayInfo

point_e/evals/npz_stream.py:15–75

Information about an array in an npz file.

@dataclass
class NumpyArrayInfo:
    """
    Information about an array in an npz file.
    """

    name: str
    dtype: np.dtype
    shape: Tuple[int]

    @classmethod
    def infos_from_first_file(cls, glob_path: str) -> Dict[str, "NumpyArrayInfo"]:
        paths, _ = _npz_paths_and_length(glob_path)
        return cls.infos_from_file(paths[0])

    @classmethod
    def infos_from_file(cls, npz_path: str) -> Dict[str, "NumpyArrayInfo"]:
        """
        Extract the info of every array in an npz file.
        """
        if not os.path.exists(npz_path):
            raise FileNotFoundError(f"batch of samples was not found: {npz_path}")
        results = {}
        with open(npz_path, "rb") as f:
            with zipfile.ZipFile(f, "r") as zip_f:
                for name in zip_f.namelist():
                    if not name.endswith(".npy"):
                        continue
                    key_name = name[: -len(".npy")]
                    with zip_f.open(name, "r") as arr_f:
                        version = np.lib.format.read_magic(arr_f)
                        if version == (1, 0):
                            header = np.lib.format.read_array_header_1_0(arr_f)
                        elif version == (2, 0):
                            header = np.lib.format.read_array_header_2_0(arr_f)
                        else:
                            raise ValueError(f"unknown numpy array version: {version}")
                        shape, _, dtype = header
                        results[key_name] = cls(name=key_name, dtype=dtype, shape=shape)
        return results

    @property
    def elem_shape(self) -> Tuple[int]:
        return self.shape[1:]

    def validate(self):
        if self.name in {"R", "G", "B"}:
            if len(self.shape) != 2:
                raise ValueError(
                    f"expecting exactly 2-D shape for '{self.name}' but got: {self.shape}"
                )
        elif self.name == "arr_0":
            if len(self.shape) < 2:
                raise ValueError(f"expecting at least 2-D shape but got: {self.shape}")
            elif len(self.shape) == 3:
                # For audio, we require continuous samples.
                if not np.issubdtype(self.dtype, np.floating):
                    raise ValueError(
                        f"invalid dtype for audio batch: {self.dtype} (expected float)"