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