SafeTensorsReader#

class vision_architectures.image_readers.safetensors_reader.SafetensorsReader(image_key, spacing_key=None, other_keys=None, add_channel_dim=True, dtype=torch.float32)[source]#

Bases: ImageReader

__init__(image_key, spacing_key=None, other_keys=None, add_channel_dim=True, dtype=torch.float32)[source]#

Reader for Safetensors image files.

Parameters:
  • iamge_key – Key to access the image tensor in the safetensors file.

  • spacing_key (Optional[str]) – Key to access the spacing tensor in the safetensors file. Leave blank if not applicable.

  • other_keys (Optional[set[str]]) – Set of keys to access other tensors in the safetensors file. Leave blank if not applicable.

  • add_channel_dim (bool) – Whether to add a channel dimension to the image tensor.

  • dtype – Desired data type for the image tensor.

verify_suffix(filename)[source]#

Ensure the file has a supported safetensors suffix.

read(filepath)[source]#

Read image data from a safetensors file.

Return type:

dict[str, Union[Tensor, Any]]

get_data(datapoint)[source]#

Extract and process image data from the datapoint.