import torch

class IdentitySubset(torch.utils.data.Dataset):
    def __init__(self, dataset, indices, id_mapping, transform=None):
        self.dataset = dataset
        self.indices = indices
        self.id_mapping = id_mapping
        self.transform = transform

    def __getitem__(self, idx):
        img, old_id = self.dataset[self.indices[idx]]

        if self.transform:
            img = self.transform(img)

        return img, self.id_mapping[old_id.item()]
    
    def __len__(self):
        return len(self.indices)