Spaces:
Running
Running
| 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 | |
| class Matterport3DDataset(BaseDataset): | |
| def __init__(self, cfg, phase, **kwargs): | |
| super(Matterport3DDataset, self).__init__( | |
| cfg=cfg, | |
| phase=phase, | |
| **kwargs) | |
| self.metric_scale = cfg.metric_scale | |
| #self.cap_range = self.depth_range # in meter | |
| def load_norm_label(self, norm_path, H, W): | |
| normal_x = cv2.imread(norm_path['x'], cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) | |
| normal_y = cv2.imread(norm_path['y'], cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) | |
| normal_z = cv2.imread(norm_path['z'], cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) | |
| raw_normal = np.array([normal_x, normal_y, normal_z]) | |
| invalid_mask = np.all(raw_normal == 0, axis=0) | |
| ego_normal = raw_normal.astype(np.float64) / 32768.0 - 1 | |
| ego2cam = np.array([[1,0,0], | |
| [0,-1,0], | |
| [0,0,-1]]) | |
| normal = (ego2cam @ ego_normal.reshape(3,-1)).reshape(ego_normal.shape) | |
| normal[:,invalid_mask] = 0 | |
| normal = normal.transpose((1,2,0)) | |
| if normal.shape[0] != H or normal.shape[1] != W: | |
| normal = cv2.resize(normal, [W,H], interpolation=cv2.INTER_NEAREST) | |
| return normal | |
| def process_depth(self, depth: np.array, rgb: np.array) -> np.array: | |
| depth[depth>65500] = 0 | |
| depth = depth / self.metric_scale | |
| return depth | |