played around with CASIA-WEB-FACE

Data.py

@@ -2,6 +2,7 @@ from torchvision import datasets, transforms, models
 import torch
 import numpy as np
 
+
 # train set transform 
 def train_transform(res):
     return transforms.Compose([
@@ -25,7 +26,6 @@ def train_transform(res):
 # test set transform
 def test_transform(res):
     return transforms.Compose([
-        # Just standard resize to 224x224
         transforms.Resize((res, res)),
         transforms.ToTensor(),
         transforms.Normalize(
@@ -67,7 +67,7 @@ def select_ids( dataset, sample_size, class_size):
     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):
+def select_top_ids(dataset, class_size):
     ids, counts = get_ids_and_counts(dataset=dataset)
 
     # sort by number of images (descending)
@@ -79,11 +79,12 @@ def select_balanced_ids(dataset, class_size):
 
 
 # split class images to train and test set.
-def get_indices(dataset, identities, split_at):
+def get_indices(dataset, identities, split_at, size = 30):
     train_indices = []
     test_indices = []
 
     #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(dataset.identity == person_id)[0].numpy()
@@ -93,6 +94,6 @@ def get_indices(dataset, identities, split_at):
         
         # split data to testing and training
         train_indices.extend(indices[:split_at])
-        test_indices.extend(indices[split_at:])
+        test_indices.extend(indices[split_at:size])
 
     return train_indices, test_indices