| 12 | @pytest.fixture |
| 13 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") |
| 14 | def test_data(): |
| 15 | N = 2 |
| 16 | xs = torch.linspace(-1, 1, N, device=device) |
| 17 | ys = torch.linspace(-1, 1, N, device=device) |
| 18 | xys = torch.stack(torch.meshgrid(xs, ys), dim=-1).reshape(-1, 2) |
| 19 | zs = torch.ones_like(xys[:, :1]) * 3 |
| 20 | means = torch.cat([xys, zs], dim=-1) |
| 21 | quats = torch.tensor([[1.0, 0.0, 0.0, 0]], device=device).repeat(len(means), 1) |
| 22 | scales = torch.ones_like(means) |
| 23 | scales[..., :2] *= 0.1 |
| 24 | opacities = torch.ones(1, len(means), device=device) * 0.5 |
| 25 | colors = torch.rand(1, len(means), 3, device=device) |
| 26 | viewmats = torch.eye(4, device=device).reshape(1, 4, 4) |
| 27 | # W, H = 24, 20 |
| 28 | W, H = 640, 480 |
| 29 | fx, fy, cx, cy = W, W, W // 2, H // 2 |
| 30 | Ks = torch.tensor( |
| 31 | [[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], device=device |
| 32 | ).reshape(1, 3, 3) |
| 33 | return { |
| 34 | "means": means, |
| 35 | "quats": quats, |
| 36 | "scales": scales, |
| 37 | "opacities": opacities, |
| 38 | "colors": colors, |
| 39 | "viewmats": viewmats, |
| 40 | "Ks": Ks, |
| 41 | "width": W, |
| 42 | "height": H, |
| 43 | } |
| 44 | |
| 45 | |
| 46 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") |