-
Notifications
You must be signed in to change notification settings - Fork 62
/
load_multiscale.py
43 lines (35 loc) · 1.63 KB
/
load_multiscale.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
import numpy as np
import os
import json
import cv2
import imageio
def _load_google_data(basedir, factor=None):
img_basedir = basedir
img_folder = 'images'
imgdir = os.path.join(img_basedir, img_folder)
imgfiles = [os.path.join(imgdir, f) for f in sorted(os.listdir(imgdir)) if f.endswith('JPG') or f.endswith('jpg') or f.endswith('png') or f.endswith('jpeg')]
sh = np.array(cv2.imread(imgfiles[0]).shape)
imgs = []
for f in imgfiles:
im = cv2.imread(f, cv2.IMREAD_UNCHANGED)
if im.shape[-1] == 3:
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
else:
im = cv2.cvtColor(im, cv2.COLOR_BGRA2RGBA)
im = cv2.resize(im, (sh[1]//factor, sh[0]//factor), interpolation=cv2.INTER_AREA)
im = im.astype(np.float32) / 255
imgs.append(im)
imgs = np.stack(imgs, -1)
imgs = np.moveaxis(imgs, -1, 0).astype(np.float32)
data = json.load(open(os.path.join(basedir, 'poses_enu.json')))
poses = np.array(data['poses'])[:, :-2].reshape([-1, 3, 5])
poses[:, :2, 4] = np.array(sh[:2]//factor).reshape([1, 2])
poses[:, 2, 4] = poses[:,2, 4] * 1./factor
scene_scaling_factor = data['scene_scale']
scene_origin = np.array(data['scene_origin'])
scale_split = data['scale_split']
return imgs, poses, scene_scaling_factor, scene_origin, scale_split
def load_multiscale_data(basedir, factor=3):
imgs, poses, scene_scaling_factor, scene_origin, scale_split = _load_google_data(basedir, factor=factor)
print('Loaded image data shape:', imgs.shape, ' hwf:', poses[0,:,-1])
return imgs, poses, scene_scaling_factor, scene_origin, scale_split