-
Notifications
You must be signed in to change notification settings - Fork 493
Expand file tree
/
Copy pathexample_raycast.py
More file actions
99 lines (73 loc) · 3.07 KB
/
example_raycast.py
File metadata and controls
99 lines (73 loc) · 3.07 KB
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#############################################################################
# Example Ray Cast
#
# Shows how to use the built-in wp.Mesh data structure and wp.mesh_query_ray()
# function to implement a basic ray-tracer.
#
##############################################################################
import os
import numpy as np
from pxr import Usd, UsdGeom
import warp as wp
import warp.examples
@wp.kernel
def draw(mesh: wp.uint64, cam_pos: wp.vec3, width: int, height: int, pixels: wp.array[wp.vec3]):
tid = wp.tid()
x = tid % width
y = tid // width
sx = 2.0 * float(x) / float(height) - 1.0
sy = 2.0 * float(y) / float(height) - 1.0
# compute view ray
ro = cam_pos
rd = wp.normalize(wp.vec3(sx, sy, -1.0))
color = wp.vec3(0.0, 0.0, 0.0)
query = wp.mesh_query_ray(mesh, ro, rd, 1.0e6)
if query.result:
color = query.normal * 0.5 + wp.vec3(0.5, 0.5, 0.5)
pixels[tid] = color
class Example:
def __init__(self, height=1024, width=1024):
self.height = height
self.width = width
self.cam_pos = (0.0, 1.0, 2.0)
asset_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
mesh_geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/root/bunny"))
points = np.array(mesh_geom.GetPointsAttr().Get())
indices = np.array(mesh_geom.GetFaceVertexIndicesAttr().Get())
self.pixels = wp.zeros(self.width * self.height, dtype=wp.vec3)
# create wp mesh
self.mesh = wp.Mesh(
points=wp.array(points, dtype=wp.vec3), velocities=None, indices=wp.array(indices, dtype=int)
)
def render(self):
with wp.ScopedTimer("render"):
wp.launch(
kernel=draw,
dim=self.width * self.height,
inputs=[self.mesh.id, self.cam_pos, self.width, self.height, self.pixels],
)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
parser.add_argument("--width", type=int, default=1024, help="Output image width in pixels.")
parser.add_argument("--height", type=int, default=1024, help="Output image height in pixels.")
parser.add_argument(
"--headless",
action="store_true",
help="Run in headless mode, suppressing the opening of any graphical windows.",
)
args = parser.parse_known_args()[0]
with wp.ScopedDevice(args.device):
example = Example(height=args.height, width=args.width)
example.render()
if not args.headless:
import matplotlib.pyplot as plt
plt.imshow(
example.pixels.numpy().reshape((example.height, example.width, 3)),
origin="lower",
interpolation="antialiased",
)
plt.show()