google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Generate training cameras.
17

18
randomly sample camera origin
19
then compute valid rotation to overlap with layout
20
"""
21
import os
22

23
from external.gsn.models.nerf_utils import get_sample_points
24
import matplotlib.pyplot as plt
25
import numpy as np
26
import torch
27
from utils import camera_util
28

29
plane_width = 256 * 0.15  # 38.4
30
nerf_far = 16
31
ymax = 0.5
32
seed = 0
33
rng = np.random.RandomState(seed)
34

35
sampled_Rts = []
36
sampled_cameras = []
37
for i in range(1000):
38
  if i % 100 == 0:
39
    print(i)
40

41
  Tx = rng.rand() * plane_width - plane_width / 2
42
  Tz = rng.rand() * plane_width - plane_width / 2
43
  Ty = rng.randn() * ymax / 3
44

45
  valid_degrees = []
46
  # find the rotations that are valid
47
  for degree in range(360):
48
    # compute world2cam matrix
49
    camera = camera_util.Camera(Tx, Ty, Tz, degree, 0.0)
50
    Rt = camera_util.pose_from_camera(camera)[None]
51
    # convert to cam2world matrix
52
    # (used FOV=90, to reproduce previous pose distribution)
53
    xyz, viewdirs, zvals, rd, ro = get_sample_points(
54
        tform_cam2world=Rt.inverse(),
55
        F=(16, -16),
56
        H=1,
57
        W=32,
58
        samples_per_ray=2,
59
        near=0,
60
        far=nerf_far / 2,
61
        perturb=False,
62
        mask=None,
63
    )
64
    if np.all(np.abs(xyz).numpy() < plane_width / 2):
65
      valid_degrees.append(degree)
66
  # sample from the valid degrees
67
  degree = rng.choice(valid_degrees) + (rng.rand() - 0.5)
68
  camera = camera_util.Camera(Tx, Ty, Tz, degree, 0.0)
69
  Rt = camera_util.pose_from_camera(camera)[None]
70

71
  # store world2cam transformation
72
  sampled_Rts.append(Rt)
73
  sampled_cameras.append(camera)
74

75
f, ax = plt.subplots()
76
for i in np.random.choice(len(sampled_Rts), 500):
77
  Rt = sampled_Rts[i]
78
  xyz, viewdirs, zvals, rd, ro = get_sample_points(
79
      tform_cam2world=Rt.inverse(),
80
      F=(16, -16),
81
      H=1,
82
      W=1,
83
      samples_per_ray=64,
84
      near=0,
85
      far=8,
86
      perturb=False,
87
      mask=None,
88
  )
89
  ax.scatter(xyz[0, 0, 0, 0], xyz[0, 0, 0, 2])
90
  ax.arrow(
91
      xyz[0, 0, 0, 0],
92
      xyz[0, 0, 0, 2],
93
      xyz[0, 0, 20, 0] - xyz[0, 0, 0, 0],
94
      xyz[0, 0, 20, 2] - xyz[0, 0, 0, 2],
95
  )
96
ax.set_xlim([-plane_width / 2, plane_width / 2])
97
ax.set_ylim([-plane_width / 2, plane_width / 2])
98
ax.set_aspect('equal', adjustable='box')
99
f.savefig('preprocessing/poses.jpg')
100

101
# save with noisy camera heights
102
os.makedirs('poses', exist_ok=True)
103
torch.save(
104
    {'Rts': torch.stack(sampled_Rts), 'cameras': sampled_cameras},
105
    f'./poses/width{plane_width}_far{nerf_far}_noisy_height.pth',
106
)
107