forked from Kai-46/SatelliteNeRF
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdata_loader.py
78 lines (64 loc) · 3.18 KB
/
data_loader.py
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
import os
import numpy as np
import imageio
imageio.plugins.freeimage.download()
from collections import OrderedDict
import json
########################################################################################################################
# camera coordinate system: x-->right, y-->down, z-->scene (opencv/colmap convention)
# poses is camera-to-world
########################################################################################################################
def read_split(basedir, split, ignore_mask):
cam_dict = json.load(open(os.path.join(basedir, split, 'cam_dict_norm.json')))
imgsizes = []
intrinsics = []
poses = []
imgfpaths = []
maskfpaths = []
for x in sorted(cam_dict.keys()):
K = np.array(cam_dict[x]['K']).reshape((4, 4))
W2C = np.array(cam_dict[x]['W2C']).reshape((4, 4))
C2W = np.linalg.inv(W2C)
W, H = cam_dict[x]['img_size']
imgsizes.append(np.array([W, H]))
intrinsics.append(K)
poses.append(C2W)
imgfpaths.append(os.path.join(basedir, split, 'image', x))
mask_fpath = os.path.join(basedir, split, 'mask', x)
if (not ignore_mask) and os.path.isfile(mask_fpath):
maskfpaths.append(mask_fpath)
else:
maskfpaths.append(None)
imgsizes = np.stack(imgsizes, axis=0)
intrinsics = np.stack(intrinsics, axis=0)
poses = np.stack(poses, axis=0)
return imgsizes, intrinsics, poses, imgfpaths, maskfpaths
def load_data(basedir, ignore_mask=True):
imgsizes, intrinsics, poses, imgfpaths, maskfpaths = read_split(basedir, split='train', ignore_mask=ignore_mask)
test_imgsizes, test_intrinsics, test_poses, test_imgfiles, test_maskfiles = read_split(basedir, split='test', ignore_mask=ignore_mask)
val_imgsizes, val_intrinsics, val_poses, val_imgfiles, val_maskfiles = read_split(basedir, split='test', ignore_mask=ignore_mask)
counts = [0] + [len(x) for x in [imgfpaths, val_imgfiles, test_imgfiles]]
counts = np.cumsum(counts)
i_split = [list(np.arange(counts[i], counts[i+1])) for i in range(3)]
imgsizes = np.concatenate([imgsizes, val_imgsizes, test_imgsizes], 0)
intrinsics = np.concatenate([intrinsics, val_intrinsics, test_intrinsics], 0)
poses = np.concatenate([poses, val_poses, test_poses], 0)
imgfpaths = imgfpaths + val_imgfiles + test_imgfiles
maskfpaths = maskfpaths + val_maskfiles + test_maskfiles
render_cams = None
data = OrderedDict([('imgfpaths', imgfpaths),
('maskfpaths', maskfpaths),
('imgsizes', imgsizes),
('intrinsics', intrinsics),
('poses', poses),
('i_train', i_split[0]),
('i_val', i_split[1]),
('i_test', i_split[2]),
('render_cams', render_cams)])
print('Data statistics:')
print('\t # of training views: ', len(data['i_train']))
print('\t # of validation views: ', len(data['i_val']))
print('\t # of test views: ', len(data['i_test']))
if data['render_cams'] is not None:
print('\t # of render cameras: ', len(data['render_cams']))
return data