optimised

sets/IdentitySubset.py

@@ -0,0 +1,34 @@
+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)