Source code for hyrax.datasets.hyrax_cifar_dataset

# ruff: noqa: D101, D102
import logging
from pathlib import Path

from PIL import Image

from .dataset_registry import HyraxDataset


[docs]
logger = logging.getLogger(__name__)




[docs]
class HyraxCifarDataset(HyraxDataset):
    """Map style CIFAR 10 dataset for Hyrax

    This utilizes the CIFAR dataset from torchvision for retrieving the dataset.
    """


[docs]
    def __init__(self, config: dict, data_location: Path = None):
        import torchvision.transforms as transforms
        from torchvision.datasets import CIFAR10


[docs]
        self.data_location = data_location



[docs]
        self.training_data = config["data_set"]["HyraxCifarDataset"]["use_training_data"]


        transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
        )


[docs]
        self.cifar = CIFAR10(
            root=self.data_location, train=self.training_data, download=True, transform=transform
        )


        n_id = len(self.cifar)

[docs]
        self.id_width = len(str(n_id))


        super().__init__(config)



[docs]
    def get_image(self, idx):
        """Get the image at the given index as a NumPy array."""
        img = self.cifar.data[idx]
        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.cifar.transform is not None:
            img = self.cifar.transform(img)

        return img.numpy()



[docs]
    def get_label(self, idx):
        """Get the label at the given index."""

        label = self.cifar.targets[idx]
        if self.cifar.target_transform is not None:
            label = self.cifar.target_transform(label)

        return label



[docs]
    def get_object_id(self, idx):
        """Get the object ID for the item as a string."""
        return f"{idx:0{self.id_width}d}"



[docs]
    def __len__(self):
        return len(self.cifar)