82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
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
|