def detect_grasps(q_out, ang_out, wid_out=None, no_grasp=1):
'''
:功能 :从抓取预测处理所得到的位置，角度，宽度映射图中提取no_grasp个最有效的抓取
:参数 q_out :int,抓取质量（位置）映射图
:参数 ang_out :int,抓取角度映射图
:参数 wid_out :int,抓取宽度映射图
:参数 no_grasp :int,想要提取的有效抓取个数
:返回 :list,包含多个grasp_cpaw对象的列表
'''
grasps_pre = []
local_max = peak_local_max(q_out,
min_distance=20,
threshold_abs=0.2,
num_peaks=no_grasp)
for grasp_point_array in local_max:
grasp_point = tuple(grasp_point_array)
grasp_angle = ang_out[grasp_point]
g = Grasp_cpaw(grasp_point, grasp_angle)
if wid_out is not None:
g.width = wid_out[grasp_point]
g.length = g.width / 2
grasps_pre.append(g)
return grasps_pre
def scale_width(gr, scale):
grasp_cpaw = Grasp_cpaw(angle=gr.angle,
center=gr.center,
length=gr.length,
width=gr.width)
grasp_cpaw.width = grasp_cpaw.width * scale
gr = grasp_cpaw.as_gr
return gr
def move_position(gr, offest):
# 先转正
angle = gr.angle
old_center = gr.center
gr.rotate(-angle, gr.center)
grasp_cpaw = Grasp_cpaw(angle=gr.angle,
center=old_center,
length=gr.length,
width=gr.width)
gr = grasp_cpaw.as_gr
# 执行平移
new_center = gr.center + offest
# 再转回来
old_center = gr.center
gr.rotate(angle, gr.center)
grasp_cpaw = Grasp_cpaw(angle=gr.angle,
center=new_center,
length=gr.length,
width=gr.width)
gr = grasp_cpaw.as_gr
return gr
def rotate_angle(gr, angle):
# 记录原始中心
old_center = gr.center
# 执行旋转
gr.rotate(angle, gr.center)
# 将抓取平移到原始的中心
grasp_cpaw = Grasp_cpaw(angle=gr.angle,
center=old_center,
length=gr.length,
width=gr.width)
gr = grasp_cpaw.as_gr
return gr
def get_patch(self,prediction,mask_height):
gt_patches = []
pre_patches = []
prob, cos_img, sin_img, width_img = prediction
# 获取batch size,后面要用
batch_size = prob.size()[0]
ang_out = (torch.atan2(sin_img, cos_img) / 2.0)
width_out= width_img * 150.0
prob_g = T.GaussianBlur(kernel_size = (3,3),sigma = 2)(prob)
ang_g = T.GaussianBlur(kernel_size = (3,3),sigma = 2)(ang_out)
width_g = T.GaussianBlur(kernel_size = (3,3),sigma = 1)(width_out)
position = torch.max(prob_g.reshape(-1,90000),dim = 1)[1]
# 求得8张图中最亮点的中心坐标
x = position % 300
y = position // 300
# 以其为中心,拓展成栅格矩阵
# # 因为上面的采样出来每个点对应五个值,所以这边就重复5次,想要更密的话可以调
# offset_x = torch.arange(-6,7,3).repeat(8*5).cuda()
# offset_y = torch.arange(-6,7,3).repeat_interleave(5).repeat(8).cuda()
# NOTE 上面这里得改,肯定要从最把握最大的中间开始检查,不行的话再检查边上啊,现在这是从边上检查的,浪费太多时间了
steps = 49
offset_x = torch.tensor([0,2,-2,4,-4,6,-6]).repeat(8*7).cuda()
offset_y = torch.tensor([0,2,-2,4,-4,6,-6]).repeat_interleave(7).repeat(8).cuda()
# 扩展x和y并与offset相加
expand_x = x.repeat_interleave(steps) + offset_x
expand_y = y.repeat_interleave(steps) + offset_y
expand_x = torch.clip(expand_x,6,293)
expand_y = torch.clip(expand_y,6,293)
indice0 = torch.arange(0,8).repeat_interleave(steps)
indice1 = indice0.new_full((len(indice0),),0)
# 索引获得每个位置对应的角度和宽度
ang = ang_g[(indice0,indice1,expand_y,expand_x)]
width = width_g[(indice0,indice1,expand_y,expand_x)]
length = width / 2
# 组合参数并分组,为后面碰撞检测做准备
boxes_params = torch.stack((expand_y,expand_x,ang,width,length)).t().reshape(8,-1,5)
# 每张图要转25个不同的角度,然后裁剪出来做检测,一共八张图,也就是说要转200次,只转图像就行了,gr不用转,然后在这个中心来进行裁剪
# 目前只能想到用for循环来做
indexes, batch_edges, batch_edges_left, batch_edges_right, batch_edges_top, batch_edges_bottom,directions = get_indexes(mask_height,boxes_params,batch_size,steps)
# 显示最终选定的抓取 NOTE 调试用
for i, image in enumerate(mask_height):
index = indexes[i]
y = boxes_params[i][index][0].cpu().data.numpy()
x = boxes_params[i][index][1].cpu().data.numpy()
angle = boxes_params[i][index][2].cpu().data.numpy()
width = boxes_params[i][index][3].cpu().data.numpy()
gr = Grasp_cpaw((y,x),angle,width/2,width)
gr_img = show_grasp(image.cpu().data.numpy()[0],gr.as_gr,50)
plt.subplot(2,4,i+1)
plt.imshow(gr_img)
plt.show()
# 对每一张图生成patch
for i in range(batch_size):
# 角度还是用位置更新之前的,因为更新后的位置角度是否合适没有经过验证的
cos_patch = cos_img[i][0][y,x].cpu().data.numpy()
sin_patch = sin_img[i][0][y,x].cpu().data.numpy()
y = int(boxes_params[i][indexes[i]][0].cpu().data.numpy())
x = int(boxes_params[i][indexes[i]][1].cpu().data.numpy())
angle = ang_g[i][0][y,x].cpu().data.numpy()
width_patch = width_img[i][0][y,x].cpu().data.numpy()
selected_edge = batch_edges[i][indexes[i]]
left_right_diff = abs(batch_edges_left[i][indexes[i]]-batch_edges_right[i][indexes[i]])
top_bottom_diff = abs(batch_edges_top[i][indexes[i]]-batch_edges_bottom[i][indexes[i]])
y,x,scale = get_patch_params(i,y,x,selected_edge,width_patch,angle,directions,left_right_diff,top_bottom_diff)
# 先裁剪出原始的预测图
pre_patch = torch.stack([prob[i][0][y-2:y+3,x-2:x+3],
cos_img[i][0][y-2:y+3,x-2:x+3],
sin_img[i][0][y-2:y+3,x-2:x+3],
width_img[i][0][y-2:y+3,x-2:x+3]])
gt_patch = torch.stack([pre_patch[0].new_full((5,5),1),
pre_patch[1].new_full((5,5),float(cos_patch)),
pre_patch[2].new_full((5,5),float(sin_patch)),
pre_patch[3].new_full((5,5),float((width_patch*scale)))])
image = mask_height[i]
index = indexes[i]
angle = boxes_params[i][index][2].cpu().data.numpy()
width = boxes_params[i][index][3].cpu().data.numpy()
old_y = int(boxes_params[i][indexes[i]][0].cpu().data.numpy())
old_x = int(boxes_params[i][indexes[i]][1].cpu().data.numpy())
gr = Grasp_cpaw((old_y,old_x),angle,width/2,width)
gr_img = show_grasp(image.cpu().data.numpy()[0],gr.as_gr,50)
plt.subplot(121)
plt.imshow(gr_img)
plt.subplot(122)
gr = Grasp_cpaw((y,x),angle,width*scale/2,width*scale)
gr_img = show_grasp(image.cpu().data.numpy()[0],gr.as_gr,50)
plt.imshow(gr_img)
plt.show()
for j in range(4):
plt.subplot(2,4,j+1)
plt.imshow(pre_patch[j].cpu().data.numpy())
for j in range(4):
plt.subplot(2,4,j+5)
plt.imshow(gt_patch[j].cpu().data.numpy())
plt.show()
pre_patches.append(pre_patch)
gt_patches.append(gt_patch)
prob_s = torch.cat((mask_height,prob_g,ang_g,width_g),dim = 0)
for i in range(32):
plt.subplot(4,8,i+1)
plt.imshow(prob_s[i][0].cpu().data.numpy())
plt.show()
return torch.stack(pre_patches),torch.stack(gt_patches)
def get_ground_truth(self, target, prediction):
gt_patches = []
pre_patches = []
gt_probs = []
probs = []
# target = ((pos_grid, gt_num, gt_points, gt_angles, gt_widths, gt_lengths), mask, rgb, idxs)
gt_data, gt_mask, gt_rgb, gt_mask_d, gt_mask_prob, idxs = target
for pos_gt, gt_num, gt_boxes, gt_center, gt_angle, gt_width, gt_length, mask, rgb, mask_d1, mask_prob, idx, pos_per_image, cos_per_image, sin_per_image, width_per_image, filter in zip(
*gt_data, gt_mask, gt_rgb, gt_mask_d, gt_mask_prob, idxs,
*prediction):
# 去除多余维度
pos_per_image = pos_per_image.squeeze()
cos_per_image = cos_per_image.squeeze()
sin_per_image = sin_per_image.squeeze()
width_per_image = width_per_image.squeeze()
# 根据位置映射图及filter计算prob分数
prob = pos_per_image * filter.squeeze().sigmoid()
prob = torch.clip(prob, min=0, max=1)
# 从预测图中反求出预测抓取 boxes, ang_img,width_img = img2boxes(cos_per_image, sin_per_image, width_per_image)
# 获取mask信息
mask = (mask / 230).to(torch.uint8).to(pos_per_image.device)
# 生成辅助采样栅格
canvas = torch.zeros((300, 300))
total = torch.sum(mask).float() # 获取mask总面积
step = (torch.sqrt(total) // 10) # 采样100个点就是sqrt(100) = 10,其他的自行设置
row = torch.arange(0, 300, max(step, 2)).to(torch.int64).unique()
col = torch.arange(0, 300, max(step, 2)).to(torch.int64).unique()
xx, yy = torch.meshgrid(row, col)
canvas[xx, yy] = 1.0
# 转移设备, 计算最终mask时用
canvas = canvas.to(pos_per_image.device)
pos_gt = pos_gt.to(pos_per_image.device)
# 找到mask及标注位置的索引值
mask_g = (mask + pos_gt + canvas) // 2
mask_index = torch.where(mask_g == 1)
# 反求所有位置的抓取框
centers, ang_img, width_img = img2boxes(cos_per_image,
sin_per_image,
width_per_image)
# index = torch.where(mask_index == quality_index)
# gt_index = gt_iou_index[index]
# NOTE 调试用图像显示
# plt.subplot(241)
# plt.title(idx.numpy())
# plt.imshow(rgb.numpy())
# plt.subplot(242)
# plt.title('grasp_label')
# grs = Grasps()
# grs.grs = [Grasp(points.cpu().data.numpy()) for points in gt_boxes]
# img = show_grasp(copy.deepcopy(rgb.numpy()),grs)
# plt.imshow(img)
# plt.subplot(243)
# plt.title('mask_d')
# plt.imshow(mask_prob.cpu().data.numpy())
# plt.subplot(244)
# plt.title('mask_g')
# plt.imshow(mask_g.cpu().data.numpy())
# plt.subplot(245)
# plt.title('pos_img')
# plt.imshow(pos_per_image.cpu().data.numpy())
# plt.subplot(246)
# plt.title('prob')
# plt.imshow(prob.cpu().data.numpy())
# 显示prob获得的最优抓取
# NOTE 调试用图像显示
prob_n = gaussian(prob.cpu().data.numpy(),
2.0,
preserve_range=True)
ang_img_n = gaussian(ang_img.cpu().data.numpy(),
2.0,
preserve_range=True)
width_img_n = gaussian(width_img.cpu().data.numpy(),
1.0,
preserve_range=True)
# prob_n = prob.cpu().data.numpy()
# ang_img_n = ang_img.cpu().data.numpy()
# width_img_n = width_img.cpu().data.numpy()
position = np.argmax(prob_n)
x = (position % 300)
y = (position // 300)
prob_ang = ang_img_n[y, x]
prob_width = width_img_n[y, x]
grasp = Grasp_cpaw((y, x), prob_ang, prob_width / 2,
prob_width).as_gr
# NOTE 调试用图像显示
# img = show_grasp(copy.deepcopy(rgb.numpy()),Grasps([grasp]))
# plt.subplot(247)
# plt.title('iou')
# plt.imshow(img)
# plt.subplot(248)
# plt.imshow(img)
# plt.show()
# NOTE 调试用图像显示
# TODO:在这下面加一个抓取合理性检测模块
#先检查一下这个区域
mask_n = copy.deepcopy(mask_d1.cpu().data.numpy())
y1 = y
x1 = x
edge = edge_check_new(
torch.from_numpy(mask_n).cuda(),
(grasp.center, grasp.angle, grasp.width, grasp.width / 2))
selected_width = edge[2]
edge_widths = []
edge_widths_left = []
edge_widths_right = []
coords = []
if selected_width < 5:
for i in range(-7, 8, 3):
for j in range(-7, 8, 3):
xo = min(x + i, 296)
yo = min(y + j, 296)
x1 = max(xo, 3)
y1 = max(yo, 3)
prob_ang = ang_img_n[y1, x1]
prob_width = width_img_n[y1, x1]
grasp = Grasp_cpaw((y1, x1), prob_ang, prob_width / 2,
prob_width).as_gr
coords.append((y1, x1))
mask_n = copy.deepcopy(mask_d1.cpu().data.numpy())
edge = edge_check_new(
torch.from_numpy(mask_n).cuda(),
(grasp.center, grasp.angle, grasp.width,
grasp.width / 2))
edge_widths_left.append(edge[0] * prob_width)
edge_widths_right.append(edge[1] * prob_width)
edge_widths.append(edge[2])
# NOTE for debug delete later
# plt.subplot(121)
# img = show_grasp(copy.deepcopy(mask_n),Grasps([grasp]))
# plt.title(edge)
# plt.imshow(img)
# plt.subplot(122)
# img = show_grasp(copy.deepcopy(rgb.numpy()),Grasps([grasp]))
# plt.title(edge)
# plt.imshow(img)
# plt.show()
# NOTE for debug delete later
if edge[2] * prob_width * 0.01 > 3:
break
# edge = edge_check_new(torch.from_numpy(mask_n).cuda(),(grasp.center,grasp.angle,grasp.width,grasp.width/2))
# print(1)
if edge[2] * prob_width * 0.01 > 3:
break
index = np.argmax(edge_widths)
# 如果找了一圈也没找到哪个无碰撞就找一个相对最好的
if np.max(edge_widths) == 0:
edge_widths_lr = (edge_widths_left + edge_widths_right)
if max(edge_widths_lr) > 0:
index = np.argmax(edge_widths_lr)
if index >= len(edge_widths):
index = index - len(edge_widths)
y1, x1 = coords[index]
selected_width = edge_widths[index]
# NOTE for debug delete later
# prob_ang = ang_img_n[y1, x1]
# prob_width = width_img_n[y1, x1]
# grasp = Grasp_cpaw((y1,x1),prob_ang,prob_width/2,prob_width).as_gr
# mask_n = copy.deepcopy(mask_d1.cpu().data.numpy())
# edge = edge_check_new(torch.from_numpy(mask_n).cuda(),(grasp.center,grasp.angle,grasp.width,grasp.width/2))
# img = show_grasp(copy.deepcopy(rgb.numpy()),Grasps([grasp]))
# plt.imshow(img)
# plt.show() # 显示最终选定的抓取
# NOTE for debug delete later
# 首先把这个抓取中心附近的区域给裁剪出来
cos = cos_per_image.cpu().data.numpy()[y1, x1]
sin = sin_per_image.cpu().data.numpy()[y1, x1]
width = width_per_image.cpu().data.numpy()[y1, x1]
slice_y, slice_x = (y1, x1) # 抓取中心x,y前面已经求出来了
pre_patch = torch.stack([
pos_per_image[slice_y - 2:slice_y + 3,
slice_x - 2:slice_x + 3],
cos_per_image[slice_y - 2:slice_y + 3,
slice_x - 2:slice_x + 3],
sin_per_image[slice_y - 2:slice_y + 3,
slice_x - 2:slice_x + 3],
width_per_image[slice_y - 2:slice_y + 3,
slice_x - 2:slice_x + 3]
])
# 为其生成对应的优化参数
scale = 1
if selected_width * width * 1.50 < 5:
scale = 1.1
elif selected_width * width * 1.50 > 10:
scale = 0.9
else:
gt_probs.append(mask_prob.to(mask.device))
probs.append(prob)
# 释放显存
torch.cuda.empty_cache()
continue
gt_patch = torch.stack([
pre_patch[0].new_full((5, 5), 1), pre_patch[1].new_full(
(5, 5), float(cos)), pre_patch[2].new_full((5, 5),
float(sin)),
pre_patch[3].new_full((5, 5), float((width * scale)))
])
# for i in range(4):
# plt.subplot(241+i)
# plt.imshow(pre_patch[i].cpu().data.numpy())
# for i in range(4):
# plt.subplot(241+4+i)
# plt.imshow(gt_patch[i].cpu().data.numpy())
# plt.show()
# plt.subplot(121)
# plt.imshow(cos_per_image.cpu().data.numpy())
# plt.subplot(122)
# plt.imshow(sin_per_image.cpu().data.numpy())
# plt.show()
pre_patches.append(pre_patch)
gt_patches.append(gt_patch)
# TODO:在这上面进行定向的参数优化生成
gt_probs.append(mask_prob.to(mask.device))
probs.append(prob)
# 释放显存
torch.cuda.empty_cache()
gt_p = torch.stack(gt_probs)
p = torch.stack(probs)
if len(pre_patches) > 0:
try:
pp = torch.stack(pre_patches)
gp = torch.stack(gt_patches)
patch_flag = 1
except:
pp = 0
gp = 0
patch_flag = 0
pass
print('堆叠错误')
else:
pp = 0
gp = 0
patch_flag = 0
# 这里返回一个prob是validate的时候用的,只有batch = 1 才有意义
return prob, gt_p, p, gp, pp, patch_flag