File size: 5,693 Bytes
3ef1661 |
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 |