File size: 2,808 Bytes
f348660 | 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 | import torch.utils.data
import numpy as np
import modules.midas.utils as utils
from PIL import Image
def load_input_image(input_image_fp):
return utils.read_image(input_image_fp)
def load_sparse_depth(input_sparse_depth_fp):
input_sparse_depth = np.array(Image.open(input_sparse_depth_fp), dtype=np.float32) / 256.0
input_sparse_depth[input_sparse_depth <= 0] = 0.0
return input_sparse_depth
class SML_dataset(torch.utils.data.Dataset):
def __init__(self,
image_paths,
radar_paths,
gt_paths,
sparse_gt_paths,
rcnet_paths,
mono_pred_paths = None,
mono_ga_paths = None,
):
self.n_sample = len(image_paths)
for paths in [image_paths, radar_paths, gt_paths, sparse_gt_paths,
rcnet_paths, mono_pred_paths, mono_ga_paths]:
if paths is not None:
assert len(paths) == self.n_sample
self.image_paths = image_paths
self.radar_paths = radar_paths
self.gt_paths = gt_paths
self.sparse_gt_paths = sparse_gt_paths
self.rcnet_paths = rcnet_paths
self.mono_pred_paths = mono_pred_paths
self.mono_ga_paths = mono_ga_paths
def __getitem__(self, index):
image = load_input_image(self.image_paths[index])
radar = load_sparse_depth(self.radar_paths[index])
gt = load_sparse_depth(self.gt_paths[index])
sparse_gt = load_sparse_depth(self.sparse_gt_paths[index])
rcnet = load_sparse_depth(self.rcnet_paths[index])
image, radar, gt, sparse_gt, rcnet = [
T.astype(np.float32)
for T in [image, radar, gt, sparse_gt, rcnet]
]
# Crop the image for ZJU dataset
if image.shape[0] == 720:
image = image[720 // 3: 720 // 4 * 3, :, :]
radar = radar[720 // 3: 720 // 4 * 3, :]
gt = gt[720 // 3: 720 // 4 * 3, :]
sparse_gt = sparse_gt[720 // 3: 720 // 4 * 3, :]
if self.mono_ga_paths is not None:
mono_pred = load_sparse_depth(self.mono_ga_paths[index])
mono_pred = mono_pred.astype(np.float32)
if mono_pred.shape[0] == 720:
mono_pred = mono_pred[720 // 3: 720 // 4 * 3, :]
else:
mono_pred = None
if self.mono_ga_paths is not None:
mono_ga = load_sparse_depth(self.mono_ga_paths[index])
mono_ga = mono_ga.astype(np.float32)
if mono_ga.shape[0] == 720:
mono_ga = mono_ga[720 // 3: 720 // 4 * 3, :]
else:
mono_ga = None
return image, mono_pred, radar, gt, sparse_gt, rcnet, mono_ga
def __len__(self):
return self.n_sample |