34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
import torch
|
|
|
|
class IdentitySubset(torch.utils.data.Dataset):
|
|
def __init__(self, dataset, indices, id_mapping, transform=None):
|
|
"""
|
|
Args:
|
|
dataset: The base dataset (CelebA or ImageFolder).
|
|
indices: List of indices belonging to the selected identities.
|
|
id_mapping: Dictionary mapping {old_label: new_label_0_to_N}.
|
|
transform: Transformations to apply to the images.
|
|
"""
|
|
self.dataset = dataset
|
|
self.indices = indices
|
|
self.id_mapping = id_mapping
|
|
self.transform = transform
|
|
|
|
def __getitem__(self, idx):
|
|
# Access the base dataset using the stored index
|
|
img, old_id = self.dataset[self.indices[idx]]
|
|
|
|
# Apply transform if provided
|
|
if self.transform:
|
|
img = self.transform(img)
|
|
|
|
# Handle Label Logic:
|
|
# CelebA returns a Tensor, ImageFolder returns an int.
|
|
# We convert to a standard Python int for the dictionary lookup.
|
|
clean_id = old_id.item() if torch.is_tensor(old_id) else old_id
|
|
|
|
# Map the original identity to our new 0 -> N-1 range
|
|
return img, self.id_mapping[clean_id]
|
|
|
|
def __len__(self):
|
|
return len(self.indices) |