def main(args):
files = Path(args.filespath).rglob('*.gz')
for f in files:
sample = du.load(str(f))
o = sample['o']
m = getattr(lie, o.lie)
K = len(sample['pi']) - 1
T = sample['theta'].shape[0]
if K == 0:
omega = np.zeros((
T,
0,
) + o.dxGm)
sample['omega'] = omega
du.save(str(f), sample)
continue
theta = sample['theta']
omega = np.stack([m.karcher(theta[:, k]) for k in range(K)])
omegaInv = np.stack([m.inv(omega[k]) for k in range(K)])
for t in range(T):
for k in range(K):
theta[t, k] = omegaInv[k] @ theta[t, k]
# re-estimate S
if K > 0: S = np.array([SED.inferSk(o, theta[:, k]) for k in range(K)])
else: S = np.zeros((0, o.dxA, o.dxA))
# re-compute LL
d = sample
o, alpha, z, pi, E, x, Q, mL, llOld = \
(d['o'], d['alpha'], d['z'], d['pi'], d['E'], d['x'],
d['Q'], d['mL'], d['ll'])
subsetIdx, dataset = (d.get('subsetIdx', None), d.get('dataset', None))
data = du.load(f'{dataset}/data')
yAll = data['y']
if subsetIdx is not None:
y = [yt[subsetIdx[t]] for t, yt in enumerate(yAll)]
else:
y = yAll
ll = SED.logJoint(o, y, z, x, theta, E, S, Q, alpha, pi, omega, mL)
# print(llOld, ll)
path, base, _ = du.fileparts(str(f))
SED.saveSample(f'{path}/{base}', o, alpha, z, pi, theta, E, S, x, Q,
omega, mL, ll, subsetIdx, dataset)
def test_optimize_omega():
sample = 'omega/se2_waving_hand/001/init.gz'
o, alpha, z, pi, theta, E, S, x, Q, omega, mL, ll, subsetIdx, datasetPath = \
SED.loadSample(sample)
data = du.load(f'{datasetPath}/data')
yAll = data['y']
T = len(z)
K = theta.shape[1]
if subsetIdx is not None:
y = [yt[subsetIdx[t]] for t, yt in enumerate(yAll)]
# omegaNew = np.zeros_like(omega)
# # for k in range(K):
# for k in range(0,1):
# yk = [ y[t][z[t]==k] for t in range(T) ]
# theta_k = theta[:,k]
# omegaNew[k] = icp.optimize_omega(o, yk, x, theta_k, E[k])
t = 2
x_t, theta_t = icp.optimize_t(o, y[t], x[t-1], omega, theta[t-1], E, S, Q)
omegaTheta_t = theta_t.copy()
for k in range(K): omegaTheta_t[k] = omega[k] @ omegaTheta_t[k]
draw.draw_t(o, y=y[t], x=x_t, theta=omegaTheta_t, E=E)
def test_optimize_all():
sample = 'omega/se2_waving_hand/001/init.gz'
o, alpha, z, pi, theta, E, S, x, Q, omega, mL, ll, subsetIdx, datasetPath = \
SED.loadSample(sample)
data = du.load(f'{datasetPath}/data')
yAll = data['y']
T = len(z)
K = theta.shape[1]
if subsetIdx is not None:
y = [yt[subsetIdx[t]] for t, yt in enumerate(yAll)]
basePath = 'optimize/se2_waving_hand/001'
def callback(s, v, cost, x, omega, theta):
z = [ SED.inferZ(o, y[t], pi, theta[t], E, x[t], omega, mL[t])
for t in range(T) ]
z, pi_, theta_, omega_, E_, S_ = SED.consolidatePartsAndResamplePi(o, z, pi,
alpha, theta, omega, E, S)
ll = SED.logJoint(o, y, z, x, theta_, E_, S_, Q, alpha, pi_, omega_, mL)
extra = dict(s=s, v=v, cost=cost)
filename = f'{basePath}/optimize-{s:05}'
SED.saveSample(filename, o, alpha, z, pi_, theta_, E_, S_, x, Q, omega_, mL,
ll, subsetIdx, datasetPath, extra=extra)
s, v, cost, x, omega, theta = icp.optimize_all(o, y, E, S, Q,
callback=callback, callbackInterval=1)
def main(args):
# load previous sample
o, alpha, z, pi, theta, E, S, x, Q, omega, mL, llInit, subsetIdx, dataset = \
SED.loadSample(args.initialSample)
# recursively make output path ignoring if it's already been created
try: os.makedirs(args.outdir)
except: pass
# load data
data = du.load(f'{dataset}/data')
yAll = data['y']
if subsetIdx is not None: y = [yt[subsetIdx[t]] for t, yt in enumerate(yAll)]
else: y = yAll
print(f'Initializing, LL: {llInit:.2f}, K: {len(pi)-1}')
ll = np.zeros(args.nSamples)
# rjmcmc moves (todo: make as args)
# pBirth, pDeath, pSwitch = (0.1, 0.1, 0.0)
# pBirth, pDeath, pSwitch = (0.0, 0.0, 0.0)
pBirth, pDeath, pSwitch = (args.pBirth, args.pDeath, args.pSwitch)
# rjmcmc proposal tracking
nBirthProp, nBirthAccept, nDeathProp, nDeathAccept = (0, 0, 0, 0)
nSwitchProp, nSwitchAccept = (0, 0)
sampleRng = range(args.firstSampleIndex, args.firstSampleIndex+args.nSamples)
for cnt, nS in enumerate(sampleRng):
du.tic()
z, pi, theta, E, S, x, Q, omega, mL, move, accept = SED.sampleRJMCMC(o, y,
alpha, z, pi, theta, E, S, x, Q, omega, mL, pBirth, pDeath, pSwitch)
ll[cnt] = SED.logJoint(o, y, z, x, theta, E, S, Q, alpha, pi, omega, mL)
# print(du.toc())
if move == 'birth':
nBirthProp += 1
if accept: nBirthAccept += 1
elif move == 'death':
nDeathProp += 1
if accept: nDeathAccept += 1
elif move == 'switch':
nSwitchProp += 1
if accept: nSwitchAccept += 1
a = '+' if accept == True else '-'
if not args.silent:
print(f'{args.outdir}, Iter {nS:05}, LL: {ll[cnt]:.2f}, K: {len(pi)-1}, Move: {move[0]}{a}')
if cnt % args.saveEvery == 0:
filename = f'{args.outdir}/sample-{nS:08}'
SED.saveSample(filename, o, alpha, z, pi, theta, E, S, x, Q, omega, mL,
ll[cnt], subsetIdx, dataset)
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 main(args):
# load previous sample
o, alpha, z, pi, theta, E, S, x, Q, omega, mL, ll, subsetIdx, dataset = \
SED.loadSample(args.sample)
K = len(pi) - 1
assert dataset
# load ground-truth labels
gtData = du.load(f'{dataset}/gtLabels')
gtLabels, gtIdx = (gtData['labels'], gtData['gtIdx'])
T, h, w = gtLabels.shape
# load dataset, take max assignments
data = du.load(f'{dataset}/data')
y = data['y']
# make mL so it is the right size
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)]
z = [
SED.inferZ(o, y[t], pi, theta[t], E, x[t], omega, mL[t], max=True)
for t in range(T)
]
# precompute image indices for se3
if o.lie == 'se3':
yy, xx = np.meshgrid(range(h), range(w), indexing='ij')
xy = np.stack((xx.flatten(), yy.flatten()), axis=1) # N x 2
# construct stacked label images for sample
sampleLabels = np.zeros((T, h, w))
for t in range(T):
if o.lie == 'se3':
idx, mask = (data['idx'][t], data['mask'][t])
# idx_t = data['idx'][t] # indexes into xy
# mask = data['mask'][t]
xyFG = xy[mask.flatten()]
xyFG = xyFG[idx] # lines up with z labels now
xs, ys = xyFG.T
else:
xs, ys = y[t].T.astype(np.int)
for k in range(K):
ztk = z[t] == k
sampleLabels[t, ys[ztk], xs[ztk]] = k + 1
# du.ViewImgs(sampleLabels)
# ip.embed()
# sys.exit()
# Have groundtruth and label images, call comparison
iou = args.iou
tp, fp, fn, ids, tilde_tp, motsa, motsp, s_motsa = \
evalSED.mots(sampleLabels, gtData, iou=iou)
name = du.fileparts(dataset)[1]
header1 = f'{name}, iou: {iou:.1f}'
header2 = 'tp & fp & fn & ids & tilde_tp & motsa & motsp & smotsa \\\\'
values = f'{tp} & {fp} & {fn} & {ids:03} & {tilde_tp:.2f} & {motsa:.2f} & {motsp:.2f} & {s_motsa:.2f}'
print(header1)
print(header2)
print(values)
def main(args):
data = du.load(f'{args.dataset_path}/data')
yAll = data['y']
T = len(yAll)
ts = range(T)
if args.maxObs > 0:
subsetIdx = [
np.random.choice(range(len(yt)),
min(args.maxObs, len(yt)),
replace=False) for yt in yAll
]
y = [yt[subsetIdx[t]] for t, yt in enumerate(yAll)]
else:
subsetIdx = [np.arange(len(yt)) for yt in yAll]
y = yAll
mL = [args.mL * np.ones(y[t].shape[0]) for t in range(T)]
if args.se3: o = SED.opts(lie='se3')
else: o = SED.opts(lie='se2')
SED.initPriorsDataDependent(o,
y,
dfQ=args.dfQ,
rotQ=args.rotQ,
dfS=args.dfS,
rotS=args.rotS,
dfE=args.dfE,
scaleE=args.scaleE,
rotX=args.rotX,
rotOmega=args.rotOmega)
x_ = SED.initXDataMeans(o, y)
Q = SED.inferQ(o, x_)
if args.tInit == -1: tInit = np.random.choice(range(T))
else: tInit = args.tInit
###
# yt = y[tInit]
# from mpl_toolkits.mplot3d import Axes3D
# fig = plt.figure()
# ax = fig.add_subplot(111, projection='3d')
# ax.scatter(yt[:,0], yt[:,1], yt[:,2], s=1)
# plt.show()
#
# ip.embed()
# sys.exit()
###
# theta_, omega0, E, S, z, pi = SED.initPartsAndAssoc(o, y[tInit:tInit+1], x_,
# args.alpha, mL, maxBreaks=args.maxBreaks, nInit=args.nInit,
# nIter=args.nIter, tInit=args.tInit, fixedBreaks=args.fixedBreaks
# )
theta_, omega0, E, S, z, pi = SED.initPartsAndAssoc(
o,
y[tInit:tInit + 1],
x_,
args.alpha,
mL,
maxBreaks=args.maxBreaks,
nInit=args.nInit,
nIter=args.nIter,
fixedBreaks=args.fixedBreaks)
K = len(pi) - 1
print(f'Initialized with {K} parts at time {tInit}')
omegaTheta0_ = np.stack([omega0[k] @ theta_[0, k] for k in range(K)])
# get omegaTheta parts and global for all time
theta = np.zeros((T, K) + o.dxGm)
# theta[tInit] = theta_[0]
theta[tInit] = omegaTheta0_
x = np.zeros((T, ) + o.dxGm)
x[tInit] = x_[0]
m = getattr(lie, o.lie)
def estimate_global_then_parts(tPrev, tNext):
# 0 -> 1, 1 -> 2, ...
yPrev, yNext = (y[tPrev], y[tNext])
xPrev = x[tPrev]
thetaPrev = theta[tPrev]
# Initialize relative body transformation from tPrev -> tNext as
# translation between observation means (with no rotation).
muDiff = np.mean(yNext, axis=0) - np.mean(yPrev, axis=0)
Q_t0 = SED.MakeRd(np.eye(o.dy), muDiff)
q_t0 = m.algi(m.logm(Q_t0))
# Estimate body frame
Q_t = icp.optimize_global(o, yNext, xPrev, thetaPrev, E, q_t=q_t0)
x[tNext] = xNext = xPrev @ Q_t
# Estimate omegaTheta with body frame
S_t = icp.optimize_local(o, yNext, xNext, thetaPrev, E, S)
for k in range(K):
theta[tNext, k] = theta[tPrev, k] @ S_t[k]
# reverse direction
# for (tInit, tInit-1) ... (1, 0)
du.tic()
for tPrev, tNext in zip(reversed(ts[:tInit]), reversed(ts[1:tInit + 1])):
print(f'Optimizing Global then Local from time {tNext:03} to {tPrev:03}. ' + \
f'Elapsed: {du.toc():.2f} seconds')
estimate_global_then_parts(tNext, tPrev)
# forward direction
for tPrev, tNext in zip(ts[tInit:-1], ts[tInit + 1:]):
print(f'Optimizing Global then Local from time {tPrev:03} to {tNext:03}. ' + \
f'Elapsed: {du.toc():.2f} seconds')
estimate_global_then_parts(tPrev, tNext)
print(f'Optimizing Global then Local for time {tInit:03}. ' + \
f'Elapsed: {du.toc():.2f} seconds')
if tInit > 0: estimate_global_then_parts(tInit - 1, tInit)
else: estimate_global_then_parts(0, 0)
# separate out omega and omega
omega = np.stack([m.karcher(theta[:, k]) for k in range(K)])
omegaInv = [m.inv(omega[k]) for k in range(K)]
for t in range(T):
for k in range(K):
theta[t, k] = omegaInv[k] @ theta[t, k]
# Re-estimate z, pi; remove unused parts (if any)
z = [[] 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])
z, pi, theta, omega, E, S = SED.consolidatePartsAndResamplePi(
o, z, pi, args.alpha, theta, omega, E, S)
# Re-estimate Q, S, E
Q = SED.inferQ(o, x)
if K > 0: S = np.array([SED.inferSk(o, theta[:, k]) for k in range(K)])
else: S = np.zeros((0, o.dxA, o.dxA))
E = SED.inferE(o, x, theta, omega, y, z)
# Compute log-likelihood
ll = SED.logJoint(o, y, z, x, theta, E, S, Q, args.alpha, pi, omega, mL)
# Save results
SED.saveSample(args.outfile, o, args.alpha, z, pi, theta, E, S, x, Q,
omega, mL, ll, subsetIdx, args.dataset_path)