目录
安装
pip install Panda3D
设置渲染宽高：
python
import panda3d.core as pd
margin = 100
screen = Tk().winfo_screenwidth() - margin, Tk().winfo_screenheight() - margin
width, height = (screen[0], int(screen[0] / 1.8)) if screen[0] / screen[1] < 1.8 else (int(screen[1] * 1.8), screen[1])
width,height=1200,800
pd.loadPrcFileData("", "win-origin {} {}".format((screen[0] + margin - width) // 2, (screen[1] + margin - height) // 2))
pd.loadPrcFileData("", "win-size {} {}".format(width, height))
pd.loadPrcFileData("", "win-fixed-size 1")渲染3d
UnderPressure/visualization.py at main · InterDigitalInc/UnderPressure · GitHub
python
"""
Copyright (c) 2022, InterDigital R&D France. All rights reserved. This source
code is made available under the license found in the LICENSE.txt at the root
directory of the repository.
"""
# UnderPressure
import anim, util
from data import FRAMERATE, TOPOLOGY, Forces
# Python
import math, time
from pathlib import Path
from tkinter import Tk
# PyTorch
import torch
# Panda3D
import panda3d.core as pd
from direct.showbase.ShowBase import ShowBase
from direct.task import Task
from direct.gui.DirectGui import DGG, DirectButton, DirectLabel, DirectSlider
def rotate(pts, axis, angle):
axis = torch.nn.functional.normalize(axis, p=2, dim=-1)
rmat = util.so3.to_SO3(angle * axis)
return torch.einsum("ij,...j->...i", rmat, pts)
class Geo:
@classmethod
def xy2xyz(cls, tensor, z=0.0, dim=-1): # [...] x 2 x [...]
assert tensor.shape[dim] == 2
x, y = tensor.unbind(dim) # [...] x [...]
return torch.stack([x, y, torch.full_like(x, z)], dim) # [...] x 3 x [...]
class GeomBuilder:
def __init__(self, name, color=None):
self.name = name
self.vdata = pd.GeomVertexData(name, pd.GeomVertexFormat.getV3n3c4(), pd.Geom.UHStatic)
self.vertexWriter = pd.GeomVertexWriter(self.vdata, "vertex")
self.normalWriter = pd.GeomVertexWriter(self.vdata, "normal")
self.colorWriter = pd.GeomVertexWriter(self.vdata, "color")
self.triangles = pd.GeomTriangles(pd.Geom.UHStatic)
self.currentRow = 0
self.color = color
def _vertex(self, point, normal, prim=None, color=None):
color = color or self.color or [1, 1, 1, 1]
self.vertexWriter.addData3(*point)
self.normalWriter.addData3(*normal)
self.colorWriter.addData4(color[0], color[1], color[2], color[3])
if prim is not None:
prim.addVertex(self.currentRow)
self.currentRow += 1
def triangle(self, p1, p2, p3, color=None, normal=None):
if normal is None:
normal = (p2 - p1).cross(p3 - p1).normalized()
self._vertex(p1, normal, self.triangles, color)
self._vertex(p2, normal, self.triangles, color)
self._vertex(p3, normal, self.triangles, color)
self.triangles.closePrimitive()
def polygon(self, pts, color=None):
triangulator = pd.Triangulator3()
for pt in pts:
triangulator.addPolygonVertex(triangulator.addVertex(*pt))
triangulator.triangulate()
for i in range(triangulator.getNumTriangles()):
self.triangle(
triangulator.getVertex(triangulator.getTriangleV0(i)),
triangulator.getVertex(triangulator.getTriangleV1(i)),
triangulator.getVertex(triangulator.getTriangleV2(i)),
color=color
)
def thick_polygon(self, pts, tickness, color=None, face_color=None, edge_color=None):
pts = torch.as_tensor(pts)
bottom, top = pts, pts + torch.as_tensor([0, 0, tickness])
self.polygon(top, face_color or color)
self.polygon(bottom, face_color or color)
for i in range(0, top.shape[0]):
j = (i + 1) % top.shape[0]
self.polygon([top[i], top[j], bottom[j], bottom[i]], edge_color or color)
def cylinder(self, p1, p2, radius, n=24, color=None):
p1, p2 = torch.as_tensor(p1), torch.as_tensor(p2)
axis = torch.nn.functional.normalize(p2 - p1, p=2, dim=-1) # 3
ex = torch.as_tensor([1.0, 0.0, 0.0])
if (axis != ex).any():
rvec = radius * torch.nn.functional.normalize(ex.cross(axis), p=2, dim=-1)
else:
rvec = torch.as_tensor([0.0, radius, 0.0])
for i in range(n):
self.triangle(
p1 + rotate(rvec, axis, i / n * math.tau),
p2 + rotate(rvec, axis, (i + 0.5) / n * math.tau),
p2 + rotate(rvec, axis, (i - 0.5) / n * math.tau),
color=color,
normal=rotate(rvec, axis, i / n * math.tau),
)
self.triangle(
p2 + rotate(rvec, axis, (i + 0.5) / n * math.tau),
p1 + rotate(rvec, axis, (i + 1) / n * math.tau),
p1 + rotate(rvec, axis, i / n * math.tau),
color=color,
normal=rotate(rvec, axis, (i + 0.5) / n * math.tau),
)
def cube(self, center, side, color=None):
return self.box(center, [side, side, side], color=color)
def box(self, center, sides, color=None):
center = torch.as_tensor(center)
sx, sy, sz = sides[0] / 2, sides[1] / 2, sides[2] / 2
vertices = torch.stack([
center + torch.as_tensor([-sx, -sy, -sz]),
center + torch.as_tensor([-sx, sy, -sz]),
center + torch.as_tensor([-sx, sy, sz]),
center + torch.as_tensor([-sx, -sy, sz]),
center + torch.as_tensor([sx, -sy, -sz]),
center + torch.as_tensor([sx, sy, -sz]),
center + torch.as_tensor([sx, sy, sz]),
center + torch.as_tensor([sx, -sy, sz]),
])
self.polygon(vertices[[0, 1, 2, 3]], color)
self.polygon(vertices[[0, 4, 5, 1]], color)
self.polygon(vertices[[1, 5, 6, 2]], color)
self.polygon(vertices[[2, 6, 7, 3]], color)
self.polygon(vertices[[3, 7, 4, 0]], color)
self.polygon(vertices[[7, 6, 5, 4]], color)
def build(self):
geom = pd.Geom(self.vdata)
geom.addPrimitive(self.triangles)
node = pd.GeomNode(self.name)
node.addGeom(geom)
return node
class AnimationLoop:
def __init__(self, length: int, fps: float, start: bool = True):
self._length, self._fps = int(length), float(fps)
self._time, self._frame = 0.0, 0.0
self._stop, self._running = False, bool(start)
self.functions = []
@property
def length(self) -> int:
return self._length
@length.setter
def length(self, length: int):
self._length = int(length)
def __len__(self) -> int:
return self.length
@property
def fps(self) -> float:
return self._fps
@fps.setter
def fps(self, value: float):
self._fps = float(value)
@property
def frame(self) -> float:
return self._frame % self.length
@property
def running(self) -> bool:
return self._running
def reset(self):
self._time = 0.0
self._frame = 0.0
def pause(self):
self._running = False
def unpause(self):
self._running = True
def stop(self):
self._stop = True
def bind(self, fn):
self.functions.append(fn)
def unbind(self, fn):
self.functions.remove(fn)
def __call__(self, state):
if self._stop:
return Task.done
if self._running:
self._frame += (state.time - self._time) * self.fps
for fn in self.functions:
fn(self.frame)
self._frame %= self.length
self._time = state.time
return Task.cont
class Animatable:
def __init__(self, loop: AnimationLoop):
self.loop = loop
@property
def loop(self):
return self._loop
@loop.setter
def loop(self, loop: AnimationLoop):
del self.loop
assert isinstance(loop, AnimationLoop)
self._loop = loop
self._loop.bind(self.update)
@loop.deleter
def loop(self):
if hasattr(self, "_loop"):
self._loop.unbind(self.update)
del self._loop
@property
def nframes(self):
return len(self.loop)
def update(self, frame: int):
prev_frame, next_frame = int(math.floor(frame)), min(int(math.ceil(frame)), self.nframes-1)
dframe = frame - prev_frame
self.__update__(frame, prev_frame, next_frame, dframe)
def __update__(self, frame: int, prev_frame: int, next_frame: int, dframe: float):
raise NotImplementedError()
class Ground(Animatable):
DEFAULT_COLOR1 = 3 * (0.4, ) + (1.0, )
DEFAULT_COLOR2 = 3 * (0.6, ) + (1.0, )
def __init__(self, parent, grid, size, origin=(0, 0), color1=DEFAULT_COLOR1, color2=DEFAULT_COLOR2, loop=None, trajectory=None):
grid, size = torch.as_tensor(grid).int(), torch.as_tensor(size)
origin = torch.as_tensor(origin)
cell = size / grid
builder = GeomBuilder("ground")
for x in range(grid[0]):
for y in range(grid[1]):
vertices = origin - 0.5*size + torch.stack([
torch.as_tensor([x, y]) * cell,
torch.as_tensor([x, y+1]) * cell,
torch.as_tensor([x+1, y+1]) * cell,
torch.as_tensor([x+1, y]) * cell,
])
color = color1 if x%2 == y%2 else color2
builder.polygon(Geo.xy2xyz(vertices), color)
self._node = parent.attachNewNode(builder.build())
if loop is not None:
super().__init__(loop)
if trajectory is not None:
self.trajectory = trajectory
@property
def trajectory(self):
return self._trajectory.clone()
@trajectory.setter
def trajectory(self, trajectory):
self._trajectory = trajectory.clone()
def __update__(self, frame: int, prev_frame: int, next_frame: int, dframe: float):
position = torch.lerp(*self.trajectory[[prev_frame, next_frame]], dframe)
self._node.setPos(*position.tolist())
class Insole(Animatable):
# geometry
INSOLE_VERTICES = torch.load(Path(__file__).parent / "geo_insoles.pth")
INSOLE_VERTICES_BBOX = torch.stack([INSOLE_VERTICES.amin(dim=0), INSOLE_VERTICES.amax(dim=0)])
INSOLE_VERTICES_LENGTH = (INSOLE_VERTICES_BBOX[1, 1] - INSOLE_VERTICES_BBOX[0, 1])
INSOLE_VERTICES_CENTER = 0.5 * (INSOLE_VERTICES_BBOX[0] + INSOLE_VERTICES_BBOX[1])
CELLS_VERTICES = torch.load(Path(__file__).parent / "geo_insole_cells.pth")
# colors
INSOLE_COLOR = (0.2, 0.2, 0.2, 1.0)
TOE_COLOR = (150/255, 48/255, 48/255, 1.0)
ARCH_COLOR = (0.5, 0.5, 0.5, 1.0)
HEEL_COLOR = (48/255, 84/255, 150/255, 1.0)
def __init__(self, parent, side, origin, length, angle, zscale=0.025, loop=None, forces=None):
self._values = torch.zeros(16)
self._side = side
self.__init_geometry__(parent, side, origin, length, angle)
self._zscale = float(zscale)
if loop is not None:
super().__init__(loop)
if forces is not None:
self.forces = forces
@classmethod
def left(cls, parent, *args, **kwargs):
return cls(parent, "left", *args, **kwargs)
@classmethod
def right(cls, parent, *args, **kwargs):
return cls(parent, "right", *args, **kwargs)
def is_left(self) -> bool:
return self._side == "left"
def is_right(self) -> bool:
return self._side == "right"
def __init_geometry__(self, parent, side, origin, length, angle):
self._origin, length = torch.as_tensor(origin), torch.as_tensor(length)
angle = (180 + torch.as_tensor(angle)) / 180 * math.pi
side = {"right": 1, "left": -1}[side]
# prepare global transform
transform_mat = torch.as_tensor([[angle.cos(), -angle.sin()], [angle.sin(), angle.cos()]]) # rotation
transform_mat = torch.matmul(transform_mat, (length / self.INSOLE_VERTICES_LENGTH).expand(2).diag()) # scale
transform_mat = torch.matmul(transform_mat, torch.as_tensor([side, 1.0]).diag()) # left
transform = lambda pts: Geo.xy2xyz(torch.einsum("ij,nj->ni", transform_mat, pts.clone() - self.INSOLE_VERTICES_CENTER))
self._base_tickness = 0.002
# insole
builders = {}
insole_vertices = transform(self.INSOLE_VERTICES)
if side == -1:
insole_vertices = insole_vertices.flip(0)
builders["insole"] = GeomBuilder("insole")
builders["insole"].thick_polygon(
self._origin + insole_vertices, tickness=self._base_tickness,
face_color=self.INSOLE_COLOR, edge_color=self.INSOLE_COLOR
)
# cells
for i, cell_vertices in enumerate(self.CELLS_VERTICES):
cell_vertices = transform(cell_vertices)
if side == -1:
cell_vertices = cell_vertices.flip(0)
cell_name = "cell{}".format(i)
cell_face_color = self.TOE_COLOR if i in Forces.FRONT_CELLS else self.HEEL_COLOR if i in Forces.BACK_CELLS else self.ARCH_COLOR
builders[cell_name] = GeomBuilder(cell_name)
builders[cell_name].thick_polygon(
cell_vertices, tickness=1.0,
face_color=cell_face_color, edge_color=cell_face_color
)
self._nodes = {key: parent.attachNewNode(builder.build()) for key, builder in builders.items()}
cells_origin = self._origin + torch.as_tensor([0.0, 0.0, self._base_tickness])
for i in range(len(self.CELLS_VERTICES)):
self._nodes["cell{}".format(i)].setPos(*cells_origin.tolist())
def __getitem__(self, index):
return self._values.__getitem__(index)
def __setitem__(self, index, value):
self._values.__setitem__(index, value)
for i in range(16):
z = self._zscale * self._values[i].item() + 1e-5
self._nodes["cell{}".format(i)].setSz(z)
@property
def forces(self):
return self._forces.clone()
@forces.setter
def forces(self, forces):
self._forces = forces.clone()
def __update__(self, frame: int, prev_frame: int, next_frame: int, dframe: float):
self[:] = torch.lerp(*self.forces[[prev_frame, next_frame]], dframe)
class InsolesPair:
def __init__(self, left, right, loop=None, forces=None):
self.left, self.right = left, right
if loop is not None:
self.left.loop = loop
self.right.loop = loop
if forces is not None:
self.forces = forces
@property
def left(self):
return self._left
@left.setter
def left(self, left):
assert isinstance(left, Insole) and left.is_left()
self._left = left
@property
def right(self):
return self._right
@right.setter
def right(self, right):
assert isinstance(right, Insole) and right.is_right()
self._right = right
def __setitem__(self, index, value):
values = self[:]
values.__setitem__(index, value)
self.left[:] = values[0]
self.right[:] = values[1]
def __getitem__(self, index):
return torch.stack([self.left[:], self.right[:]]).__getitem__(index)
@property
def forces(self):
return self.left.forces, self.right.forces
@forces.setter
def forces(self, forces):
self.left.forces = forces[0]
self.right.forces = forces[1]
def update(self, frame):
self.left.update(frame)
self.right.update(frame)
class Skeleton(Animatable):
def __init__(self, parent, skeleton, loop=None, angles=None, trajectory=None, label=""):
self._skeleton = skeleton
self._joints = []
for joint in TOPOLOGY:
node = GeomBuilder(joint)
node.cube((0, 0, 0), side=0.05, color=[0.2, 0.2, 0.2, 0.0])
self._joints.append(parent.attachNewNode(node.build()))
if loop is not None:
super().__init__(loop)
if angles is not None:
self.angles = angles
if trajectory is not None:
self.trajectory = trajectory
if label:
self._label = pd.TextNode(label)
self._label.setText(label)
self._label.setAlign(pd.TextNode.ACenter)
self._label = parent.attachNewNode(self._label)
self._label.setScale(0.1)
self._label.setPos(0.0, 0.0, 1.8)
@property
def angles(self):
return self._angles.clone()
@angles.setter
def angles(self, angles):
self._angles = angles.clone()
@property
def trajectory(self):
return self._trajectory.clone()
@trajectory.setter
def trajectory(self, trajectory):
self._trajectory = trajectory.clone()
self._label_traj = util.gma(self.trajectory, size=50, std=50/3, dim=0)
def __update__(self, frame: int, prev_frame: int, next_frame: int, dframe: float):
angles = util.SU2.slerp(*self.angles[[prev_frame, next_frame]], dframe)
positions = anim.FK(angles, self._skeleton, 0, TOPOLOGY)
if hasattr(self, "_trajectory"):
positions += torch.lerp(*self.trajectory[[prev_frame, next_frame]], dframe)
if hasattr(self, "_label") and hasattr(self, "_label_traj"):
self._label.setPos(*torch.lerp(*self._label_traj[[prev_frame, next_frame]], dframe)[:2].tolist(), 1.8)
for j, joint in enumerate(self._joints):
joint.setQuat(pd.LQuaternion(*angles[j].tolist()))
joint.setPos(*positions[j].tolist())
class MocapAndvGRFsApp(ShowBase):
def __init__(self, motions: list, vGRFs: list = [], motion_labels=None, vGRF_labels=None):
motion_labels = [None for _ in motions] if motion_labels is None else motion_labels
vGRF_labels = [None for _ in vGRFs] if vGRF_labels is None else vGRF_labels
margin = 100
screen = Tk().winfo_screenwidth() - margin, Tk().winfo_screenheight() - margin
width, height = (screen[0], int(screen[0] / 1.8)) if screen[0] / screen[1] < 1.8 else (int(screen[1] * 1.8), screen[1])
width,height=1200,800
pd.loadPrcFileData("", "win-origin {} {}".format((screen[0] + margin - width) // 2, (screen[1] + margin - height) // 2))
pd.loadPrcFileData("", "win-size {} {}".format(width, height))
pd.loadPrcFileData("", "win-fixed-size 1")
super().__init__()
self.render.setTwoSided(True)
self.disableMouse()
self.__init_lights__()
self.__init_camera__(fov=50, near=0.01, far=50, pos=[0.0, -7.0, 3.0], dir=[0.0, 1.2, -0.5])
# init animation loop
assert len({v.shape[0] for v in vGRFs} | {m[0].shape[0] for m in motions}) == 1
nframes = motions[0][0].shape[0]
self.animation_loop = AnimationLoop(length=nframes, fps=FRAMERATE, start=True)
self.taskMgr.add(self.animation_loop, "animation_loop")
# insoles
if len(vGRFs) > 0:
vGRFs_scale = 1 / torch.stack(list(vGRFs)).max()
for index, (value, label) in enumerate(zip(vGRFs, vGRF_labels)):
x = index * 0.111 - 0.111 * (len(vGRFs) - 1) / 2
self.add_insoles([
x, self.cam.getPos()[1] + 0.60 + 0.08, self.cam.getPos()[2] - 0.435],
vGRFs_scale * value, plate=True, label=label,
)
# get global trajectory
trajectories = torch.cat([motion[2] for motion in motions], dim=-2)
global_trajectory = trajectories.mean(dim=-2)
global_trajectory = util.gma(global_trajectory, size=30, std=30/3, dim=0)
global_trajectory[..., 2] = 0.0
c = 0.25
global_bbox = trajectories.amax(dim=(0, 1)) - trajectories.amin(dim=(0, 1))
global_position = global_trajectory.mean(dim=0)
ground_grid = ((global_bbox[:2] + 3) / c).ceil()
self._ground = Ground(
self.render, grid=ground_grid, size=c*ground_grid, origin=global_position[:2],
loop=self.animation_loop,
trajectory=-global_trajectory,
)
for motion, label in zip(motions, motion_labels):
angles, skeleton, trajectory = motion[0], motion[1].reshape(-1, 3), motion[2].reshape(-1, 3)
trajectory = trajectory - global_trajectory
self._skeleton = Skeleton(
self.render, skeleton,
loop=self.animation_loop, angles=angles, trajectory=trajectory,
label=label,
)
# Framerate buttons
button025 = DirectButton(text="0.25x", scale=0.05, pos=(-0.2, 0.0, 0.96), command=lambda: self.change_framerate(0.25))
button050 = DirectButton(text="0.50x", scale=0.05, pos=(0.0, 0.0, 0.96), command=lambda: self.change_framerate(0.50))
button100 = DirectButton(text="1.00x", scale=0.05, pos=(0.2, 0.0, 0.96), command=lambda: self.change_framerate(1.00))
def change_framerate(self, rate):
self.animation_loop.fps = rate * FRAMERATE
def __init_lights__(self):
def point_light(position, color, name=""):
light = pd.PointLight(name)
light.setColor(color)
light = self.render.attachNewNode(light)
light.setPos(*position)
self.render.setLight(light)
intensity = 0.6
point_light((5, 5, -5), (intensity, intensity, intensity, 1))
point_light((0, 5, -5), (intensity, intensity, intensity, 1))
point_light((5, 0, -5), (intensity, intensity, intensity, 1))
def __init_camera__(self, **kwargs):
if "near" in kwargs:
self.camLens.setNear(kwargs["near"])
if "far" in kwargs:
self.camLens.setFar(kwargs["far"])
if "fov" in kwargs:
self.camLens.setFov(kwargs["fov"])
if "pos" in kwargs:
self.cam.setPos(*kwargs["pos"])
if "target" in kwargs:
pos, target = torch.as_tensor(self.cam.getPos()), torch.as_tensor(kwargs["target"])
dir = (target - pos).tolist()
elif "dir" in kwargs:
dir = kwargs["dir"]
self.camLens.setViewVector(
*dir,
*kwargs.get("up", [0.0, 0.0, 1.0]),
)
def add_insoles(self, origin, forces, plate=False, label=None):
insoles = InsolesPair(
Insole.left(self.render, origin=[origin[0] - 0.021, origin[1], origin[2]], length=0.15, angle=10.0),
Insole.right(self.render, origin=[origin[0] + 0.021, origin[1], origin[2]], length=0.15, angle=-10.0),
loop=self.animation_loop, forces=forces.unbind(dim=-2),
)
width, depth, thickness = 0.108, 0.16, 0.01
if plate:
builder = GeomBuilder("insoles_plate")
builder.box([origin[0], origin[1], origin[2] - thickness/2], [width, depth, thickness], color=(0.8, 0.8, 0.8, 1.0))
self.render.attachNewNode(builder.build())
if isinstance(label, str):
text = pd.TextNode(label)
text.setText(label)
text.setAlign(pd.TextNode.ACenter)
text = self.render.attachNewNode(text)
text.setScale(thickness)
text.setPos(origin[0], origin[1] - depth/2 - 1e-4, origin[2] - 0.8*thickness)
if __name__ == "__main__":
import models
from data import Dataset
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument("-subj", "-subject", default="S9", help="Subject to be selected; default: S9")
parser.add_argument("-seq", "-sequence", default="WalkRandomSlow", help="Sequence to be selected; default: WalkRandomSlow")
parser.add_argument("-checkpoint", default="pretrained.tar")
args = parser.parse_args()
dataset = Dataset("{}-{}".format(args.subj, args.seq))
item = dataset[0]
angles, skeleton, trajectory = item["angles"], item["skeleton"], item["trajectory"]
vGRFs_gt = item["forces"]
positions = anim.FK(angles, skeleton, trajectory, TOPOLOGY)
checkpoint = torch.load(args.checkpoint)
model = models.DeepNetwork(state_dict=checkpoint["model"]).eval()
with torch.no_grad():
vGRFs_pred = model.vGRFs(positions.unsqueeze(0)).squeeze(0)
vGRFs_abs_error = (vGRFs_gt - vGRFs_pred).abs()
MocapAndvGRFsApp(
[(angles, skeleton, trajectory)],
[vGRFs_gt, vGRFs_pred, vGRFs_abs_error],
vGRF_labels=["Ground Truth", "Estimated", "Abs. Error"],
).run()