from torch.utils import data import os import torch import numpy as np import cv2 import random import albumentations as A pixel_transform = A.Compose([ A.SmallestMaxSize(max_size=512), A.CenterCrop(512, 512), A.Affine(scale=(0.5, 1), translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)}, rotate=(-10, 10), p=0.8), ], additional_targets={'image0': 'image', 'image1': 'image'}) hair_transform = A.Compose([ A.SmallestMaxSize(max_size=512), A.CenterCrop(512, 512), A.Affine(scale=(0.9, 1.2), rotate=(-10, 10), p=0.7)] ) # class myDataset(data.Dataset): # """Custom data.Dataset compatible with data.DataLoader.""" # def __init__(self, train_data_dir): # self.img_path = os.path.join(train_data_dir, "hair") # # self.pose_path = os.path.join(train_data_dir, "pose.npy") # # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # # self.ref_path = os.path.join(train_data_dir, "ref_hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "non-hair") # self.ref_path = os.path.join(train_data_dir, "reference") # self.lists = os.listdir(self.img_path) # self.len = len(self.lists)-10 # self.pose = np.load(self.pose_path) # #self.pose = np.random.randn(12, 4) # def __getitem__(self, index): # """Returns one data pair (source and target).""" # # seq_len, fea_dim # random_number1 = random.randrange(0, 21) # random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) # name = self.lists[index] # random_number1 = random_number1 # #* 10 # #random_number2 = random_number2 * 10 # random_number2 = random_number1 # non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') # ref_folder = os.path.join(self.ref_path, name) # files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, files[0]) # img_non_hair = cv2.imread(non_hair_path) # ref_hair = cv2.imread(ref_path) # img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) # ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) # img_non_hair = cv2.resize(img_non_hair, (512, 512)) # ref_hair = cv2.resize(ref_hair, (512, 512)) # img_non_hair = (img_non_hair / 255.0) * 2 - 1 # ref_hair = (ref_hair / 255.0) * 2 - 1 # img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) # ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) # pose1 = self.pose[random_number1] # pose1 = torch.tensor(pose1) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) # hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) # return { # 'hair_pose': pose1, # 'img_hair': img_hair, # 'bald_pose': pose2, # 'img_non_hair': img_non_hair, # 'ref_hair': ref_hair # } # def __len__(self): # return self.len class myDataset(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir): self.img_path = os.path.join(train_data_dir, "hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "no_hair") self.ref_path = os.path.join(train_data_dir, "ref_hair") self.lists = os.listdir(self.img_path) self.len = len(self.lists) self.pose = np.load(self.pose_path) def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 120) random_number2 = random.randrange(0, 120) while random_number2==random_number1: random_number2 = random.randrange(0, 120) name = self.lists[index] hair_path = os.path.join(self.img_path, name, str(random_number1)+'.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2)+'.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] ref_path = os.path.join(ref_folder, files[0]) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_hair = (img_hair/255.0)* 2 - 1 img_non_hair = (img_non_hair/255.0) ref_hair = (ref_hair/255.0)* 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose1 = self.pose[random_number1] pose1 = torch.tensor(pose1) pose2 = self.pose[random_number2] pose2 = torch.tensor(pose2) return { 'hair_pose': pose1, 'img_hair':img_hair, 'bald_pose': pose2, 'img_non_hair':img_non_hair, 'ref_hair':ref_hair } def __len__(self): return self.len # class myDataset_unet(data.Dataset): # """Custom data.Dataset compatible with data.DataLoader.""" # class myDataset_unet(data.Dataset): # """Custom data.Dataset compatible with data.DataLoader.""" # def __init__(self, train_data_dir, frame_num=6): # self.img_path = os.path.join(train_data_dir, "hair") # # self.pose_path = os.path.join(train_data_dir, "pose.npy") # # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # # self.ref_path = os.path.join(train_data_dir, "ref_hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "non-hair") # self.ref_path = os.path.join(train_data_dir, "reference") # self.lists = os.listdir(self.img_path) # self.len = len(self.lists)-10 # self.pose = np.load(self.pose_path) # self.frame_num = frame_num # #self.pose = np.random.randn(12, 4) # def __getitem__(self, index): # """Returns one data pair (source and target).""" # # seq_len, fea_dim # random_number1 = random.randrange(0, 21) # random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) # name = self.lists[index] # random_number1 = random_number1 # #* 10 # #random_number2 = random_number2 * 10 # random_number2 = random_number1 # non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') # ref_folder = os.path.join(self.ref_path, name) # ref_folder = os.path.join(self.img_path, name) # files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # #ref_path = os.path.join(ref_folder, files[0]) # ref_path = os.path.join(ref_folder, '0.jpg') # img_non_hair = cv2.imread(non_hair_path) # ref_hair = cv2.imread(ref_path) # img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) # ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) # img_non_hair = cv2.resize(img_non_hair, (512, 512)) # ref_hair = cv2.resize(ref_hair, (512, 512)) # img_non_hair = (img_non_hair / 255.0) * 2 - 1 # ref_hair = (ref_hair / 255.0) * 2 - 1 # img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) # ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) # pose1 = self.pose[random_number1] # pose1 = torch.tensor(pose1) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) # hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # # begin = random.randrange(0, 21-self.frame_num) # # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) # return { # 'hair_pose': pose1, # 'img_hair': img_hair, # 'bald_pose': pose2, # 'img_non_hair': img_non_hair, # 'ref_hair': ref_hair # } # def __len__(self): # return self.len class myDataset_unet(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir): self.img_path = os.path.join(train_data_dir, "hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists) self.pose = np.load(self.pose_path) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) name = self.lists[index] #random_number1 = random_number1 #random_number2 = random_number2 * 10 #random_number2 = random_number1 hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.img_path, name) #files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) # ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) # pose1 = self.pose[random_number1] # pose1 = torch.tensor(pose1) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) # polars = self.polars_rad[random_number1] # polars = torch.tensor(polars).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) pose = self.pose[random_number1] pose = torch.tensor(pose) return { # 'hair_pose': pose1, 'img_hair': img_hair, # 'bald_pose': pose2, # 'img_non_hair': img_non_hair, 'img_ref': ref_hair, 'pose': pose, # 'polars': polars, # 'azimuths': azimuths, } def __len__(self): return self.len-10 class myDataset_sv3d(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number3 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) while random_number3 == random_number1: random_number3 = random.randrange(0, 21) # while random_number3 == random_number1: # random_number3 = random.randrange(0, 21) name = self.lists[index] #random_number1 = random_number1 #* 10 #random_number2 = random_number2 * 10 #random_number2 = random_number1 hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') #hair_path2 = os.path.join(self.img_path, name, str(random_number3) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') #non_hair_path2 = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') #non_hair_path3 = os.path.join(self.non_hair_path, name, str(random_number3) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) # print('________') # print(files) # print('++++++++') # print(ref_folder) # print("========") # print(name) # print("********") # print(ref_path) img_hair = cv2.imread(hair_path) #img_hair2 = cv2.imread(hair_path2) img_non_hair = cv2.imread(non_hair_path) #img_non_hair2 = cv2.imread(non_hair_path2) #img_non_hair3 = cv2.imread(non_hair_path3) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) #img_non_hair2 = cv2.cvtColor(img_non_hair2, cv2.COLOR_BGR2RGB) #img_non_hair3 = cv2.cvtColor(img_non_hair3, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) #img_non_hair2 = cv2.resize(img_non_hair2, (512, 512)) #img_non_hair3 = cv2.resize(img_non_hair3, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 #img_non_hair2 = (img_non_hair2 / 255.0) * 2 - 1 #img_non_hair3 = (img_non_hair3 / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) #img_non_hair2 = torch.tensor(img_non_hair2).permute(2, 0, 1) #img_non_hair3 = torch.tensor(img_non_hair3).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) pose2 = self.pose[random_number3] pose2 = torch.tensor(pose2) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # polar = self.polars_rad[random_number1] # polar = torch.tensor(polar).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) # img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) #img_hair2 = cv2.cvtColor(img_hair2, cv2.COLOR_BGR2RGB) #img_hair2 = cv2.resize(img_hair2, (512, 512)) #img_hair2 = (img_hair2 / 255.0) * 2 - 1 #img_hair2 = torch.tensor(img_hair2).permute(2, 0, 1) # begin = random.randrange(0, 21-self.frame_num) # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) return { # 'hair_pose': pose1, 'img_hair': img_hair, #'img_hair2': img_hair2, # 'bald_pose': pose2, #'pose': pose, #'pose2': pose2, 'img_non_hair': img_non_hair, #'img_non_hair2': img_non_hair2, #'img_non_hair3': img_non_hair3, 'ref_hair': ref_hair, # 'polar': polar, # 'azimuths':azimuths, } def __len__(self): return self.len class myDataset_sv3d2(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) name = self.lists[index] #random_number1 = random_number1 #* 10 #random_number2 = random_number2 * 10 #random_number2 = random_number1 hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) # print('________') # print(files) # print('++++++++') # print(ref_folder) # print("========") # print(name) # print("********") # print(ref_path) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # polar = self.polars_rad[random_number1] # polar = torch.tensor(polar).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) # begin = random.randrange(0, 21-self.frame_num) # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) return { # 'hair_pose': pose1, 'img_hair': img_hair, # 'bald_pose': pose2, 'pose': pose, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, # 'polar': polar, # 'azimuths':azimuths, } def __len__(self): return self.len class myDataset_sv3d_temporal(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21-10) # random_number3 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) while random_number3 == random_number1: random_number3 = random.randrange(0, 21) # while random_number3 == random_number1: # random_number3 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) for i in range(10): img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(1)) hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(1)) #random_number1 = random_number1 #* 10 #random_number2 = random_number2 * 10 #random_number2 = random_number1 non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') hair_path2 = os.path.join(self.img_path, name, str(random_number3) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') non_hair_path2 = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') non_hair_path3 = os.path.join(self.non_hair_path, name, str(random_number3) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) # print('________') # print(files) # print('++++++++') # print(ref_folder) # print("========") # print(name) # print("********") # print(ref_path) img_hair = cv2.imread(hair_path) img_hair2 = cv2.imread(hair_path2) img_non_hair = cv2.imread(non_hair_path) img_non_hair2 = cv2.imread(non_hair_path2) img_non_hair3 = cv2.imread(non_hair_path3) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) img_non_hair2 = cv2.cvtColor(img_non_hair2, cv2.COLOR_BGR2RGB) img_non_hair3 = cv2.cvtColor(img_non_hair3, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) img_non_hair2 = cv2.resize(img_non_hair2, (512, 512)) img_non_hair3 = cv2.resize(img_non_hair3, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 img_non_hair2 = (img_non_hair2 / 255.0) * 2 - 1 img_non_hair3 = (img_non_hair3 / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) img_non_hair2 = torch.tensor(img_non_hair2).permute(2, 0, 1) img_non_hair3 = torch.tensor(img_non_hair3).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) pose2 = self.pose[random_number3] pose2 = torch.tensor(pose2) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # polar = self.polars_rad[random_number1] # polar = torch.tensor(polar).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) # img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) img_hair2 = cv2.cvtColor(img_hair2, cv2.COLOR_BGR2RGB) img_hair2 = cv2.resize(img_hair2, (512, 512)) img_hair2 = (img_hair2 / 255.0) * 2 - 1 img_hair2 = torch.tensor(img_hair2).permute(2, 0, 1) # begin = random.randrange(0, 21-self.frame_num) # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) return { # 'hair_pose': pose1, 'img_hair': img_hair, 'img_hair2': img_hair2, # 'bald_pose': pose2, 'pose': pose, 'pose2': pose2, 'img_non_hair': img_non_hair, 'img_non_hair2': img_non_hair2, 'img_non_hair3': img_non_hair3, 'ref_hair': ref_hair, # 'polar': polar, # 'azimuths':azimuths, } def __len__(self): return self.len class myDataset_sv3d_simple(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") # self.ref_path = os.path.join(train_data_dir, "reference") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) name = self.lists[index] #random_number1 = random_number1 #* 10 #random_number2 = random_number2 * 10 random_number2 = random_number1 hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) # print('________') # print(files) # print('++++++++') # print(ref_folder) # print("========") # print(name) # print("********") # print(ref_path) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) # hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # polar = self.polars_rad[random_number1] # polar = torch.tensor(polar).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) x = torch.tensor(self.x[random_number1]) y = torch.tensor(self.y[random_number1]) # begin = random.randrange(0, 21-self.frame_num) # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) return { # 'hair_pose': pose1, 'img_hair': img_hair, # 'bald_pose': pose2, 'pose': pose, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, # 'polar': polar, # 'azimuths':azimuths, } def __len__(self): return self.len class myDataset_sv3d_simple_ori(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) name = self.lists[index] #random_number1 = random_number1 #* 10 #random_number2 = random_number2 * 10 #random_number2 = random_number1 hair_path = os.path.join(self.img_path, name, str(random_number2) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) # print('________') # print(files) # print('++++++++') # print(ref_folder) # print("========") # print(name) # print("********") # print(ref_path) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number2] pose = torch.tensor(pose) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) hair_path = os.path.join(self.img_path, name, str(random_number2) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # polar = self.polars_rad[random_number1] # polar = torch.tensor(polar).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) x = torch.tensor(self.x[random_number2]) y = torch.tensor(self.y[random_number2]) # begin = random.randrange(0, 21-self.frame_num) # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) return { # 'hair_pose': pose1, 'img_hair': img_hair, # 'bald_pose': pose2, 'pose': pose, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, # 'polar': polar, # 'azimuths':azimuths, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21-12) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) ref_hair = self.read_img(ref_path) for i in range(12): img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) #hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') hair_path = os.path.join(self.non_hair_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal2(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement' self.img_path = os.path.join(train_data_dir, "hair") self.img_path2 = os.path.join(train_data_dir, "hair_good") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") self.pose_path2 = os.path.join(train_data_dir2, "pose.npy") self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair") self.ref_path2 = os.path.join(train_data_dir2, "reference") self.lists = os.listdir(self.img_path2) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def read_ref_img(self, path): img = cv2.imread(path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = hair_transform(image=img)['image'] # img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def reference_lists(self, reference_num, root): stacks = [] invalid = [] for i in range(12): if (reference_num-6+i)<0: invalid.append(reference_num-5+i+21) else: invalid.append((reference_num-5+i)%21) for i in range(21): if i in invalid: continue else: stacks.append(os.path.join(root, str(i)+'.jpg')) return stacks def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number = random.uniform(0, 1) if random_number<0.5: non_hair_root = self.non_hair_path2 img_path = self.img_path2 else: non_hair_root = self.non_hair_path img_path = self.img_path random_number1 = random.randrange(0, 21-12) random_number2 = random.randrange(0, 21) name = self.lists[index].split('.')[0] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') # non_hair_path = os.path.join(img_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) # files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3] + self.reference_lists(random_number2, os.path.join(self.img_path, name))[:5] files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,random.choice(files)) ref_hair = self.read_ref_img(ref_path) for i in range(12): #non_hair_path = os.path.join(img_path, name, str(random_number1+i) + '.jpg') img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') # hair_path = os.path.join(img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal3(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement' self.img_path = os.path.join(train_data_dir, "hair") self.img_path2 = os.path.join(train_data_dir, "non-hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") self.pose_path2 = os.path.join(train_data_dir2, "pose.npy") self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair") self.ref_path2 = os.path.join(train_data_dir2, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number = random.uniform(0, 1) if random_number<0.5: non_hair_root = self.non_hair_path2 img_path = self.img_path2 else: non_hair_root = self.non_hair_path img_path = self.img_path random_number1 = random.randrange(0, 21-12) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] # non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,random.choice(files)) ref_hair = self.read_img(ref_path) for i in range(12): non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1+i) + '.jpg') img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) # hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal_controlnet_without_pose(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement' self.img_path = os.path.join(train_data_dir, "hair") self.img_path2 = os.path.join(train_data_dir, "non-hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") self.pose_path2 = os.path.join(train_data_dir2, "pose.npy") self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair") self.ref_path2 = os.path.join(train_data_dir2, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number = random.uniform(0, 1) if random_number<0.5: non_hair_root = self.non_hair_path2 img_path = self.img_path2 else: non_hair_root = self.non_hair_path img_path = self.img_path random_number1 = random.randrange(0, 21-12) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,random.choice(files)) ref_hair = self.read_img(ref_path) for i in range(12): non_hair_path = os.path.join(self.non_hair_path, name, str(random_number1+i) + '.jpg') img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal_controlnet(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): train_data_dir2 = '/opt/liblibai-models/user-workspace/zyx/sky/3dhair/data/segement' self.img_path = os.path.join(train_data_dir, "hair") self.img_path2 = os.path.join(train_data_dir, "non-hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") self.pose_path2 = os.path.join(train_data_dir2, "pose.npy") self.non_hair_path2 = os.path.join(train_data_dir2, "non-hair") self.ref_path2 = os.path.join(train_data_dir2, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number = random.uniform(0, 1) if random_number<0.5: non_hair_root = self.non_hair_path2 img_path = self.img_path2 else: non_hair_root = self.non_hair_path img_path = self.img_path random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,random.choice(files)) ref_hair = self.read_img(ref_path) non_hair_path = os.path.join(img_path, name, str(random_number2) + '.jpg') img_non_hair = self.read_img(non_hair_path) hair_path = os.path.join(img_path, name, str(random_number1) + '.jpg') img_hair= self.read_img(hair_path) x = self.x[random_number1] y = self.y[random_number1] return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal_pose(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21-12) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] random_number = random.randint(0, 1) if random_number==0: img_path = self.img_path else: img_path = self.non_hair_path non_hair_path = os.path.join(img_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) ref_hair = self.read_img(ref_path) for i in range(12): img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) hair_path = os.path.join(img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal_random_reference(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21-12) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,random.choice(files)) ref_hair = self.read_img(ref_path) for i in range(12): img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_random_reference(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") self.img_path2 = os.path.join(train_data_dir, "hair_good") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") # self.lists = os.listdir(self.img_path2) self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) # while random_number2 == random_number1: # random_number2 = random.randrange(0, 21) name = self.lists[index].split('.')[0] #random_number1 = random_number1 #* 10 #random_number2 = random_number2 * 10 #random_number2 = random_number1 hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,random.choice(files)) # print('________') # print(files) # print('++++++++') # print(ref_folder) # print("========") # print(name) # print("********") # print(ref_path) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) # pose2 = self.pose[random_number2] # pose2 = torch.tensor(pose2) hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_num = [0, 2, 6, 14, 18, 21] # img_hair_stack = [] # polar = self.polars_rad[random_number1] # polar = torch.tensor(polar).unsqueeze(0) # azimuths = self.azimuths_rad[random_number1] # azimuths = torch.tensor(azimuths).unsqueeze(0) img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) x = torch.tensor(self.x[random_number1]) y = torch.tensor(self.y[random_number1]) x2 = torch.tensor(self.x[random_number2]) y2 = torch.tensor(self.y[random_number2]) # begin = random.randrange(0, 21-self.frame_num) # hair_num = [i+begin for i in range(self.frame_num)] # for i in hair_num: # img_hair = cv2.imread(hair_path) # img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) # img_hair = cv2.resize(img_hair, (512, 512)) # img_hair = (img_hair / 255.0) * 2 - 1 # img_hair = torch.tensor(img_hair).permute(2, 0, 1) # img_hair_stack.append(img_hair) # img_hair = torch.stack(img_hair_stack) return { # 'hair_pose': pose1, 'img_hair': img_hair, # 'bald_pose': pose2, 'pose': pose, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, # 'x2': x2, # 'y2': y2, # 'polar': polar, # 'azimuths':azimuths, } def __len__(self): return self.len class myDataset_sv3d_simple_random_reference_controlnet(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") self.img_path2 = os.path.join(train_data_dir, "hair_good") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "multi_reference2") # self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_ref_img(self, path): img = cv2.imread(path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = hair_transform(image=img) # img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) name = self.lists[index].split('.')[0] # hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')][:3] ref_path = os.path.join(ref_folder,random.choice(files)) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) ref_hair = hair_transform(image=ref_hair)['image'] # print(type(ref_hair)) # print(ref_hair.keys()) # ref_hair = self.read_ref_img(ref_path) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) #hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') # hair_path = os.path.join(self.non_hair_path, name, str(random_number1) + '.jpg') img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) x = torch.tensor(self.x[random_number1]) y = torch.tensor(self.y[random_number1]) x2 = torch.tensor(self.x[random_number2]) y2 = torch.tensor(self.y[random_number2]) return { # 'hair_pose': pose1, 'img_hair': img_hair, # 'bald_pose': pose2, 'pose': pose, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_random_reference_stable_hair(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21) random_number2 = random.randrange(0, 21) random_number1 = random_number2 name = self.lists[index] hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] ref_path = os.path.join(ref_folder,random.choice(files)) img_hair = cv2.imread(hair_path) img_non_hair = cv2.imread(non_hair_path) ref_hair = cv2.imread(ref_path) img_non_hair = cv2.cvtColor(img_non_hair, cv2.COLOR_BGR2RGB) ref_hair = cv2.cvtColor(ref_hair, cv2.COLOR_BGR2RGB) img_non_hair = cv2.resize(img_non_hair, (512, 512)) ref_hair = cv2.resize(ref_hair, (512, 512)) img_non_hair = (img_non_hair / 255.0) * 2 - 1 ref_hair = (ref_hair / 255.0) * 2 - 1 img_non_hair = torch.tensor(img_non_hair).permute(2, 0, 1) ref_hair = torch.tensor(ref_hair).permute(2, 0, 1) pose = self.pose[random_number1] pose = torch.tensor(pose) hair_path = os.path.join(self.img_path, name, str(random_number1) + '.jpg') img_hair = cv2.imread(hair_path) img_hair = cv2.cvtColor(img_hair, cv2.COLOR_BGR2RGB) img_hair = cv2.resize(img_hair, (512, 512)) img_hair = (img_hair / 255.0) * 2 - 1 img_hair = torch.tensor(img_hair).permute(2, 0, 1) x = torch.tensor(self.x[random_number1]) y = torch.tensor(self.y[random_number1]) return { 'img_hair': img_hair, 'pose': pose, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len class myDataset_sv3d_simple_temporal_small_squence(data.Dataset): """Custom data.Dataset compatible with data.DataLoader.""" def __init__(self, train_data_dir, frame_num=6): self.img_path = os.path.join(train_data_dir, "hair") # self.pose_path = os.path.join(train_data_dir, "pose.npy") # self.non_hair_path = os.path.join(train_data_dir, "no_hair") # self.ref_path = os.path.join(train_data_dir, "ref_hair") self.pose_path = os.path.join(train_data_dir, "pose.npy") self.non_hair_path = os.path.join(train_data_dir, "non-hair") self.ref_path = os.path.join(train_data_dir, "reference") self.lists = os.listdir(self.img_path) self.len = len(self.lists)-10 self.pose = np.load(self.pose_path) self.frame_num = frame_num #self.pose = np.random.randn(12, 4) elevations_deg = [-0.05/2*np.pi*360] * 21 azimuths_deg = np.linspace(0, 360, 21+1)[1:] % 360 Face_yaws = [0.4 * np.sin(2 * 3.14 * i / 60) for i in range(60)] for i in Face_yaws: if i<0: i = 2*np.pi+i i = i/2*np.pi*360 face_yaws = [Face_yaws[0]] for i in range(20): face_yaws.append(Face_yaws[3*i+2]) self.polars_rad = [np.deg2rad(90-e) for e in elevations_deg] self.azimuths_rad = [np.deg2rad((a) % 360) for a in azimuths_deg] self.azimuths_rad[:-1].sort() x = [0.4 * np.sin(2 * 3.14 * i / 120) for i in range(60)] y = [- 0.05 + 0.3 * np.cos(2 * 3.14 * i / 120) for i in range(60)] self.x = [x[0]] self.y = [y[0]] for i in range(20): self.x.append(x[i*3+2]) self.y.append(y[i*3+2]) def read_img(self, path): img = cv2.imread(path) img = cv2.resize(img, (512, 512)) img = (img / 255.0) * 2 - 1 img = torch.tensor(img).permute(2, 0, 1) return img def __getitem__(self, index): """Returns one data pair (source and target).""" # seq_len, fea_dim random_number1 = random.randrange(0, 21-6) random_number2 = random.randrange(0, 21) name = self.lists[index] x_stack = [] y_stack = [] img_non_hair_stack = [] img_hair_stack = [] non_hair_path = os.path.join(self.non_hair_path, name, str(random_number2) + '.jpg') ref_folder = os.path.join(self.ref_path, name) files = [f for f in os.listdir(ref_folder) if f.endswith('.jpg')] # ref_path = os.path.join(ref_folder, str(random_number2) + '.jpg') ref_path = os.path.join(ref_folder,files[0]) ref_hair = self.read_img(ref_path) for i in range(6): img_non_hair_stack.append(self.read_img(non_hair_path).unsqueeze(0)) hair_path = os.path.join(self.img_path, name, str(random_number1+i) + '.jpg') img_hair_stack.append(self.read_img(hair_path).unsqueeze(0)) x_stack.append(torch.tensor(self.x[random_number1+i]).unsqueeze(0)) y_stack.append(torch.tensor(self.y[random_number1+i]).unsqueeze(0)) img_non_hair = torch.cat(img_non_hair_stack, axis=0) img_hair = torch.cat(img_hair_stack, axis=0) x = torch.cat(x_stack, axis=0) y = torch.cat(y_stack, axis=0) return { 'img_hair': img_hair, 'img_non_hair': img_non_hair, 'ref_hair': ref_hair, 'x': x, 'y': y, } def __len__(self): return self.len if __name__ == "__main__": train_dataset = myDataset("./data") train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=1, num_workers=1, ) for epoch in range(0, len(train_dataset) + 1): for step, batch in enumerate(train_dataloader): print("batch[hair_pose]:", batch["hair_pose"]) print("batch[img_hair]:", batch["img_hair"]) print("batch[bald_pose]:", batch["bald_pose"]) print("batch[img_non_hair]:", batch["img_non_hair"]) print("batch[ref_hair]:", batch["ref_hair"])