optimised

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()