strategies tested

sets/Casia.py

@@ -0,0 +1,167 @@
+
+
+from torchvision import datasets, transforms
+from torch.utils.data import Dataset, DataLoader, Subset
+import torch
+import numpy as np
+import os
+
+# train set transform 
+def train_transform(res):
+    return transforms.Compose([
+        transforms.Resize((res, res)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(
+            brightness=0.2, 
+            contrast=0.2, 
+            saturation=0.1
+        ),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# test set transform
+def test_transform(res):
+    return transforms.Compose([
+        transforms.Resize((res, res)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# Load data using ImageFolder for CASIA-WebFace
+'''
+def get_set():
+    # This will check local cache first, then download if missing
+    print("Checking for CASIA-WebFace dataset...")
+    path = kagglehub.dataset_download("debarghamitraroy/casia-webface")
+    
+    # Kagglehub often downloads a nested structure (e.g., path/casia-webface/casia-webface)
+    # We need the folder that directly contains the identity subfolders
+    # We'll check if there's a 'casia-webface' subfolder inside the downloaded path
+    sub_path = os.path.join(path, "casia-webface")
+    final_path = sub_path if os.path.exists(sub_path) else path
+    
+    print(f"Loading dataset from: {final_path}")
+    
+    return datasets.ImageFolder(
+        root=final_path, 
+        transform=None
+    )'''
+# Load data using ImageFolder for your UNPACKED images
+def get_set():
+    # This must point to the folder created by Extractor.py
+    # NOT the kagglehub cache path
+    final_path = os.path.abspath("./data/casia-set")
+    
+    if not os.path.exists(final_path):
+        raise FileNotFoundError(
+            f"Unpacked dataset not found at {final_path}. "
+            "Please run Extractor.py first!"
+        )
+    
+    print(f"Loading unpacked CASIA dataset from: {final_path}")
+    
+    return datasets.ImageFolder(
+        root=final_path, 
+        transform=None
+    )
+
+def get_ids_and_counts(dataset):
+    # ImageFolder stores labels in .targets
+    targets = torch.tensor(dataset.targets)
+    return torch.unique(
+        input = targets,
+        return_counts=True
+    )
+
+def select_ids(dataset, sample_size, class_size):
+    ids, counts = get_ids_and_counts(dataset=dataset)
+    eligible_mask = counts >= sample_size
+    eligible_ids = ids[eligible_mask].numpy()
+
+    if len(eligible_ids) < class_size:
+        raise ValueError(
+            f"Only found {len(eligible_ids)} identities with {sample_size}+ images."
+        )
+
+    return np.random.choice(eligible_ids, class_size, replace=False)
+
+def select_balanced_ids(dataset, class_size):
+    ids, counts = get_ids_and_counts(dataset=dataset)
+    sorted_indices = torch.argsort(counts, descending=True)
+    top_ids = ids[sorted_indices][:class_size].numpy()
+    return np.array(top_ids, dtype=int)
+
+def get_indices(dataset, identities, split_at):
+    train_indices = []
+    test_indices = []
+    
+    # We convert to numpy for faster searching with np.where
+    all_targets = np.array(dataset.targets)
+
+    for person_id in identities:
+        # Get all indices for this specific person
+        indices = np.where(all_targets == person_id)[0]
+
+        # Shuffle the indices for this person
+        np.random.shuffle(indices)
+        
+        # Split data based on your split_at value
+        train_indices.extend(indices[:split_at])
+        test_indices.extend(indices[split_at:])
+
+    return train_indices, test_indices
+
+
+
+# optional function to get max amount of samples per class
+def select_top_ids(dataset, class_size):
+    ids, counts = get_ids_and_counts(dataset=dataset)
+
+    # sort by number of images (descending)
+    sorted_indices = torch.argsort(counts, descending=True)
+
+    top_ids = ids[sorted_indices][:class_size].numpy()
+
+    return np.array(top_ids, dtype=int)
+
+
+def get_forget_retain_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
+    """
+    Splits an IdentitySubset or standard Dataset into forget and retain sets 
+    based on a remapped target class index.
+    """
+    # 1. Safely extract targets whether it's a standard dataset or a Subset wrapper
+    if hasattr(dataset, 'targets'):
+        targets = dataset.targets
+    elif hasattr(dataset, 'identity'):  # Raw CelebA support
+        targets = dataset.identity
+    else:
+        # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially
+        # This guarantees we get the 0 -> (n-1) remapped labels
+        targets = [dataset[i][1] for i in range(len(dataset))]
+        
+    if not isinstance(targets, torch.Tensor):
+        targets = torch.tensor(targets)
+        
+    # 2. Generate mask indices local to this subset
+    forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
+    retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
+    
+    # 3. Create PyTorch Subsets
+    forget_subset = Subset(dataset, forget_indices)
+    retain_subset = Subset(dataset, retain_indices)
+    
+    # 4. Wrap into clean DataLoaders
+    forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
+    retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+    
+    print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
+    
+    return forget_loader, retain_loader

sets/CasiaFace.py

@@ -0,0 +1,21 @@
+import os
+from torchvision import datasets
+from torch.utils.data import Dataset
+import torch
+from .Data import Data
+
+class CasiaSet(Data):
+    def __init__(self, resolution: int = 224, sample_size = 190):
+        super().__init__(resolution = resolution, sample_size = sample_size)
+
+    def get_set(self) -> Data:
+        path_str = "./datasets/casia-set"
+        path = os.path.abspath(path_str)
+
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Unpacked dataset missing at {self.final_path}. Run Extractor.py first!")
+        print(f"Loading unpacked CASIA dataset from: {self.final_path}")
+        set = datasets.ImageFolder(root=path, transform=None)
+        # we set the target here
+        self.target = torch.tensor(set.targets)
+        return set

sets/CelebA.py

@@ -0,0 +1,20 @@
+from torchvision import datasets
+from torch.utils.data import Dataset
+import torch
+from .Data import Data
+
+class CelebA(Data):
+    def __init__(self, resolution: int = 224, sample_size = 30):
+        super().__init__(resolution, sample_size = sample_size)
+
+    def get_set(self):
+        set = datasets.CelebA(
+            root = "./data", 
+            split='all', 
+            target_type='identity', 
+            download=True, 
+            transform=None
+        )
+        # set the target first
+        self.target = set.identity
+        return set

sets/Data.py

@@ -0,0 +1,184 @@
+from torchvision import datasets, transforms
+from torch.utils.data import Dataset, DataLoader, Subset
+import torch
+import numpy as np
+import os
+
+from enum import Enum, auto
+class Set_Name(Enum):
+    CELEBA = auto()
+    CASIAFACES = auto()
+
+# train set transform 
+def train_transform(res):
+    return transforms.Compose([
+        # ResNet expects 224 x 224 res 
+        # Inception expects 299 x 299
+        transforms.Resize((res, res)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(
+            brightness=0.2, 
+            contrast=0.2, 
+            saturation=0.1
+        ),
+        transforms.ToTensor(),
+        # normalise to 
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# test set transform
+def test_transform(res):
+    return transforms.Compose([
+        transforms.Resize((res, res)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# Load data with 'identity' as target and transform it
+def get_set(set_name:Set_Name):
+    return fetch_celeb_a() if set_name == Set_Name.CELEBA else fetch_casia_faces()
+
+def fetch_celeb_a():
+    return datasets.CelebA(
+        root='./data', 
+        split='all', 
+        target_type='identity', 
+        download=True, 
+        transform=None
+    )
+
+def fetch_casia_faces():
+    # location of the data (path relative to project root)
+    final_path = os.path.abspath("./data/casia-set")
+    
+    if not os.path.exists(final_path):
+        raise FileNotFoundError(
+            f"Unpacked dataset not found at {final_path}. "
+            "Please run Extractor.py first!"
+        )
+    
+    print(f"Loading unpacked CASIA dataset from: {final_path}")
+    
+    return datasets.ImageFolder(
+        root=final_path, 
+        transform=None
+    )
+
+
+
+def get_ids_and_counts(dataset):
+
+    target = get_target(dataset=dataset)
+    return  torch.unique(
+        input = target,
+        return_counts = True
+    )
+
+
+
+# filter selected identities from dataset
+# How many classes, how many images per class
+def select_ids( dataset, sample_size, class_size):
+    ids, counts = get_ids_and_counts(dataset = dataset)
+    eligible_mask = counts >= sample_size
+    eligible_ids = ids[eligible_mask].numpy()
+
+    if len(eligible_ids) < class_size:
+        raise ValueError(
+            f"Only found {len(eligible_ids)} identities with {sample_size}+ images."
+        )
+
+    # Randomly select identities
+    return np.random.choice(eligible_ids, class_size, replace=False)
+
+# optional function to get max amount of samples per class
+def select_top_ids(dataset, class_size):
+    ids, counts = get_ids_and_counts(dataset = dataset)
+
+    # sort by number of images (descending)
+    sorted_indices = torch.argsort(counts, descending = True)
+
+    top_ids = ids[sorted_indices][:class_size].numpy()
+
+    return np.array(top_ids, dtype=int)
+
+
+def get_target(dataset):
+    """
+    Unified target extractor. 
+    Instantly reads raw dataset arrays or safely scales down to unpack wrapped Subsets.
+    """
+    if hasattr(dataset, 'identity'):
+        # celebA
+        targets = dataset.identity
+    elif hasattr(dataset, 'targets'):
+        # others
+        targets = dataset.targets
+    else:
+        # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially
+        # This guarantees we get the 0 -> (n-1) remapped labels
+        targets = [dataset[i][1] for i in range(len(dataset))]
+
+    if not isinstance(targets, torch.Tensor):
+        targets = torch.tensor(targets)
+        
+    return targets
+
+
+# split class images to train and test set.
+def get_indices(dataset, identities, split_at, size = 30):
+
+    if split_at >= size: # debug safety
+        raise ValueError(f"Split point ({split_at}) must be less than total size ({size}).")
+    
+    train_indices = []
+    test_indices = []
+
+    target = get_target(dataset=dataset)
+
+    #training_sample = int(sample_size * training_ratio)
+    np.random.seed(42)
+    for person_id in identities:
+        # Get all indices for this specific person
+        indices = torch.where(target == person_id)[0].numpy()
+
+        # Shuffle the indices for this person
+        np.random.shuffle(indices)
+        
+        # split data to testing and training
+        train_indices.extend(indices[:split_at])
+        test_indices.extend(indices[split_at:size])
+
+    return train_indices, test_indices
+
+
+
+def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
+    """
+    Splits an IdentitySubset or standard Dataset into forget and retain sets 
+    based on a remapped target class index.
+    """
+    # extract targets
+    targets = get_target(dataset=dataset)
+        
+    # mask indices local to this subset
+    forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
+    retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
+    
+    # PyTorch Subsets
+    forget_subset = Subset(dataset, forget_indices)
+    retain_subset = Subset(dataset, retain_indices)
+    
+    # DataLoaders
+    forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
+    retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+    
+    print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
+    
+    return forget_loader, retain_loader

sets/Data_OOP.py

@@ -0,0 +1,174 @@
+import torch
+import numpy as np
+from abc import ABC, abstractmethod
+from torchvision import transforms, datasets
+from torch.utils.data import Dataset, DataLoader, Subset
+
+class Data(ABC):
+    """
+    Handles image pipelines, identity filtering, indexing, and unlearning splits.
+    """
+    def __init__(self, res: int = 224, sample_size = 30, class_size = 20):
+        self.res = res
+        self.sample_size = sample_size
+        self.class_size = class_size
+        self.target = None # will have to be set in get_set()
+
+    def train_transform(self):
+        return transforms.Compose([
+            # ResNet expects 224 x 224 res 
+            # Inception expects 299 x 299
+            transforms.Resize((self.res, self.res)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.ColorJitter(
+                brightness=0.2, 
+                contrast=0.2, 
+                saturation=0.1
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406], 
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+
+
+    def test_transform(self):
+        return transforms.Compose([
+            transforms.Resize((self.res, self.res)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406], 
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+
+    @abstractmethod
+    def get_set(self)-> datasets.TorchDataset:
+        """Loads and returns the raw underlying PyTorch Dataset instance."""
+        pass
+
+    def get_targets(self) -> torch.Tensor:
+        return self.target
+
+    def get_ids_and_counts(self) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.target is None:
+            raise ValueError ("This should be called after the 'target' variable has been set.")
+        return torch.unique(
+            self.target, 
+            return_counts=True
+        )
+
+    def select_ids(self) -> np.ndarray:
+        ids, counts = self.get_ids_and_counts()
+        eligible_mask = counts >= self.sample_size
+        eligible_ids = ids[eligible_mask].numpy()
+
+        if len(eligible_ids) < self.class_size:
+            raise ValueError(
+                f"Only found {len(eligible_ids)} identities with {self.sample_size}+ images."
+            )
+
+        return np.random.choice(eligible_ids, self.class_size, replace=False)
+
+    # Function to get max amount of samples per class
+    def select_top_ids(self) -> np.ndarray:
+
+        ids, counts = self.get_ids_and_counts()
+        # sort by number of images (descending)
+        sorted_indices = torch.argsort(counts, descending=True)
+        top_ids = ids[sorted_indices][:self.class_size].numpy()
+        return np.array(top_ids, dtype=int)
+
+    def get_indices(self, identities: np.ndarray, split_at: int, max_size: int = None) -> tuple[list, list]:
+        '''train_indices = []
+        test_indices = []
+        max_size = self.sample_size if max_size is None else max_size
+        
+        # Pull raw target tensor array using concrete implementation rules
+        all_targets = np.array(self.get_targets().cpu())
+        np.random.seed(42)
+
+        for person_id in identities:
+            indices = np.where(all_targets == person_id)[0]
+            np.random.shuffle(indices)
+            
+            # Constrain total sample tracking size if requested (e.g. CelebA ceiling)
+            current_pool = indices[:max_size] if max_size else indices
+            
+            if split_at >= len(current_pool):
+                raise ValueError(f"Split point ({split_at}) exceeds slice size ({len(current_pool)}) for class {person_id}.")
+
+            train_indices.extend(current_pool[:split_at])
+            test_indices.extend(current_pool[split_at:])
+
+        return train_indices, test_indices'''
+        if split_at >= self.sample_size: # debug safety
+            raise ValueError(f"Split point ({split_at}) must be less than total size ({self.sample_size}).")
+    
+        train_indices = []
+        test_indices = []
+
+        #training_sample = int(sample_size * training_ratio)
+        np.random.seed(42)
+        target = self.get_targets()
+        for person_id in identities:
+            # Get all indices for this specific person
+            indices = torch.where(target == person_id)[0].numpy()
+
+            # Shuffle the indices for this person
+            np.random.shuffle(indices)
+            
+            # split data to testing and training
+            train_indices.extend(indices[:split_at])
+            test_indices.extend(indices[split_at:self.sample_size])
+
+        return train_indices, test_indices
+
+    @staticmethod
+    def get_unlearn_loaders(
+        dataset: Dataset, 
+        forget_class_idx: int, 
+        batch_size: int = 32
+    ) -> tuple[DataLoader, DataLoader]:
+
+        """Splits an IdentitySubset into forget/retain parts based on local class index."""
+        if hasattr(dataset, 'targets'):
+            targets = dataset.targets
+        elif hasattr(dataset, 'identity'): 
+            targets = dataset.identity
+        else:
+            targets = [dataset[i][1] for i in range(len(dataset))]
+            
+        if not isinstance(targets, torch.Tensor):
+            targets = torch.tensor(targets)
+            
+        forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
+        retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
+        
+        forget_subset = Subset(dataset, forget_indices)
+        retain_subset = Subset(dataset, retain_indices)
+        
+        forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
+        retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+        
+        print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
+        
+        return forget_loader, retain_loader
+    
+
+    @staticmethod
+    def getDataSet(set:SetType, sample_size):
+        # some test
+        if set == SetType.CASIA:
+            from sets.CasiaFace import CasiaFace
+            return CasiaFace(sample_size = sample_size)
+        if set == SetType.CELEBA:
+            from sets.CelebA import CelebA
+            return CelebA(sample_size=sample_size)
+
+
+from enum import Enum, auto
+class SetType(Enum):
+    CASIA = auto()
+    CELEBA = auto()

sets/Extractor.py

@@ -0,0 +1,131 @@
+import os
+import struct
+from tqdm import tqdm
+from collections import Counter
+import hashlib
+
+def get_top_identities_binary(rec_path, idx_path, top_n=51):
+    """
+    Pass 1: Scans the actual BINARY HEADERS in the .rec file.
+    This is the only way to be 100% sure which image belongs to whom.
+    """
+    identity_counts = Counter()
+    
+    with open(idx_path, 'r') as f:
+        offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()]
+
+    print("Pass 1: Scanning binary headers to count identities...")
+    with open(rec_path, 'rb') as f:
+        for offset in tqdm(offsets):
+            f.seek(offset)
+            header_bin = f.read(32) # Read enough for the header
+            if len(header_bin) < 32: continue
+            
+            # MXNet Header format: [Flag, Label (float), ID, ID]
+            # The label is at offset 12 (float32)
+            label = int(struct.unpack('f', header_bin[12:16])[0])
+            identity_counts[label] += 1
+
+    top_stats = identity_counts.most_common(top_n)
+    top_labels = {label for label, count in top_stats}
+
+    print(f"\nTop {top_n} Identities by Binary Label:")
+    for label, count in top_stats:
+        print(f"ID: {label:<10} | Count: {count:<10}")
+    
+    return top_labels
+
+def extract_selected_binary(rec_path, idx_path, output_dir, top_labels):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    with open(idx_path, 'r') as f:
+        offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()]
+
+    print(f"\nPass 2: Extracting verified images...")
+    
+    # NEW: Keep track of how many images we've saved for each ID
+    # to avoid overwriting files.
+    save_counters = {label: 0 for label in top_labels}
+    total_extracted = 0
+    
+    with open(rec_path, 'rb') as f:
+        for offset in tqdm(offsets):
+            f.seek(offset)
+            header_bin = f.read(32)
+            if len(header_bin) < 32: continue
+            
+            label = int(struct.unpack('f', header_bin[12:16])[0])
+            
+            if label not in top_labels:
+                continue
+            
+            # Read image content
+            _, length_flag = struct.unpack('II', header_bin[:8])
+            content_length = length_flag & ((1 << 31) - 1)
+            content = f.read(content_length)
+            
+            img_start = content.find(b'\xff\xd8')
+            if img_start == -1: continue
+
+            target_folder = os.path.join(output_dir, str(label))
+            os.makedirs(target_folder, exist_ok=True)
+            
+            # Use the counter for this specific label
+            current_count = save_counters[label]
+            img_filename = f"{current_count}.jpg"
+            img_path = os.path.join(target_folder, img_filename)
+            if(current_count > 200):
+                continue
+            
+            with open(img_path, 'wb') as img_f:
+                img_f.write(content[img_start:])
+            
+            save_counters[label] += 1
+            total_extracted += 1
+
+    print(f"\nDone! Extracted {total_extracted} total images.")
+
+
+
+def remove_duplicates(root_dir):
+    hashes = {} # hash -> first_filepath
+    duplicates_removed = 0
+
+    # Walk through every identity folder
+    for subdir, dirs, files in os.walk(root_dir):
+        for filename in tqdm(files, desc=f"Checking {os.path.basename(subdir)}"):
+            filepath = os.path.join(subdir, filename)
+            
+            # Calculate MD5 hash of the file
+            with open(filepath, 'rb') as f:
+                file_hash = hashlib.md5(f.read()).hexdigest()
+            
+            if file_hash in hashes:
+                # We've seen this image before!
+                os.remove(filepath)
+                duplicates_removed += 1
+            else:
+                hashes[file_hash] = filepath
+
+    print(f"\nClean-up complete. Removed {duplicates_removed} duplicate images.")
+
+
+'''
+if __name__ == "__main__":
+    # Point this to your unpacked Top 50 folder
+    target_dir = "./datasets/casia-set"
+    remove_duplicates(target_dir)
+'''
+if __name__ == "__main__":
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    REC = os.path.join(base_dir, 'casia', 'train.rec')
+    IDX = os.path.join(base_dir, 'casia', 'train.idx')
+    OUT = os.path.join(base_dir, 'casia-set')
+
+    # Step 1: Trust the binary, not the text file
+    top_verified_labels = get_top_identities_binary(REC, IDX, top_n=50)
+    
+    # Step 2: Extract
+    extract_selected_binary(REC, IDX, OUT, top_verified_labels)
+    

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)