def make_rigid_body(sample,
init_pos,
orientation,
linear_velocity,
angular_velocity,
device,
target_vertices=None):
# Load a triangle mesh obj file
if target_vertices is not None:
tmp_obj_file = '/tmp/obj.obj'
from blender_process import Process
p = Process(sample.obj, target_vertices, tmp_obj_file)
mesh = TriangleMesh.from_obj(tmp_obj_file)
else:
mesh = TriangleMesh.from_obj(sample.obj)
vertices = (meshutils.normalize_vertices(
mesh.vertices).unsqueeze(0).to(device))
faces = mesh.faces.unsqueeze(0).to(device)
textures = torch.cat(
(
sample.color[0] / 255 * torch.ones(
1, faces.shape[1], 2, 1, dtype=torch.float32, device=device),
sample.color[1] / 255 * torch.ones(
1, faces.shape[1], 2, 1, dtype=torch.float32, device=device),
sample.color[2] / 255 * torch.ones(
1, faces.shape[1], 2, 1, dtype=torch.float32, device=device),
),
dim=-1,
)
# (Uniform) Masses
masses = (sample.mass / vertices.shape[-2]) * torch.nn.Parameter(
torch.ones(vertices.shape[-2], dtype=vertices.dtype, device=device),
requires_grad=True,
)
# Body
body = RigidBody(
vertices[0],
position=torch.tensor(init_pos, dtype=torch.float32, device=device),
masses=masses,
orientation=torch.tensor(orientation).type(torch.float32).to(device),
friction_coefficient=sample.fric,
# linear_velocity=torch.tensor(linear_velocity).type(torch.float32).to(device)
# angular_veloctiy=torch.tensor(angular_velocity).type(torch.float32).to(device)
)
return body, vertices, faces, textures
if args.compare_every >= args.simsteps:
raise ValueError(
f"Arg --compare-every cannot be greater than or equal to {args.simsteps}."
)
# Seed RNG for repeatability
torch.manual_seed(args.seed)
# Device the tensors are all stored on.
device = "cuda:0"
# Load the template sphere mesh (for the bob)
# sphere = TriangleMesh.from_obj(args.template)
sphere = kaolin.io.obj.import_mesh(args.template)
vertices_gt = meshutils.normalize_vertices(
sphere.vertices.unsqueeze(0), scale_factor=args.radius).to(device)
faces_gt = sphere.faces.to(device).unsqueeze(0)
textures = torch.cat(
(
torch.ones(
1, faces_gt.shape[1], 2, 1, dtype=torch.float32,
device=device),
torch.zeros(
1, faces_gt.shape[1], 2, 1, dtype=torch.float32,
device=device),
torch.zeros(
1, faces_gt.shape[1], 2, 1, dtype=torch.float32,
device=device),
),
dim=-1,
)
# Initialize the renderer.
renderer = SoftRenderer(image_size=image_size,
camera_mode=camera_mode,
device=device)
renderer.set_eye_from_angles(camera_distance, elevation, azimuth)
sim_duration = 2.0 # seconds
fps = 30 # frames per second
sim_steps = int((sim_duration * fps) / 2)
# obj = get_primitive_obj(shape[0])
obj = get_primitive_obj(INT_TO_PRIMITIVE[shape[0]])
# print(obj)
mesh = TriangleMesh.from_obj(obj)
vertices = ((
meshutils.normalize_vertices(mesh.vertices) # + \
# torch.from_numpy(init_pos[i]).float().unsqueeze(0)
).to(device).unsqueeze(0))
faces = mesh.faces.unsqueeze(0).to(device)
textures = torch.cat(
(
color[0][0] / 255.0 * torch.ones(1,
faces.shape[1],
2,
1,
dtype=torch.float32,
device=device),
color[0][1] / 255.0 * torch.ones(1,
faces.shape[1],
2,
1,
device = "cuda:0"
logdir = os.path.join(args.logdir, args.expid)
if args.log:
os.makedirs(logdir, exist_ok=True)
with open(os.path.join(logdir, "args.txt"), "w") as f:
json.dump(args.__dict__, f, indent=2)
sim_dt = (1.0 / args.physics_engine_rate) / args.sim_substeps
sim_steps = int(args.sim_duration / sim_dt)
sim_time = 0.0
builder_gt = df.sim.ModelBuilder()
obj = TriangleMesh.from_obj(args.mesh)
vertices = meshutils.normalize_vertices(obj.vertices).to(device)
points = vertices.detach().cpu().numpy()
indices = list(obj.faces.numpy().reshape((-1)))
mesh = df.sim.Mesh(points, indices)
pos_gt = (0.0, 4.0, 0.0)
rot_gt = df.quat_from_axis_angle((0.0, 0.0, 1.0), math.pi * 0.3)
# rot_gt = (0.0, 0.0, 0.0, 0.0)
vel_gt = (0.0, 2.0, 0.0)
omega_gt = (0.0, 0.0, 0.0)
scale_gt = (1.0, 1.0, 1.0)
density_gt = 5.0
ke_gt = 4900.0
kd_gt = 15.0
kf_gt = 990.0
from gradsim.simulator import Simulator
from gradsim.utils import meshutils
if __name__ == "__main__":
# Device to store tensors on (MUST be CUDA-capable, for renderer to work).
device = "cuda:0"
# Output (gif) file path
outfile = Path("cache/demorestitution.gif")
# Load a body (from a triangle mesh obj file).
mesh = TriangleMesh.from_obj(Path("sampledata/cube.obj"))
# mesh = TriangleMesh.from_obj(Path("sampledata/sphere.obj"))
# mesh = TriangleMesh.from_obj(Path("sampledata/banana.obj"))
vertices = meshutils.normalize_vertices(mesh.vertices.unsqueeze(0)).to(device)
faces = mesh.faces.to(device).unsqueeze(0)
textures = torch.cat(
(
torch.ones(1, faces.shape[1], 2, 1, dtype=torch.float32, device=device),
torch.ones(1, faces.shape[1], 2, 1, dtype=torch.float32, device=device),
torch.zeros(1, faces.shape[1], 2, 1, dtype=torch.float32, device=device),
),
dim=-1,
)
masses = torch.nn.Parameter(
0.1 * torch.ones(vertices.shape[1], dtype=vertices.dtype, device=device),
requires_grad=True,
)
position = torch.tensor([0.0, 4.0, 0.0], dtype=torch.float32, device=device)
orientation = torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=torch.float32, device=device)
