File size: 1,976 Bytes
436b829 | 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 | from ppd.data.depth_estimation import Dataset as BaseDataset
from ppd.data.depth_estimation import *
import os
import json
from ppd.utils.logger import Log
class Dataset(BaseDataset):
def build_metas(self):
self.dataset_name = 'hypersim'
splits = json.load(open(self.cfg.split_path))
key_rgb = f"{self.cfg.split}_rgb_paths"
rgb_paths = splits[key_rgb]
key_dpt = f"{self.cfg.split}_dpt_paths"
dpt_paths = splits[key_dpt]
assert len(rgb_paths) == len(dpt_paths)
self.rgb_files = [os.path.join(self.cfg.data_root, rgb_path)
for rgb_path in rgb_paths]
self.depth_files = [os.path.join(
self.cfg.data_root, dpt_path) for dpt_path in dpt_paths]
def read_rgb_name(self, index):
return '__'.join(self.rgb_files[index].split('/')[-4:])
def read_depth(self, index, depth=None):
if not hasattr(self, 'depth_files'):
return None, None
Log.debug(index, self.depth_files[index])
start_time = time.time()
depth_path = self.depth_files[index]
depth = cv2.imread(depth_path, cv2.IMREAD_ANYCOLOR |
cv2.IMREAD_ANYDEPTH) / 1000.
end_time = time.time()
if end_time - start_time > 1:
Log.warn(
f'Long time to read {self.depth_files[index]}: {end_time - start_time}')
valid_mask = np.logical_and(
(depth > 0.1) & (depth < 65.0), ~np.isnan(depth)) & (~np.isinf(depth))
if valid_mask.sum() == 0:
Log.warn('No valid mask in the depth map of {}'.format(
self.depth_files[index]))
if valid_mask.sum() != 0 and np.isnan(depth).sum() != 0:
depth[np.isnan(depth)] = depth[valid_mask].max()
if valid_mask.sum() != 0 and np.isinf(depth).sum() != 0:
depth[np.isinf(depth)] = depth[valid_mask].max()
return depth, valid_mask.astype(np.uint8)
|