def quick_nrm_map(self,
dpt,
scale_to_mm,
K=intrinsic_matrix['ycb_K1'],
with_show=False):
dpt_mm = (dpt.copy() * scale_to_mm).astype(np.uint16)
nrm_map = normalSpeed.depth_normal(dpt_mm, K[0][0], K[1][1], 5, 2000,
20, False)
if with_show:
nrm_map[np.isnan(nrm_map)] = 0.0
nrm_map[np.isinf(nrm_map)] = 0.0
show_nrm = ((nrm_map[:, :, :3] + 1.0) * 127).astype(np.uint8)
return nrm_map, show_nrm
return nrm_map
def get_item(self, item_name):
if "pkl" in item_name:
data = pkl.load(open(item_name, "rb"))
dpt_mm = data['depth'] * 1000.
rgb = data['rgb']
labels = data['mask']
K = data['K']
RT = data['RT']
rnd_typ = data['rnd_typ']
if rnd_typ == "fuse":
labels = (labels == self.cls_id).astype("uint8")
else:
labels = (labels > 0).astype("uint8")
else:
with Image.open(
os.path.join(self.cls_root,
"depth/{}.png".format(item_name))) as di:
dpt_mm = np.array(di)
with Image.open(
os.path.join(self.cls_root,
"mask/{}.png".format(item_name))) as li:
labels = np.array(li)
labels = (labels > 0).astype("uint8")
with Image.open(
os.path.join(self.cls_root,
"rgb/{}.png".format(item_name))) as ri:
if self.add_noise:
ri = self.trancolor(ri)
rgb = np.array(ri)[:, :, :3]
meta = self.meta_lst[int(item_name)]
if self.cls_id == 2:
for i in range(0, len(meta)):
if meta[i]['obj_id'] == 2:
meta = meta[i]
break
else:
meta = meta[0]
R = np.resize(np.array(meta['cam_R_m2c']), (3, 3))
T = np.array(meta['cam_t_m2c']) / 1000.0
RT = np.concatenate((R, T[:, None]), axis=1)
rnd_typ = 'real'
K = self.config.intrinsic_matrix["linemod"]
cam_scale = 1000.0
if len(labels.shape) > 2:
labels = labels[:, :, 0]
rgb_labels = labels.copy()
# if self.add_noise and rnd_typ == 'render':
if self.add_noise and rnd_typ != 'real':
if rnd_typ == 'render' or self.rng.rand() < 0.8:
rgb = self.rgb_add_noise(rgb)
rgb_labels = labels.copy()
msk_dp = dpt_mm > 1e-6
rgb, dpt_mm = self.add_real_back(rgb, rgb_labels, dpt_mm,
msk_dp)
if self.rng.rand() > 0.8:
rgb = self.rgb_add_noise(rgb)
dpt_mm = dpt_mm.copy().astype(np.uint16)
nrm_map = normalSpeed.depth_normal(dpt_mm, K[0][0], K[1][1], 5, 2000,
20, False)
if self.DEBUG:
show_nrm_map = ((nrm_map + 1.0) * 127).astype(np.uint8)
imshow("nrm_map", show_nrm_map)
dpt_m = dpt_mm.astype(np.float32) / cam_scale
dpt_xyz = self.dpt_2_pcld(dpt_m, 1.0, K)
dpt_xyz[np.isnan(dpt_xyz)] = 0.0
dpt_xyz[np.isinf(dpt_xyz)] = 0.0
msk_dp = dpt_mm > 1e-6
choose = msk_dp.flatten().nonzero()[0].astype(np.uint32)
if len(choose) < 400:
return None
choose_2 = np.array([i for i in range(len(choose))])
if len(choose_2) < 400:
return None
if len(choose_2) > self.config.n_sample_points:
c_mask = np.zeros(len(choose_2), dtype=int)
c_mask[:self.config.n_sample_points] = 1
np.random.shuffle(c_mask)
choose_2 = choose_2[c_mask.nonzero()]
else:
choose_2 = np.pad(choose_2,
(0, self.config.n_sample_points - len(choose_2)),
'wrap')
choose = np.array(choose)[choose_2]
sf_idx = np.arange(choose.shape[0])
np.random.shuffle(sf_idx)
choose = choose[sf_idx]
cld = dpt_xyz.reshape(-1, 3)[choose, :]
rgb_pt = rgb.reshape(-1, 3)[choose, :].astype(np.float32)
nrm_pt = nrm_map[:, :, :3].reshape(-1, 3)[choose, :]
labels_pt = labels.flatten()[choose]
choose = np.array([choose])
cld_rgb_nrm = np.concatenate((cld, rgb_pt, nrm_pt),
axis=1).transpose(1, 0)
RTs, kp3ds, ctr3ds, cls_ids, kp_targ_ofst, ctr_targ_ofst = self.get_pose_gt_info(
cld, labels_pt, RT)
h, w = rgb_labels.shape
dpt_6c = np.concatenate((dpt_xyz, nrm_map[:, :, :3]),
axis=2).transpose(2, 0, 1)
rgb = np.transpose(rgb, (2, 0, 1)) # hwc2chw
xyz_lst = [dpt_xyz.transpose(2, 0, 1)] # c, h, w
msk_lst = [dpt_xyz[2, :, :] > 1e-8]
for i in range(3):
scale = pow(2, i + 1)
nh, nw = h // pow(2, i + 1), w // pow(2, i + 1)
ys, xs = np.mgrid[:nh, :nw]
xyz_lst.append(xyz_lst[0][:, ys * scale, xs * scale])
msk_lst.append(xyz_lst[-1][2, :, :] > 1e-8)
sr2dptxyz = {
pow(2, ii): item.reshape(3, -1).transpose(1, 0)
for ii, item in enumerate(xyz_lst)
}
rgb_ds_sr = [4, 8, 8, 8]
n_ds_layers = 4
pcld_sub_s_r = [4, 4, 4, 4]
inputs = {}
# DownSample stage
for i in range(n_ds_layers):
nei_idx = DP.knn_search(cld[None, ...], cld[None, ...],
16).astype(np.int32).squeeze(0)
sub_pts = cld[:cld.shape[0] // pcld_sub_s_r[i], :]
pool_i = nei_idx[:cld.shape[0] // pcld_sub_s_r[i], :]
up_i = DP.knn_search(sub_pts[None, ...], cld[None, ...],
1).astype(np.int32).squeeze(0)
inputs['cld_xyz%d' % i] = cld.astype(np.float32).copy()
inputs['cld_nei_idx%d' % i] = nei_idx.astype(np.int32).copy()
inputs['cld_sub_idx%d' % i] = pool_i.astype(np.int32).copy()
inputs['cld_interp_idx%d' % i] = up_i.astype(np.int32).copy()
nei_r2p = DP.knn_search(sr2dptxyz[rgb_ds_sr[i]][None, ...],
sub_pts[None, ...],
16).astype(np.int32).squeeze(0)
inputs['r2p_ds_nei_idx%d' % i] = nei_r2p.copy()
nei_p2r = DP.knn_search(sub_pts[None, ...],
sr2dptxyz[rgb_ds_sr[i]][None, ...],
1).astype(np.int32).squeeze(0)
inputs['p2r_ds_nei_idx%d' % i] = nei_p2r.copy()
cld = sub_pts
n_up_layers = 3
rgb_up_sr = [4, 2, 2]
for i in range(n_up_layers):
r2p_nei = DP.knn_search(
sr2dptxyz[rgb_up_sr[i]][None, ...],
inputs['cld_xyz%d' % (n_ds_layers - i - 1)][None, ...],
16).astype(np.int32).squeeze(0)
inputs['r2p_up_nei_idx%d' % i] = r2p_nei.copy()
p2r_nei = DP.knn_search(
inputs['cld_xyz%d' % (n_ds_layers - i - 1)][None, ...],
sr2dptxyz[rgb_up_sr[i]][None,
...], 1).astype(np.int32).squeeze(0)
inputs['p2r_up_nei_idx%d' % i] = p2r_nei.copy()
show_rgb = rgb.transpose(1, 2, 0).copy()[:, :, ::-1]
if self.DEBUG:
for ip, xyz in enumerate(xyz_lst):
pcld = xyz.reshape(3, -1).transpose(1, 0)
p2ds = self.bs_utils.project_p3d(pcld, cam_scale, K)
srgb = self.bs_utils.paste_p2ds(show_rgb.copy(), p2ds,
(0, 0, 255))
# imshow("rz_pcld_%d" % ip, srgb)
p2ds = self.bs_utils.project_p3d(inputs['cld_xyz%d' % ip],
cam_scale, K)
srgb1 = self.bs_utils.paste_p2ds(show_rgb.copy(), p2ds,
(0, 0, 255))
# imshow("rz_pcld_%d_rnd" % ip, srgb1)
# print(
# "kp3ds:", kp3ds.shape, kp3ds, "\n",
# "kp3ds.mean:", np.mean(kp3ds, axis=0), "\n",
# "ctr3ds:", ctr3ds.shape, ctr3ds, "\n",
# "cls_ids:", cls_ids, "\n",
# "labels.unique:", np.unique(labels),
# )
item_dict = dict(
rgb=rgb.astype(np.uint8), # [c, h, w]
cld_rgb_nrm=cld_rgb_nrm.astype(np.float32), # [9, npts]
choose=choose.astype(np.int32), # [1, npts]
labels=labels_pt.astype(np.int32), # [npts]
rgb_labels=rgb_labels.astype(np.int32), # [h, w]
dpt_map_m=dpt_m.astype(np.float32), # [h, w]
RTs=RTs.astype(np.float32),
kp_targ_ofst=kp_targ_ofst.astype(np.float32),
ctr_targ_ofst=ctr_targ_ofst.astype(np.float32),
cls_ids=cls_ids.astype(np.int32),
ctr_3ds=ctr3ds.astype(np.float32),
kp_3ds=kp3ds.astype(np.float32),
)
item_dict.update(inputs)
if self.DEBUG:
extra_d = dict(
dpt_xyz_nrm=dpt_6c.astype(np.float32), # [6, h, w]
cam_scale=np.array([cam_scale]).astype(np.float32),
K=K.astype(np.float32),
)
item_dict.update(extra_d)
item_dict['normal_map'] = nrm_map[:, :, :3].astype(np.float32)
return item_dict
def get_item(self, item_name):
with Image.open(os.path.join(self.root, item_name+'-depth.png')) as di:
dpt_um = np.array(di)
with Image.open(os.path.join(self.root, item_name+'-label.png')) as li:
labels = np.array(li)
rgb_labels = labels.copy()
meta = scio.loadmat(os.path.join(self.root, item_name+'-meta.mat'))
if item_name[:8] != 'data_syn' and int(item_name[5:9]) >= 60:
K = config.intrinsic_matrix['ycb_K2']
else:
K = config.intrinsic_matrix['ycb_K1']
with Image.open(os.path.join(self.root, item_name+'-color.png')) as ri:
if self.add_noise:
ri = self.trancolor(ri)
rgb = np.array(ri)[:, :, :3]
rnd_typ = 'syn' if 'syn' in item_name else 'real'
cam_scale = meta['factor_depth'].astype(np.float32)[0][0]
msk_dp = dpt_um > 1e-6
if self.add_noise and rnd_typ == 'syn':
rgb = self.rgb_add_noise(rgb)
rgb, dpt_um = self.add_real_back(rgb, rgb_labels, dpt_um, msk_dp)
if self.rng.rand() > 0.8:
rgb = self.rgb_add_noise(rgb)
dpt_um = bs_utils.fill_missing(dpt_um, cam_scale, 1)
msk_dp = dpt_um > 1e-6
dpt_mm = (dpt_um.copy()/10).astype(np.uint16)
nrm_map = normalSpeed.depth_normal(
dpt_mm, K[0][0], K[1][1], 5, 2000, 20, False
)
if self.debug:
show_nrm_map = ((nrm_map + 1.0) * 127).astype(np.uint8)
imshow("nrm_map", show_nrm_map)
dpt_m = dpt_um.astype(np.float32) / cam_scale
dpt_xyz = self.dpt_2_pcld(dpt_m, 1.0, K)
choose = msk_dp.flatten().nonzero()[0].astype(np.uint32)
if len(choose) < 400:
return None
choose_2 = np.array([i for i in range(len(choose))])
if len(choose_2) < 400:
return None
if len(choose_2) > config.n_sample_points:
c_mask = np.zeros(len(choose_2), dtype=int)
c_mask[:config.n_sample_points] = 1
np.random.shuffle(c_mask)
choose_2 = choose_2[c_mask.nonzero()]
else:
choose_2 = np.pad(choose_2, (0, config.n_sample_points-len(choose_2)), 'wrap')
choose = np.array(choose)[choose_2]
sf_idx = np.arange(choose.shape[0])
np.random.shuffle(sf_idx)
choose = choose[sf_idx]
cld = dpt_xyz.reshape(-1, 3)[choose, :]
rgb_pt = rgb.reshape(-1, 3)[choose, :].astype(np.float32)
nrm_pt = nrm_map[:, :, :3].reshape(-1, 3)[choose, :]
labels_pt = labels.flatten()[choose]
choose = np.array([choose])
cld_rgb_nrm = np.concatenate((cld, rgb_pt, nrm_pt), axis=1).transpose(1, 0)
cls_id_lst = meta['cls_indexes'].flatten().astype(np.uint32)
RTs, kp3ds, ctr3ds, cls_ids, kp_targ_ofst, ctr_targ_ofst = self.get_pose_gt_info(
cld, labels_pt, cls_id_lst, meta
)
h, w = rgb_labels.shape
dpt_6c = np.concatenate((dpt_xyz, nrm_map[:, :, :3]), axis=2).transpose(2, 0, 1)
rgb = np.transpose(rgb, (2, 0, 1)) # hwc2chw
xyz_lst = [dpt_xyz.transpose(2, 0, 1)] # c, h, w
msk_lst = [dpt_xyz[2, :, :] > 1e-8]
for i in range(3):
scale = pow(2, i+1)
nh, nw = h // pow(2, i+1), w // pow(2, i+1)
ys, xs = np.mgrid[:nh, :nw]
xyz_lst.append(xyz_lst[0][:, ys*scale, xs*scale])
msk_lst.append(xyz_lst[-1][2, :, :] > 1e-8)
sr2dptxyz = {
pow(2, ii): item.reshape(3, -1).transpose(1, 0) for ii, item in enumerate(xyz_lst)
}
sr2msk = {
pow(2, ii): item.reshape(-1) for ii, item in enumerate(msk_lst)
}
rgb_ds_sr = [4, 8, 8, 8]
n_ds_layers = 4
pcld_sub_s_r = [4, 4, 4, 4]
inputs = {}
# DownSample stage
for i in range(n_ds_layers):
nei_idx = DP.knn_search(
cld[None, ...], cld[None, ...], 16
).astype(np.int32).squeeze(0)
sub_pts = cld[:cld.shape[0] // pcld_sub_s_r[i], :]
pool_i = nei_idx[:cld.shape[0] // pcld_sub_s_r[i], :]
up_i = DP.knn_search(
sub_pts[None, ...], cld[None, ...], 1
).astype(np.int32).squeeze(0)
inputs['cld_xyz%d'%i] = cld.astype(np.float32).copy()
inputs['cld_nei_idx%d'%i] = nei_idx.astype(np.int32).copy()
inputs['cld_sub_idx%d'%i] = pool_i.astype(np.int32).copy()
inputs['cld_interp_idx%d'%i] = up_i.astype(np.int32).copy()
nei_r2p = DP.knn_search(
sr2dptxyz[rgb_ds_sr[i]][None, ...], sub_pts[None, ...], 16
).astype(np.int32).squeeze(0)
inputs['r2p_ds_nei_idx%d'%i] = nei_r2p.copy()
nei_p2r = DP.knn_search(
sub_pts[None, ...], sr2dptxyz[rgb_ds_sr[i]][None, ...], 1
).astype(np.int32).squeeze(0)
inputs['p2r_ds_nei_idx%d'%i] = nei_p2r.copy()
cld = sub_pts
n_up_layers = 3
rgb_up_sr = [4, 2, 2]
for i in range(n_up_layers):
r2p_nei = DP.knn_search(
sr2dptxyz[rgb_up_sr[i]][None, ...],
inputs['cld_xyz%d'%(n_ds_layers-i-1)][None, ...], 16
).astype(np.int32).squeeze(0)
inputs['r2p_up_nei_idx%d'%i] = r2p_nei.copy()
p2r_nei = DP.knn_search(
inputs['cld_xyz%d'%(n_ds_layers-i-1)][None, ...],
sr2dptxyz[rgb_up_sr[i]][None, ...], 1
).astype(np.int32).squeeze(0)
inputs['p2r_up_nei_idx%d'%i] = p2r_nei.copy()
show_rgb = rgb.transpose(1, 2, 0).copy()[:, :, ::-1]
if self.debug:
for ip, xyz in enumerate(xyz_lst):
pcld = xyz.reshape(3, -1).transpose(1, 0)
p2ds = bs_utils.project_p3d(pcld, cam_scale, K)
print(show_rgb.shape, pcld.shape)
srgb = bs_utils.paste_p2ds(show_rgb.copy(), p2ds, (0, 0, 255))
imshow("rz_pcld_%d" % ip, srgb)
p2ds = bs_utils.project_p3d(inputs['cld_xyz%d'%ip], cam_scale, K)
srgb1 = bs_utils.paste_p2ds(show_rgb.copy(), p2ds, (0, 0, 255))
imshow("rz_pcld_%d_rnd" % ip, srgb1)
item_dict = dict(
rgb=rgb.astype(np.uint8), # [c, h, w]
cld_rgb_nrm=cld_rgb_nrm.astype(np.float32), # [9, npts]
choose=choose.astype(np.int32), # [1, npts]
labels=labels_pt.astype(np.int32), # [npts]
rgb_labels=rgb_labels.astype(np.int32), # [h, w]
dpt_map_m=dpt_m.astype(np.float32), # [h, w]
RTs=RTs.astype(np.float32),
kp_targ_ofst=kp_targ_ofst.astype(np.float32),
ctr_targ_ofst=ctr_targ_ofst.astype(np.float32),
cls_ids=cls_ids.astype(np.int32),
ctr_3ds=ctr3ds.astype(np.float32),
kp_3ds=kp3ds.astype(np.float32),
)
item_dict.update(inputs)
if self.debug:
extra_d = dict(
dpt_xyz_nrm=dpt_6c.astype(np.float32), # [6, h, w]
cam_scale=np.array([cam_scale]).astype(np.float32),
K=K.astype(np.float32),
)
item_dict.update(extra_d)
item_dict['normal_map'] = nrm_map[:, :, :3].astype(np.float32)
return item_dict
distance_threshold = 2000
difference_threshold = 20
point_into_surface = False
color = np.load("examplePicture/color.npy")
depth = np.load("examplePicture/depth.npy")
depth = depth.astype(np.uint16)
t1 = time.time()
"""
The coordinate of depth and normal is in cv coordinate:
- x is horizontal
- y is down (to align to the actual pixel coordinates used in digital images)
- right-handed: positive z look-at direction
"""
normals_map_out = normalSpeed.depth_normal(depth, fx, fy, k_size, distance_threshold, difference_threshold, point_into_surface)
# If the normal point out the surface, the z of normal should be negetive.
print("normals_map_out z mean:", normals_map_out[:, :, 2].mean())
# If the normal point into the surface, the z of normal should be positive, as z of depth.
normals_map_in = normalSpeed.depth_normal(depth, fx, fy, k_size, distance_threshold, difference_threshold, True)
print("normals_map_in z mean:", normals_map_in[:, :, 2].mean())
imshow('norm_out', norm2bgr(normals_map_out))
imwrite("examplePicture/normal_out.jpg", norm2bgr(normals_map_out))
imshow('norm_in', norm2bgr(normals_map_in))
imwrite("examplePicture/normal_in.jpg", norm2bgr(normals_map_in))
imshow('depth', depth2show(depth))
imwrite('examplePicture/depth_view.jpg', depth2show(depth))