def summ_unp(self, name, unp, occ):
if self.save_this:
B, C, D, H, W = list(unp.shape)
occ = occ.repeat(1, C, 1, 1, 1)
unp = utils_basic.reduce_masked_mean(unp, occ, dim=3)
# unp = [unp.transpose(1, 2) for unp in unp] #rotate 90 degree counter-clockwise
self.summ_rgb(name=name, ims=unp, blacken_zeros=True)
def summ_unps(self, name, unps, occs):
if self.save_this:
unps = torch.stack(unps, dim=1)
occs = torch.stack(occs, dim=1)
B, S, C, D, H, W = list(unps.shape)
occs = occs.repeat(1, 1, C, 1, 1, 1)
unps = utils_basic.reduce_masked_mean(unps, occs, dim=4)
unps = torch.unbind(unps, dim=1) #should be S x B x W x D x C
# unps = [unp.transpose(1, 2) for unp in unps] #rotate 90 degree counter-clockwise
self.summ_rgbs(name=name, ims=unps, blacken_zeros=True)
def summ_feat(self, name, feat, valid=None, pca=True):
if self.save_this:
if feat.ndim == 5: #B x C x D x H x W
if valid is None:
feat = torch.mean(feat, dim=3)
else:
feat = utils_basic.reduce_masked_mean(feat, valid, dim=3)
B, C, D, W = list(feat.shape)
if not pca:
feat = torch.mean(torch.abs(feat), dim=1, keepdims=True)
# feat is B x 1 x D x W
self.summ_oned(name=name, im=feat, norm=True)
else:
feat_pca = get_feat_pca(feat)
self.summ_rgb(name, feat_pca)
def summ_feats(self, name, feats, valids=None, pca=True):
if self.save_this:
feats = torch.stack(feats, dim=1)
# feats leads with B x S x C
if feats.ndim == 6:
# feats is B x S x C x D x H x W
if valids is None:
feats = torch.mean(feats, dim=4)
else:
valids = torch.stack(valids, dim=1)
valids = valids.repeat(1, 1, feats.size()[2], 1, 1, 1)
feats = utils_basic.reduce_masked_mean(feats,
valids,
dim=4)
B, S, C, D, W = list(feats.size()) # seq is 2 for me here
if not pca:
# feats leads with B x S x C
feats = torch.mean(torch.abs(feats), dim=2, keepdims=True)
# feats leads with B x S x 1
# feats is B x S x D x W
feats = torch.unbind(feats, dim=1)
# feats is a len=S list, each element of shape B x W x D
# # make "forward" point up, and make "right" point right
# feats = [feat.transpose(1, 2) for feat in feats]
self.summ_oneds(name=name, ims=feats, norm=True)
else: #pca
__p = lambda x: utils_basic.pack_seqdim(x, B)
__u = lambda x: utils_basic.unpack_seqdim(x, B)
feats_ = __p(feats)
feats_pca_ = get_feat_pca(feats_)
feats_pca = __u(feats_pca_)
self.summ_rgbs(name=name, ims=torch.unbind(feats_pca, dim=1))
def get_gt_flow(obj_lrtlist_camRs,
obj_scorelist,
camRs_T_camXs,
occR,
K=2,
occ_only=True,
mod='',
vis=True,
summ_writer=None):
# this constructs the flow field according to the given
# box trajectories (obj_lrtlist_camRs) (collected from a moving camR)
# and egomotion (encoded in camRs_T_camXs)
# (so they do not take into account egomotion)
# so, we first generate the flow for all the objects,
# then in the background, put the ego flow
N, B, S, D = list(obj_lrtlist_camRs.shape)
assert (S == 2) # as a flow util, this expects S=2
B, _, Z, Y, X = list(occR.shape)
flows = []
masks = []
for k in range(K):
obj_masklistR0 = utils_vox.assemble_padded_obj_masklist(
obj_lrtlist_camRs[k, :, 0:1],
obj_scorelist[k, :, 0:1],
Z,
Y,
X,
coeff=1.0)
# this is B x 1(N) x 1(C) x Z x Y x Z
# obj_masklistR0 = obj_masklistR0.squeeze(1)
# this is B x 1 x Z x Y x X
obj_mask0 = obj_masklistR0.squeeze(1)
# this is B x 1 x Z x Y x X
camR_T_cam0 = camRs_T_camXs[:, 0]
camR_T_cam1 = camRs_T_camXs[:, 1]
cam0_T_camR = utils_geom.safe_inverse(camR_T_cam0)
cam1_T_camR = utils_geom.safe_inverse(camR_T_cam1)
# camR0_T_camR1 = camR0_T_camRs[:,1]
# camR1_T_camR0 = utils_geom.safe_inverse(camR0_T_camR1)
# obj_masklistA1 = utils_vox.apply_4x4_to_vox(camR1_T_camR0, obj_masklistA0)
# if vis and (summ_writer is not None):
# summ_writer.summ_occ('flow/obj%d_maskA0' % k, obj_masklistA0)
# summ_writer.summ_occ('flow/obj%d_maskA1' % k, obj_masklistA1)
if vis and (summ_writer is not None):
# summ_writer.summ_occ('flow/obj%d_mask0' % k, obj_mask0)
summ_writer.summ_oned('flow/obj%d_mask0' % k,
torch.mean(obj_mask0, 3))
_, ref_T_objs_list = utils_geom.split_lrtlist(obj_lrtlist_camRs[k])
# this is B x S x 4 x 4
ref_T_obj0 = ref_T_objs_list[:, 0]
ref_T_obj1 = ref_T_objs_list[:, 1]
obj0_T_ref = utils_geom.safe_inverse(ref_T_obj0)
obj1_T_ref = utils_geom.safe_inverse(ref_T_obj1)
# these are B x 4 x 4
mem_T_ref = utils_vox.get_mem_T_ref(B, Z, Y, X)
ref_T_mem = utils_vox.get_ref_T_mem(B, Z, Y, X)
ref1_T_ref0 = utils_basic.matmul2(ref_T_obj1, obj0_T_ref)
cam1_T_cam0 = utils_basic.matmul3(cam1_T_camR, ref1_T_ref0,
camR_T_cam0)
mem1_T_mem0 = utils_basic.matmul3(mem_T_ref, cam1_T_cam0, ref_T_mem)
xyz_mem0 = utils_basic.gridcloud3D(B, Z, Y, X)
xyz_mem1 = utils_geom.apply_4x4(mem1_T_mem0, xyz_mem0)
xyz_mem0 = xyz_mem0.reshape(B, Z, Y, X, 3)
xyz_mem1 = xyz_mem1.reshape(B, Z, Y, X, 3)
# only use these displaced points within the obj mask
# obj_mask03 = obj_mask0.view(B, Z, Y, X, 1).repeat(1, 1, 1, 1, 3)
obj_mask0 = obj_mask0.view(B, Z, Y, X, 1)
# # xyz_mem1[(obj_mask03 < 1.0).bool()] = xyz_mem0
# cond = (obj_mask03 < 1.0).float()
cond = (obj_mask0 > 0.0).float()
xyz_mem1 = cond * xyz_mem1 + (1.0 - cond) * xyz_mem0
flow = xyz_mem1 - xyz_mem0
flow = flow.permute(0, 4, 1, 2, 3)
obj_mask0 = obj_mask0.permute(0, 4, 1, 2, 3)
# flow is centered on frame0, so we use occ0 to mask it
if occ_only:
flow = flow * occR
# if vis and k==0:
if vis:
summ_writer.summ_3D_flow('flow/gt_%d' % k, flow, clip=4.0)
masks.append(obj_mask0)
flows.append(flow)
camR_T_cam0 = camRs_T_camXs[:, 0]
camR_T_cam1 = camRs_T_camXs[:, 1]
cam0_T_camR = utils_geom.safe_inverse(camR_T_cam0)
cam1_T_camR = utils_geom.safe_inverse(camR_T_cam1)
mem_T_ref = utils_vox.get_mem_T_ref(B, Z, Y, X)
ref_T_mem = utils_vox.get_ref_T_mem(B, Z, Y, X)
cam1_T_cam0 = utils_basic.matmul2(cam1_T_camR, camR_T_cam0)
mem1_T_mem0 = utils_basic.matmul3(mem_T_ref, cam1_T_cam0, ref_T_mem)
xyz_mem0 = utils_basic.gridcloud3D(B, Z, Y, X)
xyz_mem1 = utils_geom.apply_4x4(mem1_T_mem0, xyz_mem0)
xyz_mem0 = xyz_mem0.reshape(B, Z, Y, X, 3)
xyz_mem1 = xyz_mem1.reshape(B, Z, Y, X, 3)
flow = xyz_mem1 - xyz_mem0
flow = flow.permute(0, 4, 1, 2, 3)
if occ_only:
flow = flow * occXs[:, 0]
bkg_flow = flow
# allow zero motion in the bkg
any_mask = torch.max(torch.stack(masks, axis=0), axis=0)[0]
masks.append(1.0 - any_mask)
flows.append(bkg_flow)
flows = torch.stack(flows, axis=0)
masks = torch.stack(masks, axis=0)
masks = masks.repeat(1, 1, 3, 1, 1, 1)
flow = utils_basic.reduce_masked_mean(flows, masks, dim=0)
if vis:
summ_writer.summ_3D_flow('flow/gt_complete', flow, clip=4.0)
# flow is shaped B x 3 x D x H x W
return flow
def get_unps_vis(unps, occs, dim=5):
B, S, C, D, H, W = list(unps.shape)
occs = occs.repeat(1, 1, C, 1, 1, 1)
unps = utils_basic.reduce_masked_mean(unps, occs, dim=dim)
# unps is B x S x C x D x W
return unps