(basedir, half_res=False, testskip=1)
| 35 | |
| 36 | |
| 37 | def load_blender_data(basedir, half_res=False, testskip=1): |
| 38 | splits = ['train', 'holdout', 'val', 'test'] |
| 39 | metas = {} |
| 40 | for s in splits: |
| 41 | with open(os.path.join(basedir, 'transforms_{}.json'.format(s)), 'r') as fp: |
| 42 | metas[s] = json.load(fp) |
| 43 | |
| 44 | all_imgs = [] |
| 45 | all_poses = [] |
| 46 | counts = [0] |
| 47 | for s in splits: |
| 48 | meta = metas[s] |
| 49 | imgs = [] |
| 50 | poses = [] |
| 51 | if s=='train' or testskip==0: |
| 52 | skip = 1 |
| 53 | elif s=='holdout': |
| 54 | skip = 1 |
| 55 | else: |
| 56 | skip = testskip |
| 57 | |
| 58 | for frame in meta['frames'][::skip]: |
| 59 | fname = os.path.join(basedir, frame['file_path'] + '.png') |
| 60 | imgs.append(imageio.imread(fname)) |
| 61 | poses.append(np.array(frame['transform_matrix'])) |
| 62 | imgs = (np.array(imgs) / 255.).astype(np.float32) # keep all 4 channels (RGBA) |
| 63 | poses = np.array(poses).astype(np.float32) |
| 64 | counts.append(counts[-1] + imgs.shape[0]) |
| 65 | all_imgs.append(imgs) |
| 66 | all_poses.append(poses) |
| 67 | |
| 68 | i_split = [np.arange(counts[i], counts[i+1]) for i in range(4)] |
| 69 | |
| 70 | imgs = np.concatenate(all_imgs, 0) |
| 71 | poses = np.concatenate(all_poses, 0) |
| 72 | |
| 73 | H, W = imgs[0].shape[:2] |
| 74 | camera_angle_x = float(meta['camera_angle_x']) |
| 75 | focal = .5 * W / np.tan(.5 * camera_angle_x) |
| 76 | |
| 77 | render_poses = torch.stack([pose_spherical(angle, -30.0, 4.0) for angle in np.linspace(-180,180,40+1)[:-1]], 0) |
| 78 | |
| 79 | if half_res: |
| 80 | H = H//2 |
| 81 | W = W//2 |
| 82 | focal = focal/2. |
| 83 | |
| 84 | imgs_half_res = np.zeros((imgs.shape[0], H, W, 4))\ |
| 85 | |
| 86 | for i, img in enumerate(imgs): |
| 87 | imgs_half_res[i] = cv2.resize(img, (H, W), interpolation=cv2.INTER_AREA) |
| 88 | imgs = imgs_half_res |
| 89 | # imgs = tf.image.resize_area(imgs, [400, 400]).numpy() |
| 90 | |
| 91 | |
| 92 | return imgs, poses, render_poses, [H, W, focal], i_split |
| 93 | |
| 94 |
no test coverage detected