def testRotationBasic():
# general parameters
o = SED.opts(lie='se2')
m = getattr(lie, o.lie)
mRot = getattr(lie, o.lieRot)
# T, K = ( 10000, 4 )
T, K = ( 10, 4 )
T = 10
y, z, x, omega, theta, Q, S, E = generateSyntheticK4(o, T)
K = S.shape[0]
Strans = S[:,:o.dy,:o.dy]
Srot = S[:,o.dy:,o.dy:]
thetaR = theta.copy()
for t in range(1,T-1):
for k in range(K):
R_cur, d_cur = m.Rt(theta[t,k])
zeroAlgebra = np.zeros(o.dxA)
# test inference here
R_prev, d_prev = m.Rt(theta[t-1,k])
R_next, d_next = m.Rt(theta[t+1,k])
lhs = x[t] @ omega[k]
rhs = np.eye(o.dy+1)
y_tk = y[t][z[t]==k]
R_U = SED.sampleRotation(o, y_tk, d_cur, lhs, rhs, E[k], theta[t-1,k],
S[k], nextU=theta[t+1,k], controlled=True)
thetaR[t,k,:o.dy,:o.dy] = R_U
colorsA = du.diffcolors(K, bgCols=[[0,0,0],[1,1,1]], alpha=0.5)
colors = du.diffcolors(K, bgCols=[[0,0,0],[1,1,1]])
def show(t):
# plot x
T_world_object = x[t]
m.plot(T_world_object, np.tile([0,0,0], [2,1]), l=10.0)
# plot theta[t,k]
for k in range(K):
T_world_part = x[t] @ omega[k] @ theta[t,k]
m.plot(T_world_part, np.tile(colorsA[k], [2,1]), l=10.0)
plt.scatter(*y[t][k].T, color=colors[k], s=1, alpha=0.25)
# plot thetaR[t,k]
for k in range(K):
T_world_part = x[t] @ omega[k] @ thetaR[t,k]
m.plot(T_world_part, np.tile(colors[k], [2,1]), l=10.0)
plt.xlim(-100, 100)
plt.ylim(-100, 100)
du.ViewPlots(range(1,T-1), show)
plt.show()
def show(o, y, x, omega, theta, E, z, t):
o = SED.opts(lie='se2')
m = getattr(lie, o.lie)
mRot = getattr(lie, o.lieRot)
K = E.shape[0]
colors = du.diffcolors(K, bgCols=[[0,0,0],[1,1,1]])
# plot x
T_world_object = x[t]
m.plot(T_world_object, np.tile([0,0,0], [2,1]), l=10.0)
# plot x[t] @ omega[k] @ theta[t,k] in world coordinates
for k in range(K):
T_world_part = x[t] @ omega[k] @ theta[t,k]
m.plot(T_world_part, np.tile(colors[k], [2,1]), l=10.0)
yMu = SED.TransformPointsNonHomog(T_world_part, o.zeroObs)
R = T_world_part[:-1,:-1]
ySig = R.dot(E[k]).dot(R.T)
plt.plot( *du.stats.Gauss2DPoints(yMu, ySig, deviations=2.0), color=colors[k])
plt.scatter(*y[t].T, s=1, color='k')
plt.gca().set_aspect('equal', 'box')
plt.xlim(-50, 50)
plt.ylim(-50, 50)
plt.title(f'{t:04}')
def testControlledWalkCov():
o = SED.opts(lie='se2')
m = getattr(lie, o.lie)
mRot = getattr(lie, o.lieRot)
T = 10000
y, z, x, omega, theta, Q, S, E = generateSyntheticK4(o, T, Nk=2)
Strans = S[:,:o.dy,:o.dy]
Srot = S[:,o.dy:,o.dy:]
K = S.shape[0]
plt.figure(figsize=(12,8))
colors = du.diffcolors(K, bgCols=[[0,0,0],[1,1,1]])
for k in range(K):
plt.subplot(2,2,k+1)
d_k = theta[:,k,:-1,-1]
omega_trans = omega[k,:-1,-1]
plt.scatter(*d_k.T, color=colors[k], s=1)
plt.plot(*du.stats.Gauss2DPoints(o.zeroObs, Strans[k]), color=colors[k])
plt.title(f'Part {k+1}')
plt.gca().set_aspect('equal', 'box')
# plt.plot(*du.stats.Gauss2DPoints(omega_trans, Strans), color=colors[k])
path = 'reparam/controlled'
plt.savefig(f'{path}/cov.png', dpi=300, bbox_inches='tight')
plt.show()
def main(args):
if args.sampleIdx is not None:
samples = du.GetFilePaths(f'{args.resultPath}', 'gz')
o, alpha, z, pi, theta, E, S, x, Q, omega, mL, ll, subsetIdx, datasetPath = \
SED.loadSample(samples[int(args.sampleIdx)])
else:
o, alpha, z, pi, theta, E, S, x, Q, omega, mL, ll, subsetIdx, datasetPath = \
SED.loadSample(args.resultPath)
data = du.load(f'{datasetPath}/data')
yAll = data['y']
m = getattr(lie, o.lie)
T = len(z)
K = theta.shape[1]
if args.drawMesh:
meshFiles = du.GetFilePaths(f'{datasetPath}/mesh', 'obj')[:T]
mesh = du.Parfor(tm.load, meshFiles)
y = [ mesh[t].vertices for t in range(T) ]
mL_const = np.mean([np.mean(mL[t]) for t in range(T)])
mL = [ mL_const*np.ones(y[t].shape[0]) for t in range(T) ]
for t in range(T):
z[t] = SED.inferZ(o, y[t], pi, theta[t], E, x[t], omega, mL[t], max=args.maxZ)
elif args.no_resampleZ:
if subsetIdx is not None:
y = [yt[subsetIdx[t]] for t, yt in enumerate(yAll)]
else:
y = yAll
else:
# don't use subsetIdx
y = yAll
mL_const = np.mean([np.mean(mL[t]) for t in range(T)])
mL = [ mL_const*np.ones(y[t].shape[0]) for t in range(T) ]
if args.maxZ: max=True
else: max=False
for t in range(T):
z[t] = SED.inferZ(o, y[t], pi, theta[t], E, x[t], omega, mL[t], max=args.maxZ)
# omegaTheta
if args.omega:
theta = np.tile(omega, (T,1,1,1))
else:
for t in range(T):
for k in range(K):
theta[t,k] = omega[k] @ theta[t,k]
if args.noE: E = None
if args.noTheta: theta = None
if args.noZ: z = None
noX = args.noX
if args.noTitle: title = [ f'' for t in range(T) ]
else: title = [ f'{t:05}' for t in range(T) ]
if args.decimate > 0:
nPts = args.decimate
print('decimate')
for t in range(T):
decimateIdx = np.random.choice(range(len(y[t])), nPts)
y[t] = y[t][decimateIdx]
if z is not None: z[t] = z[t][decimateIdx]
if args.wiggle:
from scipy.stats import multivariate_normal as mvn
for t in range(T):
y[t] += mvn.rvs(np.zeros(3), args.wiggle_eps*np.eye(3), size=y[t].shape[0])
# if theta is not None: theta[0,0][1,3] += 0.2
if args.save:
try: os.makedirs(args.save)
except: pass
fnames = [ f'{args.save}/img-{t:05}.png' for t in range(T) ]
else:
fnames = None
def getImgs(path):
imgPaths = du.GetImgPaths(imgPath)
if len(imgPaths) > 100:
du.imread(imgPaths[0]); imgs = du.ParforT(du.imread, imgPaths)
else:
imgs = du.For(du.imread, imgPaths)
return imgs
# three cases
# se2, se3 + draw2d, se3
if o.lie == 'se2':
imgPath = f'{datasetPath}/imgs'
if isdir(imgPath): imgs = getImgs(imgPath)
else: imgs = None
scenes_or_none = drawSED.draw(o, y=y, x=x, theta=theta, E=E, img=imgs, z=z,
filename=fnames, noX=noX, title=title)
if not args.save: plt.show()
elif o.lie == 'se3' and not args.draw2d and not args.drawMesh:
scenes = drawSED.draw(o, y=y, x=x, theta=theta, E=E, z=z, noX=noX, title=title)
transform = tmu.CameraFromScenes(scenes)
# set transform as 1.25 of min z. This is specific to se3_marmoset for now
# multiply instead of divide because maybe camera is looking backwards?
for scene in scenes:
transform_t = scene.camera.transform.copy()
transform_t[2,3] = transform[2,3] * 1.25
scene.camera.transform = transform_t
if args.save:
for t in range(T): tmu.save_render(scenes[t], fnames[t])
else:
if args.single_frame:
t = 0
tmu.show_scene_with_bindings(scenes[t], res=[1920,1080])
else:
for t in range(T): tmu.show_scene_with_bindings(scenes[t], res=[1920,1080])
elif o.lie == 'se3' and args.drawMesh:
meshFiles = du.GetFilePaths(f'{datasetPath}/mesh', 'obj')[:T]
mesh = du.Parfor(tm.load, meshFiles)
mL_const = np.mean([np.mean(mL[t]) for t in range(T)])
zCols = (255*du.diffcolors(100, bgCols=[[1,1,1],[0,0,0]], alpha=1.0)).astype(np.uint8)
zColsFloat = du.diffcolors(100, bgCols=[[1,1,1],[0,0,0]], alpha=1.0)
for t in range(T):
mesh[t].visual.vertex_colors = zCols[z[t]]
# draw but don't render scenes just to get camera
scenes = drawSED.draw(o, y=y)
transform = tmu.CameraFromScenes(scenes)
# same colors as drawSED
for t in range(T):
scene = tm.scene.Scene()
scene.add_geometry(mesh[t])
# scene = mesh[t]
transform_t = scene.camera.transform.copy()
transform_t[2,3] = transform[2,3] * 1.5
transform_t[2,2] = transform[2,2]
scene.camera.transform = transform_t
if not args.noX:
scene.add_geometry(tmu.MakeAxes(0.2, x[t], np.tile([0, 0, 0, 255],
[4,1]).astype(np.int), minor=0.01))
if not args.orbit:
if args.save:
tmu.save_render(scene, fnames[t], res=[1920,1080])
else:
scene.show()
else:
nCams = 8
cams = tmu.MakeCamerasOrbit(scene, nCams)
for i in range(nCams):
print('Time {t}, Cam {i}')
cam = cams[i] @ transform_t
scene.camera.transform = cam
fname = f'{args.save}/camera-{i:02}-img-{t:05}.png'
tmu.save_render(scene, fname, res=[1920,1080])
# re-associate to mesh vertices
elif o.lie == 'se3' and args.draw2d:
imgPath = f'{datasetPath}/rgb'
assert isdir(imgPath)
imgs = getImgs(imgPath)
yImg = [ ]
for t in range(T):
mask = data['mask'][t]
# get ordered image indices
h, w = imgs[0].shape[:2]
yy, xx = np.meshgrid(range(h), range(w), indexing='ij')
xy = np.stack((xx.flatten(),yy.flatten()), axis=1) # N x 2
xyFG = xy[mask.flatten()]
idx_t = data['idx'][t] # indexes into xy
if subsetIdx is not None and args.no_resampleZ:
subsetIdx_t = subsetIdx[t]
idx_t = idx_t[subsetIdx_t]
# ip.embed()
xyFG = xyFG[ idx_t ]
yImg.append(xyFG)
# make fake options with se2 for display
oImg = SED.opts(lie='se2')
drawSED.draw(oImg, y=yImg, z=z, img=imgs, filename=fnames, noX=noX, title=title)
if not args.save: plt.show()
def testTranslationBasic():
# general parameters
o = SED.opts(lie='se2')
m = getattr(lie, o.lie)
mRot = getattr(lie, o.lieRot)
T = 10
y, z, x, omega, theta, Q, S, E = generateSyntheticK4(o, T)
K = S.shape[0]
Strans = S[:,:o.dy,:o.dy]
Srot = S[:,o.dy:,o.dy:]
# t, k = (5, 2)
for t in range(1,T-1):
for k in range(K):
R_cur, d_cur = m.Rt(theta[t,k])
zeroAlgebra = np.zeros(o.dxA)
# test inference here
R_prev, d_prev = m.Rt(theta[t-1,k])
R_next, d_next = m.Rt(theta[t+1,k])
def thetat_fwdBack(v):
ll = 0.0
# v is a translation; we'll already have a rotation so make the full
V = SED.MakeRd(R_cur, v)
# past
phi = mRot.algi(mRot.logm(R_prev.T @ R_cur))
phi = np.atleast_1d(phi)
m = o.Bi @ (v - o.A @ d_prev)
val = np.concatenate((m, phi))
ll = mvn.logpdf(val, zeroAlgebra, S[k])
# future
phi_ = mRot.algi(mRot.logm(R_cur.T @ R_next))
phi_ = np.atleast_1d(phi_)
m_ = o.Bi @ (d_next - o.A @ v)
val_ = np.concatenate((m_, phi_))
ll += mvn.logpdf(val_, zeroAlgebra, S[k])
return -ll
mu, Sigma = SED.thetaTranslationDynamicsPosterior(o, R_cur, theta[t-1,k],
S[k], nextU=theta[t+1,k])
mu_noFuture, Sigma_noFuture = SED.thetaTranslationDynamicsPosterior(o, R_cur, theta[t-1,k],
S[k])
# print(f'd_cur:\n {d_cur}\n mu:\n {mu}\n Sigma:\n {Sigma}\n S:\n{S[k]}')
g = nd.Gradient(thetat_fwdBack)
map_est = so.minimize(thetat_fwdBack, np.zeros(o.dy), method='BFGS',
jac=g)
norm = np.linalg.norm(map_est.x - mu)
print(f'norm: {norm:.6f}')
assert norm < 1e-2, \
f'SED.test5 bad, norm {norm:.6f}'
plt.scatter(*d_cur, color='g')
plt.plot(*du.stats.Gauss2DPoints(mu, Sigma), color='b')
plt.plot(*du.stats.Gauss2DPoints(mu_noFuture, Sigma_noFuture), color='r')
plt.scatter(*mu, color='b')
plt.scatter(*mu_noFuture, color='r')
plt.show()
sys.exit()
colors = du.diffcolors(K, bgCols=[[0,0,0],[1,1,1]])
def show(t):
# plot x
T_world_object = x[t]
m.plot(T_world_object, np.tile([0,0,0], [2,1]), l=10.0)
# plot theta[t,k]
for k in range(K):
T_world_part = x[t] @ omega[k] @ theta[t,k]
m.plot(T_world_part, np.tile(colors[k], [2,1]), l=10.0)
plt.xlim(-100, 100)
plt.ylim(-100, 100)
path = 'reparam/controlled'
y[t][idx] = igp.TransformPointsNonHomog(T_world_part, ytk_part)
z[t][idx] = k
# randomly permute y_t, z_t
perm = np.random.permutation(range(Nt[t]))
y[t] = y[t][perm]
z[t] = z[t][perm]
# pi
pi = np.ones(K) / K
mins = np.array([np.min(y[t], axis=0) for t in range(T)])
maxs = np.array([np.max(y[t], axis=0) for t in range(T)])
xlim = np.array([np.min(mins[:, 0]), np.max(maxs[:, 0])])
ylim = np.array([np.min(mins[:, 1]), np.max(maxs[:, 1])])
cols = du.diffcolors(2, bgCols=[[1, 1, 1], [0, 0, 0]])
# # uncomment if we want to save new dataset
# du.save('%s/data' % path, {
# 'y': y, 'z': z,
# 'H_Q': H_Q, 'H_x': H_x,
# 'H_S': H_S, 'H_theta': H_theta,
# 'H_E': H_E,
# 'Q': Q, 'S': S, 'E': E,
# 'x': x, 'theta': theta, 'z': z,
# 'pi': pi
# })
oFake = igp.opts(lie='se2')
# filename = [ f'{path}/images/img-{t:05}.jpg' for t in range(T) ]
# title = [ f'Frame {t:05}' for t in range(T) ]
def draw_t_SE2(o, **kwargs):
""" Draw or save model drawing for given time t.
INPUT
o (Namespace): options
KEYWORD INPUT
y (ndarray, [N, m.n]): observations
x (ndarray, dxGm): latent global state
z (ndarray, [N_t,]): Associations
zCols (ndarray, [K,3/4]): Association Colors
theta (ndarray, [K_est, dxGm]): latent part state
E (ndarray, [K_est, dy, dy]): observation noise extent
isObserved (ndarray, [K,]): boolean array of whether parts are observed
xlim: min, max x plot limits
ylim: min, max y plot limits
title: plot title name
filename: save path if desired, else returns figure reference
style: string for modifying plot style. Currently ignored.
reverseY: boolean, default False
noX: boolean, default False
img (ndarray, [nY, nX, 3]): image to draw on
linestyle (str): linestyle for x, theta arrows
deviations (float): deviations to draw E at
"""
assert o.lie == 'se2'
m = getattr(lie, o.lie)
zCols = kwargs.get('zCols', None)
if zCols is None:
zCols = du.diffcolors(100, bgCols=[[1, 1, 1], [0, 0, 0]])
reverseY = kwargs.get('reverseY', False)
noX = kwargs.get('noX', False)
linestyle = kwargs.get('linestyle', '-')
deviations = kwargs.get('deviations', 1.25)
img = kwargs.get('img', None)
if img is None:
# fillStyle = kwargs.get('fillStyle', 'full')
import matplotlib
marker = matplotlib.markers.MarkerStyle('o', 'full')
y = kwargs.get('y', None)
if y is not None:
z = kwargs.get('z', None)
if z is None:
plt.scatter(y[:, 0], y[:, 1], color='k', s=0.1, zorder=0.1)
else:
assoc = z >= 0
plt.scatter(y[assoc, 0],
y[assoc, 1],
color=zCols[z[assoc]],
s=0.1,
zorder=0.1,
marker=marker)
noAssoc = np.logical_not(assoc)
plt.scatter(y[noAssoc, 0],
y[noAssoc, 1],
color='gray',
alpha=0.5,
s=0.1,
zorder=0.1,
marker=marker)
else:
# assume y are foreground indices
y = kwargs.get('y', None)
if y is not None:
z = kwargs.get('z', None)
if z is not None:
yInt = y.astype(np.int)
if zCols.shape[1] == 3:
zColsA = np.concatenate((zCols, 0.5 * np.ones(
(zCols.shape[0], 1))),
axis=1)
else:
zColsA = zCols
for k in np.unique(z[z >= 0]):
yk = yInt[z == k]
img = du.DrawOnImage(img, (yk[:, 1], yk[:, 0]), zColsA[k])
gray = np.array([128, 128, 128, 0.5])
yNoAssoc = yInt[z == -1]
img = du.DrawOnImage(img, (yNoAssoc[:, 1], yNoAssoc[:, 0]),
gray)
plt.imshow(img)
# both x, xTrue are black, need to handle linestyles
x = kwargs.get('x', None)
if x is not None and not noX:
m.plot(x, np.tile([0, 0, 0], [2, 1]), l=10.0, linestyle=linestyle)
theta = kwargs.get('theta', None)
if theta is not None and x is not None:
K = theta.shape[0]
isObserved = kwargs.get('isObserved', np.ones(K, dtype=np.bool))
for k in range(K):
if not isObserved[k]: continue
T_world_part = x.dot(theta[k])
m.plot(T_world_part,
np.tile(zCols[k], [2, 1]),
l=10.0,
linestyle=linestyle)
zero = np.zeros(2)
E = kwargs.get('E', None)
if E is not None and theta is not None and x is not None:
K = theta.shape[0]
for k in range(K):
if not isObserved[k]: continue
T_world_part = x.dot(theta[k])
yMu = SED.TransformPointsNonHomog(T_world_part, zero)
R = T_world_part[:-1, :-1]
ySig = R.dot(E[k]).dot(R.T)
plt.plot(*du.stats.Gauss2DPoints(yMu, ySig, deviations=deviations),
c=zCols[k],
linestyle=linestyle)
ax = plt.gca()
xlim = kwargs.get('xlim', None)
if xlim is not None: ax.set_xlim(xlim)
elif img is not None:
xlim = (0, img.shape[1])
ax.set_xlim(xlim)
ylim = kwargs.get('ylim', None)
if ylim is not None: ax.set_ylim(ylim)
elif img is not None:
ylim = (img.shape[0], 0)
ax.set_ylim(ylim)
plt.gca().set_aspect('equal', 'box')
plt.gca().set_xticks([])
plt.gca().set_yticks([])
plt.grid(color=[.5, .5, .5], alpha=0.25)
title = kwargs.get('title', None)
if title is not None: plt.title(title)
if reverseY: plt.gca().invert_yaxis()
filename = kwargs.get('filename', None)
if filename is not None:
plt.savefig(filename, dpi=300, bbox_inches='tight')
plt.close()
else:
return plt.gcf()
def draw_t_SE3(o, **kwargs):
""" Draw or save model drawing for given time t.
INPUT
o (Namespace): options
KEYWORD INPUT
y (ndarray, [N, m.n]): observations
x (ndarray, [dxGf,]): latent global state
z (ndarray, [N_t,]): Associations
zCols (ndarray, [K,3/4]): Association Colors
xTrue (ndarray, [dx,]): true latent global state
theta (ndarray, [K_est, dt]): latent part state
thetaTrue (ndarray, [K, dt]): latent part state
E (ndarray, [K_est, dy]): observation noise extent
ETrue (ndarray, [K, dy]): true observation noise extent
xlim: min, max x plot limits
ylim: min, max y plot limits
title: plot title name
filename: save path if desired, else returns figure reference
style: string for modifying plot style. Currently ignored.
noX: bool, default False
"""
assert o.lie == 'se3'
m = getattr(lie, o.lie)
zCols = kwargs.get('zCols', None)
if zCols is None:
zCols = du.diffcolors(100, bgCols=[[1, 1, 1], [0, 0, 0]])
# zCols[1] = zCols[7]
# zCols[4] = zCols[11]
# zColsInt = (zCols*255).astype(np.int)
scene = tm.scene.Scene()
y = kwargs.get('y', None)
if y is not None:
z = kwargs.get('z', None)
if z is None:
pc = tmu.ConstructPointCloud(y)
scene.add_geometry(pc)
else:
instIdx = z >= 0
if np.sum(instIdx) > 0:
zPart = z[instIdx]
pc1 = tmu.ConstructPointCloud(y[instIdx], zCols[zPart])
scene.add_geometry(pc1)
noInstIdx = z == -1
nNonInstantiated = np.sum(noInstIdx)
if nNonInstantiated > 0:
grayCols = np.tile(0.5 * np.ones(3), [nNonInstantiated, 1])
if nNonInstantiated == 1 and len(y[noInstIdx].shape) == 1:
yNon = y[noInstIdx][np.newaxis]
else:
yNon = y[noInstIdx]
pc2 = tmu.ConstructPointCloud(yNon, grayCols)
scene.add_geometry(pc2)
# pc = tmu.ConstructPointCloud(y, zCols[z])
# scene.add_geometry(pc)
# todo: add part majorAxix (60 / 20 / 2 / 0.2) as a parameter
# else: marmoset needs ~60 and articulated meshes need ~0.2
x = kwargs.get('x', None)
noX = kwargs.get('noX', None)
if x is not None and not noX:
scene.add_geometry(
tmu.MakeAxes(0.1,
du.asShape(x, o.dxGm),
np.tile([0, 0, 0, 255], [4, 1]).astype(np.int),
minor=0.01))
# scene.add_geometry(tmu.MakeAxes(0.2, du.asShape(x, o.dxGm),
# np.tile([0, 0, 0, 255], [4,1]).astype(np.int), minor=0.01))
# scene.add_geometry(tmu.MakeAxes(2.0, du.asShape(x, o.dxGm),
# np.tile([0, 0, 0, 255], [4,1]).astype(np.int), minor=0.01))
# scene.add_geometry(tmu.MakeAxes(60.0, du.asShape(x, o.dxGm),
# np.tile([0, 0, 0, 255], [4,1]).astype(np.int), minor=0.01))
theta = kwargs.get('theta', None)
if theta is not None and x is not None:
K = theta.shape[0]
for k in range(K):
c = zCols[k]
if len(c) == 3: c = np.concatenate((c, np.ones(1)))
c = np.tile((c * 255).astype(np.uint8), [4, 1])
T_world_part = du.asShape(x,
o.dxGm).dot(du.asShape(theta[k], o.dxGm))
scene.add_geometry(tmu.MakeAxes(0.1, T_world_part, c, minor=0.01))
# scene.add_geometry(tmu.MakeAxes(0.2, T_world_part, c, minor=0.01))
# scene.add_geometry(tmu.MakeAxes(20.0, T_world_part, c, minor=0.01))
# scene.add_geometry(tmu.MakeAxes(60.0, T_world_part, c, minor=0.01))
# scene.add_geometry(tmu.MakeAxes(2.0, T_world_part, c, minor=0.01))
E = kwargs.get('E', None)
if E is not None and theta is not None and x is not None:
K = theta.shape[0]
for k in range(K):
c = zCols[k]
if len(c) == 3: c = np.concatenate((c, 0.25 * np.ones(1)))
c = (c * 255).astype(np.uint8)
T_world_part = du.asShape(x,
o.dxGm).dot(du.asShape(theta[k], o.dxGm))
# TODO: eigvals not sorted, this is probably bad idea
# ell = tmu.MakeEllipsoid(T_world_part,
# np.sqrt(np.linalg.eigvals(E[k]))*2, c)
ell = tmu.MakeEllipsoid(T_world_part,
np.sqrt(np.linalg.eigvals(E[k])) * 3, c)
# setattr(ell, 'wire', True)
scene.add_geometry(ell)
filename = kwargs.get('filename', None)
if filename is not None: tmu.save_render(scene, filename)
else: return scene