File size: 4,817 Bytes
2bd5c7a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import os.path
from data.base_dataset import BaseDataset, get_params, get_transform
from data.image_folder import make_dataset
from PIL import Image
import torch
import json
import numpy as np
import os.path as osp
from PIL import ImageDraw
class AlignedDataset(BaseDataset):
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.diction={}
if opt.isTrain or opt.use_encoded_image:
dir_A = '_A' if self.opt.label_nc == 0 else '_label'
self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A)
self.A_paths = sorted(make_dataset(self.dir_A))
self.fine_height=256
self.fine_width=192
self.radius=5
dir_B = '_B' if self.opt.label_nc == 0 else '_img'
self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B)
self.B_paths = sorted(make_dataset(self.dir_B))
self.dataset_size = len(self.A_paths)
if opt.isTrain or opt.use_encoded_image:
dir_E = '_edge'
self.dir_E = os.path.join(opt.dataroot, opt.phase + dir_E)
self.E_paths = sorted(make_dataset(self.dir_E))
if opt.isTrain or opt.use_encoded_image:
dir_C = '_color'
self.dir_C = os.path.join(opt.dataroot, opt.phase + dir_C)
self.C_paths = sorted(make_dataset(self.dir_C))
def __getitem__(self, index):
A_path = self.A_paths[index]
A = Image.open(A_path).convert('L')
params = get_params(self.opt, A.size)
if self.opt.label_nc == 0:
transform_A = get_transform(self.opt, params)
A_tensor = transform_A(A.convert('RGB'))
else:
transform_A = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
A_tensor = transform_A(A) * 255.0
B_path = self.B_paths[index]
B = Image.open(B_path).convert('RGB')
transform_B = get_transform(self.opt, params)
B_tensor = transform_B(B)
C_path = self.C_paths[index]
C = Image.open(C_path).convert('RGB')
C_tensor = transform_B(C)
E_path = self.E_paths[index]
E = Image.open(E_path).convert('L')
E_tensor = transform_A(E)
index_un = np.random.randint(14221)
C_un_path = self.C_paths[index_un]
C_un = Image.open(C_un_path).convert('RGB')
C_un_tensor = transform_B(C_un)
E_un_path = self.E_paths[index_un]
E_un = Image.open(E_un_path).convert('L')
E_un_tensor = transform_A(E_un)
pose_name =B_path.replace('.png', '_keypoints.json').replace('.jpg','_keypoints.json').replace('train_img','train_pose')
with open(osp.join(pose_name), 'r') as f:
pose_label = json.load(f)
try:
pose_data = pose_label['people'][0]['pose_keypoints']
except IndexError:
pose_data = [0 for i in range(54)]
pose_data = np.array(pose_data)
pose_data = pose_data.reshape((-1,3))
point_num = pose_data.shape[0]
pose_map = torch.zeros(point_num, self.fine_height, self.fine_width)
r = self.radius
im_pose = Image.new('L', (self.fine_width, self.fine_height))
pose_draw = ImageDraw.Draw(im_pose)
for i in range(point_num):
one_map = Image.new('L', (self.fine_width, self.fine_height))
draw = ImageDraw.Draw(one_map)
pointx = pose_data[i,0]
pointy = pose_data[i,1]
if pointx > 1 and pointy > 1:
draw.rectangle((pointx-r, pointy-r, pointx+r, pointy+r), 'white', 'white')
pose_draw.rectangle((pointx-r, pointy-r, pointx+r, pointy+r), 'white', 'white')
one_map = transform_B(one_map.convert('RGB'))
pose_map[i] = one_map[0]
P_tensor=pose_map
densepose_name = B_path.replace('.png', '.npy').replace('.jpg','.npy').replace('train_img','train_densepose')
dense_mask = np.load(densepose_name).astype(np.float32)
dense_mask = transform_A(dense_mask)
if self.opt.isTrain:
input_dict = { 'label': A_tensor, 'image': B_tensor, 'path': A_path, 'img_path': B_path ,'color_path': C_path,'color_un_path': C_un_path,
'edge': E_tensor, 'color': C_tensor, 'edge_un': E_un_tensor, 'color_un': C_un_tensor, 'pose':P_tensor, 'densepose':dense_mask
}
return input_dict
def __len__(self):
return len(self.A_paths) // (self.opt.batchSize * self.opt.num_gpus) * (self.opt.batchSize * self.opt.num_gpus)
def name(self):
return 'AlignedDataset'
|