-
Notifications
You must be signed in to change notification settings - Fork 121
/
Copy pathhypersim_dataset.py
141 lines (124 loc) · 5.56 KB
/
hypersim_dataset.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
import os
import json
import torch
import torchvision.transforms as transforms
import os.path
import numpy as np
import cv2
from PIL import Image
from torch.utils.data import Dataset
import random
from .__base_dataset__ import BaseDataset
import h5py
def creat_uv_mesh(H, W):
y, x = np.meshgrid(np.arange(0, H, dtype=np.float), np.arange(0, W, dtype=np.float), indexing='ij')
meshgrid = np.stack((x,y))
ones = np.ones((1,H*W), dtype=np.float)
xy = meshgrid.reshape(2, -1)
return np.concatenate([xy, ones], axis=0)
class HypersimDataset(BaseDataset):
def __init__(self, cfg, phase, **kwargs):
super(HypersimDataset, self).__init__(
cfg=cfg,
phase=phase,
**kwargs)
self.metric_scale = cfg.metric_scale
#self.cap_range = self.depth_range # in meter
# init uv
# meshgrid for depth reprojection
self.xy = creat_uv_mesh(768, 1024)
def load_batch(self, meta_data, data_path):
curr_intrinsic = meta_data['cam_in']
# load rgb/depth
curr_rgb, curr_depth = self.load_rgb_depth(data_path['rgb_path'], data_path['depth_path'])
# get semantic labels
curr_sem = self.load_sem_label(data_path['sem_path'], curr_depth)
# create camera model
curr_cam_model = self.create_cam_model(curr_rgb.shape[0], curr_rgb.shape[1], curr_intrinsic)
# get normal labels
curr_normal = self.load_norm_label(data_path['normal_path'], H=curr_rgb.shape[0], W=curr_rgb.shape[1], depth=curr_depth, K=curr_intrinsic) # !!! this is diff of BaseDataset
# get depth mask
depth_mask = self.load_depth_valid_mask(data_path['depth_mask_path'])
curr_depth[~depth_mask] = -1
data_batch = dict(
curr_rgb = curr_rgb,
curr_depth = curr_depth,
curr_sem = curr_sem,
curr_normal = curr_normal,
curr_cam_model=curr_cam_model,
)
return data_batch
def load_data_path(self, meta_data):
# 'rgbs': {'rgb_color': 'Hypersim/data/ai_001_001/images/scene_cam_00_final_preview/frame.0008.color.jpg',
# 'rgb_gamma': 'Hypersim/data/ai_001_001/images/scene_cam_00_final_preview/frame.0008.gamma.jpg',
# 'rgb_tonemap': 'Hypersim/data/ai_001_001/images/scene_cam_00_final_preview/frame.0008.tonemap.jpg',
# 'rgb_raw': 'Hypersim/data/ai_001_001/images/scene_cam_00_final_hdf5/frame.0008.color.hdf5'}
meta_data['rgb'] = meta_data['rgbs']['rgb_color'] # this is diff of BaseDataset
curr_rgb_path = os.path.join(self.data_root, meta_data['rgb'])
curr_depth_path = os.path.join(self.depth_root, meta_data['depth'])
curr_sem_path = os.path.join(self.sem_root, meta_data['sem']) \
if self.sem_root is not None and ('sem' in meta_data) and (meta_data['sem'] is not None) \
else None
curr_norm_path = os.path.join(self.norm_root, meta_data['normal']) \
if ('normal' in meta_data) and (meta_data['normal'] is not None) and (self.norm_root is not None) \
else None
curr_depth_mask_path = os.path.join(self.depth_mask_root, meta_data['depth_mask']) \
if self.depth_mask_root is not None and ('depth_mask' in meta_data) and (meta_data['depth_mask'] is not None) \
else None
data_path=dict(
rgb_path=curr_rgb_path,
depth_path=curr_depth_path,
sem_path=curr_sem_path,
normal_path=curr_norm_path,
depth_mask_path=curr_depth_mask_path,
)
return data_path
def load_rgb_depth(self, rgb_path: str, depth_path: str):
"""
Load the rgb and depth map with the paths.
"""
rgb = self.load_data(rgb_path, is_rgb_img=True)
if rgb is None:
self.logger.info(f'>>>>{rgb_path} has errors.')
# depth = self.load_data(depth_path)
with h5py.File(depth_path, "r") as f: depth = f["dataset"][:]
np.nan_to_num(depth, copy=False, nan=0) # fill nan in gt
if depth is None:
self.logger.info(f'{depth_path} has errors.')
depth = depth.astype(np.float)
depth = self.process_depth(depth, rgb)
return rgb, depth
def load_norm_label(self, norm_path, H, W, depth, K):
with h5py.File(norm_path, "r") as f:
normal = f["dataset"][:]
np.nan_to_num(normal, copy=False, nan=0)
normal[:,:,1:] *= -1
normal = normal.astype(np.float)
return self.align_normal(normal, depth, K, H, W)
def process_depth(self, depth: np.array, rgb: np.array) -> np.array:
depth[depth>60000] = 0
depth = depth / self.metric_scale
return depth
def align_normal(self, normal, depth, K, H, W):
'''
Orientation of surface normals in hypersim is not always consistent
see https://github.com/apple/ml-hypersim/issues/26
'''
# inv K
K = np.array([[K[0], 0 ,K[2]],
[0, K[1], K[3]],
[0, 0, 1]])
inv_K = np.linalg.inv(K)
# reprojection depth to camera points
if H == 768 and W == 1024:
xy = self.xy
else:
print('img size no-equal 768x1024')
xy = creat_uv_mesh(H, W)
points = np.matmul(inv_K[:3, :3], xy).reshape(3, H, W)
points = depth * points
points = points.transpose((1,2,0))
# align normal
orient_mask = np.sum(normal * points, axis=2) > 0
normal[orient_mask] *= -1
return normal