File size: 4,609 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 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 | 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
def load_in_extrinsic(path):
with open(path, 'r') as f:
lines = f.readlines()
lines = [line.strip() for line in lines]
lines = [line.split(' ') for line in lines]
lines = [[float(x) for x in line] for line in lines]
intrinsic = np.array(lines[0]).reshape(3,3)
extrinsic_ = np.array(lines[1]).reshape(3,4)
extrinsic = np.eye(4)
extrinsic[:3,:3] = extrinsic_[:3,:3]
extrinsic[:3,3] = extrinsic_[:3,3]
return intrinsic, extrinsic
def get_baseline_focal(left_path, right_path):
intrinsic0, extrinsic0 = load_in_extrinsic(left_path)
intrinsic1, extrinsic1 = load_in_extrinsic(right_path)
cam_extrinsic = extrinsic1 @ np.linalg.inv(extrinsic0) # ref:https://github.com/fabiotosi92/SMD-Nets/issues/8
baseline = abs(cam_extrinsic[0, 3])
return baseline, (intrinsic0[0,0] + intrinsic0[1,1]) / 2
def get_depth(disp, focal, baseline):
"""
get depth from reference frame disparity and camera intrinsics
"""
return focal * baseline / disp
class Dataset(BaseDataset):
def build_metas(self):
self.focal = 1920
self.dataset_name = 'unrealstereo4k'
# Auto-discover scene folders (mirror layout uses 00000..0000N).
if "scenes" in self.cfg:
folder_names = list(self.cfg.scenes)
else:
folder_names = sorted(
d for d in os.listdir(self.cfg.data_root)
if os.path.isdir(os.path.join(self.cfg.data_root, d, 'Image0'))
)
self.rgb_files = []
self.depth_files = []
# The fabiotosi92/UnrealStereo4K mirror ships .jpg under Image0/; the
# original release uses .png. Accept either.
exts = ('.png', '.jpg', '.jpeg')
for folder_name in folder_names:
folder_path = os.path.join(self.cfg.data_root, folder_name, 'Image0')
if not os.path.isdir(folder_path):
continue
for file_name in os.listdir(folder_path):
base, ext = os.path.splitext(file_name)
if ext.lower() not in exts:
continue
rgb_path = os.path.join(folder_path, file_name)
dpt_path = os.path.join(self.cfg.data_root, folder_name, 'Disp0', base + '.npy')
if os.path.isfile(rgb_path) and os.path.isfile(dpt_path):
self.rgb_files.append(rgb_path)
self.depth_files.append(dpt_path)
assert len(self.rgb_files) == len(self.depth_files), \
f"{len(self.rgb_files)} rgbs vs {len(self.depth_files)} depths"
def read_depth(self, index, depth=None):
disp = np.load(self.depth_files[index])
invalid_mask = np.isnan(disp) | (disp == 0)
disp[invalid_mask] = 1e-4
left_extrin_path = self.depth_files[index].replace('Disp0', 'Extrinsics0').replace('.npy', '.txt')
right_extrin_path = self.depth_files[index].replace('Disp0', 'Extrinsics1').replace('.npy', '.txt')
if os.path.isfile(left_extrin_path) and os.path.isfile(right_extrin_path):
baseline, focal = get_baseline_focal(left_extrin_path, right_extrin_path)
else:
# Mirrors that ship only Image0/Image1/Disp0/Disp1 (e.g. fabiotosi92's HF
# mirror) lack the Extrinsics .txt files. Fall back to the published
# SMD-Nets defaults — UnrealStereo4K is rendered with a 30 cm baseline
# at 3840×2160 with focal length 1920.
baseline, focal = 0.3, 1920.0
depth = get_depth(disp, focal, baseline).astype(np.float32)
min_val, max_val = 0.1, 80
tiankong_mask = depth > 200.
valid_mask = np.logical_and(
depth > 0.1, ~np.isnan(depth)) & (~np.isinf(depth))
valid_mask = np.logical_and(valid_mask, depth < max_val)
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()
depth = np.clip(depth, min_val, max_val)
depth[tiankong_mask] = depth.max() + 1.0
valid_mask = np.logical_or(valid_mask, tiankong_mask)
return depth, valid_mask.astype(np.uint8)
|