Initial commit

Data.py

@@ -0,0 +1,81 @@
+from torchvision import datasets, transforms, models
+import torch
+import numpy as np
+
+# transform images to size
+def transform(res):
+    return transforms.Compose([
+        # ResNet expects 224 x 224 res 
+        transforms.Resize((res, res)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        # normalise to 
+        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(res):
+    return datasets.CelebA(
+        root='./data', 
+        split='all', 
+        target_type='identity', 
+        download=True, 
+        transform=transform(res)
+    )
+
+
+def get_ids_and_counts(dataset):
+    return  torch.unique(
+        dataset.identity, 
+        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 50 identities
+    return np.random.choice(eligible_ids, class_size, replace=False)
+
+# optional function to get max amount of samples per class
+def select_balanced_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)
+
+
+# split class images to train and test set.
+def get_indices(dataset, identities, split_at):
+    train_indices = []
+    test_indices = []
+
+    #training_sample = int(sample_size * training_ratio)
+
+    for person_id in identities:
+        # Get all indices for this specific person
+        indices = torch.where(dataset.identity == 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:])
+
+    return train_indices, test_indices