175 lines
6.2 KiB
Python
175 lines
6.2 KiB
Python
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()
|