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