keypoint.py
3.23 KB
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
import os.path
import torchvision.transforms as transforms
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
import PIL
import random
import pandas as pd
import numpy as np
import torch
class KeyDataset(BaseDataset):
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_P = os.path.join(opt.dataroot, opt.phase) #person images
self.dir_K = os.path.join(opt.dataroot, opt.phase + 'K') #keypoints
self.init_categories(opt.pairLst)
self.transform = get_transform(opt)
def init_categories(self, pairLst):
pairs_file_train = pd.read_csv(pairLst)
self.size = len(pairs_file_train)
self.pairs = []
print('Loading data pairs ...')
for i in range(self.size):
pair = [pairs_file_train.iloc[i]['from'], pairs_file_train.iloc[i]['to']]
self.pairs.append(pair)
print('Loading data pairs finished ...')
def __getitem__(self, index):
if self.opt.phase == 'train':
index = random.randint(0, self.size-1)
P1_name, P2_name = self.pairs[index]
P1_path = os.path.join(self.dir_P, P1_name) # person 1
BP1_path = os.path.join(self.dir_K, P1_name + '.npy') # bone of person 1
# person 2 and its bone
P2_path = os.path.join(self.dir_P, P2_name) # person 2
BP2_path = os.path.join(self.dir_K, P2_name + '.npy') # bone of person 2
P1_img = Image.open(P1_path).convert('RGB')
P2_img = Image.open(P2_path).convert('RGB')
BP1_img = np.load(BP1_path) # h, w, c
BP2_img = np.load(BP2_path)
# use flip
if self.opt.phase == 'train' and self.opt.use_flip:
# print ('use_flip ...')
flip_random = random.uniform(0,1)
if flip_random > 0.5:
# print('fliped ...')
P1_img = P1_img.transpose(Image.FLIP_LEFT_RIGHT)
P2_img = P2_img.transpose(Image.FLIP_LEFT_RIGHT)
BP1_img = np.array(BP1_img[:, ::-1, :]) # flip
BP2_img = np.array(BP2_img[:, ::-1, :]) # flip
BP1 = torch.from_numpy(BP1_img).float() #h, w, c
BP1 = BP1.transpose(2, 0) #c,w,h
BP1 = BP1.transpose(2, 1) #c,h,w
BP2 = torch.from_numpy(BP2_img).float()
BP2 = BP2.transpose(2, 0) #c,w,h
BP2 = BP2.transpose(2, 1) #c,h,w
P1 = self.transform(P1_img)
P2 = self.transform(P2_img)
else:
BP1 = torch.from_numpy(BP1_img).float() #h, w, c
BP1 = BP1.transpose(2, 0) #c,w,h
BP1 = BP1.transpose(2, 1) #c,h,w
BP2 = torch.from_numpy(BP2_img).float()
BP2 = BP2.transpose(2, 0) #c,w,h
BP2 = BP2.transpose(2, 1) #c,h,w
P1 = self.transform(P1_img)
P2 = self.transform(P2_img)
return {'P1': P1, 'BP1': BP1, 'P2': P2, 'BP2': BP2,
'P1_path': P1_name, 'P2_path': P2_name}
def __len__(self):
if self.opt.phase == 'train':
return 4000
elif self.opt.phase == 'test':
return self.size
def name(self):
return 'KeyDataset'