CelebA_class.py 3.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. import torch
  2. from torch.utils.data import Dataset
  3. import os
  4. import numpy as np
  5. import random
  6. from torchvision import transforms
  7. from PIL import Image
  8. import cv2
  9. class FaceDataSet(Dataset):
  10. def __init__(self, dataset_path, batch_size):
  11. super(FaceDataSet, self).__init__()
  12. '''picture_dir_list = []
  13. for i in range(self.people_num):
  14. picture_dir_list.append('/data/home/renwangchen/vgg_align_224/'+self.people_list[i])
  15. self.people_pic_list = []
  16. for i in range(self.people_num):
  17. pic_list = os.listdir(picture_dir_list[i])
  18. person_pic_list = []
  19. for j in range(len(pic_list)):
  20. pic_dir = os.path.join(picture_dir_list[i], pic_list[j])
  21. person_pic_list.append(pic_dir)
  22. self.people_pic_list.append(person_pic_list)'''
  23. pic_dir = '/data/home/renwangchen/CelebA_224/'
  24. latent_dir = '/data/home/renwangchen/CelebA_latent/'
  25. tmp_list = os.listdir(pic_dir)
  26. self.pic_list = []
  27. self.latent_list = []
  28. for i in range(len(tmp_list)):
  29. self.pic_list.append(pic_dir + tmp_list[i])
  30. self.latent_list.append(latent_dir + tmp_list[i][:-3] + 'npy')
  31. self.pic_list = self.pic_list[:29984]
  32. '''for i in range(29984):
  33. print(self.pic_list[i])'''
  34. self.latent_list = self.latent_list[:29984]
  35. self.people_num = len(self.pic_list)
  36. self.type = 1
  37. self.bs = batch_size
  38. self.count = 0
  39. self.transformer = transforms.Compose([
  40. transforms.ToTensor(),
  41. transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  42. ])
  43. def __getitem__(self, index):
  44. p1 = random.randint(0, self.people_num - 1)
  45. p2 = p1
  46. if self.type == 0:
  47. # load pictures from the same folder
  48. pass
  49. else:
  50. # load pictures from different folders
  51. p2 = p1
  52. while p2 == p1:
  53. p2 = random.randint(0, self.people_num - 1)
  54. pic_id_dir = self.pic_list[p1]
  55. pic_att_dir = self.pic_list[p2]
  56. latent_id_dir = self.latent_list[p1]
  57. latent_att_dir = self.latent_list[p2]
  58. img_id = Image.open(pic_id_dir).convert('RGB')
  59. img_id = self.transformer(img_id)
  60. latent_id = np.load(latent_id_dir)
  61. latent_id = latent_id / np.linalg.norm(latent_id)
  62. latent_id = torch.from_numpy(latent_id)
  63. img_att = Image.open(pic_att_dir).convert('RGB')
  64. img_att = self.transformer(img_att)
  65. latent_att = np.load(latent_att_dir)
  66. latent_att = latent_att / np.linalg.norm(latent_att)
  67. latent_att = torch.from_numpy(latent_att)
  68. self.count += 1
  69. data_type = self.type
  70. if self.count == self.bs:
  71. self.type = 1 - self.type
  72. self.count = 0
  73. return img_id, img_att, latent_id, latent_att, data_type
  74. def __len__(self):
  75. return len(self.pic_list)