def test():
import sys
sys.path.append('../')
from my_folder import MyImageFolder
from neupeak.utils import webcv2 as cv2
minibatch = 8
theta = torch.from_numpy(np.array(
[[0.5, 0, 0],
[0, 0.5, 0]]
))
theta = theta.expand(minibatch, 2,3).float()
img_size = 224
data_path = '/unsullied/sharefs/chenyaolin/cylfile/data_2/train'
train_transform = transforms.Compose(
[transforms.Scale((448,448)),
transforms.ToTensor()
]
)
data = MyImageFolder(data_path, train_transform, data_cached= True)
imgs = torch.utils.data.DataLoader(data, batch_size= minibatch, shuffle = False)
stn = STN((img_size, img_size))
stn = torch.nn.DataParallel(stn, device_ids=list(range(8)))
stn.cuda()
for batch in imgs:
transformed = stn(batch[0].cuda(), theta)
o = batch[0]
t = transformed
for i in range(minibatch):
# from IPython import embed
# embed()
cv2.imshow('o',np.transpose( o[i].cpu().data.numpy()*255, (1,2,0)))
cv2.imshow('transformed', np.transpose(t[i].cpu().data.numpy()*255,(1,2,0)))
cv2.waitKey(0)
def vis_input(prefix, image, seg, lm):
image = np.asarray(image)
seg = (255 / seg.max() * seg).astype(np.uint8)
lm = lm.reshape(-1, 2).round().astype(np.int32)[:, ::-1]
for point in lm:
image = cv2.circle(image, tuple(point), 2, (255, 0, 0))
webcv2.imshow(prefix + "-image", image[..., ::-1])
webcv2.imshow(prefix + "-seg", seg)
webcv2.waitKey()
def output(vis, fname):
if args.show:
print(fname)
webcv2.imshow("window", vis.get_image()[:, :, ::-1])
webcv2.waitKey()
else:
filepath = os.path.join(dirname, fname)
print("Saving to {} ...".format(filepath))
vis.save(filepath)
def main():
# config.mini_batch_size = 1
global DEBUG
DEBUG = True
ds = {}
ds['train'] = Dataset('train', DEBUG=True)
# ds['val'] = Dataset('validation')
ds['test'] = Dataset('test', DEBUG=True)
idx = dict(
train=0,
val=0,
test=0
)
while True:
# for cat in ['val', 'test']:
# for cat in ['train']:
for cat in ['test']:
datum = ds[cat].__getitem__(idx[cat])
idx[cat] += 1
K = datum['K']
cam_scale = datum['cam_scale']
rgb = datum['rgb'].transpose(1, 2, 0)[...,::-1].copy()# [...,::-1].copy()
for i in range(22):
pcld = datum['cld_rgb_nrm'][:3, :].transpose(1, 0).copy()
p2ds = bs_utils.project_p3d(pcld, cam_scale, K)
# rgb = bs_utils.draw_p2ds(rgb, p2ds)
kp3d = datum['kp_3ds'][i]
if kp3d.sum() < 1e-6:
break
kp_2ds = bs_utils.project_p3d(kp3d, cam_scale, K)
rgb = bs_utils.draw_p2ds(
rgb, kp_2ds, 3, bs_utils.get_label_color(datum['cls_ids'][i][0], mode=1)
)
ctr3d = datum['ctr_3ds'][i]
ctr_2ds = bs_utils.project_p3d(ctr3d[None, :], cam_scale, K)
rgb = bs_utils.draw_p2ds(
rgb, ctr_2ds, 4, (0, 0, 255)
)
imshow('{}_rgb'.format(cat), rgb)
cmd = waitKey(0)
if cmd == ord('q'):
exit()
else:
continue
def test():
while True:
a = np.random.rand(1000, 2)
ta = torch.from_numpy(a.astype(np.float32)).cuda()
ms = MeanShiftTorch(0.05)
ctr, _ = ms.fit(ta)
a_idx = (a * 480).astype("uint8")
show_a = np.zeros((480, 480, 3), dtype="uint8")
show_a[a_idx[:, 0], a_idx[:, 1], :] = np.array([255, 255, 255])
ctr = (ctr.cpu().numpy() * 480).astype("uint8")
show_a = cv2.circle(show_a, (ctr[1], ctr[0]), 3, (0, 0, 255), -1)
ms_cpu = MeanShift(bandwidth=0.05, n_jobs=8)
ms_cpu.fit(a)
clus_ctrs = np.array(ms_cpu.cluster_centers_)
clus_labels = ms_cpu.labels_
ctr = (clus_ctrs[0] * 480).astype("uint8")
show_a = cv2.circle(show_a, (ctr[1], ctr[0]), 3, (255, 0, 0), -1)
imshow('show_a', show_a)
waitKey(0)
print(clus_ctrs[0])
def save_fuse_data(output_dir, idx, fuse_img, fuse_mask, fuse_depth,
fuse_begins, t_pose, cls):
cls_id = lm_obj_dict[cls]
if (fuse_mask == cls_id).sum() < 20:
return None
os.makedirs(output_dir, exist_ok=True)
fuse_mask = fuse_mask.astype(np.uint8)
data = {}
data['rgb'] = fuse_img
data['mask'] = fuse_mask
data['depth'] = fuse_depth.astype(np.float32) / 1000.0
data['K'] = Intrinsic_matrix['linemod']
data['RT'] = t_pose
data['cls_typ'] = cls
data['rnd_typ'] = 'fuse'
data['begins'] = fuse_begins
if DEBUG:
imshow("rgb", fuse_img[:, :, ::-1])
imshow("depth", (fuse_depth / fuse_depth.max() * 255).astype('uint8'))
imshow("label", (fuse_mask / fuse_mask.max() * 255).astype("uint8"))
waitKey(0)
sv_pth = os.path.join(output_dir, "{}.pkl".format(idx))
pickle.dump(data, open(sv_pth, 'wb'))
sv_pth = os.path.abspath(sv_pth)
return sv_pth
def test2():
sv_ptn = '/data/workspace/3D_Point_Det/config/ycb.onestage.rs14.nofarflatFocalls/train_log/eval_result/051_large_clamp/mask_res_pic/{}sv_info_1.pkl'
for i in range(2000):
data = pkl.load(open(sv_ptn.format(i), 'rb'))
all_p3ds = data['p3ds']
for cls_id in data['gt_cls_ids'][0]:
if cls_id == 0:
break
p3ds = all_p3ds[np.where(data['labels'] == cls_id)[0], :]
show_img = np.zeros((480, 640, 3), dtype="uint8")
p2ds = my_utils.project_p3d(p3ds, 1.0)
show_img[p2ds[:, 1], p2ds[:, 0], :] = np.array([255, 255, 255])
gpu_label = np.zeros((480, 640, 3), dtype="uint8")
cpu_label = gpu_label.copy()
p3ds_cu = torch.from_numpy(p3ds).cuda()
ms_gpu = MeanShiftTorch(0.05)
start = time.time()
ctr, labels = ms_gpu.fit(p3ds_cu)
ctr = ctr.cpu().numpy().reshape(1, 3)
labels = labels.cpu().numpy()
p2ds_gt_lb = p2ds[np.where(labels==1)[0], :]
gpu_label[p2ds_gt_lb[:, 1], p2ds_gt_lb[:, 0], :] = np.array(
[255, 255, 255]
)
end = time.time()
print("gpu time:\t", end - start)
ctr_2d = my_utils.project_p3d(ctr, 1.0)
show_img = cv2.circle(
show_img, (ctr_2d[0][0], ctr_2d[0][1]), 3, (0, 0, 255), -1
)
ms_cpu = MeanShift(
bandwidth=0.05, n_jobs=40
)
start = time.time()
ms_cpu.fit(p3ds)
end = time.time()
print("sklearn cpu time:\t", end - start)
clus_ctrs = np.array(ms_cpu.cluster_centers_)
clus_labels = ms_cpu.labels_
ctr_2d = my_utils.project_p3d(clus_ctrs[0].reshape(1, 3), 1.0)
show_img = cv2.circle(
show_img, (ctr_2d[0][0], ctr_2d[0][1]), 3, (255, 0, 0), -1
)
p2ds_gt_lb = p2ds[np.where(clus_labels==0)[0], :]
cpu_label[p2ds_gt_lb[:, 1], p2ds_gt_lb[:, 0], :] = np.array(
[255, 255, 255]
)
imshow('show_img', show_img)
imshow('gpu', gpu_label)
imshow('cpu', cpu_label)
waitKey(0)
def complete_dpt(nid_p):
nid_lst = my_utils.read_lines(nid_p)
# fill_type = 'fast'
cnt = 0
import random
# random.shuffle(nid_lst)
with concurrent.futures.ProcessPoolExecutor(15) as executor:
for info in executor.map(get_one_show, nid_lst):
print(np.min(info['final_dpt']), np.max(info['final_dpt']))
show_depth('ori_dpth', info['ori_dpt'])
show_depth('cmplt_dpth', info['final_dpt'])
imshow('rgb', info['rgb'])
imshow('nrm_map', info['nrm_map'])
imshow('nrm_map_final', info['nrm_map_final'])
if cnt == 0:
cmd = waitKey(0)
# cnt += 1
else:
cmd = waitKey(2)
def main(makeup_dir="faces/makeup",
non_makeup_dir="faces/no_makeup",
img_size=256,
speed=False):
makeup_dir = smart_path(makeup_dir)
non_makeup_dir = smart_path(non_makeup_dir)
makeup_paths = list(makeup_dir.glob("*"))
non_makeup_paths = list(non_makeup_dir.glob("*"))
assert len(makeup_paths) > 0
assert len(non_makeup_paths) > 0
random = np.random.RandomState(seed=0)
while True:
makeup_path = random.choice(makeup_paths)
ant_makeup_path = random.choice(makeup_paths)
non_makeup_path = random.choice(non_makeup_paths)
makeup_image = load_image(makeup_path)
non_makeup_image = load_image(non_makeup_path)
if makeup_image is None or non_makeup_image is None:
continue
transferred_image = solver.test(
*preprocess(Image.fromarray(non_makeup_image)),
*preprocess(Image.fromarray(makeup_image)))
if speed:
input_1 = preprocess(Image.fromarray(non_makeup_image))
input_2 = preprocess(Image.fromarray(makeup_image))
start = time.time()
for _ in tqdm(range(100)):
transferred_image = solver.test(*input_1, *input_2)
print("inference time", time.time() - start)
webcv2.imshow("source", non_makeup_image[..., ::-1])
webcv2.imshow("reference", makeup_image[..., ::-1])
webcv2.imshow("result", np.array(transferred_image)[..., ::-1])
webcv2.waitKey()
def test():
while True:
# a = np.random.rand(20000, 2)
n_clus = 5
n_samples = 100
bw = 10
centroids = np.random.uniform(0, 480, (n_clus, 2))
slices = [
np.random.multivariate_normal(centroids[i], np.diag([50., 50.]),
n_samples + i * 100)
for i in range(n_clus)
]
a = np.concatenate(slices).astype(np.float32)
print("npts:", a.shape)
ta = torch.from_numpy(a.astype(np.float32)).cuda()
a_idx = (a / a.max() * 480).astype("uint8")
show_a = np.zeros((480, 480, 3), dtype="uint8")
show_a[a_idx[:, 0], a_idx[:, 1], :] = np.array([255, 255, 255])
ms = MeanShiftTorch(bw)
ctr, lb = ms.fit(ta)
ctr = (ctr.cpu().numpy() / a.max() * 480).astype("uint8")
show_a_one = cv2.circle(show_a.copy(), (ctr[1], ctr[0]), 3,
(0, 0, 255), -1)
ctr_lst, lb, n_in_lst = ms.fit_multi_clus(ta)
print(ctr_lst, n_in_lst)
show_a_multi = show_a.copy()
for ctr in ctr_lst:
ctr = (ctr.cpu().numpy() / a.max() * 480).astype("uint8")
show_a_multi = cv2.circle(show_a_multi, (ctr[1], ctr[0]), 3,
(0, 0, 255), -1)
def get_color(cls_id, n_obj=30):
mul_col = 255 * 255 * 255 // n_obj * cls_id
r, g, b = mul_col // 255 // 255, (mul_col //
255) % 255, mul_col % 255
bgr = (int(r), int(g), int(b))
return bgr
show_ca = np.zeros((480, 480, 3), dtype="uint8")
print(lb.unique())
n_clus = lb.max()
for cls in range(1, n_clus + 1):
inl = a_idx[lb.cpu().numpy() == cls, :]
show_ca[inl[:, 0], inl[:, 1], :] = np.array(
list(get_color(cls, n_obj=n_clus + 1)))
# ms_cpu = MeanShift(
# bandwidth=bw, n_jobs=8
# )
# ms_cpu.fit(a)
# clus_ctrs = np.array(ms_cpu.cluster_centers_)
# clus_labels = ms_cpu.labels_
# ctr = (clus_ctrs[0] / a.max() * 480).astype("uint8")
# show_a = cv2.circle(show_a, (ctr[1], ctr[0]), 3, (255, 0, 0), -1)
# imshow("show_b", show_b)
imshow('show_a_one', show_a_one)
imshow("show_a_multi", show_a_multi)
imshow('show_ca', show_ca)
waitKey(0)
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 gen_one_zbuf_render(args, meshc, RT):
if args.extractor == 'SIFT':
extractor = cv2.xfeatures2d.SIFT_create()
else: # use orb
extractor = cv2.ORB_create()
h, w = args.h, args.w
if type(args.K) == list:
K = np.array(args.K).reshape(3, 3)
R, T = RT[:3, :3], RT[:3, 3]
new_xyz = meshc['xyz'].copy()
new_xyz = np.dot(new_xyz, R.T) + T
p2ds = np.dot(new_xyz.copy(), K.T)
p2ds = p2ds[:, :2] / p2ds[:, 2:]
p2ds = np.require(p2ds.flatten(), 'float32', 'C')
zs = np.require(new_xyz[:, 2].copy(), 'float32', 'C')
zbuf = np.require(np.zeros(h*w), 'float32', 'C')
rbuf = np.require(np.zeros(h*w), 'int32', 'C')
gbuf = np.require(np.zeros(h*w), 'int32', 'C')
bbuf = np.require(np.zeros(h*w), 'int32', 'C')
RENDERER.rgbzbuffer(
ct.c_int(h),
ct.c_int(w),
p2ds.ctypes.data_as(ct.c_void_p),
new_xyz.ctypes.data_as(ct.c_void_p),
zs.ctypes.data_as(ct.c_void_p),
meshc['r'].ctypes.data_as(ct.c_void_p),
meshc['g'].ctypes.data_as(ct.c_void_p),
meshc['b'].ctypes.data_as(ct.c_void_p),
ct.c_int(meshc['n_face']),
meshc['face'].ctypes.data_as(ct.c_void_p),
zbuf.ctypes.data_as(ct.c_void_p),
rbuf.ctypes.data_as(ct.c_void_p),
gbuf.ctypes.data_as(ct.c_void_p),
bbuf.ctypes.data_as(ct.c_void_p),
)
zbuf.resize((h, w))
msk = (zbuf > 1e-8).astype('uint8')
if len(np.where(msk.flatten() > 0)[0]) < 500:
return None
zbuf *= msk.astype(zbuf.dtype) # * 1000.0
bbuf.resize((h, w)), rbuf.resize((h, w)), gbuf.resize((h, w))
bgr = np.concatenate((bbuf[:, :, None], gbuf[:, :, None], rbuf[:, :, None]), axis=2)
bgr = bgr.astype('uint8')
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
if args.vis:
imshow("bgr", bgr.astype("uint8"))
show_zbuf = zbuf.copy()
min_d, max_d = show_zbuf[show_zbuf > 0].min(), show_zbuf.max()
show_zbuf[show_zbuf > 0] = (show_zbuf[show_zbuf > 0] - min_d) / (max_d - min_d) * 255
show_zbuf = show_zbuf.astype(np.uint8)
imshow("dpt", show_zbuf)
show_msk = (msk / msk.max() * 255).astype("uint8")
imshow("msk", show_msk)
waitKey(0)
data = {}
data['depth'] = zbuf
data['rgb'] = rgb
data['mask'] = msk
data['K'] = K
data['RT'] = RT
data['cls_typ'] = args.obj_name
data['rnd_typ'] = 'render'
kps, des = extractor.detectAndCompute(bgr, None)
kp_xys = np.array([kp.pt for kp in kps]).astype(np.int32)
kp_idxs = (kp_xys[:, 1], kp_xys[:, 0])
dpt_xyz = img_pcld_utils.dpt_2_cld(zbuf, 1.0, K)
kp_x = dpt_xyz[:, :, 0][kp_idxs][..., None]
kp_y = dpt_xyz[:, :, 1][kp_idxs][..., None]
kp_z = dpt_xyz[:, :, 2][kp_idxs][..., None]
kp_xyz = np.concatenate((kp_x, kp_y, kp_z), axis=1)
# filter by dpt (pcld)
kp_xyz, msk = img_pcld_utils.filter_pcld(kp_xyz)
kps = [kp for kp, valid in zip(kps, msk) if valid] # kps[msk]
des = des[msk, :]
# 6D pose of object in cv camer coordinate system
# transform to object coordinate system
kp_xyz = (kp_xyz - RT[:3, 3]).dot(RT[:3, :3])
dpt_xyz = dpt_xyz[dpt_xyz[:, :, 2] > 0, :]
dpt_pcld = (dpt_xyz.reshape(-1, 3) - RT[:3, 3]).dot(RT[:3, :3])
data['kp_xyz'] = kp_xyz
data['dpt_pcld'] = dpt_pcld
return data
def cal_frame_poses_lm(
pcld, mask, ctr_of, pred_kp_of, use_ctr, n_cls, use_ctr_clus_flter, obj_id,
debug=False
):
"""
Calculates pose parameters by 3D keypoints & center points voting to build
the 3D-3D corresponding then use least-squares fitting to get the pose parameters.
"""
n_kps, n_pts, _ = pred_kp_of.size()
pred_ctr = pcld - ctr_of[0]
pred_kp = pcld.view(1, n_pts, 3).repeat(n_kps, 1, 1) - pred_kp_of
radius = 0.08
if use_ctr:
cls_kps = torch.zeros(n_cls, n_kps+1, 3).cuda()
else:
cls_kps = torch.zeros(n_cls, n_kps, 3).cuda()
pred_pose_lst = []
cls_id = 1
cls_msk = mask == cls_id
if cls_msk.sum() < 1:
pred_pose_lst.append(np.identity(4)[:3, :])
else:
cls_voted_kps = pred_kp[:, cls_msk, :]
ms = MeanShiftTorch(bandwidth=radius)
ctr, ctr_labels = ms.fit(pred_ctr[cls_msk, :])
if ctr_labels.sum() < 1:
ctr_labels[0] = 1
if use_ctr:
cls_kps[cls_id, n_kps, :] = ctr
if use_ctr_clus_flter:
in_pred_kp = cls_voted_kps[:, ctr_labels, :]
else:
in_pred_kp = cls_voted_kps
for ikp, kps3d in enumerate(in_pred_kp):
cls_kps[cls_id, ikp, :], _ = ms.fit(kps3d)
# visualize
if debug:
show_kp_img = np.zeros((480, 640, 3), np.uint8)
kp_2ds = bs_utils.project_p3d(
cls_kps[cls_id].cpu().numpy(), 1000.0, K='linemod'
)
# print("cls_id = ", cls_id)
# print("kp3d:", cls_kps[cls_id])
# print("kp2d:", kp_2ds, "\n")
color = (0, 0, 255) # bs_utils.get_label_color(cls_id.item())
show_kp_img = bs_utils.draw_p2ds(show_kp_img, kp_2ds, r=3, color=color)
imshow("kp: cls_id=%d" % cls_id, show_kp_img)
waitKey(0)
mesh_kps = bs_utils_lm.get_kps(obj_id, ds_type="linemod")
if use_ctr:
mesh_ctr = bs_utils_lm.get_ctr(obj_id, ds_type="linemod").reshape(1, 3)
mesh_kps = np.concatenate((mesh_kps, mesh_ctr), axis=0)
# mesh_kps = torch.from_numpy(mesh_kps.astype(np.float32)).cuda()
pred_RT = best_fit_transform(
mesh_kps,
cls_kps[cls_id].squeeze().contiguous().cpu().numpy()
)
pred_pose_lst.append(pred_RT)
return pred_pose_lst
def cal_view_pred_pose(model, data, epoch=0, obj_id=-1):
model.eval()
with torch.set_grad_enabled(False):
cu_dt = [item.to("cuda", non_blocking=True) for item in data]
# rgb:[1,3,480,640] pcld:[1,122881,3]
rgb, pcld, cld_rgb_nrm, choose, kp_targ_ofst, ctr_targ_ofst, \
cls_ids, rts, labels, kp_3ds, ctr_3ds = cu_dt
pred_kp_of, pred_rgbd_seg, pred_ctr_of = model(
cld_rgb_nrm, rgb, choose
)
_, classes_rgbd = torch.max(pred_rgbd_seg, -1)
# ----------------------------------------------------------------------------------------------------------
# 3D visualization
from copy import deepcopy
def vis3d(cld_rgb_nrm):
cld_rgb_nrm = deepcopy(cld_rgb_nrm)
np_xyz = cld_rgb_nrm.cpu().numpy()[0][:, :3]
np_rgb = cld_rgb_nrm.cpu().numpy()[0][:, 3:6] / 255.
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(np_xyz)
pcd.colors = o3d.utility.Vector3dVector(np_rgb)
o3d.visualization.draw_geometries([pcd])
# vis3d(cld_rgb_nrm)
# ----------------------------------------------------------------------------------------------------------
if args.dataset == "ycb":
pred_cls_ids, pred_pose_lst = cal_frame_poses(
pcld[0], classes_rgbd[0], pred_ctr_of[0], pred_kp_of[0], True, # classes_rgbd[0]应该是mask才对
config.n_objects, True
)
else:
pred_pose_lst = cal_frame_poses_lm(
pcld[0], classes_rgbd[0], pred_ctr_of[0], pred_kp_of[0], True,
config.n_objects, False, obj_id
)
pred_cls_ids = np.array([[1]])
np_rgb = rgb.cpu().numpy().astype("uint8")[0].transpose(1, 2, 0).copy()
if args.dataset == "ycb":
np_rgb = np_rgb[:, :, ::-1].copy()
ori_rgb = np_rgb.copy()
for cls_id in cls_ids[0].cpu().numpy():
idx = np.where(pred_cls_ids == cls_id)[0]
if len(idx) == 0:
continue
pose = pred_pose_lst[idx[0]]
if args.dataset == "ycb":
obj_id = int(cls_id[0])
mesh_pts = bs_utils.get_pointxyz(obj_id, ds_type=args.dataset).copy()
mesh_pts = np.dot(mesh_pts, pose[:, :3].T) + pose[:, 3]
if args.dataset == "ycb":
K = config.intrinsic_matrix["ycb_K1"]
else:
K = config.intrinsic_matrix["linemod"]
# ----------------------------------------------------------------------------------------------------------
# 3D visualization
from copy import deepcopy
def vis3d(cld_rgb_nrm):
cld_rgb_nrm = deepcopy(cld_rgb_nrm)
np_xyz = cld_rgb_nrm.cpu().numpy()[0][:, :3]
np_rgb = cld_rgb_nrm.cpu().numpy()[0][:, 3:6] / 255.
np_rgb[:, [0, 2]] = np_rgb[:, [2, 0]]
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(np_xyz)
pcd.colors = o3d.utility.Vector3dVector(np_rgb)
obj_model = o3d.geometry.PointCloud()
obj_model.points = o3d.utility.Vector3dVector(mesh_pts)
obj_model.paint_uniform_color(color=[0, 0, 1])
o3d.visualization.draw_geometries([pcd, obj_model])
vis3d(cld_rgb_nrm)
# ----------------------------------------------------------------------------------------------------------
mesh_p2ds = bs_utils.project_p3d(mesh_pts, 1.0, K)
color = bs_utils.get_label_color(obj_id, n_obj=22, mode=1)
np_rgb = bs_utils.draw_p2ds(np_rgb, mesh_p2ds, color=color)
vis_dir = os.path.join(config.log_eval_dir, "pose_vis")
ensure_fd(vis_dir)
f_pth = os.path.join(vis_dir, "{}.jpg".format(epoch))
# --------------------------------------------------------------------------------------------------------------
# 2D visualization
# cv2.imwrite(f_pth, np_rgb)
imshow("projected_pose_rgb", np_rgb)
# imshow("ori_rgb", ori_rgb)
waitKey(0)
# --------------------------------------------------------------------------------------------------------------
if epoch == 0:
print("\n\nResults saved in {}".format(vis_dir))
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
def cal_frame_poses(
pcld, mask, ctr_of, pred_kp_of, use_ctr, n_cls, use_ctr_clus_flter,
gt_kps, gt_ctrs, debug=False, kp_type='farthest'
):
"""
Calculates pose parameters by 3D keypoints & center points voting to build
the 3D-3D corresponding then use least-squares fitting to get the pose parameters.
"""
n_kps, n_pts, _ = pred_kp_of.size()
pred_ctr = pcld - ctr_of[0]
pred_kp = pcld.view(1, n_pts, 3).repeat(n_kps, 1, 1) - pred_kp_of
radius = 0.08
if use_ctr:
cls_kps = torch.zeros(n_cls, n_kps+1, 3).cuda()
else:
cls_kps = torch.zeros(n_cls, n_kps, 3).cuda()
# Use center clustering filter to improve the predicted mask.
pred_cls_ids = np.unique(mask[mask > 0].contiguous().cpu().numpy())
if use_ctr_clus_flter:
ctrs = []
for icls, cls_id in enumerate(pred_cls_ids):
cls_msk = (mask == cls_id)
ms = MeanShiftTorch(bandwidth=radius)
ctr, ctr_labels = ms.fit(pred_ctr[cls_msk, :])
ctrs.append(ctr.detach().contiguous().cpu().numpy())
try:
ctrs = torch.from_numpy(np.array(ctrs).astype(np.float32)).cuda()
n_ctrs, _ = ctrs.size()
pred_ctr_rp = pred_ctr.view(n_pts, 1, 3).repeat(1, n_ctrs, 1)
ctrs_rp = ctrs.view(1, n_ctrs, 3).repeat(n_pts, 1, 1)
ctr_dis = torch.norm((pred_ctr_rp - ctrs_rp), dim=2)
min_dis, min_idx = torch.min(ctr_dis, dim=1)
msk_closest_ctr = torch.LongTensor(pred_cls_ids).cuda()[min_idx]
new_msk = mask.clone()
for cls_id in pred_cls_ids:
if cls_id == 0:
break
min_msk = min_dis < config.ycb_r_lst[cls_id-1] * 0.8
update_msk = (mask > 0) & (msk_closest_ctr == cls_id) & min_msk
new_msk[update_msk] = msk_closest_ctr[update_msk]
mask = new_msk
except Exception:
pass
# 3D keypoints voting and least squares fitting for pose parameters estimation.
pred_pose_lst = []
pred_kps_lst = []
for icls, cls_id in enumerate(pred_cls_ids):
if cls_id == 0:
break
cls_msk = mask == cls_id
if cls_msk.sum() < 1:
pred_pose_lst.append(np.identity(4)[:3, :])
pred_kps_lst.append(np.zeros((n_kps+1, 3)))
continue
cls_voted_kps = pred_kp[:, cls_msk, :]
ms = MeanShiftTorch(bandwidth=radius)
ctr, ctr_labels = ms.fit(pred_ctr[cls_msk, :])
if ctr_labels.sum() < 1:
ctr_labels[0] = 1
if use_ctr:
cls_kps[cls_id, n_kps, :] = ctr
if use_ctr_clus_flter:
in_pred_kp = cls_voted_kps[:, ctr_labels, :]
else:
in_pred_kp = cls_voted_kps
for ikp, kps3d in enumerate(in_pred_kp):
cls_kps[cls_id, ikp, :], _ = ms.fit(kps3d)
# visualize
if debug:
show_kp_img = np.zeros((480, 640, 3), np.uint8)
kp_2ds = bs_utils.project_p3d(cls_kps[cls_id].cpu().numpy(), 1000.0)
color = bs_utils.get_label_color(cls_id.item())
show_kp_img = bs_utils.draw_p2ds(show_kp_img, kp_2ds, r=3, color=color)
imshow("kp: cls_id=%d" % cls_id, show_kp_img)
waitKey(0)
# Get mesh keypoint & center point in the object coordinate system.
# If you use your own objects, check that you load them correctly.
mesh_kps = bs_utils.get_kps(cls_lst[cls_id-1], kp_type=kp_type)
if use_ctr:
mesh_ctr = bs_utils.get_ctr(cls_lst[cls_id-1]).reshape(1, 3)
mesh_kps = np.concatenate((mesh_kps, mesh_ctr), axis=0)
pred_kpc = cls_kps[cls_id].squeeze().contiguous().cpu().numpy()
pred_RT = best_fit_transform(mesh_kps, pred_kpc)
pred_kps_lst.append(pred_kpc)
pred_pose_lst.append(pred_RT)
return (pred_cls_ids, pred_pose_lst, pred_kps_lst)
def gen_pack_zbuf_render(self):
pth_lst = []
for idx, RT in tqdm(enumerate(self.RT_lst)):
h, w = self.h, self.w
R, T = RT[:, :3], RT[:, 3]
K = self.K
new_xyz = self.xyz.copy()
new_xyz = np.dot(new_xyz, R.T) + T
p2ds = np.dot(new_xyz.copy(), K.T)
p2ds = p2ds[:, :2] / p2ds[:, 2:]
p2ds = np.require(p2ds.flatten(), 'float32', 'C')
zs = np.require(new_xyz[:, 2].copy(), 'float32', 'C')
zbuf = np.require(np.zeros(h * w), 'float32', 'C')
rbuf = np.require(np.zeros(h * w), 'int32', 'C')
gbuf = np.require(np.zeros(h * w), 'int32', 'C')
bbuf = np.require(np.zeros(h * w), 'int32', 'C')
xyzs = np.require(new_xyz.flatten(), 'float32', 'C')
self.dll.rgbzbuffer(
ct.c_int(h),
ct.c_int(w),
p2ds.ctypes.data_as(ct.c_void_p),
new_xyz.ctypes.data_as(ct.c_void_p),
zs.ctypes.data_as(ct.c_void_p),
self.r.ctypes.data_as(ct.c_void_p),
self.g.ctypes.data_as(ct.c_void_p),
self.b.ctypes.data_as(ct.c_void_p),
ct.c_int(self.n_face),
self.face.ctypes.data_as(ct.c_void_p),
zbuf.ctypes.data_as(ct.c_void_p),
rbuf.ctypes.data_as(ct.c_void_p),
gbuf.ctypes.data_as(ct.c_void_p),
bbuf.ctypes.data_as(ct.c_void_p),
)
zbuf.resize((h, w))
msk = (zbuf > -1e8).astype('uint8')
if len(np.where(msk.flatten() > 0)[0]) < 500:
continue
zbuf *= msk.astype(zbuf.dtype) # * 1000.0
bbuf.resize((h, w)), rbuf.resize((h, w)), gbuf.resize((h, w))
bgr = np.concatenate(
(bbuf[:, :, None], gbuf[:, :, None], rbuf[:, :, None]), axis=2)
bgr = bgr.astype('uint8')
bg = None
len_bg_lst = len(self.bg_img_pth_lst)
while bg is None or len(bg.shape) < 3:
bg_pth = self.bg_img_pth_lst[randint(0, len_bg_lst - 1)]
bg = cv2.imread(bg_pth)
if len(bg.shape) < 3:
continue
bg_h, bg_w, _ = bg.shape
if bg_h < h or bg_w < w:
bg = None
continue
else:
sh = sw = 0
sh, sw = randint(0, bg_h - h), randint(0, bg_w - w)
if sh + h > bg_h:
sh -= 1
if sw + w > bg_w:
sw -= 1
bg = bg[sh:sh + h, sw:sw + w, :]
msk_3c = np.repeat(msk[:, :, None], 3, axis=2)
bgr = bg * (msk_3c <= 0).astype(bg.dtype) + bgr * (msk_3c).astype(
bg.dtype)
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
data = {}
data['depth'] = zbuf
data['rgb'] = rgb
data['mask'] = msk
data['K'] = self.K
data['RT'] = RT
data['cls_typ'] = self.cls_type
data['rnd_typ'] = 'render'
data_str = pkl.dumps(data)
sv_pth = os.path.join(self.render_dir, "{}.pkl".format(idx))
if DEBUG:
imshow("rgb", rgb[:, :, ::-1].astype("uint8"))
imshow("depth", (zbuf / zbuf.max() * 255).astype("uint8"))
imshow("mask", (msk / msk.max() * 255).astype("uint8"))
waitKey(0)
pkl.dump(data, open(sv_pth, "wb"))
pth_lst.append(os.path.abspath(sv_pth))
plst_pth = os.path.join(self.render_dir, "file_list.txt")
with open(plst_pth, 'w') as of:
for pth in pth_lst:
print(pth, file=of)
from dataloder import get_loader
from psgan.inference import Inference
from setup import setup_config, setup_argparser
import numpy as np
import neupeak.utils.webcv2 as cv2
args = setup_argparser().parse_args()
config = setup_config(args)
loader = get_loader(config)
inference = Inference(config)
for source_input, reference_input in loader:
ret = inference.solver.test(*source_input, *reference_input)
source = (source_input[0].squeeze(0).squeeze(0).numpy().transpose(1, 2, 0)
+ 1) / 2
reference = (reference_input[0].squeeze(0).squeeze(0).numpy().transpose(
1, 2, 0) + 1) / 2
mask_s = (source_input[1][0, :, 0].numpy().transpose(1, 2, 0) * 255 +
0.5).astype(np.uint8)
mask_r = (
reference_input[1][0, :, 0].squeeze(2).numpy().transpose(1, 2, 0) * 255
+ 0.5).astype(np.uint8)
cv2.imshow("source", source[..., ::-1])
cv2.imshow("reference", reference[..., ::-1])
cv2.imshow("mask_s", mask_s)
cv2.imshow("mask_r", mask_r)
cv2.imshow("ret", np.asarray(ret)[..., ::-1])
cv2.waitKey()
def show_depth(name, dpt):
dpt = (dpt / np.max(dpt) * 255).astype(np.uint8)
imshow(name, dpt)
def gen_pack_zbuf_render(self):
pth_lst = []
for idx, RT in tqdm(enumerate(self.RT_lst)):
h, w = self.h, self.w
R, T = RT[:, :3], RT[:, 3]
K = self.K
new_xyz = self.xyz.copy()
new_xyz = np.dot(new_xyz, R.T) + T
p2ds = np.dot(new_xyz.copy(), K.T)
p2ds = p2ds[:, :2] / p2ds[:, 2:]
p2ds = np.require(p2ds.flatten(), 'float32', 'C')
zs = np.require(new_xyz[:, 2].copy(), 'float32', 'C')
zbuf = np.require(np.zeros(h * w), 'float32', 'C')
rbuf = np.require(np.zeros(h * w), 'int32', 'C')
gbuf = np.require(np.zeros(h * w), 'int32', 'C')
bbuf = np.require(np.zeros(h * w), 'int32', 'C')
self.dll.rgbzbuffer(
ct.c_int(h),
ct.c_int(w),
p2ds.ctypes.data_as(ct.c_void_p),
new_xyz.ctypes.data_as(ct.c_void_p),
zs.ctypes.data_as(ct.c_void_p),
self.r.ctypes.data_as(ct.c_void_p),
self.g.ctypes.data_as(ct.c_void_p),
self.b.ctypes.data_as(ct.c_void_p),
ct.c_int(self.n_face),
self.face.ctypes.data_as(ct.c_void_p),
zbuf.ctypes.data_as(ct.c_void_p),
rbuf.ctypes.data_as(ct.c_void_p),
gbuf.ctypes.data_as(ct.c_void_p),
bbuf.ctypes.data_as(ct.c_void_p),
)
zbuf.resize((h, w))
msk = (zbuf > 1e-8).astype('uint8')
if len(np.where(msk.flatten() > 0)[0]) < 500:
continue
zbuf *= msk.astype(zbuf.dtype) # * 1000.0
bbuf.resize((h, w)), rbuf.resize((h, w)), gbuf.resize((h, w))
bgr = np.concatenate(
(bbuf[:, :, None], gbuf[:, :, None], rbuf[:, :, None]), axis=2)
bgr = bgr.astype('uint8')
bg = None
len_bg_lst = len(self.bg_img_pth_lst)
while bg is None or len(bg.shape) < 3:
bg_pth = self.bg_img_pth_lst[randint(0, len_bg_lst - 1)]
bg = cv2.imread(bg_pth)
if len(bg.shape) < 3:
bg = None
continue
bg_h, bg_w, _ = bg.shape
if bg_h < h:
new_w = int(float(h) / bg_h * bg_w)
bg = cv2.resize(bg, (new_w, h))
bg_h, bg_w, _ = bg.shape
if bg_w < w:
new_h = int(float(w) / bg_w * bg_h)
bg = cv2.resize(bg, (w, new_h))
bg_h, bg_w, _ = bg.shape
if bg_h > h:
sh = randint(0, bg_h - h)
bg = bg[sh:sh + h, :, :]
bg_h, bg_w, _ = bg.shape
if bg_w > w:
sw = randint(0, bg_w - w)
bg = bg[:, sw:sw + w, :]
msk_3c = np.repeat(msk[:, :, None], 3, axis=2)
bgr = bg * (msk_3c <= 0).astype(bg.dtype) + bgr * (msk_3c).astype(
bg.dtype)
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
if args.vis:
try:
from neupeak.utils.webcv2 import imshow, waitKey
except ImportError:
from cv2 import imshow, waitKey
imshow("bgr", bgr.astype("uint8"))
show_zbuf = zbuf.copy()
min_d, max_d = show_zbuf[show_zbuf > 0].min(), show_zbuf.max()
show_zbuf[show_zbuf > 0] = (show_zbuf[show_zbuf > 0] -
min_d) / (max_d - min_d) * 255
show_zbuf = show_zbuf.astype(np.uint8)
imshow("dpt", show_zbuf)
show_msk = (msk / msk.max() * 255).astype("uint8")
imshow("msk", show_msk)
waitKey(0)
data = {}
data['depth'] = zbuf
data['rgb'] = rgb
data['mask'] = msk
data['K'] = self.K
data['RT'] = RT
data['cls_typ'] = self.cls_type
data['rnd_typ'] = 'render'
sv_pth = os.path.join(self.render_dir, "{}.pkl".format(idx))
if DEBUG:
imshow("rgb", rgb[:, :, ::-1].astype("uint8"))
imshow("depth", (zbuf / zbuf.max() * 255).astype("uint8"))
imshow("mask", (msk / msk.max() * 255).astype("uint8"))
waitKey(0)
pkl.dump(data, open(sv_pth, "wb"))
pth_lst.append(os.path.abspath(sv_pth))
plst_pth = os.path.join(self.render_dir, "file_list.txt")
with open(plst_pth, 'w') as of:
for pth in pth_lst:
print(pth, file=of)
def cal_view_pred_pose(model, data, epoch=0, obj_id=-1):
model.eval()
with torch.set_grad_enabled(False):
cu_dt = {}
# device = torch.device('cuda:{}'.format(args.local_rank))
for key in data.keys():
if data[key].dtype in [np.float32, np.uint8]:
cu_dt[key] = torch.from_numpy(data[key].astype(
np.float32)).cuda()
elif data[key].dtype in [np.int32, np.uint32]:
cu_dt[key] = torch.LongTensor(data[key].astype(
np.int32)).cuda()
elif data[key].dtype in [torch.uint8, torch.float32]:
cu_dt[key] = data[key].float().cuda()
elif data[key].dtype in [torch.int32, torch.int16]:
cu_dt[key] = data[key].long().cuda()
end_points = model(cu_dt)
_, classes_rgbd = torch.max(end_points['pred_rgbd_segs'], 1)
pcld = cu_dt['cld_rgb_nrm'][:, :3, :].permute(0, 2, 1).contiguous()
if args.dataset == "ycb":
pred_cls_ids, pred_pose_lst, _ = cal_frame_poses(
pcld[0], classes_rgbd[0], end_points['pred_ctr_ofs'][0],
end_points['pred_kp_ofs'][0], True, config.n_objects, True,
None, None)
else:
pred_pose_lst = cal_frame_poses_lm(pcld[0], classes_rgbd[0],
end_points['pred_ctr_ofs'][0],
end_points['pred_kp_ofs'][0],
True, config.n_objects, False,
obj_id)
pred_cls_ids = np.array([[1]])
np_rgb = cu_dt['rgb'].cpu().numpy().astype("uint8")[0].transpose(
1, 2, 0).copy()
if args.dataset == "ycb":
np_rgb = np_rgb[:, :, ::-1].copy()
ori_rgb = np_rgb.copy()
for cls_id in cu_dt['cls_ids'][0].cpu().numpy():
idx = np.where(pred_cls_ids == cls_id)[0]
if len(idx) == 0:
continue
pose = pred_pose_lst[idx[0]]
if args.dataset == "ycb":
obj_id = int(cls_id[0])
mesh_pts = bs_utils.get_pointxyz(obj_id,
ds_type=args.dataset).copy()
mesh_pts = np.dot(mesh_pts, pose[:, :3].T) + pose[:, 3]
if args.dataset == "ycb":
K = config.intrinsic_matrix["ycb_K1"]
else:
K = config.intrinsic_matrix["linemod"]
mesh_p2ds = bs_utils.project_p3d(mesh_pts, 1.0, K)
color = bs_utils.get_label_color(obj_id, n_obj=22, mode=2)
np_rgb = bs_utils.draw_p2ds(np_rgb, mesh_p2ds, color=color)
vis_dir = os.path.join(config.log_eval_dir, "pose_vis")
ensure_fd(vis_dir)
f_pth = os.path.join(vis_dir, "{}.jpg".format(epoch))
if args.dataset == 'ycb':
bgr = np_rgb
ori_bgr = ori_rgb
else:
bgr = np_rgb[:, :, ::-1]
ori_bgr = ori_rgb[:, :, ::-1]
cv2.imwrite(f_pth, bgr)
if args.show:
imshow("projected_pose_rgb", bgr)
imshow("original_rgb", ori_bgr)
waitKey()
if epoch == 0:
print("\n\nResults saved in {}".format(vis_dir))
