DatasetLoader.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. import torch
  4. import numpy
  5. import random
  6. import pdb
  7. import glob
  8. import os
  9. from torch.utils.data import Dataset, DataLoader
  10. from PIL import Image
  11. def round_down(num, divisor):
  12. return num - (num%divisor)
  13. def worker_init_fn(worker_id):
  14. numpy.random.seed(numpy.random.get_state()[1][0] + worker_id)
  15. class meta_loader(Dataset):
  16. def __init__(self, train_path, train_ext, transform):
  17. ## Read Training Files
  18. files = glob.glob('%s/*/*.%s'%(train_path,train_ext))
  19. ## Make a mapping from Class Name to Class Number
  20. dictkeys = list(set([x.split('/')[-2] for x in files]))
  21. dictkeys.sort()
  22. dictkeys = { key : ii for ii, key in enumerate(dictkeys) }
  23. self.transform = transform
  24. self.label_dict = {}
  25. self.data_list = []
  26. self.data_label = []
  27. for lidx, file in enumerate(files):
  28. speaker_name = file.split('/')[-2]
  29. speaker_label = dictkeys[speaker_name];
  30. if not (speaker_label in self.label_dict):
  31. self.label_dict[speaker_label] = [];
  32. self.label_dict[speaker_label].append(lidx);
  33. self.data_label.append(speaker_label)
  34. self.data_list.append(file)
  35. print('%d files from %d classes found.'%(len(self.data_list),len(self.label_dict)))
  36. def __getitem__(self, indices):
  37. feat = []
  38. for index in indices:
  39. feat.append(self.transform(Image.open(self.data_list[index])));
  40. feat = numpy.stack(feat, axis=0)
  41. return torch.FloatTensor(feat), self.data_label[index]
  42. def __len__(self):
  43. return len(self.data_list)
  44. class test_dataset_loader(Dataset):
  45. def __init__(self, test_list, test_path, transform, **kwargs):
  46. self.test_path = test_path
  47. self.data_list = test_list
  48. self.transform = transform
  49. def __getitem__(self, index):
  50. img = Image.open(os.path.join(self.test_path, self.data_list[index]))
  51. return self.transform(img), self.data_list[index]
  52. def __len__(self):
  53. return len(self.data_list)
  54. class meta_sampler(torch.utils.data.Sampler):
  55. def __init__(self, data_source, nPerClass, max_img_per_cls, batch_size):
  56. self.label_dict = data_source.label_dict
  57. self.nPerClass = nPerClass
  58. self.max_img_per_cls = max_img_per_cls;
  59. self.batch_size = batch_size;
  60. def __iter__(self):
  61. ## Get a list of identities
  62. dictkeys = list(self.label_dict.keys());
  63. dictkeys.sort()
  64. lol = lambda lst, sz: [lst[i:i+sz] for i in range(0, len(lst), sz)]
  65. flattened_list = []
  66. flattened_label = []
  67. ## Data for each class
  68. for findex, key in enumerate(dictkeys):
  69. data = self.label_dict[key]
  70. numSeg = round_down(min(len(data),self.max_img_per_cls),self.nPerClass)
  71. rp = lol(numpy.random.permutation(len(data))[:numSeg],self.nPerClass)
  72. flattened_label.extend([findex] * (len(rp)))
  73. for indices in rp:
  74. flattened_list.append([data[i] for i in indices])
  75. ## Data in random order
  76. mixid = numpy.random.permutation(len(flattened_label))
  77. mixlabel = []
  78. mixmap = []
  79. ## Prevent two pairs of the same speaker in the same batch
  80. for ii in mixid:
  81. startbatch = len(mixlabel) - len(mixlabel) % self.batch_size
  82. if flattened_label[ii] not in mixlabel[startbatch:]:
  83. mixlabel.append(flattened_label[ii])
  84. mixmap.append(ii)
  85. return iter([flattened_list[i] for i in mixmap])
  86. def __len__(self):
  87. return len(self.data_source)
  88. def get_data_loader(batch_size, max_img_per_cls, nDataLoaderThread, nPerClass, train_path, train_ext, transform, **kwargs):
  89. train_dataset = meta_loader(train_path, train_ext, transform)
  90. train_sampler = meta_sampler(train_dataset, nPerClass, max_img_per_cls, batch_size)
  91. train_loader = torch.utils.data.DataLoader(
  92. train_dataset,
  93. batch_size=batch_size,
  94. num_workers=nDataLoaderThread,
  95. sampler=train_sampler,
  96. pin_memory=False,
  97. worker_init_fn=worker_init_fn,
  98. drop_last=True,
  99. )
  100. return train_loader