20 lines
532 B
Python
20 lines
532 B
Python
|
|
import torch
|
|
|
|
class IdentitySubset(torch.utils.data.Dataset):
|
|
def __init__(self, dataset, indices, id_mapping, transform=None):
|
|
self.dataset = dataset
|
|
self.indices = indices
|
|
self.id_mapping = id_mapping
|
|
self.transform = transform
|
|
|
|
def __getitem__(self, idx):
|
|
img, old_id = self.dataset[self.indices[idx]]
|
|
|
|
if self.transform:
|
|
img = self.transform(img)
|
|
|
|
return img, self.id_mapping[old_id.item()]
|
|
|
|
def __len__(self):
|
|
return len(self.indices) |