74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
import cv2
|
|
import h5py
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
from torchvision.transforms import Compose
|
|
|
|
from dataset.transform import Resize, NormalizeImage, PrepareForNet, Crop
|
|
|
|
|
|
def hypersim_distance_to_depth(npyDistance):
|
|
intWidth, intHeight, fltFocal = 1024, 768, 886.81
|
|
|
|
npyImageplaneX = np.linspace((-0.5 * intWidth) + 0.5, (0.5 * intWidth) - 0.5, intWidth).reshape(
|
|
1, intWidth).repeat(intHeight, 0).astype(np.float32)[:, :, None]
|
|
npyImageplaneY = np.linspace((-0.5 * intHeight) + 0.5, (0.5 * intHeight) - 0.5,
|
|
intHeight).reshape(intHeight, 1).repeat(intWidth, 1).astype(np.float32)[:, :, None]
|
|
npyImageplaneZ = np.full([intHeight, intWidth, 1], fltFocal, np.float32)
|
|
npyImageplane = np.concatenate(
|
|
[npyImageplaneX, npyImageplaneY, npyImageplaneZ], 2)
|
|
|
|
npyDepth = npyDistance / np.linalg.norm(npyImageplane, 2, 2) * fltFocal
|
|
return npyDepth
|
|
|
|
|
|
class Hypersim(Dataset):
|
|
def __init__(self, filelist_path, mode, size=(518, 518)):
|
|
|
|
self.mode = mode
|
|
self.size = size
|
|
|
|
with open(filelist_path, 'r') as f:
|
|
self.filelist = f.read().splitlines()
|
|
|
|
net_w, net_h = size
|
|
self.transform = Compose([
|
|
Resize(
|
|
width=net_w,
|
|
height=net_h,
|
|
resize_target=True if mode == 'train' else False,
|
|
keep_aspect_ratio=True,
|
|
ensure_multiple_of=14,
|
|
resize_method='lower_bound',
|
|
image_interpolation_method=cv2.INTER_CUBIC,
|
|
),
|
|
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
PrepareForNet(),
|
|
] + ([Crop(size[0])] if self.mode == 'train' else []))
|
|
|
|
def __getitem__(self, item):
|
|
img_path = self.filelist[item].split(' ')[0]
|
|
depth_path = self.filelist[item].split(' ')[1]
|
|
|
|
image = cv2.imread(img_path)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
|
|
|
|
depth_fd = h5py.File(depth_path, "r")
|
|
distance_meters = np.array(depth_fd['dataset'])
|
|
depth = hypersim_distance_to_depth(distance_meters)
|
|
|
|
sample = self.transform({'image': image, 'depth': depth})
|
|
|
|
sample['image'] = torch.from_numpy(sample['image'])
|
|
sample['depth'] = torch.from_numpy(sample['depth'])
|
|
|
|
sample['valid_mask'] = (torch.isnan(sample['depth']) == 0)
|
|
sample['depth'][sample['valid_mask'] == 0] = 0
|
|
|
|
sample['image_path'] = self.filelist[item].split(' ')[0]
|
|
|
|
return sample
|
|
|
|
def __len__(self):
|
|
return len(self.filelist) |