-
Notifications
You must be signed in to change notification settings - Fork 343
/
EvalDataset.py
executable file
·166 lines (143 loc) · 6.23 KB
/
EvalDataset.py
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
from torch.utils.data import Dataset
import numpy as np
import os
import random
import torchvision.transforms as transforms
from PIL import Image, ImageOps
import cv2
import torch
from PIL.ImageFilter import GaussianBlur
import trimesh
import cv2
class EvalDataset(Dataset):
@staticmethod
def modify_commandline_options(parser):
return parser
def __init__(self, opt, root=None):
self.opt = opt
self.projection_mode = 'orthogonal'
# Path setup
self.root = self.opt.dataroot
if root is not None:
self.root = root
self.RENDER = os.path.join(self.root, 'RENDER')
self.MASK = os.path.join(self.root, 'MASK')
self.PARAM = os.path.join(self.root, 'PARAM')
self.OBJ = os.path.join(self.root, 'GEO', 'OBJ')
self.phase = 'val'
self.load_size = self.opt.loadSize
self.num_views = self.opt.num_views
self.max_view_angle = 360
self.interval = 1
self.subjects = self.get_subjects()
# PIL to tensor
self.to_tensor = transforms.Compose([
transforms.Resize(self.load_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
def get_subjects(self):
var_file = os.path.join(self.root, 'val.txt')
if os.path.exists(var_file):
var_subjects = np.loadtxt(var_file, dtype=str)
return sorted(list(var_subjects))
all_subjects = os.listdir(self.RENDER)
return sorted(list(all_subjects))
def __len__(self):
return len(self.subjects) * self.max_view_angle // self.interval
def get_render(self, subject, num_views, view_id=None, random_sample=False):
'''
Return the render data
:param subject: subject name
:param num_views: how many views to return
:param view_id: the first view_id. If None, select a random one.
:return:
'img': [num_views, C, W, H] images
'calib': [num_views, 4, 4] calibration matrix
'extrinsic': [num_views, 4, 4] extrinsic matrix
'mask': [num_views, 1, W, H] masks
'''
# For now we only have pitch = 00. Hard code it here
pitch = 0
# Select a random view_id from self.max_view_angle if not given
if view_id is None:
view_id = np.random.randint(self.max_view_angle)
# The ids are an even distribution of num_views around view_id
view_ids = [(view_id + self.max_view_angle // num_views * offset) % self.max_view_angle
for offset in range(num_views)]
if random_sample:
view_ids = np.random.choice(self.max_view_angle, num_views, replace=False)
calib_list = []
render_list = []
mask_list = []
extrinsic_list = []
for vid in view_ids:
param_path = os.path.join(self.PARAM, subject, '%d_%02d.npy' % (vid, pitch))
render_path = os.path.join(self.RENDER, subject, '%d_%02d.jpg' % (vid, pitch))
mask_path = os.path.join(self.MASK, subject, '%d_%02d.png' % (vid, pitch))
# loading calibration data
param = np.load(param_path)
# pixel unit / world unit
ortho_ratio = param.item().get('ortho_ratio')
# world unit / model unit
scale = param.item().get('scale')
# camera center world coordinate
center = param.item().get('center')
# model rotation
R = param.item().get('R')
translate = -np.matmul(R, center).reshape(3, 1)
extrinsic = np.concatenate([R, translate], axis=1)
extrinsic = np.concatenate([extrinsic, np.array([0, 0, 0, 1]).reshape(1, 4)], 0)
# Match camera space to image pixel space
scale_intrinsic = np.identity(4)
scale_intrinsic[0, 0] = scale / ortho_ratio
scale_intrinsic[1, 1] = -scale / ortho_ratio
scale_intrinsic[2, 2] = -scale / ortho_ratio
# Match image pixel space to image uv space
uv_intrinsic = np.identity(4)
uv_intrinsic[0, 0] = 1.0 / float(self.opt.loadSize // 2)
uv_intrinsic[1, 1] = 1.0 / float(self.opt.loadSize // 2)
uv_intrinsic[2, 2] = 1.0 / float(self.opt.loadSize // 2)
# Transform under image pixel space
trans_intrinsic = np.identity(4)
mask = Image.open(mask_path).convert('L')
render = Image.open(render_path).convert('RGB')
intrinsic = np.matmul(trans_intrinsic, np.matmul(uv_intrinsic, scale_intrinsic))
calib = torch.Tensor(np.matmul(intrinsic, extrinsic)).float()
extrinsic = torch.Tensor(extrinsic).float()
mask = transforms.Resize(self.load_size)(mask)
mask = transforms.ToTensor()(mask).float()
mask_list.append(mask)
render = self.to_tensor(render)
render = mask.expand_as(render) * render
render_list.append(render)
calib_list.append(calib)
extrinsic_list.append(extrinsic)
return {
'img': torch.stack(render_list, dim=0),
'calib': torch.stack(calib_list, dim=0),
'extrinsic': torch.stack(extrinsic_list, dim=0),
'mask': torch.stack(mask_list, dim=0)
}
def get_item(self, index):
# In case of a missing file or IO error, switch to a random sample instead
try:
sid = index % len(self.subjects)
vid = (index // len(self.subjects)) * self.interval
# name of the subject 'rp_xxxx_xxx'
subject = self.subjects[sid]
res = {
'name': subject,
'mesh_path': os.path.join(self.OBJ, subject + '.obj'),
'sid': sid,
'vid': vid,
}
render_data = self.get_render(subject, num_views=self.num_views, view_id=vid,
random_sample=self.opt.random_multiview)
res.update(render_data)
return res
except Exception as e:
print(e)
return self.get_item(index=random.randint(0, self.__len__() - 1))
def __getitem__(self, index):
return self.get_item(index)