def getLocalPointsV2(pc, contacts, centers, local_th=0.011, local_pn=100):
pc = pc.squeeze(0) * torch.FloatTensor([0.22 / 2, 0.22 / 2, 0.22]).to(
pc.device)
point_num = pc.size(0)
contacts = contacts.squeeze(0)
con_num = contacts.size(0)
centers = centers.squeeze(0)
interval = 1e3
num = con_num // interval + 1
num = int(num)
delta = (con_num + num - 1) // num
local_points_list = []
for i in range(num):
s = delta * i
e = delta * (i + 1)
if i + 1 == num:
e = max(e, con_num)
# print(con_num, num, e)
contacts_ = contacts[s:e]
centers_ = centers[s:e]
con_cen_vec = nn.functional.normalize(contacts_ - centers_, dim=-1)
con_pc_dist = dist_matrix_torch(contacts_, pc)
con_pc_vec = contacts_.view(-1, 1, 3) - pc.view(1, -1, 3)
con_pc_vec /= con_pc_dist.view(-1, point_num, 1)
dist_ = con_pc_dist.unsqueeze(1) * (1.0 + torch.abs(
con_cen_vec.unsqueeze(1).matmul(con_pc_vec.transpose(1, 2))))
dist = dist_.to('cuda:0')
local_points_ = pu.matrix_k_min(local_th, local_pn, dist).long()
local_points_list.append(local_points_.view(1, -1, local_pn))
del (dist_, dist, local_points_, con_cen_vec, con_pc_dist, con_pc_vec)
local_points = torch.cat(local_points_list, 1)
return local_points
def getTestProposals(obj_pc, grids, contact_index, radius=0.022*np.sqrt(3), \
local_th=0.011, local_pn=100, return_time=False):
st = time.time()
contact_index = contact_index.squeeze(0)
grids = grids.squeeze(0)
obj_pc = obj_pc.squeeze(0) * torch.FloatTensor([0.22 / 2, 0.22 / 2, 0.22
]).to(obj_pc.device)
# print('pc max-min:', obj_pc.max(0), obj_pc.min(0))
contact = obj_pc[contact_index]
# dist_matrix = dist_matrix_torch(center, grids)
point_dist = dist_matrix_torch(contact, obj_pc)
con_num = contact_index.size(0)
grid_num = grids.size(0)
pn_num = obj_pc.size(0)
contact_exp = contact.view(-1, 1, 3).expand(-1, grid_num, -1)
grids_exp = grids.view(1, -1, 3).expand(con_num, -1, -1)
pairs_ = torch.stack([contact_exp, grids_exp], 2).view(-1, 2,
3).unsqueeze(0)
pairs = pairs_.cpu()
del pairs_
pg_vec = contact_exp - grids_exp
pg_vec = pg_vec / torch.sqrt(torch.sum(pg_vec**2, -1, keepdim=True))
obj_pc_exp = obj_pc.view(1, -1, 3) #.expand(con_num, -1, -1)
pp_vec = obj_pc_exp - contact.view(-1, 1, 3)
point_dist_view = point_dist.view(con_num, -1, 1)
pp_vec = pp_vec / point_dist_view
del (obj_pc_exp, contact_exp, grids_exp, point_dist)
data_num = grid_num * con_num * pn_num
num = data_num // 7e8 + 1
num = int(num)
# print(data_num, num)
delta = (con_num + num - 1) // num
local_points_list = []
t1 = time.time() - st
st = time.time()
for i in range(num):
s = delta * i
e = delta * (i + 1)
if i + 1 == num:
e = max(e, con_num)
# print(data_num, num, e)
dist_ = point_dist_view[s:e].transpose(1, 2) * (
1.0 + torch.abs(pg_vec[s:e].matmul(pp_vec[s:e].transpose(1, 2))))
dist = dist_.to('cuda:0')
local_points = pu.matrix_k_min(local_th, local_pn, dist).long()
local_points_list.append(local_points.cpu())
del (dist_, dist)
t2 = time.time() - st
t = t1 + t2 / num
local_points = torch.cat(local_points_list, 0).view(1, -1, local_pn).long()
assert local_points.size(1) == con_num * grid_num, local_points.size(1)
del (pg_vec, pp_vec)
if not return_time:
return pairs, local_points
else:
return pairs, local_points, t
def getTestProposalsV3(obj_pc,
grids,
contact_index,
local_th=0.011,
local_pn=100,
grid_th=0.0425 + 0.022 * np.sqrt(3)):
contact_index = contact_index.squeeze(0)
grids = grids.squeeze(0)
grid_num = grids.size(0)
obj_pc = obj_pc.squeeze(0) * torch.FloatTensor([0.22 / 2, 0.22 / 2, 0.22
]).to(obj_pc.device)
contact = obj_pc[contact_index]
point_dist = dist_matrix_torch(contact, obj_pc)
contact_grid_dist = dist_matrix_torch(contact, grids)
# contact_grid_dist = contact_grid_dist_.to('cuda:0')
# del (contact_grid_dist_)
con_num = contact_index.size(0)
pn_num = obj_pc.size(0)
local_points_list = []
pairs_list = []
for i in range(grid_num):
g = grids[i]
c_idx = torch.nonzero(contact_grid_dist[:, i] < grid_th).view(-1)
if len(c_idx) == 0:
continue
cons = contact[c_idx] # (c num, 3)
point_dist_ = point_dist[c_idx].unsqueeze(-1) # (c num, pc num, 1)
pp_vec = obj_pc.view(1, -1, 3) - cons.view(-1, 1,
3) # (c num, pc num, 3)
pp_vec = pp_vec / point_dist_
pg_vec = cons - g # (c num, 3)
pg_vec = pg_vec / torch.sqrt(torch.sum(pg_vec**2, -1, keepdim=True))
dist_ = point_dist_.transpose(1, 2) * (1.0 + torch.abs(
pg_vec.unsqueeze(1).matmul(pp_vec.transpose(1, 2))))
dist = dist_.to('cuda:0')
local_points = pu.matrix_k_min(local_th, local_pn, dist).long()
local_points_list.append(local_points)
del (dist_, dist)
pair = torch.stack([cons, g.view(1, -1).expand(cons.size(0), -1)], 1)
pairs_list.append(pair)
pairs_ = torch.cat(pairs_list, 0)
pairs = pairs_.cpu().unsqueeze(0)
# print('pairs size', pairs.size())
local_points = torch.cat(local_points_list, 0).view(1, -1, local_pn).long()
del (pg_vec, pp_vec, pairs_, point_dist_, point_dist)
assert pairs.size(1) == local_points.size(1)
return pairs, local_points
def getLocalPoints(pc, contacts, centers, local_th=0.011, local_pn=100):
pc = pc.squeeze(0) * torch.FloatTensor([0.22 / 2, 0.22 / 2, 0.22]).to(
pc.device)
point_num = pc.size(0)
contacts = contacts.squeeze(0)
con_num = contacts.size(0)
centers = centers.squeeze(0)
con_cen_vec = nn.functional.normalize(contacts - centers, dim=-1)
con_pc_dist = dist_matrix_torch(contacts, pc)
con_pc_vec = contacts.view(-1, 1, 3) - pc.view(1, -1, 3)
con_pc_vec /= con_pc_dist.view(con_num, point_num, 1)
dist_ = con_pc_dist.unsqueeze(1) * (1.0 + torch.abs(
con_cen_vec.unsqueeze(1).matmul(con_pc_vec.transpose(1, 2))))
dist = dist_.to('cuda:0')
local_points = pu.matrix_k_min(local_th, local_pn, dist).long()
# print(dist.size(), local_points.size())
# print(dist.device, local_points.device)
del (dist_, dist)
return local_points.view(1, -1, local_pn)
def getTestProposalsV2(obj_pc,
grids,
contact_index,
radius=0.022 * np.sqrt(3),
local_th=0.011,
local_pn=100,
grid_th=0.0425 + 0.022 * np.sqrt(3),
grid_num=150):
contact_index = contact_index.squeeze(0)
grids = grids.squeeze(0)
obj_pc = obj_pc.squeeze(0) * torch.FloatTensor([0.22 / 2, 0.22 / 2, 0.22
]).to(obj_pc.device)
contact = obj_pc[contact_index]
point_dist = dist_matrix_torch(contact, obj_pc)
contact_grid_dist_ = dist_matrix_torch(contact, grids)
contact_grid_dist = contact_grid_dist_.to('cuda:0')
del (contact_grid_dist_)
st = time.time()
grids_idx_ = pu.matrix_k_min(grid_th, grid_num,
contact_grid_dist.unsqueeze(0))
print('grids select time: ', time.time() - st)
grids_idx = grids_idx_[0].to(obj_pc.device)
grids_idx = grids_idx.unsqueeze(-1).expand(-1, -1, 3).long()
del (contact_grid_dist, grids_idx_)
con_num = contact_index.size(0)
# grid_num = grids.size(0)
pn_num = obj_pc.size(0)
contact_exp = contact.view(-1, 1, 3).expand(-1, grid_num, -1)
grids_exp = grids.view(1, -1, 3).expand(con_num, -1, -1)
grids_exp = torch.gather(grids_exp, 1, grids_idx)
pairs_ = torch.stack([contact_exp, grids_exp], 2).view(-1, 2,
3).unsqueeze(0)
pairs = pairs_.cpu()
del pairs_
pg_vec = contact_exp - grids_exp
pg_vec = pg_vec / torch.sqrt(torch.sum(pg_vec**2, -1, keepdim=True))
obj_pc_exp = obj_pc.view(1, -1, 3) #.expand(con_num, -1, -1)
pp_vec = obj_pc_exp - contact.view(-1, 1, 3)
point_dist_view = point_dist.view(con_num, -1, 1)
pp_vec = pp_vec / point_dist_view
del (obj_pc_exp, contact_exp, grids_exp, point_dist)
data_num = grid_num * con_num * pn_num
num = data_num // 6e8 + 1
# num = int(num)
num = int(num)
# print(data_num, num)
delta = (con_num + num - 1) // num
local_points_list = []
for i in range(num):
s = delta * i
e = delta * (i + 1)
if i + 1 == num:
e = max(e, con_num)
# print(data_num, num, e)
dist_ = point_dist_view[s:e].transpose(1, 2) * (
1.0 + torch.abs(pg_vec[s:e].matmul(pp_vec[s:e].transpose(1, 2))))
dist = dist_.to('cuda:0')
local_points = pu.matrix_k_min(local_th, local_pn, dist).long()
local_points_list.append(local_points)
del (dist_, dist)
local_points = torch.cat(local_points_list, 0).view(1, -1, local_pn).long()
assert local_points.size(1) == con_num * grid_num, local_points.size(1)
del (pg_vec, pp_vec)
return pairs, local_points
def getProposals(obj_pc,
grids,
center,
index,
scores,
data_index,
radius=0.022 * np.sqrt(3),
local_th=0.011,
local_pn=100):
center = center.squeeze(0)
index = index.squeeze(0)
scores = scores.squeeze(0)
grids = grids.squeeze(0)
obj_pc = obj_pc.squeeze(0) * torch.FloatTensor([0.22 / 2, 0.22 / 2, 0.22
]).to(obj_pc.device)
contact = obj_pc[index]
cent_grid_dist_matrix = dist_matrix_torch(center, grids)
point_dist = dist_matrix_torch(contact, obj_pc)
con_num = index.size(0)
grid_num = grids.size(0)
pn_num = obj_pc.size(0)
# get contact-grid pairs
contact_exp = contact.view(-1, 1, 3).expand(-1, grid_num, -1)
grids_exp = grids.view(1, -1, 3).expand(con_num, -1, -1)
pairs_ = torch.stack([contact_exp, grids_exp], 2).view(-1, 2,
3).unsqueeze(0)
pairs = pairs_.cpu()
del pairs_
# get positive and negative proposals
select = (cent_grid_dist_matrix < radius).float()
posi_prop_idx_ = torch.nonzero(select.view(-1)).view(-1)
nega_prop_idx_ = torch.nonzero(select.view(-1) == 0).view(-1)
posi_prop_idx, nega_prop_idx = posi_prop_idx_.cpu(), nega_prop_idx_.cpu()
offsets_ = (grids_exp - center.view(-1, 1, 3)).view(1, -1, 3) / radius
offsets = (offsets_ * select.view(1, -1, 1)).cpu()
del offsets_
# get proposals scores
scores_all_ = scores.view(-1, 1) * select
scores_all_ = select * scores_all_
scores_all = scores_all_.view(1, -1).cpu()
posi_prop_scores = scores_all_.view(-1)[posi_prop_idx_]
posi_idx_ = torch.nonzero(posi_prop_scores).view(
-1) # positive proposals associated to positive grasps
nega_idx_ = torch.nonzero(posi_prop_scores == 0).view(
-1) # positive proposals associated to negative grasps
posi_idx, nega_idx = posi_idx_.cpu(), nega_idx_.cpu()
anti_label = select.view(1, -1).cpu() # proposals labels
del (scores_all_, posi_prop_idx_, nega_prop_idx_, posi_idx_, nega_idx_)
# get local points
pg_vec = contact_exp - grids_exp # vectors from grids to contacts
pg_vec = pg_vec / torch.sqrt(torch.sum(pg_vec**2, -1, keepdim=True))
obj_pc_exp = obj_pc.view(1, -1, 3) #.expand(con_num, -1, -1)
pp_vec = obj_pc_exp - contact.view(
-1, 1, 3) # vectors from contacts to other points
point_dist_view = point_dist.view(con_num, -1, 1)
pp_vec = pp_vec / point_dist_view
del (obj_pc_exp, contact_exp, grids_exp, cent_grid_dist_matrix, point_dist)
data_num = grid_num * con_num * pn_num
num = data_num // 5e8 + 1
num = int(num)
delta = (con_num + num - 1) // num
local_points_list = []
# in case of out of memory
for i in range(num):
s = delta * i
e = delta * (i + 1)
if i + 1 == num:
e = max(e, con_num)
dist_ = point_dist_view[s:e].transpose(1, 2) * (
1.0 + torch.abs(pg_vec[s:e].matmul(pp_vec[s:e].transpose(1, 2))))
dist = dist_.to('cuda:0')
local_points = pu.matrix_k_min(local_th, local_pn, dist).long()
local_points_list.append(local_points)
del (dist_, dist)
local_points = torch.cat(local_points_list, 0).view(1, -1, local_pn).long()
assert local_points.size(1) == con_num * grid_num, local_points.size(1)
data_index = data_index.new(con_num, grid_num).zero_() + data_index.view(
-1, 1)
data_index = data_index.view(1, -1)
del (select, pg_vec, pp_vec)
return pairs, scores_all, offsets, local_points, data_index, anti_label, posi_prop_idx, nega_prop_idx, posi_idx, nega_idx